# Lecture 10 · 推理

> **CS336: Language Modeling from Scratch** · Stanford · Spring 2025
> 📅 May 1 · 💻 [可执行讲义](https://github.com/stanford-cs336/spring2025-lectures/blob/main/lecture_10.py) · 讲者 Percy Liang

---

## 承上启下

**上一讲（[[Lecture 9 · 缩放定律 I：基础]]）**：讨论的是训练前的预算分配问题——给定计算预算 $C$，如何在参数量 $N$ 和数据量 $D$ 之间分配才能最小化训练 loss。训练侧的核心指标是 loss、FLOPs、batch size、学习率与 scaling curve。

**本讲（推理，inference）**：视角从「训练一个模型花多少钱」切换到「用一个模型花多少钱」。训练是一次性成本，推理是随每次调用反复发生的成本——产品一旦跑起来，推理账单往往远超训练账单。更关键的是，推理 workload 的性质与训练完全不同：训练可以在整个序列维度上并行，把计算组织成大矩阵乘法；而自回归生成必须一个 token 一个 token 地串行进行，瓶颈从「训练算力」转为「生成阶段的内存带宽（memory bandwidth）+ 动态调度」。分析工具正是 [[Lecture 5 · GPU]] 的 roofline 模型与 [[Lecture 2 · PyTorch 与资源核算]] 的 FLOPs 核算方法。

**与后续的关联**：[[Lecture 12 · 评估]] 中跑 benchmark、[[Lecture 16 · 对齐 II：RLVR]] 中生成 RL rollout，本质上都是大规模推理 workload，推理效率直接决定评估与后训练的成本。下一讲 [[Lecture 11 · 缩放定律 II：细节]] 则回到训练侧，看工业界如何在实践中拟合与使用 scaling laws。

> [!abstract] 本讲一句话总览
> 推理分为 compute-bound 的 prefill 和 memory-bound 的 generation 两个阶段。KV cache 让生成的计算量从 $O(T^3)$ 降到 $O(T^2)$，但也把整个服务系统变成「KV cache 显存管理问题」；GQA/MLA/CLA/局部注意力、量化/剪枝/蒸馏、投机解码、continuous batching 与 PagedAttention，全都是围绕这个瓶颈做文章。

---

## 1. 推理 workload：训练之后的长期账单

推理无处不在，而且形态各异：

- **聊天机器人**：对延迟极其敏感，用户盯着屏幕等第一个字出现。
- **代码补全**：要求比聊天更低的延迟——补全必须在程序员打下一个字符之前出来。
- **批量数据处理**：如用 LLM 清洗/标注海量文档，不在乎单条延迟，只在乎总吞吐和总成本。
- **模型评估**：跑一遍 [[Lecture 12 · 评估]] 里的 benchmark 就是一次大规模批量推理。
- **Test-time compute**：o1/R1 式「先思考再回答」用更长的生成换更好的答案，直接把推理成本乘了一个数量级。
- **RL 训练中的 rollout 生成**：[[Lecture 16 · 对齐 II：RLVR]] 中策略模型要不断采样输出算 reward，推理慢会直接卡住训练循环。

这些场景对「快」的定义并不相同，因此需要先把指标拆清楚。

### 1.1 三个指标

| 指标 | 含义 | 关注场景 |
| :--- | :--- | :--- |
| **TTFT**（time-to-first-token） | 用户等到第一个 token 的时间，由 prefill 决定 | 交互体验 |
| **Latency**（每 token 延迟） | 生成阶段每个 token 的间隔，由 decode 决定 | 聊天/代码补全 |
| **Throughput**（吞吐） | 系统整体每秒产出的 tokens/s | 批处理、平台服务 |

三者经常互相冲突：

- **小 batch**：单请求延迟低，但 GPU 利用率差，每 token 成本高；
- **大 batch**：吞吐高、成本低，但请求要排队，TTFT 和每 token 延迟都变差；
- prefill 与 decode 两个阶段的最优 batch 策略也不同（见第 4 节）。

> [!tip] 直觉
> 推理优化没有「唯一的快」。聊天产品优化的是 TTFT 和每 token 延迟（人的阅读速度约 10 tokens/s，比这快就够用），批处理平台优化的是每百万 token 的成本。同一套硬件、同一个模型，为这两类目标设计出的 serving 系统完全不同——所以先问「为谁优化」，再谈技术。

---

## 2. 推理的算术强度

### 2.1 训练 vs 生成

回顾 [[Lecture 5 · GPU]]：**算术强度（arithmetic intensity）**定义为每从显存搬运 1 字节数据所做的浮点运算数：

$$\text{intensity} = \frac{\text{FLOPs}}{\text{bytes transferred}}$$

硬件也有一个对应的比值。以 H100 为例：bf16 算力峰值约 $989\ \text{TFLOP/s}$，HBM 带宽约 $3.35\ \text{TB/s}$，两者之比约

$$\frac{989 \times 10^{12}}{3.35 \times 10^{12}} \approx 295\ \text{FLOPs/byte}$$

这就是 roofline 的「屋脊点」：算术强度高于 295 时是 compute-bound（算力受限），低于 295 时是 memory-bound（带宽受限），此时算力再强也只能空转等数据。

**训练**：输入完整序列，可以在序列维并行，前向/反向都是大矩阵乘法，算术强度轻松超过屋脊点，是 compute-bound。

**生成**：每步只产出一个新 token，且必须等上一个 token 出来才能算下一个。大量操作退化为 matrix-vector 乘法，读权重和 KV cache 的搬运成本主导一切，是 memory-bound。

### 2.2 Matrix-vector 的问题

考察 MLP 层的一次矩阵乘 $Y = XW$，其中 $X \in \mathbb{R}^{B \times D}$（$B$ 为 batch 中的 token 数，$D$ 为隐藏维度），$W \in \mathbb{R}^{D \times F}$（$F$ 为 MLP 中间维度）：

- FLOPs：$2BDF$（每个输出元素一次乘加）；
- 数据搬运（bf16，每元素 2 字节）：$2(BD + DF + BF)$ 字节。

当 $D, F \gg B$ 时，分母由权重 $DF$ 主导，于是

$$\text{intensity} = \frac{2BDF}{2(BD + DF + BF)} \approx B$$

**算术强度约等于 batch 中的 token 数**。生成时若 $B = 1$：为了产出一个 token 的计算，要把整个 $W$ 从 HBM 读一遍，算术强度约为 1 FLOP/byte，离屋脊点 295 差了两个数量级——GPU 的算力利用率不到 1%。

> [!tip] 直觉
> 把权重从 HBM 搬到片上的成本是固定的「过路费」。batch 里的 token 越多，这笔过路费被摊得越薄；$B=1$ 时一个 token 独自付全款，自然亏到极点。这就是为什么「换更快的 GPU」不一定让单请求生成按比例变快——瓶颈不是算力，而是把参数和 KV cache 从 HBM 搬出来的带宽。

### 2.3 Attention 的算术强度：batch 帮不上忙

MLP 的强度随 $B$ 增长，但 attention 不同。设一次前向处理 $T$ 个 query token、历史共 $S$ 个 KV token，单头维度 $H$：

- FLOPs：$QK^\top$ 与 $AV$ 两次矩阵乘，约 $4STH$；
- 搬运：读 $K, V$（各 $S \times H$）与 $Q$（$T \times H$）、写输出，约 $4H(S+T)$ 字节量级。

$$\text{intensity}_{\text{attn}} \approx \frac{4STH}{4H(S+T)} = \frac{ST}{S+T}$$

代入两个阶段：

- **Prefill**（$T = S$，整段 prompt 一起算）：强度 $\approx S/2$。$S = 1024$ 时约 512，超过屋脊点，compute-bound；
- **Decode**（$T = 1$，一次一个 token）：强度 $\approx S/(S+1) < 1$，深度 memory-bound。

注意这个公式里**没有 $B$**：每条序列都有自己私有的 KV cache，batch 内不同请求之间无法共享读取。增大 batch 能摊薄 MLP 读权重的成本，却对 attention 读 KV cache 毫无帮助。

> [!warning] 易错点
> 「加大 batch 就能解决 memory-bound」只对 MLP/权重部分成立。长上下文场景下 KV cache 的读取量随 $B \times S$ 线性增长，batch 越大总搬运量越大——attention 部分的算术强度始终小于 1。这也是长上下文推理格外昂贵的根本原因。

---

## 3. KV Cache：复用前缀计算

### 3.1 朴素自回归推理

生成第 $t$ 个 token 时，如果把前面所有 token 都重新喂进 Transformer 完整算一遍前向：第 $t$ 步 attention 计算量为 $O(t^2)$，累加 $T$ 步：

$$\sum_{t=1}^{T} O(t^2) = O(T^3)$$

生成 $T$ 个 token 要 $O(T^3)$ FLOPs——大量计算在重复：前缀 token 的 K/V 表示每一步都被原样重算了一遍。

### 3.2 KV Cache

**KV cache** 的做法：对每层、每个 KV 头、每个历史 token，把算过的 key/value 向量缓存下来：

$$\text{KV cache 形状} \approx [\,B,\ S,\ L,\ K,\ H\,]$$

其中 $B$ 为 batch size，$S$ 为序列长度，$L$ 为层数，$K$ 为 KV 头数，$H$ 为单头维度。生成新 token 时：

- 新 token 只计算自己的 Q/K/V，并把 K/V 追加进 cache；
- 旧 token 的 K/V 直接从 cache 读，不再重算；
- attention 用新 Q attend 到全部历史 K/V。

于是每步 attention 降为 $O(t)$，总量降为 $O(T^2)$；MLP 每步只处理 1 个 token，总量 $O(T)$。代价是显存里多了一块随序列增长的状态。

由此推理天然分成两个阶段：

| 阶段 | 输入 | 特征 | 瓶颈 |
| :--- | :--- | :--- | :--- |
| **Prefill** | prompt 全部 token | 序列内并行，一次前向填满 cache | compute-bound |
| **Generation / Decode** | 每次 1 个 token | 自回归串行 | memory-bound |

> [!tip] 直觉
> KV cache 是「用空间换时间」的教科书案例：把 $O(T^3)$ 的重复计算换成 $O(T^2)$ 计算 + $O(T)$ 的显存状态。但这笔交易把问题的性质变了——推理系统的核心资源从「FLOPs」变成「装 KV cache 的显存 + 读 KV cache 的带宽」，后面所有优化（GQA、MLA、PagedAttention）都是在管理这块状态。

### 3.3 KV Cache 的内存成本

每个请求的 KV cache 随上下文长度线性增长：

$$\text{bytes} = \underbrace{B}_{\text{batch}} \times \underbrace{S}_{\text{序列长}} \times \underbrace{L}_{\text{层数}} \times \underbrace{K}_{\text{KV 头数}} \times \underbrace{H}_{\text{头维度}} \times \underbrace{2}_{K\text{ 与 }V} \times \underbrace{2}_{\text{bf16 字节数}}$$

> [!example] 算一算：Llama 2 13B 在 H100 上
> 配置：$L=40$ 层，$K=40$ 个 KV 头（标准 MHA），$H=128$，$D=5120$，$F=13824$，词表 32000。
> - 每 token 的 KV cache：$40 \times 40 \times 128 \times 2 \times 2 = 819{,}200$ 字节 $\approx 0.8\ \text{MB}$；
> - $S=1024$ 时每条序列约 $0.8\ \text{GB}$；
> - $B=64$ 时 KV cache 共约 $54\ \text{GB}$，加上 bf16 权重 $26\ \text{GB}$，恰好逼近 H100 的 80 GB 显存上限；
> - $B=256$ 直接 OOM。**限制 batch size（也就限制了吞吐）的不是算力，而是 KV cache 的显存占用。**

---

## 4. Throughput 与 Latency 的权衡

batch size 增大时：

- MLP 权重被 batch 内所有请求共享读取，每 token 的搬运成本被摊薄，throughput 提升；
- 但单个请求要与更多请求共享算力和带宽，且可能排队，latency 和 TTFT 变差；
- KV cache 总量随 batch 线性增长，显存上限反过来限制 batch 的天花板（见 3.3 的算例）。

常见策略：

- **prefill 与 decode 分离调度**：prefill 是 compute-bound 的大矩阵乘，decode 是 memory-bound 的小步快跑，两者混在一个 batch 里会互相拖累（一个长 prompt 的 prefill 会卡住所有人的 decode）。分开调度甚至分开部署（disaggregated serving）能让各自都跑在合适的 regime；
- **decode 用尽量大的 batch**：提高吞吐，直到 KV cache 吃满显存；
- **多副本部署**：复制 $M$ 份模型实例，latency 不变，throughput 乘 $M$——用钱换吞吐，简单有效；
- **模型/KV cache 分片**：结合 [[Lecture 7 · 并行 I：基础]] 的张量并行等手段把模型切到多卡，更复杂，但能服务单卡装不下的大模型，也顺带扩了 KV cache 的可用显存。

> [!tip] 直觉
> Throughput 和 latency 的冲突本质上是「公交车 vs 出租车」：公交车（大 batch）人均成本低但要等发车，出租车（小 batch）即叫即走但贵。serving 系统的工作就是根据 SLA 在两者之间找运营方案，而 KV cache 显存决定了「公交车最多能坐多少人」。

---

## 5. 降低 KV Cache：架构捷径

以下都是**有损（lossy）**捷径：改动模型架构来缩小 KV cache，可能付出一点准确率代价，换取更大的 batch 与更快的生成。

### 5.1 GQA / MQA

标准 **MHA**（multi-head attention）中 KV 头数等于 query 头数（$K = N$）；**MQA**（multi-query attention）走到另一个极端，所有 query 头共享 1 组 KV（$K = 1$）；**GQA**（grouped-query attention）取中间值，让每组 query 头共享一组 KV（$1 < K < N$）。

KV cache 的缩减比例正是头数之比：

$$\text{缩减倍数} = \frac{N}{K}$$

> [!example] GQA 让 batch 翻 4 倍
> Llama 2 13B 若把 $K = 40$ 改为 $K = 8$（每 5 个 query 头共享一组 KV），KV cache 缩小 5 倍：$B=64$ 时从约 54 GB 降到约 11 GB，于是 batch 可以开到 256，吞吐大幅提升。Llama 3 系列已默认采用 GQA（见 [[Lecture 3 · 架构与超参数]]）。

GQA 是当前主流折中：实验表明适度分组几乎不掉点，比 MQA 表达力好，比 MHA 省数倍 cache。

### 5.2 MLA（Multi-head Latent Attention）

DeepSeek-V2/V3 使用的思路：不直接缓存完整的 K/V，而是训练一个低秩压缩——把每个 token 的隐状态投影到一个低维 latent 向量 $c_t$ 缓存起来，需要时再用上投影矩阵恢复出各头的 K/V。

> [!example] DeepSeek-V2 的压缩幅度
> 128 头 × 128 维的完整 K/V 每 token 每层本需 $2 \times 16384$ 维；MLA 只缓存 $c_t \in \mathbb{R}^{512}$ 外加 64 维解耦的 RoPE key，共 **576 维**——缓存量降到约 1/57，与只分 2~3 组的 GQA 相当，但论文报告效果反而优于 MHA。

优点：

- KV cache 极大降低，generation 阶段要读的字节数同步降低；
- 因为「读 cache」正是 decode 的瓶颈，这直接转化为更大的 batch 和更高的吞吐，对 MoE 大模型（参见 [[Lecture 4 · 混合专家模型 MoE]]）的推理尤其重要。

难点：

- **与 RoPE 的兼容需要特殊处理**：MLA 的效率依赖「把 K 的上投影矩阵吸收进 Q 的投影」这一代数技巧，而 RoPE 给 K 乘上了随位置变化的旋转矩阵，夹在中间导致无法吸收。DeepSeek 的解法是解耦（decoupled）RoPE：额外留出 64 维专门携带位置信息，其余维度不做旋转；
- 架构复杂度明显上升，且必须从预训练阶段就采用，无法像量化那样对现成模型事后加装。

### 5.3 Cross-Layer Attention / Local Attention

**CLA（cross-layer attention）**：相邻的多层共享同一份 K/V——GQA 是「跨头共享」，CLA 是「跨层共享」，两者正交、可叠加。实验显示它能改善「准确率 vs KV cache 大小」的帕累托前沿。

**Local attention（滑动窗口注意力）**：每个 token 只 attend 最近 $W$ 个 token，KV cache 只需保留窗口内的条目，**与总上下文长度无关**。代价是长程信息会丢失、准确率受损，因此实践中与少数全局 attention 层交错使用——例如 Character.AI 的部署是每 6 层局部才配 1 层全局。

> [!tip] 直觉
> 5.1–5.3 的所有花样共享同一个目标函数：**在尽量不掉准确率的前提下，减少生成每个 token 时必须从 HBM 读出的 KV 字节数**。方向有三个正交的轴——跨头共享（GQA/MQA）、跨层共享（CLA）、压缩表示（MLA）、截断历史（local attention）——工业模型往往组合使用。

---

## 6. Transformer 之外：为推理重设架构

如果把「KV cache 随 $T$ 线性增长」看作 Transformer 的原罪，那么更激进的路线是换掉架构本身。

### 6.1 State Space Models

**SSM（state space model）**的核心思路：用一个**固定大小**的状态向量递归地总结全部历史，生成时只更新这个状态，而不是维护 $O(T)$ 的 KV cache——推理成本与上下文长度解耦。

代表工作：

- **S4**：结构化状态空间模型的开山之作，线性时不变系统，可用卷积并行训练；
- **Mamba**：引入输入依赖（selective）的状态更新参数，弥补了 SSM 在内容选择上的短板，在 1B 规模上与 Transformer 打平；
- **Jamba**：AI21 的混合架构，Transformer 层与 Mamba 层按 1:7 交错，再叠加 MoE——用少量 attention 层保住检索能力，用大量 Mamba 层省 cache；
- **BASED / MiniMax-01 等线性注意力变体**：线性注意力 + 局部注意力，再保留少量完整 attention 层。

优势是 $O(1)$ 的状态大小，长上下文推理内存和带宽都极省。问题在于固定大小的状态是有损压缩：**associative recall**（如「前文提过的某个电话号码是多少」这类精确检索）任务上 Transformer 仍显著更强。因此纯替代不容易，混合架构是当前的务实答案。

### 6.2 Diffusion LM

扩散式文本生成不做自回归：从一个全是噪声/掩码的序列出发，**并行地**对整个序列做多步迭代精化（refine）。理论上每步都能同时更新所有位置的 token，摆脱「一步一个 token」的串行瓶颈——Inception Labs 的 Mercury 声称在代码生成上比同档自回归模型快数倍。

挑战：离散文本上的扩散在数学上比连续图像更别扭，生成质量、可控性与训练配方都还在探索期，尚未撼动自回归的主流地位。

> [!tip] 直觉
> SSM 与 diffusion 分别攻击自回归 Transformer 的两个软肋：SSM 砍掉 $O(T)$ 的 KV cache（状态太大），diffusion 砍掉逐 token 的串行依赖（步数太多）。它们都是「推理成本倒逼架构设计」的例证——架构选择越来越不由训练效果单独决定。

---

## 7. 量化、剪枝与蒸馏

另一类有损捷径：不改架构，直接压缩模型本身。生成是 memory-bound 的，**权重字节数减半 ≈ 读权重的时间减半**，所以压缩直接换来速度。

### 7.1 Quantization

| 格式 | 字节 | 用途 |
| :--- | ---: | :--- |
| FP32 | 4 | 训练中的 optimizer state |
| BF16 / FP16 | 2 | 常规训练与推理默认 |
| FP8 | 1 | 高效推理；训练「艺高人胆大」也可用 |
| INT8 | 1 | 高效推理 |
| INT4 | 0.5 | 极限压缩 |

**PTQ（post-training quantization，训练后量化）**：用少量校准数据估计每个通道/分组的 scale 与 zero-point，把权重映射到低比特整数网格，不需要重新训练。**QAT（quantization-aware training，量化感知训练）**：训练时就模拟量化误差，效果更好但要动训练流程，大规模场景不总现实。

**LLM.int8() 的关键发现**：模型规模超过约 6.7B 后，激活中会系统性出现少数**离群（outlier）特征通道**，其数值大到让普通 int8 量化的网格完全失真。解法是混合精度分解——把约 0.1% 的 outlier 通道单独用 fp16 计算，其余走 int8。代价是推理反而比 fp16 慢 15~23%（分解本身有开销），它换来的是显存减半而非速度。

**AWQ（activation-aware weight quantization）**：观察到权重的重要性由激活分布决定，保护约 0.1%~1% 的关键权重（通过按通道缩放而非真的保留 fp16），其余压到 int3/int4。fp16 → int3 可得约 4 倍显存缩减与 3.2 倍实测加速。

> [!warning] 易错点
> 量化的评估不能只看 perplexity——PPL 对少数关键能力的退化不敏感（这与 [[Lecture 12 · 评估]] 讲的「PPL 不等于下游能力」是同一件事）。低比特量化后应在下游任务上复测，尤其是长推理链任务，误差会沿生成步数累积。

### 7.2 Pruning + Distillation

**剪枝（pruning）**：按重要性评分删除整个层、attention 头或隐藏维度，直接得到一个更小、更快的模型（结构化剪枝，硬件友好）。剪完的模型通常明显掉点。

**蒸馏（distillation）**：用原模型（teacher）的输出分布做监督信号，训练/修复剪枝后的小模型（student），把丢掉的能力找补回来。

两种得到小模型的 recipe：

| 路线 | 步骤 |
| :--- | :--- |
| From scratch | 设计快架构 → 从头预训练 |
| Prune + Distill | 取大模型 → 重要性评分并剪枝 → 用 teacher 蒸馏修复 |

NVIDIA 的 Minitron 系列与 Llama 3.2 的 1B/3B 都走第二条路线：其总成本可比从头训练低一个数量级，因为大模型已经付过的「学费」通过蒸馏被继承了下来。

---

## 8. Speculative Decoding：用小模型猜，大模型验

这是**无损（lossless）**加速：输出分布与目标模型逐 token 精确一致。

出发点是 decode 阶段的一个不对称性：**生成慢，验证快**。生成 $K$ 个 token 要 $K$ 次串行前向；但给定 $K$ 个候选 token，用一次并行前向就能算出它们每个位置的概率（像 prefill 一样并行）。又因为 decode 是 memory-bound，一次前向处理 1 个还是 $K$ 个 token，读权重的时间几乎一样——**验证近乎免费**。

**投机采样（speculative sampling）**流程（讲义记号：draft 模型为 $p$，target 模型为 $q$）：

1. 便宜的 draft 模型自回归猜 $K$ 个 token（如 $K=4$）；
2. target 模型**一次并行前向**算出这些位置上的概率；
3. 对每个候选 token $x \sim p$，以概率

$$\alpha(x) = \min\!\left(1,\ \frac{q(x)}{p(x)}\right)$$

   接受；一旦拒绝，从**残差分布**

$$r(x) = \frac{\max\bigl(0,\ q(x) - p(x)\bigr)}{\sum_{x'} \max\bigl(0,\ q(x') - p(x')\bigr)}$$

   中重新采样一个 token 并停止本轮；
4. 这个修正的 rejection sampling 保证最终每个 token 都精确服从 $q$。

> [!example] 两词表上的正确性证明
> 设词表只有 $\{A, B\}$，且 draft 高估了 $A$：$p(A) > q(A)$（于是 $p(B) < q(B)$）。
> - 输出 $A$ 的概率 $= p(A) \cdot \frac{q(A)}{p(A)} = q(A)$ ✓
> - 输出 $B$ 的概率 $= \underbrace{p(B) \cdot 1}_{\text{直接接受}} + \underbrace{p(A)\bigl(1 - \tfrac{q(A)}{p(A)}\bigr)}_{\text{拒绝 }A\text{ 后重采样，残差全在 }B} = p(B) + p(A) - q(A) = 1 - q(A) = q(B)$ ✓
> 无论 draft 猜得多差，输出分布都严格等于 target——draft 质量只影响**速度**（接受率），不影响**正确性**。

```mermaid
sequenceDiagram
    participant Draft as "Draft 模型 p"
    participant Target as "Target 模型 q"
    Draft->>Draft: 串行猜 token1..token4
    Draft->>Target: 提交候选序列
    Target->>Target: 一次前向并行算各位置概率
    Target-->>Draft: 按 min(1, q/p) 接受前 k 个，拒绝处从残差重采样
```

关键性质与实践：

- **exact sampling**：与直接从 target 采样不可区分，无需任何质量评估；
- draft 越接近 target，接受率越高、加速越明显。典型搭配是同家族小模型：70B target + 8B draft、8B target + 1B draft；
- draft 与 target 必须共用同一个 tokenizer（参见 [[Lecture 1 · 概览与分词]]），否则「逐 token 验证」无从谈起。

扩展：

- **Medusa**：不用独立 draft 模型，在 target 模型上加多个解码头一次并行猜多个后续 token，省掉 draft 的前向；
- **EAGLE**：用 target 模型最后一层的 hidden features（而非 token）做自回归 draft，猜测质量更高、接受率更好。

> [!tip] 直觉
> 投机解码把 memory-bound 阶段「算力大量闲置」这个坏消息变成好消息：反正读一遍权重的时间是固定的，不如让这次前向多验证几个 token，把闲置的 FLOPs 用起来。它在概念上与 rejection sampling 完全同源——用一个便宜的 proposal 分布逼近昂贵的目标分布，错了就修正。

---

## 9. 动态请求：Continuous Batching 与 PagedAttention

前面假设 batch 是静态给定的；线上真实流量是 **ragged workload**（参差不齐的负载）：

- 请求到达时间不同；
- prompt 长度不同；
- 生成长度不同且事先未知（有人生成 10 个 token 就停，有人生成 1000 个）；
- 大量请求共享同一个 system prompt 前缀；
- 同一 prompt 可能要采样多个候选（best-of-n）。

静态 batching 下，整个 batch 必须等最慢的请求生成完才能换人，GPU 大量空转。

### 9.1 Continuous Batching

Orca 提出两个核心机制：

- **Iteration-level scheduling（迭代级调度）**：调度的粒度不是「请求」而是「一步 decode」。每个 decode step 结束后重新组 batch——完成的请求立即腾位，新请求随时插入，不必等整个 batch 收工；
- **Selective batching（选择性批处理）**：不同请求序列长度不同，无法拼成规整张量。解法是把所有序列的 token 拼接（flatten）成一个大 batch 跑非 attention 算子（MLP、LayerNorm 等对 token 独立的运算），只在 attention 处按序列拆开分别计算——因为只有 attention 关心序列边界。

### 9.2 PagedAttention

传统 KV cache 分配按「最坏情况」给每个请求预留最大长度的连续显存，造成严重浪费：内部碎片（预留了但没用到）加外部碎片（空闲块不连续没法用），实测有效利用率可能只有 20%~40%。

**PagedAttention**（vLLM 的核心）把操作系统的虚拟内存分页思想搬进 KV cache 管理：

- 把每个请求的 KV cache 切成固定大小的 **block**（如每块 16 个 token），按需分配；
- 逻辑上连续、物理上不连续，靠 block table 做「页表」翻译；
- 天然支持 **prefix sharing**：共享同一 system prompt 的请求映射到同一批物理 block，只存一份；
- 分叉采样（一个 prompt 出多个候选）用 **copy-on-write**：共享的 block 被写入时才复制。

```mermaid
graph LR
    REQ["请求的逻辑 KV cache"] --> B1["block 1"]
    REQ --> B2["block 2"]
    REQ --> B3["block 3"]
    B1 -.物理映射.-> HBM1["HBM page A"]
    B2 -.物理映射.-> HBM2["HBM page Q"]
    B3 -.物理映射.-> HBM3["HBM page C"]
```

vLLM 还叠加了一系列系统优化：把 block 读取融合进 attention kernel（fused kernel）、FlashAttention / FlashDecoding（参见 [[Lecture 6 · Kernel 与 Triton]]）、CUDA graphs 消减 kernel launch 开销。综合下来相对此前系统有 2~4 倍吞吐提升。

> [!tip] 直觉
> KV cache 管理面对的问题——变长、动态到达、共享前缀、碎片化——和操作系统管理进程内存的问题一模一样，所以答案也一样：分页 + 页表 + 写时复制。Serving 的进步很多不是 ML 创新，而是把 OS/数据库社区五十年的老智慧搬进 GPU 显存。

---

## 总结

```mermaid
mindmap
  root((推理))
    Metrics
      TTFT
      latency
      throughput
    Workload
      prefill compute-bound
      generation memory-bound
      KV cache
    Architecture
      GQA
      MLA
      CLA
      local attention
      SSM
    Compression
      quantization
      pruning
      distillation
    Speedups
      speculative decoding
      Medusa
      EAGLE
    Serving
      continuous batching
      PagedAttention
      prefix sharing
```

**关键要点**：

1. **推理和训练是两种 workload**
   - 训练可在序列维并行、compute-bound；生成串行、算术强度约等于 batch size，深度 memory-bound。H100 屋脊点约 295 FLOPs/byte，$B=1$ 的生成只有约 1。

2. **KV cache 是推理系统的核心状态**
   - 它把 $O(T^3)$ 计算降到 $O(T^2)$，但吃显存和带宽；限制 batch（吞吐）的往往是 KV cache 而非算力。

3. **Prefill 和 generation 要分开优化**
   - prefill 像训练（大矩阵乘），generation 像反复搬运权重和 cache；两阶段混合调度会互相拖累。

4. **GQA/MLA/CLA/local attention 都是在减 KV cache**
   - 跨头共享、压缩表示、跨层共享、截断历史四条正交的轴；架构设计越来越受推理成本驱动，SSM 与 diffusion 是更激进的重设。

5. **量化/剪枝/蒸馏用少量精度换字节数**
   - memory-bound 下权重减半近似读取时间减半；outlier 是量化的头号敌人（LLM.int8()、AWQ 都在处理它）。

6. **Speculative decoding 是 lossless shortcut**
   - 小模型猜、大模型并行验，修正的 rejection sampling 保证输出分布精确等于 target；draft 质量只影响速度不影响正确性。

7. **Serving 是系统问题**
   - continuous batching（迭代级调度 + 选择性批处理）和 PagedAttention（分页 + 前缀共享 + 写时复制）是把 OS 思想搬进 LLM serving。

**下一讲预告（[[Lecture 11 · 缩放定律 II：细节]]）**：

推理说明了模型训练后的成本；下一讲回到训练前的模型设计，深入看 CerebrasGPT、MiniCPM、DeepSeek 等模型如何在实践中拟合 scaling laws、选择 batch size 和学习率，并讨论 muP 是否真能让超参数跨规模迁移。

---

## 复习自测

> [!question]- Q1：为什么在 H100 上 $B=1$ 的自回归生成必然是 memory-bound？给出量化论证。
> MLP 层矩阵乘的算术强度约等于 batch 中的 token 数，$B=1$ 时约为 1 FLOP/byte；而 H100 的屋脊点是 $989\ \text{TFLOP/s} \div 3.35\ \text{TB/s} \approx 295$ FLOPs/byte。1 远小于 295，说明时间花在从 HBM 搬权重上，算力利用率不足 1%——瓶颈是带宽而非算力。

> [!question]- Q2：算一算：Llama 2 13B（$L=40, K=40, H=128$，bf16），$S=1024$、$B=64$ 时 KV cache 多大？这说明什么？
> 每 token：$40 \times 40 \times 128 \times 2 \times 2 \approx 0.8$ MB；每序列（1024 token）约 0.8 GB；$B=64$ 时约 54 GB。加上 26 GB 权重正好逼近 H100 的 80 GB。说明限制 batch size（进而限制吞吐）的是 KV cache 显存，这正是 GQA/MLA 等「减 cache」架构的动机。

> [!question]- Q3：为什么增大 batch 能提升 MLP 的吞吐，却帮不了 attention？
> MLP 的权重 $W$ 是所有请求共享的，batch 越大，读一遍 $W$ 的固定成本被摊得越薄（强度 $\approx B$）。而每条序列的 KV cache 是私有的，batch 内无法共享读取——attention 的算术强度 $\approx ST/(S+T)$，公式里根本没有 $B$；decode 时（$T=1$）恒小于 1，始终 memory-bound。

> [!question]- Q4：投机解码为什么是无损的？为什么「验证 $K$ 个 token」不比「生成 1 个 token」慢多少？
> 无损性来自修正的 rejection sampling：对 draft 提议的 $x \sim p$ 以 $\min(1, q(x)/p(x))$ 接受，拒绝时从残差分布 $r(x) \propto \max(0, q(x)-p(x))$ 重采样，可证输出边际分布精确等于 $q$。验证便宜有两个原因：K 个候选位置的概率可以在一次前向中并行算出（类似 prefill）；且 decode 是 memory-bound，一次前向的耗时由读权重主导，处理 1 个还是 $K$ 个 token 时间几乎相同。

> [!question]- Q5：MLA 与 RoPE 为什么冲突？DeepSeek 如何解决？
> MLA 的推理效率依赖把 K 的上投影矩阵「吸收」进 Q 的投影，从而只需缓存低维 latent。但 RoPE 在 K 上乘了随位置变化的旋转矩阵，卡在两个投影之间使吸收不再成立。DeepSeek 的解法是解耦 RoPE：额外留出 64 维专门做旋转、携带位置信息，其余 512 维 latent 不做旋转，两部分拼接后参与 attention。

---

## 参考资料

- 💻 [lecture_10.py（官方可执行讲义）](https://github.com/stanford-cs336/spring2025-lectures/blob/main/lecture_10.py)
- 📖 [Scaling Book: Inference](https://jax-ml.github.io/scaling-book/inference/)
- 📄 [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://arxiv.org/pdf/2305.13245.pdf)
- 📄 [DeepSeek-V2](https://arxiv.org/abs/2405.04434)
- 📄 [Reducing Transformer Key-Value Cache Size with Cross-Layer Attention](https://arxiv.org/abs/2405.12981)
- 📄 [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://arxiv.org/abs/2312.00752)
- 📄 [LLM.int8()](https://arxiv.org/abs/2208.07339)
- 📄 [AWQ](https://arxiv.org/abs/2306.00978)
- 📄 [Compact Language Models via Pruning and Knowledge Distillation (Minitron)](https://arxiv.org/abs/2407.14679)
- 📄 [Fast Inference from Transformers via Speculative Decoding](https://arxiv.org/abs/2211.17192)
- 📄 [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://arxiv.org/abs/2401.10774)
- 📄 [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://arxiv.org/abs/2401.15077)
- 📄 [Orca: A Distributed Serving System for Transformer-Based Generative Models](https://www.usenix.org/conference/osdi22/presentation/yu)
- 📄 [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://arxiv.org/pdf/2309.06180.pdf)
- 🧰 [vLLM](https://github.com/vllm-project/vllm)