Оптимизация инференса LLM через TensorRT-LLM

Оптимизация инференса LLM с помощью TensorRT-LLM

TensorRT-LLM — NVIDIA-библиотека для максимально эффективного LLM-инференса на NVIDIA GPU. Если vLLM — удобный production сервер, то TensorRT-LLM — низкоуровневый движок для максимальной производительности на NVIDIA оборудовании. Ускорение: 2–4x по сравнению с vLLM на тех же GPU.

Архитектура и принцип работы

TensorRT-LLM компилирует модель в оптимизированный TensorRT движок:

1. Graph compilation: граф модели компилируется с учётом конкретного GPU (архитектура, VRAM, тензорные ядра)
2. Kernel fusion: несколько операций объединяются в один CUDA-kernel (LayerNorm + Linear, Flash Attention)
3. Quantization: FP8, INT8, INT4 с точными calibration методами
4. In-flight batching: наиболее продвинутая реализация continuous batching

Установка и компиляция модели

# Установка через Docker (рекомендуется)
docker pull nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3

# Или pip
pip install tensorrt-llm --extra-index-url https://pypi.nvidia.com

import tensorrt_llm
from tensorrt_llm.builder import BuildConfig, build_model
from tensorrt_llm.models import LLaMAForCausalLM

# Загрузка HuggingFace модели
hf_model_path = "meta-llama/Llama-3-8b-instruct"

# Конфигурация компиляции
build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_output_len=512,
    max_beam_width=1,               # greedy decoding
    strongly_typed=True,
    plugin_config={
        "gpt_attention_plugin": "float16",
        "gemm_plugin": "float16",
        "rmsnorm_quantization_plugin": False,
        "use_paged_context_fmha": True,    # PagedAttention
        "use_fp8_context_fmha": False,
    }
)

# Компиляция занимает 5-30 минут в зависимости от модели и GPU
engine = build_model(
    model=LLaMAForCausalLM.from_hugging_face(hf_model_path),
    build_config=build_config
)
engine.save("./llama3-8b-engine/")

FP8 Quantization на H100

H100 имеет аппаратную поддержку FP8 — наибольший прирост производительности:

from tensorrt_llm.quantization import QuantAlgo

build_config_fp8 = BuildConfig(
    max_batch_size=128,
    max_input_len=4096,
    max_output_len=1024,
    quant_config=QuantConfig(
        quant_algo=QuantAlgo.FP8,
        kv_cache_quant_algo=QuantAlgo.FP8,  # KV-кеш тоже в FP8
    ),
    plugin_config={
        "use_fp8_context_fmha": True,  # Flash Attention в FP8
        "gemm_plugin": "float16",
    }
)

FP8 на H100: примерно 2x прирост throughput по сравнению с BF16, деградация качества < 0.5% на стандартных бенчмарках.

Интеграция с Triton Inference Server

TensorRT-LLM нативно интегрируется с NVIDIA Triton:

# Структура для Triton
model_repository/
├── ensemble/
│   └── config.pbtxt
├── preprocessing/      # токенизация
│   ├── config.pbtxt
│   └── 1/model.py
├── tensorrt_llm/       # TRT-LLM движок
│   ├── config.pbtxt
│   └── 1/
│       ├── model.engine
│       └── config.json
└── postprocessing/     # детокенизация
    ├── config.pbtxt
    └── 1/model.py

# tensorrt_llm/config.pbtxt
name: "tensorrt_llm"
backend: "tensorrtllm"
max_batch_size: 128

parameters {
  key: "max_beam_width"
  value: { string_value: "1" }
}
parameters {
  key: "executor_worker_path"
  value: { string_value: "/opt/tritonserver/backends/tensorrtllm/trtllmExecutorWorker" }
}
parameters {
  key: "decoding_mode"
  value: { string_value: "top_p_top_k" }
}

Multi-GPU с Tensor Parallelism

# LLaMA-70B на 4xH100
build_config_tp4 = BuildConfig(
    max_batch_size=64,
    max_input_len=8192,
    max_output_len=2048,
    auto_parallel_config=AutoParallelConfig(
        world_size=4,
        gpus_per_node=4,
        shards_along_head=4,       # tensor parallelism
    )
)

# Запуск mpirun для multi-GPU
# mpirun -n 4 python run_inference.py

Сравнение с vLLM

Рекомендация: vLLM для большинства production use cases. TensorRT-LLM — когда нужно максимально использовать NVIDIA GPU (высоконагруженные сервисы, cost optimization на cloud GPU).

Сроки внедрения

День 1–3: Установка TRT-LLM, компиляция первой модели, измерение baseline метрик

Неделя 1–2: Подбор оптимальных параметров компиляции, квантизация, интеграция с Triton

Неделя 3–4: Load testing, monitoring, production deployment

Месяц 2: Оптимизация для конкретных use cases (latency vs throughput), multi-model deployment