TensorRT-LLM

# TensorRT-LLM 설치 및 모델 변환
# pip install tensorrt-llm

from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.hlapi import BuildConfig
import torch

# 1. 간단한 사용법 (High-Level API)
# Hugging Face 모델 직접 로드 및 최적화
llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    tensor_parallel_size=1,  # GPU 수
    dtype="float16"
)

# 샘플링 파라미터 설정
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256
)

# 추론 실행
prompts = [
    "인공지능의 미래에 대해 설명해주세요.",
    "Python으로 퀵소트 알고리즘을 작성해주세요."
]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(f"입력: {output.prompt}")
    print(f"출력: {output.outputs[0].text}")
    print("-" * 50)


# 2. 고급 사용법 - 엔진 빌드 및 최적화
from tensorrt_llm.builder import Builder
from tensorrt_llm.models import LLaMAForCausalLM

# 빌드 설정 (INT8 양자화 포함)
build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_output_len=512,
    max_beam_width=1,
    strongly_typed=True,
    builder_opt_level=5,  # 최대 최적화
)

# INT8 양자화 설정
build_config.plugin_config.update({
    "use_paged_context_fmha": True,
    "use_fp8_context_fmha": False,  # H100에서 True
    "gemm_plugin": "auto",
    "gpt_attention_plugin": "auto"
})

# 엔진 빌드
builder = Builder()
engine = builder.build_engine(
    model_dir="./llama-7b-hf",
    build_config=build_config,
    output_dir="./llama-7b-engine"
)

print("TensorRT-LLM 엔진 빌드 완료!")


# 3. Triton Inference Server 배포 설정
triton_config = """
name: "llama_tensorrt_llm"
backend: "tensorrtllm"
max_batch_size: 64

model_transaction_policy {
  decoupled: True
}

input [
  {
    name: "text_input"
    data_type: TYPE_STRING
    dims: [ 1 ]
  },
  {
    name: "max_tokens"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }
]

output [
  {
    name: "text_output"
    data_type: TYPE_STRING
    dims: [ -1 ]
  }
]

instance_group [
  {
    count: 1
    kind: KIND_GPU
  }
]
"""
print(triton_config)