TensorRT-LLM 2026：NVIDIA 推理加速终极方案

# Blackwell B200 上的 FP4 推理
import tensorrt_llm as trtllm

# 编译模型为 FP4 精度
builder = trtllm.Builder()
config = builder.create_builder_config(
    precision="fp4",                    # FP4 量化
    plugin_config=trtllm.PluginConfig(
        attention_plugin=True,
        nccl_plugin=True,
        gemm_plugin=True,
        rmsnorm_plugin=True,
        # Blackwell 专属
        transformer_engine=True,
        moe_plugin=True,
    ),
    max_batch_size=256,
    max_input_len=32768,
    max_output_len=4096,
    max_num_tokens=8192,
    use_paged_context_fmha=True,        # Paged Attention
    use_context_fmha=True,              # Flash Attention
    multiple_profiles=True,             # 多优化 Profile
    tensor_parallel=8,                  # 8 路张量并行
    pipeline_parallel=1,
)

# 编译（AOT，提前编译为优化引擎）
engine = builder.build(
    model_dir="Qwen/Qwen3-72B-Instruct",
    config=config,
    output_dir="engines/qwen3-72b-fp4"
)

# 编译后的引擎不可移植，绑定特定 GPU 架构
# 但性能比 JIT 方案高 30-60%

# 步骤 1：从 HuggingFace 模型编译 TensorRT 引擎
import tensorrt_llm as trtllm
from tensorrt_llm.models import QWenForCausalLM

# 加载模型配置
model_config = trtllm.ModelConfig.from_pretrained(
    "Qwen/Qwen3-72B-Instruct"
)

# 编译配置
build_config = trtllm.BuildConfig(
    max_input_len=32768,
    max_output_len=4096,
    max_batch_size=128,
    max_num_tokens=8192,
    opt_batch_size=32,
    opt_input_len=4096,
    
    # 精度配置
    precision="fp8",                    # fp4/fp8/fp16
    int8_kv_cache=True,                # KV Cache 量化
    
    # 插件配置
    plugin_config=trtllm.PluginConfig(
        paged_kv_cache=True,
        attention_plugin=True,
        gemm_plugin=True,
        nccl_plugin=True,
        rmsnorm_plugin=True,
        rotary_plugin=True,
        remove_input_padding=True,     # 移除 padding 优化
    ),
    
    # 并行配置
    tensor_parallel=4,
    pipeline_parallel=1,
    
    # 高级优化
    use_fused_mlp=True,                # MLP 算子融合
    use_fused_qkv=True,                # QKV 融合
    use_dynamic_shape=True,            # 动态形状
    weight_sparsity=True,              # 权重稀疏化
)

# 编译引擎
builder = trtllm.Builder()
engine = builder.build_model(
    model_config=model_config,
    build_config=build_config,
    output_dir="engines/qwen3-72b-fp8-tp4"
)

# 使用 Triton Inference Server 部署
# 模型仓库结构
models/
└── qwen3-72b/
    ├── config.pbtxt
    ├── 1/
    │   └── model.py
    └── engines/
        └── qwen3-72b-fp8-tp4/
            ├── rank0.engine
            ├── rank1.engine
            ├── rank2.engine
            └── rank3.engine

# config.pbtxt
name: "qwen3-72b"
backend: "tensorrtllm"
max_batch_size: 128
input [
  {
    name: "input_ids"
    data_type: TYPE_INT32
    dims: [ -1 ]
  },
  {
    name: "input_lengths"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }
]
output [
  {
    name: "output_ids"
    data_type: TYPE_INT32
    dims: [ -1, -1 ]
  }
]
dynamic_batching {
  preferred_batch_size: [ 4, 8, 16, 32 ]
  max_queue_delay_microseconds: 100000
}
instance_group [
  {
    count: 1
    kind: KIND_GPU
  }
]
parameters [
  {
    key: "tensorrt_llm_model_dir"
    value: { string_value: "/models/qwen3-72b/engines/qwen3-72b-fp8-tp4" }
  },
  {
    key: "max_output_len"
    value: { string_value: "4096" }
  },
  {
    key: "temperature"
    value: { string_value: "0.7" }
  },
  {
    key: "top_p"
    value: { string_value: "0.9" }
  }
]

# 启动 Triton Server
docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v /models:/models \
  nvcr.io/nvidia/tritonserver:25.06-py3 \
  tritonserver --model-repository=/models \
  --backend-directory=/opt/tritonserver/backends \
  --log-verbose=1

import tritonclient.grpc as grpcclient
import numpy as np

class TensorRTLLMClient:
    def __init__(self, url="localhost:8001"):
        self.client = grpcclient.InferenceServerClient(url)
    
    def generate(self, prompt: str, max_tokens: int = 512, 
                 temperature: float = 0.7, stream: bool = True):
        # Tokenize
        input_ids = self.tokenizer.encode(prompt)
        
        inputs = [
            grpcclient.InferInput("input_ids", [1, len(input_ids)], "INT32"),
            grpcclient.InferInput("input_lengths", [1, 1], "INT32"),
        ]
        inputs[0].set_data_from_numpy(np.array([input_ids], dtype=np.int32))
        inputs[1].set_data_from_numpy(np.array([[len(input_ids)]], dtype=np.int32))
        
        outputs = [grpcclient.InferRequestedOutput("output_ids")]
        
        # 推理
        result = self.client.infer(
            model_name="qwen3-72b",
            inputs=inputs,
            outputs=outputs
        )
        
        output_ids = result.as_numpy("output_ids")
        # Detokenize
        return self.tokenizer.decode(output_ids[0])

# 张量并行 + 流水线并行配置
config = trtllm.BuildConfig(
    tensor_parallel=4,      # 4 路张量并行
    pipeline_parallel=2,    # 2 路流水线并行
    # 总共 8 GPU
    
    # 专家并行（MoE 模型）
    moe_config=trtllm.MoEConfig(
        num_experts=256,
        expert_parallel_size=8,
        moe_plugin=True,
    ),
)