2. TensorRT-LLM 基础

在深入动态批处理调优之前，必须先建立对 TensorRT-LLM 整体框架的清晰认知。TensorRT-LLM 并非一个简单的模型加载器，而是一套面向大语言模型推理的编译型执行引擎。它将用户定义的模型结构（通常来自 PyTorch 或 HuggingFace）转换为高度优化的 TensorRT 引擎，并在运行时通过一组精心设计的组件管理 GPU 资源、KV Cache 与请求调度。理解其分层架构与数据流，是后续所有调优动作（如 inflight batching、KV Cache 复用、显存池调整）的前提。

2.1 整体架构：从模型定义到引擎执行

TensorRT-LLM 的架构可以划分为三个逻辑层次：前端模型定义层、编译优化层与运行时执行层。前端层负责将模型描述为统一的 Module 结构；编译层通过 trtllm-build 命令将模型图转换为可执行的引擎文件；运行时层则加载引擎，并通过 GptSession 与 Executor 对外提供推理服务。下图展示了这一核心架构及推理请求的主要流转路径。

① 前端模型定义层 HuggingFace 模型 Llama / Qwen / GPT TensorRT-LLM Module DecoderLayer / Attention 权重绑定与转换 FP16 / INT8 / FP8 配置参数 max_batch / max_len ② 编译优化层 (trtllm-build) 图优化与层融合 QKV融合 / MLP融合 Kernel 自动调优 Tactic / 算子选择 内存池规划 KV Cache / Workspace 引擎 .engine ③ 运行时执行层 (Runtime) Executor 请求队列管理 动态批调度 KV Cache 管理 (inflight batching) GptSession 上下文阶段 生成阶段 状态跟踪 beam search TensorRT 引擎 CUDA Graph 融合 Kernel 多流执行 量化推理 输出 Token 序列 logits 性能指标 请求输入: prompt + 采样参数 + 请求ID

上图清晰地展示了三个层次之间的依赖关系。前端层将模型结构、权重与配置参数统一交给编译层；编译层经过图优化、算子融合与内存规划后，产出包含可执行 Kernel 序列的引擎文件；运行时层中的 Executor 负责接收外部请求，将其放入队列，并依据当前 GPU 资源状态决定哪些请求可以立即执行。理解这一分层有助于定位性能瓶颈：例如，若吞吐量低但 GPU 利用率不高，问题往往出在运行时调度策略；若 GPU 利用率高但单请求延迟大，则需检查编译层的 Kernel 选择是否最优。

2.2 核心组件：Executor、GptSession 与 KV Cache Manager

在运行时层中，有三个组件对动态批处理调优至关重要。它们各自承担不同职责，且相互协作方式直接影响系统整体吞吐。

  • Executor（执行器）：对外提供 enqueue_request() 与 await_responses() 接口。它内部维护一个请求优先级队列，并周期性触发调度决策。Executor 是动态批处理（inflight batching）策略的实际执行者，它决定哪些请求的 prefill 阶段与 decode 阶段可以交错执行。
  • GptSession（会话管理器）：负责跟踪每个请求的生成状态，包括当前已生成的 token 数、KV Cache 的引用位置、采样参数（temperature、top_p）等。GptSession 将请求拆分为“上下文阶段（prefill）”与“生成阶段（decode）”，并协调引擎执行。
  • KV Cache Manager（缓存管理器）：这是显存利用效率的核心。它维护一个可复用的显存块池，每个块存储特定层、特定 token 位置的 Key 和 Value 张量。动态批处理要求 KV Cache 能够按需分配与释放，因此该管理器采用分页式管理（类似操作系统虚拟内存），将逻辑 KV Cache 映射到物理显存块。
调优视角：在 TensorRT-LLM 中，max_num_tokens 参数决定了 KV Cache 管理器可分配的最大 token 总数。该值设置过小会导致显存利用率不足，设置过大则可能因预留过多显存而限制并发批大小。通常建议根据实际请求长度分布动态调整，而非简单取最大值。

2.3 推理流程：从请求进入到 Token 输出

一次完整的推理请求在 TensorRT-LLM 内部经历以下阶段。理解该流程有助于识别动态批处理插入点。

  1. 请求接收与校验：Executor 接收包含 prompt、采样参数、请求 ID 的请求对象，并检查其合法性（如序列长度是否超过 max_seq_len）。
  2. KV Cache 分配：KV Cache Manager 根据请求的预估长度（通常由 max_new_tokens 决定）分配逻辑块，并映射到物理显存。若显存不足，请求将进入等待队列。
  3. Prefill（上下文处理）：引擎对 prompt 进行并行处理，生成第一个 token 的 logits，并填充 KV Cache。此阶段计算密集，GPU 利用率高。
  4. Decode（自回归生成）：逐 token 生成。每步仅处理一个 token，但需要读取全部历史 KV Cache。此阶段访存密集，GPU 利用率通常较低。
  5. 采样与终止判断：根据 logits 与采样参数选择下一个 token。若达到 max_new_tokens 或触发 EOS，则结束生成。
  6. 资源释放与响应返回：释放该请求占用的 KV Cache 块，将完整生成序列返回给调用方。

动态批处理（inflight batching）的核心思想是：在 decode 阶段，当某些请求因等待采样或显存不足而暂停时，插入其他请求的 prefill 阶段，从而填补 GPU 计算空隙。这要求 Executor 能够精细控制每个请求在 prefill 与 decode 之间的切换时机。

2.4 典型场景与配置示例

以下展示一个最小化的 TensorRT-LLM 引擎构建与推理配置示例，帮助读者将上述组件与具体参数对应起来。

# 1. 构建引擎（编译层）
trtllm-build \
  --model_dir ./llama-7b-hf \
  --output_dir ./engine \
  --max_batch_size 64 \
  --max_input_len 2048 \
  --max_seq_len 4096 \
  --max_num_tokens 8192 \
  --dtype float16 \
  --use_inflight_batching

# 2. 运行时初始化（Python API）
from tensorrt_llm import LLM, SamplingParams

llm = LLM(
    engine_dir="./engine",
    max_batch_size=64,          # 与构建时一致
    max_num_tokens=8192,        # KV Cache 池大小
    enable_inflight_batching=True
)

# 3. 发送请求
prompts = ["解释动态批处理原理", "写一段快速排序代码"]
outputs = llm.generate(prompts, SamplingParams(max_new_tokens=256))

在上述配置中，max_num_tokens=8192 是动态批处理调优的关键旋钮。它表示 KV Cache 池最多容纳 8192 个 token 位置。若请求平均长度为 512 token，则理论上可同时处理约 16 个请求（8192/512）。但实际中，由于 prefill 与 decode 交错，有效并发数会动态变化。

2.5 常见误区与复盘

误区 现象 正确理解
认为 max_batch_size 越大越好 显存溢出或 KV Cache 分配失败 batch 大小受限于 KV Cache 池与计算资源，需与 max_num_tokens 联动调整
忽略 prefill 与 decode 的资源差异 GPU 利用率波动大，存在明显气泡 prefill 计算密集，decode 访存密集，动态批处理正是为了混合两者以填满 GPU
将 KV Cache 视为固定显存开销 长请求导致后续请求排队 KV Cache 是分页管理的，请求结束后立即释放，可被后续请求复用
复盘提示：在调优实践中，一个常见错误是仅调整 max_batch_size 而忽略 max_num_tokens。实际上，后者决定了 KV Cache 池的总容量，是限制并发请求数量的更直接因素。建议先通过 profiling 工具观察实际 KV Cache 峰值使用量，再反推合理的 max_num_tokens 值，避免盲目设置。
← 上一章 📖 返回目录 下一章 →