2. 大模型推理基础

在深入 TensorRT-LLM 的部署优化之前，必须先厘清大模型推理的底层计算逻辑。当前主流的大语言模型（如 GPT、LLaMA、Qwen 等）均基于 Transformer 解码器架构，其推理过程并非一次性生成完整文本，而是通过自回归方式逐 token 预测。理解这一过程的核心矛盾——计算密集与访存密集的交替，以及 KV Cache 如何缓解该矛盾，是后续掌握连续批处理、PagedAttention、量化与算子融合的前提。

2.1 Transformer 解码器：推理时的数据流视角

训练时的 Transformer 是并行处理整个序列的，但推理时只能逐个 token 进行。一个典型的解码器层（Decoder Layer）在推理单个 token 时，其内部数据流如下：

  1. 输入嵌入：当前 token 的 ID 经 Embedding 表映射为向量，并与位置编码相加。
  2. 自注意力（Self-Attention）：将输入向量分别乘以三个权重矩阵 W_Q、W_K、W_V，得到 Query、Key、Value。用 Query 与当前序列中所有 token 的 Key 做点积，经 Softmax 归一化后加权求和 Value。
  3. 前馈网络（FFN）：注意力输出经过两层线性变换（通常中间层有 4 倍隐藏层宽度）和激活函数（如 SwiGLU）。
  4. 归一化与残差连接：每个子层前后均有 LayerNorm 和残差连接，最终输出当前 token 的隐状态向量。

该向量经最后的 LM Head（输出投影层）映射为词表大小的 logits，取最大值或按概率采样得到下一个 token。随后，这个新 token 被拼接到输入序列末尾，进入下一轮推理。

关键认知：在推理第 t 个 token 时，自注意力需要与前面 t-1 个 token 的 Key 和 Value 进行交互。若每次重新计算所有历史 token 的 K/V，计算量将随序列长度呈平方级增长，这是完全不可接受的。

Transformer 解码器推理数据流（单 token 步进） 输入 Token ID Embedding + 位置编码 解码器层（Decoder Layer）× N 自注意力（Self-Attention） Q = x·W_Q K = x·W_K V = x·W_V 当前 Q 历史 K (KV Cache) 历史 V (KV Cache) Q·K^T → Softmax → × V 加权聚合 注意力输出 残差连接 前馈网络（FFN） 线性层 → 激活 → 线性层 隐状态 LM Head：隐状态 → 词表 Logits → Softmax → 采样下一个 Token 新 Token 拼接至输入序列，进入下一轮推理（自回归循环） 新 K/V 追加到缓存 （供后续 token 使用）

图中红色虚线回路揭示了自回归推理的本质：当前 token 的输出依赖于之前所有 token 的 K/V 信息。如果不做缓存，每生成一个新 token 都要重新计算整个序列的 K/V，时间复杂度为 O(n²)，而 KV Cache 将这部分历史信息存储下来，使单步生成的时间复杂度降为 O(n)（仅需计算当前 token 的 Q/K/V，并与缓存的 K/V 做注意力）。

2.2 自回归生成：Prefill 与 Decode 两个阶段

一次完整的推理请求被划分为两个计算特征截然不同的阶段：

阶段 处理内容 计算特征 瓶颈类型
Prefill（预填充） 处理用户输入的整个 Prompt（如 512 个 token） 并行计算所有输入 token 的 K/V，并生成第一个输出 token 计算密集（矩阵乘法规模大）
Decode（解码） 逐 token 生成后续输出（每个 step 仅 1 个 token） 每次只计算 1 个 token 的 Q，与缓存中所有 K/V 做注意力 访存密集（权重和 KV Cache 的读取量远大于计算量）

以 LLaMA-7B 为例，假设隐藏层维度为 4096，词表大小为 32000。在 Decode 阶段生成 1 个 token 时，需要读取约 14GB 的模型权重（7B 参数 × 2 字节），但实际浮点运算量仅约 14GFLOPs。按现代 GPU 的算力（如 A100 的 312 TFLOPS FP16）和显存带宽（2TB/s）估算：

计算时间 ≈ 14GFLOPs / 312TFLOPS ≈ 0.045 ms
访存时间 ≈ 14GB / 2TB/s ≈ 7 ms

访存时间比计算时间高出两个数量级！这就是为什么 Decode 阶段被称为访存密集（Memory-Bound）。TensorRT-LLM 的许多优化（如算子融合、量化、In-flight Batching）本质上都是在缓解访存压力或提高访存效率。

工程启示：Prefill 阶段适合用高算力 GPU 批量并行处理；Decode 阶段则受限于显存带宽。因此，在部署时通常将两者分开调度，或使用连续批处理（Continuous Batching）让多个请求的 Prefill 和 Decode 交错执行，以填满 GPU 的计算和访存流水线。

2.3 KV Cache 的显存占用与动态性

KV Cache 并非固定大小的缓冲区，它随序列长度和并发请求数动态增长。其显存占用公式为：

KV Cache 显存 = 2（K 和 V）× 层数 × 隐藏维度 × 序列长度 × 精度字节数 × 并发请求数

以 LLaMA-7B（32 层，隐藏维度 4096）为例，使用 FP16（2 字节）推理，单个请求的 KV Cache 占用为：

2 × 32 × 4096 × 序列长度 × 2 字节
= 2 × 32 × 4096 × 2 = 524,288 字节 / token ≈ 0.5 MB / token

当序列长度达到 2048 时，单请求 KV Cache 约 1GB；若并发 16 个请求，则仅 KV Cache 就需要 16GB 显存。相比之下，模型权重本身约 14GB（FP16）。这意味着长序列和高峰并发下，KV Cache 可能占据超过一半的显存，成为部署时的主要瓶颈之一。

常见误区：许多初学者认为 KV Cache 是“缓存”所以不重要，或者认为它只是加速的小技巧。实际上，KV Cache 是推理正确性的必要组成部分——它保存了历史 token 的 Key 和 Value，使得注意力计算能够正确进行。删除或错误管理 KV Cache 会导致生成结果错误或显存溢出（OOM）。

2.4 案例：KV Cache 的分配与释放过程

在 TensorRT-LLM 中，KV Cache 的管理是核心调度逻辑。以下是一个典型请求的生命周期：

  1. 请求到达：系统根据当前显存余量，为该请求分配一块连续的 KV Cache 空间（或按块分配）。
  2. Prefill 阶段：计算 Prompt 中所有 token 的 K/V，写入分配的缓存空间。
  3. Decode 阶段：每生成一个新 token，将其 K/V 追加到缓存末尾；同时读取全部缓存用于注意力计算。
  4. 请求结束：释放该请求占用的 KV Cache 空间，供其他请求使用。

TensorRT-LLM 引入了类似操作系统虚拟内存的PagedAttention机制，将 KV Cache 划分为固定大小的块（Block），通过块表（Block Table）管理逻辑地址到物理地址的映射。这样即使物理显存中存在碎片，也能通过非连续块完成缓存分配，大幅提升显存利用率。

核心要点回顾：

  • Transformer 解码器推理 = 自回归逐 token 生成，分为 Prefill（计算密集）和 Decode（访存密集）两阶段。
  • KV Cache 存储历史 token 的 Key/Value，避免重复计算，是推理性能的关键。
  • KV Cache 显存占用随序列长度和并发数线性增长，是部署规划的重要参数。
  • TensorRT-LLM 通过 PagedAttention 等机制实现 KV Cache 的高效动态管理。
← 上一章 📖 返回目录 下一章 →