Lecture 10：推理优化

为什么推理如此重要#

大语言模型的训练是一次性的固定成本——虽然昂贵，但一旦完成就是完成了。而推理（inference）是持续性的重复成本，每天都在发生。据估计，OpenAI 每天产出约 8.6 万亿（trillion）个 token。作为对比，DeepSeek-V4 的整个训练过程使用了约 14.8 万亿 token——也就是说，不到四天的推理量就等于 DeepSeek-V4 的全部训练数据量。

推理的应用场景远不止对话：

• 交互式应用：AI 助手、聊天机器人、代码补全
• Agent 工作流：agent 需要大量推理来做内部思考和工具调用，其中大部分生成的 token 并不直接呈现给人类阅读——真正的计算开销体现在 token 生成量上
• 批量数据处理：大规模文档处理、数据标注
• 评估与训练内循环：需要生成式评估的场景，以及强化学习中的 rollout 生成

在传统的聊天场景中，推理速度在某个阈值之上就”够用”了，因为人类阅读速度有限。但在 agentic 场景下，不存在这样的上限——问题越复杂，能从更快推理中获得的价值就越大。

推理生态方面，商业 API 提供商（OpenAI、Anthropic 等）需要服务自己的模型；开放权重模型的推理服务商同样活跃。开源社区中，vLLM 是目前最主流的推理框架，SGLang 在 agentic 负载上有特别的优势，TensorRT-LLM（NVIDIA）追求极致速度但适用范围较窄，llama.cpp 则用于 CPU 推理。

三个核心度量指标#

衡量推理”快不快”有三个维度，它们适用于不同场景且存在内在张力：

1. Time to First Token (TTFT)：从用户发出请求到第一个 token 返回的等待时间。本质上是 prefill 阶段的耗时。对交互式应用至关重要——用户在等待期间完全无反馈。

2. Latency（延迟）：单个请求的 token 生成速度（秒/token）。同样关乎交互体验，决定了用户看到文本”流出”的速度。

3. Throughput（吞吐量）：系统级的总 token 产出速率（token/秒），衡量的是跨多个请求的整体处理能力。对批量处理和高并发场景至关重要——此时单个请求的延迟不重要，重要的是整体任务完成时间。

延迟和吞吐量看似正相关，但实际上存在根本性的权衡：提升吞吐量的主要手段是增大 batch size，而增大 batch size 会恶化延迟。这个权衡贯穿整节课。

符号约定#

分析基于 Google 的 Scaling Book（Transformers 章节），采用类似 einops 的张量符号：

张量乘法中，维度用颜色编码：红色维度是收缩维度（出现在两个操作数中，从结果中消失），黑色维度只出现在一个操作数中并保留到结果中，蓝色维度是批处理维度（出现在两个操作数中且保留到结果中，不做收缩）。

Transformer Block 的张量运算#

一个 Transformer block 的完整计算可以精确地用张量维度来描述。

Attention 部分：输入 $X$（维度 $BTD$）分别经过 $W_Q$（$D \times NH$）、$W_K$（$D \times KH$）、$W_V$（$D \times KH$）投影，得到 $Q$（$BTNH$）、$K$（$BSKH$）、$V$（$BSKH$）。注意力矩阵的计算中，$B$ 作为蓝色批处理维度出现在两个操作数中——这一点对后续的算术强度分析至关重要。

MLP 部分：包含 gate 投影（$W_{\text{in1}}$: $D \times F$）、up 投影（$W_{\text{in2}}$: $D \times F$）和 down 投影（$W_{\text{out}}$: $F \times D$），加上 GeLU 激活。MLP 是标准的矩阵乘法，各序列之间完全独立。

关键区别在于：训练时 $T = S$（看到全部 token，可以并行处理序列维度），而推理生成时 $T = 1$（一次只生成一个 token），$S$ 是已生成的序列长度。

矩阵乘法的算术强度推导#

以一个简单的矩阵乘法 $X_{B \times D} \times W_{D \times F}$ 为例进行分析：

1. 读取：$X$ 从 HBM 读取 $2BD$ 字节（BF16），$W$ 读取 $2DF$ 字节
2. 计算：FLOPs = $2BDF$（矩阵乘法）
3. 写回：结果 $Y_{B \times F}$ 写回 $2BF$ 字节

$\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{Bytes Transferred}} = \frac{2BDF}{2BD + 2DF + 2BF}$

当 $D, F \gg B$ 时（即模型维度远大于 batch size），分母中 $2DF$ 项占主导，算术强度近似为 $B$。

这意味着对于 H100（加速器强度 $\approx 295$），当 $B \geq 295$ 时才是 compute-bound。当 $B = 1$ 时，算术强度为 1——这正是推理的典型场景：你读取了一整个 $D \times F$ 的权重矩阵，但只做了 $2DF$ 次 FLOPs，权重没有被充分复用。

MLP 层的算术强度#

MLP 层由三个矩阵乘法组成（gate、up、down），但分析逻辑完全相同：

• Prefill 阶段：intensity $\approx B \cdot S$（batch size × 序列长度），只要有足够长的序列或足够大的 batch 就能 compute-bound
• Generation 阶段：$T = 1$，intensity $\approx B$——需要足够大的 batch 来达到 compute-bound

MLP 的算术强度能随 $B$ 提升，因为所有序列共享同一组 MLP 权重：读一次权重，可以为 $B$ 个序列复用。

Attention 层的算术强度#

Attention 层的情况截然不同。计算注意力时需要读取 $Q$（$BTD$）、$K$（$BSD$）、$V$（$BSD$），其中 $K$ 和 $V$ 都依赖于 $B$——每个序列有自己独立的 KV cache。

$\text{Attention Intensity} = \frac{S \cdot T}{S + T}$

• Prefill（$T = S$）：intensity $= S/2$，只要序列足够长就没问题
• Generation（$T = 1$）：intensity $= S/(S+1) < 1$

Generation 阶段 attention 的算术强度恒小于 1，无论 batch size 多大都无法改善。这是因为每个序列有自己的 KV cache（$Q$、$K$、$V$ 都依赖于 $B$），增大 batch 相当于做更多独立的 matmul——每个 matmul 的算术强度仍然很低，本质上退化为一系列点积运算。

核心结论：Prefill 是 compute-bound，Generation 是 memory-bound。Generation 阶段 attention 层的算术强度 < 1 是推理的根本瓶颈，这是 Transformer 架构的固有特性，无法通过增大 batch 来解决。

朴素自回归推理的代价#

最朴素的推理方式是把 Transformer 视为黑盒：输入一个序列，输出下一个 token 的分布，采样一个 token 后拼接到序列末尾，再重新输入整个序列。这种方式每生成一个 token 都需要对整个序列做一次完整的前向传播，其中注意力计算是 $O(T^2)$，生成 $T$ 个 token 的总代价是 $O(T^3)$。

但这里有大量冗余计算：由于 causal mask 的存在，已有 token 的 key 和 value 不会因为追加新 token 而改变。例如生成序列 “never going to give you up” 时，“never” 的 key/value 向量在整个过程中保持不变。（注意这只对因果 Transformer 成立——如果是双向注意力，追加新 token 会改变所有位置的激活。）

KV Cache 机制#

利用上述因果性，可以缓存已计算过的 key 和 value 向量，避免重复计算。推理因此被分为两个阶段：

Prefill 阶段：接收完整的 prompt，一次性计算所有 token 的 key 和 value 并存入 KV cache。这一步与训练类似，可以并行处理整个序列，是 compute-bound 的。同时产出第一个生成 token 的 logits。

Generation 阶段：逐 token 生成。每步只需处理一个新 token：计算它的 query，与 KV cache 中所有已有的 key 做注意力，然后将新 token 的 key/value 追加到 cache 中。不再需要重新计算之前所有 token 的 key 和 value。

KV cache 的形式化定义：对每个序列（共 $B$ 个），每一层、每个 KV head 都存储一个 $H$ 维的向量，key 和 value 各一份。总大小为：

$\text{KV Cache Size} = B \times S \times L \times K \times H \times 2 \times 2 \text{ bytes}$

其中最后两个 2 分别是 key+value 两份，以及 BF16 的 2 字节。KV cache 的大小与序列长度 $S$ 和 batch size $B$ 成正比——当 batch size 足够大时，KV cache 的内存占用可以超过模型参数本身。

从算术强度到实际性能指标#

既然推理是 memory-bound 的，那么性能分析就可以大幅简化：延迟主要取决于需要搬运的内存量。假设通信和计算可以重叠，瓶颈就是 HBM 的读写吞吐。

需要占用内存（也需要搬运）的主要有两部分：

1. 模型参数：大小固定，$\text{param\_size} = 2 \times \text{num\_params}$ 字节（BF16）
2. KV Cache：随 $B$ 和 $S$ 线性增长，$\text{kv\_cache\_per\_seq} = S \times L \times K \times H \times 2 \times 2$

总内存为 $\text{memory} = B \times \text{kv\_cache\_per\_seq} + \text{param\_size}$。

由此推导：

$\text{Latency} = \frac{\text{memory}}{\text{memory\_bandwidth}} \quad \text{(seconds/token)}$

$\text{Throughput} = \frac{B}{\text{Latency}} \quad \text{(tokens/second)}$

吞吐量在分子中有 $B$，因为一个 batch 同时产出 $B$ 个 token。

Llama-2 13B 在 H100 上的实例#

以 Llama-2 13B 为例（$S=1024$, $D=5120$, $F=13824$, $N=40$, $K=40$, $H=128$, $L=40$, $V=32000$），H100 内存带宽 3.35 TB/s：

• 参数量：13,015,449,600（约 130 亿，与标称一致）
• 内存：$838{,}860{,}800 \times B + 26{,}030{,}899{,}200$（线性函数）

Batch Size 的根本权衡#

从数值中可以清楚地看到延迟和吞吐量的反向关系：

• 增大 batch size 恶化延迟：因为 KV cache 随 $B$ 线性增长，每步需要读写更多内存。每个请求都必须等待整个 batch 处理完毕——类比”等公交”：延迟高（要等），但吞吐量好（一次运很多人）。
• 增大 batch size 提升吞吐量：因为模型参数被所有序列共享，读取一次权重可以服务 $B$ 个序列。但这个提升有递减效应——当 KV cache 的内存开销主导时，增大 $B$ 带来的边际吞吐增长越来越小。
• 内存墙：batch size 不能无限增大，因为最终会超出 GPU 显存。即使用 B200，也只是把上限推高而已。

因此在实际部署中，小 batch size 适合低延迟场景（交互对话），大 batch size 适合高吞吐场景（批量处理）。

并行化与 TTFT#

• 简单并行：部署 $M$ 个模型副本，延迟不变，吞吐量提升 $M$ 倍
• 模型并行：分片模型和 KV cache 到多卡上，可以降低延迟但引入通信开销（详见 Scaling Book 推理章节）
• TTFT：本质上就是 prefill 的时间。更小的 batch size → 更快的 TTFT；更大的 batch size → 更高的吞吐量。两者在 batch size 这个旋钮上是对立的。

Grouped Query Attention (GQA)#

GQA 是 Multi-Head Attention (MHA) 和 Multi-Query Attention (MQA) 之间的折中方案。

在标准 MHA 中，$K = N$（key/value head 数等于 query head 数），每个 query head 有自己独立的 key 和 value。MQA 走到另一个极端：$K = 1$，所有 query head 共享同一组 key/value，KV cache 大幅减小但精度损失严重。GQA 取中间值：保持 $N$ 个 query head，但只用 $K$ 个 key/value head，每 $N/K$ 个 query head 共享一组 key/value。

GQA 将 KV cache 缩小 $N/K$ 倍。回到 Llama-2 13B 的例子，将 $K$ 从 40 降到 8（1:5 的比例）：

GQA 带来两重好处：一是直接提速（内存减少 → 延迟和吞吐量都改善，此时两者并不对立）；二是释放显存空间，允许使用更大的 batch size 来进一步提升吞吐量。GQA 原始论文（Ainslie+ 2023）显示精度损失很小——但后来 DeepSeek 的论文指出 GQA 在某些场景下的精度损失比最初声称的更显著。

Multi-Latent Attention (MLA)#

DeepSeek-V2 提出的 MLA 采取了更激进的压缩策略。传统方法中，每个 token 的 key/value 向量维度是 $NH$（约等于模型维度 $D$）。MLA 将激活先投影到一个低维空间（$C$ 维），然后在需要时再展开回 key/value。

DeepSeek-V2 将维度从 16,384 压缩到 512——压缩比超过 30 倍。在 KV cache 中只需存储这个 $C$ 维的压缩表示，需要 key/value 时再通过投影矩阵”物化”出来。

一个技术细节是 MLA 与 RoPE 不直接兼容（RoPE 直接操作 key/value 向量），DeepSeek 的解决方案是为处理位置编码单独保留一小部分未压缩的 key。

DeepSeek 的对比实验表明，MLA 在精度上甚至略好于 MHA（至少不逊色），同时 KV cache 大幅缩小。不过也显示了 GQA 的精度损失比 GQA 原始论文声称的更明显——这提醒我们，单一论文的 ablation 结论要谨慎看待。

Cross-Layer Attention (CLA)#

CLA 的思路是在层间共享 KV cache（而 GQA 是在 head 间共享）。标准做法是每一层独立计算自己的 key/value，CLA 让部分层直接复用前面层的 KV cache，不再独立计算。

实验表明 CLA 在给定 KV cache 预算下能取得更好的精度-效率 Pareto 前沿。

滑动窗口注意力与混合架构#

另一个自然的想法是限制注意力范围：每个 token 只关注最近的 $k$ 个 token（滑动窗口），而不是全部历史。这使得 KV cache 大小独立于序列总长度，对长序列场景尤其有利。

由于 Transformer 有多层，信息可以通过逐层传播覆盖比窗口大小更远的上下文——有效感受野随层数线性增长。

注意力模式有多种变体：(a) 全连接 $n^2$ 注意力，(b) 滑动窗口注意力，(c) 空洞滑动窗口（dilated），(d) 全局+滑动窗口混合。

纯滑动窗口仍然会损失精度（表达能力下降），实践中的解法是混合架构：部分层用全注意力、部分层用局部注意力，兼顾精度和效率。

DeepSeek-V4 更进一步，引入了 Compressed Sparse Attention：先将每 $M$ 个连续 token 压缩为一个 token（compressed attention），然后通过轻量级的 query-key 计算（“Lightning Indexer”）选出最重要的 token 子集（selected compressed KV entries），最终只对这个子集做完整注意力，再结合滑动窗口的局部注意力。

线性注意力与状态空间模型#

关于 linear attention（线性注意力），它的核心是将全部历史压缩为一个固定大小的状态向量，彻底消除 KV cache 对序列长度的依赖。Mamba、Delta Net 等是更强的变体——它们在压缩的同时试图保留更多信息。

线性注意力与滑动窗口的权衡：线性注意力更擅长捕获全局的”宽泛摘要”，而滑动窗口更擅长高分辨率的局部关注。两者并非互斥——最先进的混合架构会组合全注意力、滑动窗口注意力和线性注意力，让不同层负责不同粒度的信息。

但核心问题无法回避：如果你有一个非常长的上下文并且需要精确检索其中的细节（needle-in-a-haystack），将全部历史压缩到一个小向量中必然会丢失信息。没有免费午餐。

量化#

量化的核心思想：降低数值精度。更少的位数 → 更小的内存占用 → 更快的读写 → 更低的延迟和更高的吞吐量。当然代价是潜在的精度损失。

各种精度格式的定位：

• FP32（4 字节）：训练时用于参数和优化器状态
• BF16（2 字节）：推理的默认格式
• FP8（1 字节，e4m3 格式，范围 [-240, 240]）：H100 上可用，甚至可以用于训练
• INT8（1 字节，范围 [-128, 127]）：比 FP8 便宜但精度更低，仅用于推理
• INT4（0.5 字节，范围 [-8, 7]）：更便宜，精度更低

模型剪枝#

剪枝是一种更”粗暴”但有效的方法：直接从大模型中移除部分组件（层、head、隐藏单元），得到一个更小的模型，然后用少量数据微调修复。

NVIDIA 的方法（Muralidharan+ 2024）分五步迭代：

1. 从训练好的 LLM 出发
2. 估计重要性：在校准集（约 1024 个样本）上运行前向传播，观察各组件（embedding、attention head、MLP 通道）的激活幅度
3. 排序：按重要性对所有组件排名
4. 裁剪：移除低重要性的组件
5. 蒸馏修复：用原始大模型作为 teacher，在目标数据上微调裁剪后的模型

他们成功地将 15B 参数的模型压缩到更小的尺寸，精度损失很小，而且修复所需的训练量远少于从头训练。

方法论总结#

减小推理复杂度的思路可以归纳为两条路径：

1. 设计更快的模型架构 → 从头训练：GQA、MLA、CLA、滑动窗口注意力等都属于这类
2. 从已有大模型中提取更快的版本 → 修复：量化、剪枝，然后通过蒸馏恢复精度。架构可能不同（比如剪枝后的模型结构与原始不一致），但通过初始化自好模型的部分权重来加速收敛

两条路径的共同目标是减少参数量或 KV cache 大小，最终都需要验证精度是否可接受。

核心洞察：验证比生成快#

这个方法的基础是一个关键的不对称性：

• 生成（generation）必须一次一个 token，是 memory-bound 的，很慢
• 验证（verification / prefill）可以并行处理一批 token，是 compute-bound 的，很快

给定一个序列，判断它有多好（计算每个 token 的概率）远比逐 token 生成这个序列要快。

Draft-Target 架构#

推测解码利用这个不对称性，引入两个模型：

• Draft 模型 $p$：小而快，用于猜测接下来的 $K$ 个 token
• Target 模型 $q$：大而精确，用于并行验证 draft 模型的猜测

工作流程：

1. Draft 模型自回归地生成 $K$ 个候选 token $\tilde{x}_1, \ldots, \tilde{x}_K$（虽然是逐 token 的，但 draft 模型很小，速度可接受）
2. Target 模型并行地对这 $K$ 个 token 计算 logits（一次 prefill，高效）
3. 对每个候选 token，以概率 $\min\left(1, \frac{q(\tilde{x}_t | \text{context})}{p(\tilde{x}_t | \text{context})}\right)$ 决定是否接受：
   • $q > p$ 时大概率接受（target 模型认为这是好选择）
   • $q < p$ 时按比例接受（draft 模型过于自信的地方打折）
4. 一旦某个 token 被拒绝，从修正后的残差分布 $(q - p)_+$ 中采样一个替代 token，并终止本轮
5. 如果所有 $K$ 个 token 都被接受，额外从 target 模型采样一个 bonus token

数学保证#

这本质上是改进的拒绝采样（modified rejection sampling）。标准拒绝采样在 reject 时什么都得不到，但推测解码保证每轮至少产出一个精确来自 target 分布的样本。可以严格证明，推测解码输出的 token 分布与直接从 target 模型采样完全一致——这是无损的。

实际效果#

Draft token 数 $K$ 的选择存在 sweet spot：

• $K$ 太小：target 模型端的批处理不够高效，加速倍率有限
• $K$ 太大：后面的 token 被拒绝的概率累积增高，浪费 draft 模型的计算

实验表明最优 $K$ 通常在 4-8 左右。

Draft 模型的选择至关重要：它需要足够小以保证速度优势，但又要足够接近 target 模型以提高接受率。实践中通常通过蒸馏来训练 draft 模型——这与前面提到的剪枝+蒸馏思路相通。也就是说，如果你用前面的方法（GQA、量化、剪枝）得到了一个精度下降但速度更快的模型，它至少可以作为一个不错的 draft 模型，然后用 target 模型来”兜底”。

推测解码已经发展出了大量变体和改进，是一个活跃的研究方向。

Continuous Batching (Orca)#

传统的 static batching 要求一个 batch 中所有序列同时开始、同时结束。但实际中请求是交错到达的，每个请求的生成长度也不同。如果强制等所有序列都结束才处理下一个 batch，会造成大量空闲等待。

Continuous Batching（Orca 系统提出）的核心思想是迭代级调度（iteration-level scheduling）：

1. 每一步解码时，为 batch 中所有序列各生成一个 token
2. 某个序列完成（遇到 EOS）后立即从 batch 中移除
3. 新到达的请求随时插入 batch，无需等待当前 batch 全部完成

这样 batch 是动态变化的——老序列不断退出，新请求不断加入，GPU 始终保持忙碌。

一个技术难点是不等长序列的高效处理。注意力计算无法回避长度差异（每个序列有自己的 KV cache），但 MLP 层不依赖序列长度——可以将所有序列的当前 token 拼接成一个”超级序列”统一处理。

PagedAttention (vLLM)#

KV cache 的内存管理引入了类似操作系统磁盘管理的经典问题——碎片化。

内部碎片：由于不知道序列会生成多长，必须为每个请求预分配最大长度（如 1024 token）的 KV cache 缓冲区。如果实际只用了 200 个 token，剩余 824 个位置的内存就被浪费了。

外部碎片：多个请求的 KV cache 散布在内存中，留下一些太小而无法被有效利用的空隙。

PagedAttention（vLLM 论文的核心创新）借用了操作系统虚拟内存的分页机制来解决这个问题：

将 KV cache 划分为固定大小的块（block），每个块容纳若干 token 的 key/value（例如 4 个 token 一块）。逻辑上连续的 token 在物理内存中可以不连续存放——只要维护一个块表（page table）记录映射关系即可。

这带来几个重要好处：

1. 消除碎片：不需要预分配连续的大块内存，按需分配固定大小的块
2. 前缀共享：如果多个请求使用相同的 system prompt，它们可以共享同一块 KV cache——只存一份，通过块表指向同一个物理块。这对高并发场景下的 system prompt 重用意义重大
3. Copy-on-Write：如果从同一个 prompt 生成多个不同的 response（如 best-of-N 采样），它们共享 prompt 的 KV cache，直到某个位置产生不同的 token 时才拷贝并分裂对应的块——在此之前的共享前缀零成本

例如，两个请求都以 “Four score and seven years ago our fathers brought forth” 开头：它们在逻辑上各自维护 Block 0-3，但物理上指向相同的存储位置。只有当生成内容开始分叉时，才会分配新的物理块。