TensorRT-LLM 基础：架构、核心组件与推理流程

在深入内存管理与显存优化之前，必须先建立对 TensorRT-LLM 整体框架的清晰认知。只有理解了推理引擎如何组织计算图、如何管理 KV Cache、如何调度 GPU 资源，才能准确判断显存消耗在何处、瓶颈出现在哪个环节。本节将拆解 TensorRT-LLM 的运行时架构、核心组件以及一次完整推理请求的生命周期。

2.1 整体架构：从模型权重到可执行引擎

TensorRT-LLM 的架构可以划分为 编译期（Build Time） 与 运行期（Runtime） 两大阶段。编译期负责将 PyTorch 等框架的模型权重转换为高度优化的 TensorRT 引擎（Plan 文件），运行期则加载该引擎并执行推理。这种“一次编译、多次执行”的设计，使得运行时开销极小，但也意味着显存规划必须在编译期就通过配置项确定下来。

核心编译流程：

# 伪代码示意：构建引擎的核心步骤
from tensorrt_llm import Builder, ModelConfig, EngineConfig

builder = Builder()
model_config = ModelConfig(model_dir="llama-7b-hf")  # 加载 HuggingFace 权重
engine_config = EngineConfig(
    max_batch_size=8,
    max_seq_len=4096,
    max_num_tokens=8192,      # 影响激活显存上限
    kv_cache_config={"free_gpu_memory_fraction": 0.3}  # 预留 KV Cache 比例
)
engine = builder.build_engine(model_config, engine_config)
engine.save("llama7b_trt_engine.plan")

编译期最关键的决定是 静态显存预算。TensorRT-LLM 在构建引擎时，会根据 max_batch_size、max_seq_len、max_num_tokens 等参数预先分配所有层的激活内存缓冲区。这意味着即使实际推理时 batch size 很小，显存占用也可能接近编译时设定的上限。理解这一点，是后续显存调优的基石。

编译期 (Build Time) PyTorch 权重 (HuggingFace 格式) 引擎配置 max_batch_size max_seq_len / KV Cache 策略 TensorRT-LLM Builder 层融合 / 量化 / 图优化 TensorRT 引擎 (.plan) 包含权重 + 优化后的计算图 + 内存池规划 运行期 (Runtime) 引擎加载器 反序列化 .plan 文件 加载 运行时配置 KV Cache 比例 / 显存限制 并行策略 (TP/PP) GptSession 执行器 管理 KV Cache / 调度请求 推理循环 (Decode Loop) Prefill → Decode 迭代 动态 KV Cache 分配与释放 GPU 显存池 权重区 | 激活区 | KV Cache 区 编译期静态规划 + 运行期动态管理 静态规划：显存上限在此确定 动态执行：KV Cache 按需分配

2.2 核心组件：GptSession 与 KV Cache 管理器

运行期最核心的组件是 GptSession。它负责维护整个推理会话的状态，包括输入 token 的序列信息、当前生成进度、以及最重要的——KV Cache 的分配与释放。GptSession 内部通过一个显存池（Memory Pool）来管理 KV Cache，该池的大小在引擎加载时根据配置确定。

KV Cache 管理器的工作机制可以概括为：

  • 按需分配：每个序列（Sequence）在 Prefill 阶段后，根据实际序列长度向缓存池申请对应大小的 KV 块。
  • 分块管理：缓存池被划分为固定大小的块（Block），每个块可存储固定数量 token 的 KV 值。这种分块设计减少了碎片化，并支持 PagedAttention 式的非连续存储。
  • 主动释放：当序列生成结束（遇到 EOS 或达到最大长度）时，其占用的 KV 块立即归还给池子，供其他序列使用。

显存优化关键点：KV Cache 是推理过程中最大的动态显存消费者。对于一个 7B 模型，当 batch size 为 8、序列长度为 2048 时，KV Cache 可能占用超过 4GB 显存。因此，kv_cache_config 中的 free_gpu_memory_fraction 参数直接决定了 KV Cache 的上限，设置不当会导致显存不足（OOM）或利用率低下。

2.3 推理流程：一次请求的完整生命周期

理解一次推理请求如何从文本输入变为 token 输出，有助于定位显存峰值出现的阶段。整个流程分为 Prefill（预填充） 和 Decode（解码） 两个阶段。

阶段 处理方式 显存特征 计算特征
Prefill 并行处理输入序列的所有 token 激活内存峰值高（所有 token 的中间激活） 计算密集，GPU 利用率高
Decode 逐个生成 token，每步只处理 1 个新 token 激活内存低，但 KV Cache 持续增长 访存密集，受显存带宽限制

一个典型的推理请求流程如下：

  1. 请求接入：客户端发送输入文本，TensorRT-LLM 将其 tokenize 为 token ID 序列。
  2. Prefill 阶段：GptSession 将整个输入序列送入引擎，并行计算所有 token 的注意力结果。此时，模型权重 + 输入激活 + 初始 KV Cache 同时驻留显存，形成第一个显存峰值。
  3. KV Cache 分配：Prefill 完成后，根据输入序列长度，向缓存池申请对应的 KV 块，存储所有层的 K 和 V 矩阵。
  4. Decode 循环：进入自回归生成循环。每步只输入上一步生成的 1 个 token，更新 KV Cache（追加新 token 的 KV 值），并采样出下一个 token。此阶段显存占用相对平稳，但 KV Cache 随生成长度线性增长。
  5. 终止与释放：当生成达到 max_new_tokens 或遇到 EOS token，序列结束。GptSession 释放该序列占用的所有 KV 块，返回结果。

常见误区：初学者常误以为显存峰值出现在 Decode 阶段，因为此时序列最长。实际上，Prefill 阶段的激活内存峰值往往更高，尤其是处理长提示词（如 4K 以上上下文）时。激活内存与输入 token 数成正比，而 KV Cache 增长是渐进的。因此，优化显存时需同时关注 max_num_tokens（限制 Prefill 批处理量）和 KV Cache 复用策略。

2.4 案例复盘：一个 7B 模型的显存分配实例

假设我们在一张 24GB 显存的 GPU（如 RTX 4090）上部署 LLaMA-7B 模型，使用 FP16 精度。通过 TensorRT-LLM 的日志，可以观察到显存分配情况：

[TensorRT-LLM] Engine loaded in 2.34s
[TensorRT-LLM] Total weights memory: 13.5 GB
[TensorRT-LLM] Activation memory (max): 1.8 GB
[TensorRT-LLM] KV Cache pool size: 6.2 GB (free_gpu_memory_fraction=0.30)
[TensorRT-LLM] Remaining GPU memory: 2.5 GB (用于运行时上下文等)

从这个日志中可以解读出关键信息：

  • 权重固定占用 13.5GB，这是模型参数本身的大小（7B × 2 字节），无法通过运行时优化减少，只能通过量化（如 INT4/INT8）降低。
  • 激活内存上限 1.8GB，由编译期的 max_num_tokens 决定。如果实际请求的 batch size 和序列长度远小于上限，这部分显存会被浪费。
  • KV Cache 池 6.2GB，由 free_gpu_memory_fraction=0.30 控制，即预留 30% 的显存（约 7.2GB）给 KV Cache，但实际只用了 6.2GB（因为总显存中还需扣除权重和激活）。

如果此时并发请求数增加，导致 KV Cache 池耗尽，TensorRT-LLM 会返回类似 Out of Memory: KV Cache block exhausted 的错误。此时需要调整 free_gpu_memory_fraction 或减少 max_batch_size，而不是盲目增加并发。

复盘结论：TensorRT-LLM 的显存管理是“编译期规划 + 运行期调度”的组合。编译期决定了显存使用的上限和分区比例，运行期则在固定池内动态分配。优化显存的第一步，是理解当前配置下权重、激活、KV Cache 三者的比例关系，再针对瓶颈进行调整。

← 上一章 📖 返回目录 下一章 →