第 10 讲：Inference 优化
训练一个大语言模型是一次性投入，而推理（inference）是持续的成本——每一次用户查询都需要消耗计算资源。对于部署在生产环境中的模型，推理效率直接决定了服务成本和用户体验。本讲将系统介绍 LLM inference 的核心优化技术，从 KV Cache 到 Quantization，从 Speculative Decoding 到 PagedAttention。
Autoregressive Generation 流程
LLM 的文本生成是 autoregressive（自回归）的——每次只生成一个 token，然后将其作为输入的一部分来生成下一个 token。这意味着生成 T 个 token 需要 T 次模型推理。
一次完整的生成过程分为两个阶段：
Prefill 阶段
处理用户输入的 prompt（长度为 P 个 tokens）。这个阶段类似训练时的 forward pass——一次性处理整个 prompt，产生所有 position 的 hidden states 和 attention 中间结果。
- 计算特征：大量矩阵乘法，序列长度 P 可能很长
- 计算量：≈ 2NP FLOPs（N 为参数量，P 为 prompt 长度）
- 瓶颈：通常是 Compute-bound（GPU 算力利用率高）
Decode 阶段
逐个生成新 token。每一步只有 1 个新 token 作为输入（sequence length = 1），但需要 attend to 所有之前的 tokens。
- 计算特征：每步的矩阵乘法中，一个维度为 1（当前 token），另一个维度很大（模型参数）
- 计算量：每步 ≈ 2N FLOPs
- 瓶颈：通常是 Memory-bound（受限于从显存读取模型权重的带宽）
Prefill vs Decode 的根本区别
Prefill 可以并行处理所有 prompt tokens，GPU 利用率高；Decode 每步只处理 1 个 token，大部分时间花在从显存读取权重上，计算单元严重空闲。这是 LLM inference 优化的核心矛盾——Decode 阶段的 arithmetic intensity 极低。
KV Cache 原理
在 Decode 阶段，每生成一个新 token，attention 层需要计算新 token 与所有之前 tokens 的 attention。如果每次都重新计算所有 tokens 的 K 和 V，计算量会随序列长度平方增长。
KV Cache 的思想很简单：缓存之前所有 token 的 Key 和 Value 向量，每步只计算新 token 的 K、V，然后追加到缓存中。
工作流程
- Prefill：计算 prompt 中所有 token 的 K、V，存入 KV Cache
- Decode step t：
- 只计算新 token 的 Q、K、V
- 将新的 K、V 追加到 KV Cache
- 用新 token 的 Q 与 cache 中所有的 K 做 attention score 计算
- 用 attention weights 和 cache 中所有的 V 做加权求和
KV Cache 的内存占用
KV Cache 的大小取决于模型架构和序列长度：
其中因子 2 代表 K 和 V 两个缓存。用模型超参数表示：
具体例子：
| 模型 | 层数 L | d_model | 序列长度 S | KV Cache (FP16) |
|---|---|---|---|---|
| LLaMA-7B | 32 | 4096 | 2048 | 1.0 GB |
| LLaMA-7B | 32 | 4096 | 128K | 64 GB |
| LLaMA-70B | 80 | 8192 | 2048 | 5.0 GB |
| LLaMA-70B | 80 | 8192 | 128K | 320 GB |
KV Cache 是推理的内存瓶颈
在长序列和大 batch 场景下，KV Cache 的内存占用可能远超模型参数本身。例如 LLaMA-70B 在 128K context 下，KV Cache 需要 320GB（FP16），而模型参数只有 140GB。KV Cache 的优化是 inference 系统的核心挑战。
KV Cache 优化
Grouped Query Attention (GQA)
GQA 是 LLaMA-2 引入的架构改进。在标准的 Multi-Head Attention (MHA) 中，每个 head 有独立的 Q、K、V 投影。GQA 将 heads 分成若干组，同一组内的 heads 共享 K 和 V，但各自有独立的 Q。
GQA: n_kv_heads = n_heads / group_size（多个 head 共享 K、V）
MQA: n_kv_heads = 1（所有 head 共享一组 K、V）
| 方法 | KV Heads | KV Cache 大小 | 质量影响 |
|---|---|---|---|
| MHA | = n_heads | 1× (baseline) | 无 |
| GQA (group=8) | = n_heads/8 | 1/8× | 极小 |
| MQA | = 1 | 1/n_heads× | 略有下降 |
GQA 在几乎不损失模型质量的情况下，将 KV Cache 大小减少到原来的 1/group_size。LLaMA-2 70B 使用 GQA 将 KV Cache 减小了 8 倍。
Multi-Query Attention (MQA)
MQA 是 GQA 的极端情况——所有 query heads 共享同一组 K 和 V。它在 2019 年由 Shazeer 提出，是最激进的 KV Cache 压缩方式。
MQA 对推理速度的提升显著（KV Cache 可能减少 32 倍或更多），但在某些任务上会有质量损失。GQA 是 MHA 和 MQA 之间的折中方案。
Quantization 量化
Quantization 通过降低数值精度来减少模型大小和加速推理。在推理时，我们不需要训练时的高精度——将 FP16 权重量化到 INT8 甚至 INT4 通常对质量影响很小。
基本概念
量化将连续的浮点值映射到离散的整数值：
x_dequantized = (x_quantized - zero_point) × scale
其中 scale 和 zero_point 是量化参数，决定了浮点数到整数的映射范围。
常见量化方法
| 方法 | 精度 | 大小压缩 | 特点 |
|---|---|---|---|
| FP16 (baseline) | 16-bit | 1× | 标准推理精度 |
| INT8 | 8-bit | 2× | 几乎无质量损失 |
| INT4 (GPTQ) | 4-bit | 4× | 轻微质量损失 |
| INT4 (AWQ) | 4-bit | 4× | 保护显著权重 |
| INT3/INT2 | 3/2-bit | 5-8× | 较大质量损失 |
GPTQ
GPTQ（Generative Pre-Trained Transformer Quantization）是一种 Post-Training Quantization (PTQ) 方法，即在训练完成后进行量化。它的核心思想来自 Optimal Brain Quantization (OBQ)：
- 逐列量化权重矩阵
- 量化一列后，用 Hessian 信息补偿其他列的值，最小化量化引入的输出误差
- 使用一小组校准数据（calibration data，通常 128 条）来估计 Hessian
GPTQ 可以在几分钟内将一个大模型量化到 INT4，且质量损失极小。
AWQ (Activation-Aware Weight Quantization)
AWQ 的关键观察是：不是所有权重都同等重要。与大 activation magnitude 对应的权重通道（channels）对模型输出的影响更大。AWQ 的策略：
- 通过少量校准数据，观察哪些权重通道对应较大的 activation
- 对这些显著通道（salient channels）进行特殊处理——在量化前先缩放这些通道的权重，使它们在量化后误差更小
- 在推理时，通过相应的反缩放来保持等价性
Quantization 对推理速度的影响
Quantization 的加速来自两方面：
- 减少显存带宽需求：INT4 权重只有 FP16 的 1/4 大小，从显存读取更快。对于 memory-bound 的 decode 阶段，这是最主要的加速来源
- 更高效的计算：INT8/INT4 矩阵乘法的吞吐量通常是 FP16 的 2-4 倍（利用 Tensor Core 的 INT8 路径）
Speculative Decoding
Speculative Decoding（推测性解码）是一种不改变模型输出分布、但能显著加速生成的技术。
核心思想：Draft + Verify
- Draft：使用一个小而快的草稿模型（draft model）快速生成 K 个候选 token
- Verify：用大的目标模型（target model）一次性并行验证这 K 个 token
- 如果前 k 个 token 都被接受（k ≤ K），则采纳它们；否则从第一个被拒绝的位置重新生成
为什么能加速？
Decode 阶段是 memory-bound 的——每步都需要从显存加载全部模型权重，但只处理 1 个 token。Verify 阶段一次处理 K 个 token，也只需要加载一次权重（像 prefill 一样并行计算），实际增加的时间很少。
当 draft model 的预测与 target model 比较接近时（acceptance rate 高），每个 verify step 可以推进多个 token，实现 2-3× 的加速。
接受/拒绝准则
关键的数学保证：Speculative Decoding 不改变目标模型的输出分布。接受第 i 个 token 的概率为：
如果 draft model 的概率 p_draft 比 target 高（即 draft 过于自信），则按比例拒绝。被拒绝时，从修正后的分布中重新采样。这个过程保证了最终的输出分布与单独使用 target model 完全一致。
Draft Model 的选择
Draft model 可以是同系列的小模型（如用 LLaMA-7B 作为 LLaMA-70B 的 draft），也可以是通过 pruning/distillation 得到的模型。甚至可以用基于 n-gram 的简单模型。关键是 draft model 要足够快，且与 target model 的分布尽量接近。
Continuous Batching
在服务多用户请求时，传统的 Static Batching 将多个请求打包成一个 batch 一起处理，但必须等 batch 中最长的请求完成后才能释放所有请求。
Continuous Batching（也叫 Iteration-Level Batching）的改进：
- 每个 decode step 结束后，检查哪些请求已经完成（生成了 EOS token）
- 完成的请求立即释放，空出的 slot 可以插入新的请求
- 请求可以在任意 step 加入和离开 batch
Continuous Batching 显著提高了 GPU 利用率和吞吐量（throughput），是生产推理系统的标准做法。
PagedAttention 与 vLLM
PagedAttention 是 vLLM 系统的核心创新，灵感来自操作系统的虚拟内存管理。
问题：KV Cache 的内存浪费
传统 KV Cache 实现为每个请求预分配一块连续的显存，大小为 max_seq_len。但实际生成的长度通常远小于最大值，导致大量内存碎片（internal fragmentation）。此外，不同请求的 KV Cache 长度不同，也会产生外部碎片（external fragmentation）。
在 vLLM 的测试中，传统方案的内存浪费高达 60-80%。
PagedAttention 的解决方案
将 KV Cache 分成固定大小的 block（类似操作系统的页 page），使用 block table 来管理逻辑到物理地址的映射：
- KV Cache 不需要存储在连续的显存中
- 每个 block 包含固定数量 token（如 16 个）的 K、V 向量
- 新 token 产生时，只需分配新的 block，不需要移动已有数据
- 请求完成时，归还所有 block 到空闲池（free pool）
PagedAttention 实现了接近零的内存浪费（浪费仅来自最后一个 block 的未满部分），大幅提高了可同时服务的请求数量。
额外优势：共享前缀
PagedAttention 还支持 prefix sharing：如果多个请求有共同的前缀（如相同的 system prompt），它们可以共享前缀的 KV Cache blocks，进一步节省内存。这通过 copy-on-write 机制实现。
Inference 的 Roofline 分析
Roofline Model 帮助我们判断一个计算任务是 compute-bound 还是 memory-bound。
Arithmetic Intensity
将 AI 与硬件的 ridge point（计算峰值 / 内存带宽）比较：
- 如果 AI < ridge point：Memory-bound（受限于内存带宽）
- 如果 AI > ridge point：Compute-bound（受限于计算能力）
H100 的 Ridge Point
| 精度 | Peak FLOPS | HBM Bandwidth | Ridge Point |
|---|---|---|---|
| BF16 Tensor Core | 989 TFLOPS | 3.35 TB/s | ~295 FLOPs/byte |
| INT8 Tensor Core | 1979 TFLOPS | 3.35 TB/s | ~591 FLOPs/byte |
Decode 阶段的 Roofline 分析
在 decode 阶段，对于一个线性层 y = xW（batch=1）：
- FLOPs：2 × d_in × d_out（一次乘加）
- Bytes：2 × d_in × d_out（读取 FP16 权重）+ 少量输入/输出
- AI ≈ 1 FLOP/byte
而 H100 的 ridge point 是 295——AI 只有 ridge point 的 1/295。这意味着 decode 阶段的计算单元 99.7% 的时间都在等待数据。
Batching 的作用
Batch size = B 时：
当 B ≥ 295 时（对于 BF16），decode 才变成 compute-bound。但实际中受限于 KV Cache 的显存占用，往往很难达到如此大的 batch size。这就是推理优化的核心矛盾。
推理优化的本质
推理优化可以归结为两个目标：（1）减少要加载的数据量——Quantization（减少权重大小）、GQA/MQA（减少 KV Cache）；（2）增加每次加载后的计算量——Batching（增加 batch size）、Speculative Decoding（一次 verify 多个 token）。两者都在提高 arithmetic intensity，使计算更接近 compute-bound。
推理优化技术总览
| 技术 | 目标 | 机制 | 加速倍数 |
|---|---|---|---|
| KV Cache | 避免重复计算 | 缓存 K、V 向量 | 序列长度倍 |
| GQA/MQA | 减少 KV Cache | 共享 K、V heads | 间接（更大 batch） |
| INT8 Quantization | 减少模型大小 | 权重量化到 8-bit | ~1.5-2× |
| INT4 Quantization | 进一步压缩 | 权重量化到 4-bit | ~2-3× |
| Speculative Decoding | 减少 decode 步数 | Draft + Verify | ~2-3× |
| Continuous Batching | 提高 GPU 利用率 | 动态 batch 管理 | ~2-4× throughput |
| PagedAttention | 减少内存浪费 | KV Cache 分页管理 | ~2-4× throughput |
| Flash Attention | 减少 HBM 访问 | Tiled attention + fusion | ~2-4× attention |
小结
- LLM 推理分 Prefill（compute-bound）和 Decode（memory-bound）两个阶段
- KV Cache 避免重复计算，但在长序列/大 batch 下成为内存瓶颈
- GQA/MQA 通过共享 K、V heads 减少 KV Cache 大小
- Quantization（GPTQ、AWQ）将权重压缩到 INT8/INT4，减少显存和带宽需求
- Speculative Decoding 使用 draft model 加速生成，且不改变输出分布
- Continuous Batching 和 PagedAttention 分别从调度和内存管理层面提高吞吐量
- 推理优化的本质是提高 arithmetic intensity，让 GPU 从 memory-bound 走向 compute-bound