系统第 10 讲 / 共 17 讲

第 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 的大小取决于模型架构和序列长度：

KV Cache Size = 2 × n_layers × n_heads × d_head × seq_len × bytes_per_element

其中因子 2 代表 K 和 V 两个缓存。用模型超参数表示：

KV Cache Size = 2 × L × d_model × S × dtype_size

具体例子：

模型	层数 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。

MHA: n_kv_heads = n_heads（每个 head 独立的 K、V）
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_quantized = round(x / scale) + zero_point
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 一样并行计算），实际增加的时间很少。

Speedup ≈ E[accepted_tokens + 1] / (T_draft × K + T_verify)

当 draft model 的预测与 target model 比较接近时（acceptance rate 高），每个 verify step 可以推进多个 token，实现 2-3× 的加速。

接受/拒绝准则

关键的数学保证：Speculative Decoding 不改变目标模型的输出分布。接受第 i 个 token 的概率为：

P(accept) = min(1, p_target(x_i) / p_draft(x_i))

如果 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

Arithmetic Intensity (AI) = FLOPs / Bytes Accessed

将 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 时：

AI = 2 × B × d_in × d_out / (2 × d_in × d_out + ...) ≈ 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