TensorRT-LLM 推理优化

import tensorrt as trt

# 创建 Logger
logger = trt.Logger(trt.Logger.WARNING)

# 创建 Builder
builder = trt.Builder(logger)

# 关键属性
# builder.max_batch_size  -- 最大批次大小（已弃用，建议使用动态维度）
# builder.max_workspace_size  -- 最大工作空间大小（已弃用）

# 创建网络定义，flags=1 表示启用显式批次维度
network = builder.create_network(1)

# 添加输入
input_tensor = network.add_input(
    name="input_ids",           # 张量名称
    dtype=trt.int32,            # 数据类型
    shape=(1, -1)               # 形状，-1 表示动态维度
)

# 添加全连接层
fc_layer = network.add_multiply(input_tensor, weight_tensor)

# 标记输出
network.mark_output(fc_layer.get_output(0))

# ONNX Parser
parser = trt.OnnxParser(network, logger)
success = parser.parse_from_file("model.onnx")

# 检查解析结果
if not success:
    for i in range(parser.num_errors):
        print(parser.get_error(i))

from tensorrt_llm import LLM, ModelConfig

# 方式一：使用 LLM 高层 API（推荐，0.12+ 版本）
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",   # HuggingFace 模型ID或本地路径
    tensor_parallel_size=2,               # 张量并行度
    pipeline_parallel_size=1,             # 流水线并行度
    quantization=None,                    # 量化配置，如 "int8-weight-only"
    max_batch_size=32,                    # 最大批次大小
    max_seq_len=4096,                     # 最大序列长度
    kv_cache_type="paged",               # KV Cache 类型: "paged" | "contiguous"
    trust_remote_code=True,              # 是否信任远程代码
)

# 方式二：使用 ModelConfig + 从权重构建
config = ModelConfig(
    model_name="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=2,
    max_batch_size=32,
    max_seq_len=4096,
)

from tensorrt_llm.commands import build

# 使用命令行构建引擎
# trtllm-build --model_dir ./llama-7b-hf \
#              --dtype float16 \
#              --tp_size 2 \
#              --output_dir ./engine_output \
#              --max_batch_size 32 \
#              --max_input_len 1024 \
#              --max_seq_len 4096

# 使用 Python API 构建引擎
from tensorrt_llm.builder import Builder, BuildConfig

builder = Builder()
build_config = BuildConfig(
    max_batch_size=32,       # 最大批次大小
    max_input_len=1024,      # 最大输入长度
    max_seq_len=4096,        # 最大总序列长度
    max_num_tokens=4096,     # 单次最大token数（用于 Inflight Batching）
    max_beam_width=1,        # 最大beam宽度
    tp_size=2,               # 张量并行度
    pp_size=1,               # 流水线并行度
    dtype="float16",         # 计算精度
)

# 构建并保存引擎
engine = builder.build(model_config, build_config)
engine.save("./engine_output")

from tensorrt_llm.runtime import GenerationSession, ModelConfig

# 创建模型配置
model_config = ModelConfig(
    max_batch_size=32,
    max_seq_len=4096,
    num_heads=32,
    num_kv_heads=32,        # GQA 时可小于 num_heads
    hidden_size=4096,
    vocab_size=32000,
    num_layers=32,
    dtype="float16",
)

# 创建生成会话
session = GenerationSession(
    model_config=model_config,
    engine_path="./engine_output",       # 引擎路径
    max_batch_size=32,
)

# 执行推理
output = session.generate(
    input_ids=input_tensor,              # 输入 token IDs，形状 [batch, seq_len]
    sampling_config=sampling_config,     # 采样配置
    max_new_tokens=256,                  # 最大生成 token 数
)

from tensorrt_llm.runtime import SamplingConfig

sampling_config = SamplingConfig(
    end_id=2,                    # EOS token ID
    pad_id=0,                    # PAD token ID
    num_beams=1,                 # Beam Search 宽度（1=贪心）
    temperature=0.7,             # 采样温度，越高越随机
    top_k=50,                    # Top-K 采样：仅从概率最高的K个token中采样
    top_p=0.9,                   # Top-P (nucleus) 采样：从累积概率<=P的token中采样
    length_penalty=1.0,          # 长度惩罚（仅 Beam Search）
    repetition_penalty=1.1,      # 重复惩罚
    min_tokens=1,                # 最小生成长度
)

# 不同采样策略的典型配置
greedy_config = SamplingConfig(end_id=2, pad_id=0, num_beams=1)
creative_config = SamplingConfig(end_id=2, pad_id=0, temperature=0.9, top_p=0.95)
beam_config = SamplingConfig(end_id=2, pad_id=0, num_beams=4, length_penalty=0.6)

import torch
from tensorrt_llm import LLM, SamplingParams

# 第一步：初始化 LLM（自动下载模型 + 构建引擎）
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",   # HuggingFace 模型 ID
    tensor_parallel_size=1,               # 单卡推理
    max_seq_len=4096,                     # 最大序列长度
    kv_cache_type="paged",               # 使用 Paged KV Cache
)

# 第二步：配置采样参数
sampling_params = SamplingParams(
    max_tokens=512,           # 最大生成 token 数
    temperature=0.7,          # 采样温度
    top_p=0.9,               # Top-P 采样
    stop_token_ids=[151645],  # 停止 token（Qwen2 的 <|im_end|>）
)

# 第三步：执行推理
prompts = [
    "请解释什么是TensorRT-LLM？",
    "用Python写一个快速排序算法。",
]
outputs = llm.generate(prompts, sampling_params)

# 第四步：查看结果
for output in outputs:
    print(f"Prompt: {output.prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}")
    print(f"Tokens: {len(output.outputs[0].token_ids)}")
    print("---")

# 第五步：释放资源
del llm
torch.cuda.empty_cache()

# === 离线阶段：构建引擎 ===
from tensorrt_llm.commands import build

# 方式一：命令行（推荐生产环境使用）
# trtllm-build \
#   --model_dir /models/llama-7b-hf \
#   --dtype float16 \
#   --tp_size 2 \
#   --max_batch_size 64 \
#   --max_input_len 2048 \
#   --max_seq_len 4096 \
#   --output_dir /engines/llama-7b-tp2-fp16

# 方式二：Python API
from tensorrt_llm import LLM

llm = LLM(
    model="/models/llama-7b-hf",
    tensor_parallel_size=2,
    max_batch_size=64,
    max_seq_len=4096,
)
llm.save("/engines/llama-7b-tp2-fp16")
del llm

# === 在线阶段：加载引擎并推理 ===
from tensorrt_llm import LLM, SamplingParams

# 直接加载预构建的引擎
llm = LLM(
    model="/engines/llama-7b-tp2-fp16",   # 引擎目录路径
    tensor_parallel_size=2,
)

sampling_params = SamplingParams(max_tokens=256, temperature=0.8)
output = llm.generate("Hello, world!", sampling_params)
print(output[0].outputs[0].text)

from tensorrt_llm import LLM, SamplingParams

# Weight-Only INT8 量化
llm_wo_int8 = LLM(
    model="meta-llama/Llama-2-7b-hf",
    quantization="int8-weight-only",    # 仅权重量化为 INT8
    tensor_parallel_size=1,
    max_seq_len=4096,
)
# 权重从 FP16 (2 bytes/param) 压缩到 INT8 (1 byte/param)，显存减半

# SmoothQuant INT8 量化（权重+激活均量化）
llm_sq_int8 = LLM(
    model="meta-llama/Llama-2-7b-hf",
    quantization="int8-smoothquant",    # SmoothQuant: 权重+激活 INT8
    tensor_parallel_size=1,
    max_seq_len=4096,
)
# 激活也量化为 INT8，计算吞吐更高

# FP8 量化（需要 H100/H200 等 Hopper 架构 GPU）
llm_fp8 = LLM(
    model="meta-llama/Llama-2-7b-hf",
    quantization="fp8",                 # FP8 E4M3 格式
    tensor_parallel_size=1,
    max_seq_len=4096,
)

# 推理
sampling_params = SamplingParams(max_tokens=128, temperature=0.7)
output = llm_wo_int8.generate("解释量子计算的基本原理", sampling_params)
print(output[0].outputs[0].text)

from tensorrt_llm import LLM, SamplingParams

# 2-GPU 张量并行
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,              # 4 卡张量并行
    max_seq_len=4096,
    kv_cache_type="paged",
)

# 推理（自动分配到 4 张 GPU）
sampling_params = SamplingParams(max_tokens=256, temperature=0.8)
outputs = llm.generate(["你好", "Hello"], sampling_params)

for out in outputs:
    print(out.outputs[0].text)

import tensorrt as trt
from tensorrt_llm import Builder, Network
from tensorrt_llm.layers import Attention, MLP
from tensorrt_llm.module import Module

class CustomTransformerBlock(Module):
    def __init__(self, hidden_size, num_heads, intermediate_size):
        super().__init__()
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads

        # QKV 投影层（融合为一个 GEMM）
        self.qkv_proj = ColumnLinear(hidden_size, hidden_size * 3, bias=False)
        # 输出投影层
        self.o_proj = RowLinear(hidden_size, hidden_size, bias=False)
        # MLP 层
        self.mlp = MLP(hidden_size, intermediate_size, bias=False)
        # Layer Norm
        self.input_layernorm = LayerNorm(hidden_size)
        self.post_layernorm = LayerNorm(hidden_size)

    def forward(self, hidden_states, attention_mask, kv_cache):
        # Pre-LN Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        # QKV 融合投影
        qkv = self.qkv_proj(hidden_states)  # [batch, seq, 3*hidden]

        # 分离 Q, K, V
        q, k, v = split(qkv, self.hidden_size, dim=-1)

        # Attention 计算
        attn_output = attention(q, k, v, attention_mask, kv_cache)

        # 输出投影 + 残差连接
        hidden_states = self.o_proj(attn_output) + residual

        # Pre-LN MLP
        residual = hidden_states
        hidden_states = self.post_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states) + residual

        return hidden_states

# 计算模型 + KV Cache 显存需求
# 示例：Llama-2-7B, FP16, max_seq_len=4096, max_batch_size=32

# 模型权重: 7B × 2 bytes = 14 GB
# KV Cache:  2 × num_layers × num_kv_heads × head_dim × max_seq_len × max_batch_size × 2 bytes
#          = 2 × 32 × 32 × 128 × 4096 × 32 × 2 bytes ≈ 68.7 GB

# 这说明 KV Cache 可能比模型权重本身占用更多显存！
# 使用 PagedKVCache + INT8 量化 KV Cache 可大幅降低显存需求