3、TensorRT-LLM基础：模型加载、推理流程、引擎构建与序列化

3.2 方法流程：引擎构建与序列化的完整链路

TensorRT-LLM的引擎构建遵循“权重准备 → 网络定义 → 引擎编译 → 序列化”四步流程。首先，通过from_pretrained加载HuggingFace格式的模型权重，并调用save_checkpoint将权重转换为TensorRT-LLM的checkpoint格式。这一步骤会完成权重重排与格式统一，例如将FP32权重转换为FP16或INT8的量化格式。

随后，在build阶段，TensorRT-LLM读取checkpoint，通过模型定义文件（如LLaMAForCausalLM）构建TensorRT网络。此时可以指定并行方式（张量并行、流水线并行）、量化策略（FP8、INT4-AWQ等）、最大序列长度、批处理大小等编译选项。引擎编译完成后，通过engine_buffer = engine.serialize()得到二进制引擎文件，该文件可直接落盘保存，供后续推理加载。

# 典型构建流程（LLaMA-7B 示例）
from tensorrt_llm.models import LLaMAForCausalLM

# 1. 加载HF权重并转换为TRT-LLM checkpoint
llama = LLaMAForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    dtype="float16",
    tensor_parallel_size=2,   # 张量并行度
    quant_config={"quant_algo": "FP8"}
)
llama.save_checkpoint("./llama7b_trtllm_ckpt")

# 2. 构建引擎
from tensorrt_llm.builder import build_engine
engine = build_engine(
    model_dir="./llama7b_trtllm_ckpt",
    max_batch_size=8,
    max_seq_len=4096,
    use_fused_mlp=True,
    use_gpt_attention_plugin="float16"
)

# 3. 序列化保存
with open("llama7b_fp8.engine", "wb") as f:
    f.write(engine.serialize())

3.3 推理流程：从引擎加载到Token生成的内部机制

推理阶段，TensorRT-LLM采用运行时（Runtime）与会话（Session）两层抽象。首先通过EngineBuffer从文件加载引擎，然后创建GenerationSession。每次推理请求，需要准备输入张量：input_ids（形状[batch, seq_len]）、position_ids、attention_mask等。会话内部维护KV Cache，用于缓存历史token的Key和Value，避免重复计算。

生成过程采用自回归循环：每次前向传播只生成下一个token，然后将新token追加到输入序列，更新KV Cache，再继续下一次前向。TensorRT-LLM通过max_new_tokens控制生成长度，通过end_id（如</s>）判断终止。整个循环在GPU上以同步批处理方式执行，多个请求可以共享同一个引擎实例，通过in_flight_batching实现动态批处理。

# 推理调用示例
from tensorrt_llm.runtime import GenerationSession

session = GenerationSession(engine_buffer)
input_ids = tokenizer(["今天天气不错"], return_tensors="pt").input_ids.cuda()

output = session.generate(
    input_ids,
    max_new_tokens=128,
    temperature=0.7,
    top_k=50,
    top_p=0.9,
    end_id=tokenizer.eos_token_id
)
generated_text = tokenizer.decode(output[0])

# 步骤1：安装与导入
pip install tensorrt-llm transformers

# 步骤2：权重转换
from tensorrt_llm.models import GPT2LMHeadModel
model = GPT2LMHeadModel.from_pretrained("gpt2", dtype="float16")
model.save_checkpoint("./gpt2_ckpt")

# 步骤3：构建引擎（使用默认配置）
from tensorrt_llm.builder import build_engine
engine = build_engine(
    model_dir="./gpt2_ckpt",
    max_batch_size=4,
    max_seq_len=1024,
    use_gpt_attention_plugin="float16"
)
with open("gpt2.engine", "wb") as f:
    f.write(engine.serialize())

# 步骤4：加载并推理
from tensorrt_llm.runtime import GenerationSession
from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
session = GenerationSession(engine_buffer=open("gpt2.engine","rb").read())

prompt = "The future of AI is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
output = session.generate(input_ids, max_new_tokens=50, temperature=0.8)
print(tokenizer.decode(output[0]))