大模型推理基础：Transformer架构、自回归生成与KV Cache原理

在深入对比 TensorRT-LLM 与 vLLM 的性能之前，必须先建立大模型推理的底层认知。当前主流的大语言模型（如 GPT、LLaMA、Qwen）几乎都基于 Transformer 解码器架构。推理过程并非一次性生成完整回答，而是逐 token 迭代的自回归生成。理解这一过程，尤其是 KV Cache 的引入与优化，是后续分析推理引擎性能差异的关键前提。

一、Transformer 解码器：推理时的核心组件

Transformer 解码器在推理时，其核心计算模块可简化为三个部分：自注意力（Self-Attention）、前馈网络（FFN）以及层归一化与残差连接。其中，自注意力机制负责建模输入序列中 token 之间的依赖关系，是计算量最大的部分之一。

自注意力的计算过程可表示为以下公式：

Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V

其中，Q（查询）、K（键）、V（值）由输入向量与对应的权重矩阵相乘得到。在推理阶段，对于当前生成的 token，其 Q 向量需要与序列中所有 token 的 K 向量做点积，以计算注意力权重。这意味着，如果不做缓存，每次生成新 token 时，都需要重新计算历史 token 的 K 和 V 向量，造成大量冗余计算。

关键认知：推理引擎的优化重点，往往集中在如何高效管理注意力机制所需的中间数据，而非仅仅提升矩阵乘法的速度。

二、自回归生成：逐 Token 的迭代循环

自回归生成是指模型将已生成的 token 序列作为输入，预测下一个 token的概率分布，然后选取（或采样）一个 token 加入序列，再将该序列作为新的输入，重复此过程，直到生成结束符或达到最大长度。

其基本流程如下：

  1. 输入提示词：将用户输入的文本转换为 token ID 序列。
  2. 前向传播：将整个 token 序列输入模型，计算最后一个位置输出的概率分布。
  3. 采样/解码：根据概率分布选择下一个 token（贪心、top-k、top-p 等策略）。
  4. 拼接序列：将新 token 拼接到原序列末尾。
  5. 重复步骤 2-4，直到满足停止条件。

这种串行依赖关系决定了推理的延迟与序列长度成正比。每一步生成都依赖上一步的输出，无法像训练阶段那样对序列进行并行化处理。因此，降低单步延迟和提高吞吐量成为推理引擎的核心目标。

三、KV Cache：消除重复计算的基石

观察自回归过程可以发现，在生成第 N 个 token 时，前 N-1 个 token 的 K 和 V 向量已经计算过。若将其丢弃，下一步需要重新计算，计算复杂度为 O(N²)。KV Cache 的核心思想是用显存换算力：将历史 token 的 K、V 向量缓存下来，供后续步骤直接使用。

引入 KV Cache 后，单步生成的计算复杂度降为 O(N)，但显存占用随序列长度线性增长。下图展示了有无 KV Cache 时，自回归生成第 4 个 token 的计算差异：

无 KV Cache（重复计算） Token1 Token2 Token3 新Token 重新计算所有 K、V 计算量 ∝ 序列长度 N 延迟高，显存占用低 有 KV Cache（高效推理） K1,V1 | K2,V2 | K3,V3 新Token 仅计算新 Token 的 K、V 计算量 ∝ 1（常数） 延迟低，显存占用高 直接复用 KV Cache 的权衡 优势 • 避免重复矩阵乘法，大幅降低单步延迟 • 支持更高吞吐量的并行处理 代价 • 显存占用随序列长度线性增长 • 长序列场景下可能触发显存不足（OOM）

上图清晰展示了两种方式的差异。现代推理引擎（包括 TensorRT-LLM 和 vLLM）均默认启用 KV Cache，并在此基础上进行显存管理优化。

四、KV Cache 的显存占用估算

KV Cache 的显存占用并非固定值，它由模型参数和推理配置共同决定。计算公式如下：

显存占用（字节） = 2（K 和 V） × 层数 × 注意力头数 × 头维度 × 序列长度 × 批次大小 × 数据类型大小

以 LLaMA-7B 模型为例，假设使用 FP16 精度（每个元素占 2 字节），其关键参数为：层数 32、注意力头数 32、头维度 128。当批次大小为 1、序列长度为 2048 时，KV Cache 占用约为：

2 × 32 × 32 × 128 × 2048 × 1 × 2 = 1,073,741,824 字节 ≈ 1 GB

这意味着，仅处理一个 2048 token 的序列，就需要约 1GB 显存来存储 KV Cache。若批次大小提升至 8，则需 8GB。下表展示了不同配置下的显存需求：

模型规模 层数 头数 × 头维度 序列长度 批次大小 KV Cache 占用（FP16）
LLaMA-7B 32 32 × 128 2048 1 ≈ 1 GB
LLaMA-13B 40 40 × 128 2048 4 ≈ 6.4 GB
LLaMA-70B 80 64 × 128 4096 8 ≈ 40 GB

从表中可以看出，KV Cache 的显存开销可能超过模型权重本身。例如 LLaMA-70B 的权重在 FP16 下约 140GB，而 KV Cache 在长序列、大批次下同样可能达到数十 GB。因此，KV Cache 的管理策略直接决定了推理引擎能否在有限显存下支持更大的并发和更长的上下文。

五、常见误区与优化方向

误区一：认为 KV Cache 越大越好。实际上，KV Cache 占用过多显存会挤压模型权重和激活值的存储空间，导致批次大小被迫减小，反而降低整体吞吐量。推理引擎需要在 KV Cache 容量与并发度之间寻找平衡。
误区二：忽略 KV Cache 的碎片化问题。在动态批次（continuous batching）场景下，不同请求的序列长度不同，KV Cache 的分配和释放会产生大量显存碎片。vLLM 提出的 PagedAttention 正是为了解决这一问题，将 KV Cache 划分为固定大小的块，按需分配，类似操作系统中的虚拟内存分页。

TensorRT-LLM 则采用了不同的优化策略，包括：

  • KV Cache 量化：将 FP16 的 KV Cache 量化为 INT8 或 FP8，减少显存占用，但需注意精度损失。
  • 多轮对话的 KV Cache 复用：对于前缀相同的请求（如系统提示词），复用已计算的 KV Cache，避免重复计算。
  • 显存池管理：预分配大块显存，通过内部分配器管理 KV Cache 的申请与释放，减少碎片。
← 上一章 📖 返回目录 下一章 →