CS336 · 第 10 课
训练阶段追求稳定地消化大批量 token；推理阶段面对的是动态请求、自回归生成和持续增长的 KV cache。同一套模型权重没有改变，但 workload 改变了，因此性能瓶颈和优化目标也随之改变。
Spring 2026 官方来源：Lecture 10 trace · 课程日程
本页是中文学习导读；讲次主题、教师与材料版本以 Stanford 官方页面和讲义为准。
前面主要讨论如何高效训练模型；推理面对的是另一种 workload：请求到达时间不固定，生成必须逐 token 推进，KV cache 会随上下文增长，而且延迟与吞吐量往往互相牵制。
因此本讲不把“推理”视为一次更小的训练 forward，而是回答：prefill 和 decode 各在搬什么、算什么，系统怎样在用户延迟、总吞吐和显存容量之间做权衡？
prefill / decode workload → arithmetic intensity → 权重与 KV traffic → 压缩数据（GQA / MLA / quantization）→ 减少必须执行的步骤（speculative decoding）→ 用 continuous batching 与 PagedAttention 管理动态请求。
Training
tokens → forward → loss → backward → gradients → optimizer → updated weights
memory = weights + activations + gradients + optimizer states + buffers
Inference
prompt → prefill → KV cache → decode one token → append new K/V → repeat
memory = weights + KV cache + temporary activations
(no gradients, no optimizer states, no backward)推理不是“训练减掉 optimizer”这么简单：自回归请求新增了随 batch、层数和 context 增长的 KV cache，也新增了 TTFT、inter-token latency 与动态调度问题。
| 指标 | 定义 | 重要性 |
|---|---|---|
| TTFT | Time-to-first-token | 用户等待时间 |
| Latency | 秒/token | 生成速度 |
| Throughput | tokens/second | 系统容量 |
| 阶段 | 并行形态 | 更可靠的 mental model |
|---|---|---|
| Prefill | 一次处理 prompt 中的大量 token，容易形成较大的 GEMM | 并行度高，因而更容易进入 compute-bound 区域；是否真的 compute-bound 仍取决于 shape、batch、精度和 kernel |
| Decode MLP | 每个请求一次生成一个 token，但可把多个请求组成 concurrent batch | batching 提高权重 reuse 和 arithmetic intensity，因此并非永远 memory-bound |
| Decode attention | 新 query 必须读取随 context 增长的 KV cache | 当 \(T=1\) 时 arithmetic intensity 很低，通常强烈受 memory bandwidth 限制 |
“decode 是 memory-bound”只是入口，不是终点。需要把 MLP 的权重复用与 attention 的 KV cache traffic 分开分析，才能知道 batching 对哪里有效。
# MLP 前向传播（SwiGLU）:
# 读 X (B×T×D), W_gate (D×F), W_up (D×F), W_down (F×D)
# 计算: U = X @ W_gate, G = X @ W_up, Y = Swish(G)*U @ W_down
# 写 Y (B×T×D)
FLOPs = 6 × B × T × D × F
Bytes = 4BTD + 6DF + 4BTD + 4BTD + 4BTF + 4BTD
≈ 4BTD + 4BTF + 6DF (忽略写入)
# 假设 B*T << D, F:
算术强度 ≈ B*T
# Prefill: B*T 很大 → compute-bound
# Generation: B*T = B（batch size）
# B=1: 强度 = 1 → memory-bound
# B=100: 强度 = 100 → 接近 compute-bound# Attention（FlashAttention 优化后）:
# 读 Q (B×T×D), K (B×S×D), V (B×S×D)
# 计算: Q @ K^T, softmax, @ V
# 写 Y (B×T×D)
FLOPs = 4 × B × S × T × D
Bytes = 4BSD + 4BTD + 4BTD ≈ 4BSD + 4BTD
算术强度 = S*T / (S + T)
# Prefill: T = S → 强度 = S/2 → 大 S 时 compute-bound
# Generation: T = 1 → 强度 ≈ 1 → memory-bound
# 在这个一阶模型里 B 同时缩放 FLOPs 与 KV bytes，因此不会像 MLP 那样
# 直接提高权重 reuse；真实实现仍受 paged layout、并行度和 kernel 影响。核心洞察：
Llama 2 13B 在 H100 上的理论性能:
# Batch size = 1:
# 内存: 参数 26GB (bf16) + KV cache ~0.5GB
# 延迟: ~13 ms/token（受内存带宽限制）
# 吞吐量: ~77 tokens/s
# Batch size = 64:
# 内存: 26GB + 32GB KV cache
# 延迟: ~20 ms/token（更慢）
# 吞吐量: ~3200 tokens/s（更高！）
# 权衡: 大 batch → 高吞吐、高延迟
# 小 batch → 低延迟、低吞吐这些数字只用于展示 batch 带来的方向性权衡，不应当作特定 serving engine、精度或软件版本的性能承诺。
设已有 prefix token \(x_1,\ldots,x_n\)。在某一 attention layer，它们已经产生：
Q₁...Qₙ, K₁...Kₙ, V₁...Vₙ生成新 token \(x_{n+1}\) 时，本层先计算 \(Q_{new},K_{new},V_{new}\)。新的 query 必须与所有允许看到的历史 keys 比较，并用 attention weights 读取历史 values：
\[ a=\operatorname{softmax}\left(\frac{Q_{new}[K_1;\ldots;K_n;K_{new}]^\top}{\sqrt{d_h}}\right),\qquad y_{new}=a[V_1;\ldots;V_n;V_{new}] \]所以历史 \(K_1\ldots K_n\) 和 \(V_1\ldots V_n\) 在未来每一步还会被读取。若不保存，每生成一个 token 都要让整个 prefix 重新经过各层来重算它们。
past Q_i
→ 只在“token i 作为 query”那一次计算自己的 attention output
→ 未来 token 有自己的 Q_new，不再查询 Q_i
→ discard
past K_i / V_i
→ 每个未来 Q_new 都可能匹配 K_i 并读取 V_i
→ cacheKV cache 是由未来数据依赖推出来的，不是因为 K/V 比 Q 特殊或更容易存。
| 阶段 | toy 场景 | 做什么 | 主要输出 |
|---|---|---|---|
| Prefill | 一次输入 5000 prompt tokens | 并行处理整个 prompt；causal mask 保持语义 | 每层 5000 个历史 K/V + 首个 next-token logits |
| Decode | 每轮新增 1 token | 只算新 token 的层输出，同时读取历史 KV | 新 logits，并向每层 cache 追加一个 K/V |
Prefill 的较大 token 维通常更容易形成高效 GEMM，因而常更 compute-heavy；decode 尤其是长上下文 attention，每步计算一个小 query 却读取大量历史 KV，常更 bandwidth-sensitive。这些是 workload-dependent 判断，不是无条件定律。
对 MHA，若每层有 \(n_{kv}\) 个 KV heads、每 head 维度 \(d_h\)、缓存 dtype 每元素 \(s\) bytes，则每个 sequence 的 cache：
\[ M_{KV}=2\times L\times S\times n_{kv}\times d_h\times s \]其中 2 来自 K 与 V，\(L\) 是层数，\(S\) 是当前 cache token 数。若 batch 中每个请求各有 cache，再乘或逐请求求和；padding、block metadata 与 allocator overhead 另算。
toy MHA: n_kv=64, d_h=128, bf16 s=2 bytes
per token per layer = 2 × 64 × 128 × 2 = 32 KiB
S=100,000, L=80:
32 KiB × 100,000 × 80 ≈ 244 GiB（十进制约 262 GB）原始数字若把 KB、GB 与 KiB、GiB 混用会看似得到“256GB”。资源账本必须声明单位制；方向结论不变：长 context 与大 concurrent batch 会让 KV 成为主导显存对象。
“不用 cache，显存就更省，所以一定更快。”显存少了，但每一步重算 prefix；生成 \(N\) 个 token 时会重复执行大量历史计算。
“cache 全部 Q/K/V 更完整。”历史 Q 没有未来消费者，只增加约 50% cache 元素而不避免重算。
“decode 永远 memory-bound。”大 batch 可显著提高 MLP 权重复用；短 context 或不同 kernel 下瓶颈可能变化。应把 decode MLP 与 decode attention 分开。
GQA 把 KV cache 缩小 n_heads / n_kv_heads 倍
Llama 2 70B: 64 Q heads, 8 KV heads → 8× 缩小
DeepSeek-V2 的 MLA:
不存完整的 K, V 向量
存一个低秩压缩的 latent 向量 c (C=512 维)
需要时通过上投影恢复
KV cache: 512 + 64 (for RoPE) = 576 维
vs MHA: 64 × 128 × 2 = 16,384 维
压缩比: 28×！
跨层共享 KV:
GQA 是共享 KV across heads
CLA 是共享 KV across layers
例如: 每 4 层共享一组 KV
KV cache 缩小 4×
部分层用 local attention（只关注附近 token）。若实现会淘汰窗口外 K/V，那么这些 local layers 的 cache 不再随完整历史 context 无限增长，而被 attention window size 上界约束；它仍然线性依赖 window size。
混合策略（character.ai）:
- 每 6 层有 1 层 global attention
- 其余 5 层用 local attention
- 有效信息传播范围会随 window 与层数扩大，但不等于单层可直接访问完整历史
许多逐 token decode workload 对权重与 KV 的 memory traffic 很敏感
减少表示大小 → 减少数据搬运 → 可能降低延迟并提高吞吐
量化级别:
fp32 (4 bytes) → bf16 (2 bytes) → fp8 (1 byte) → int4 (0.5 byte)
int4 相对 fp32 可把权重表示压缩到约 1/8；实际加速取决于 kernel、量化元数据、batch、瓶颈位置和硬件支持，不能直接等同于 8×。
# 训练后量化: 不需要重新训练
# 简单量化（per-tensor）:
scale = max(abs(x)) / 127 # int8
x_quantized = round(x / scale)
# 问题: outlier 会导致精度大幅下降
# LLM 中有 ~0.1% 的 outlier（值特别大的元素）Dettmers et al. (2022):
1. 对大多数元素做 int8 量化
2. 对 outlier 保持 fp16
3. 矩阵乘法: outlier 用 fp16, 其余用 int8
效果: 几乎无损，但 15-23% 更慢（因为混合精度）
Lin et al. (2023) 的关键洞察：权重是否难以量化，不能只看权重自身；输入 activation 的分布会放大量化误差。AWQ 用一小批校准 activation 识别 salient channels，再搜索每个输入通道的缩放因子。
方法：
1. 统计 activation，定位对输出更敏感的输入通道
2. 对 activation 与 weight 做数学等价的逆向缩放：一边乘 \(s\)，另一边除 \(s\)
3. 把缩放后的权重统一量化到低比特，而不是永久保留一小撮 FP16 权重
Implementation-specific：压缩率与端到端加速取决于 bit width、group size、kernel、batch 与硬件。论文中的单个速度数字不是所有 workload 的保证。
剪枝试图移除对输出影响较小的权重、通道、attention heads 或整个层。非结构化稀疏只减少非零元素，若硬件/库不能跳过稀疏位置，wall-clock 不一定变快；结构化剪枝更容易映射到 dense kernel，但通常更伤精度。
Spring 2026 官方 trace 给出的 recipe 是：先定义更快的 student 架构，再用原模型的输出/中间分布初始化或训练它，最后通过蒸馏修复精度。蒸馏可以与 GQA、local attention、较浅网络或更窄 hidden size 组合。
loss = alpha * cross_entropy(student_logits, labels) \
+ (1 - alpha) * kl_div(student_logits / T, teacher_logits / T) * T**2
# 关键权衡：教师生成数据有成本；student 的架构必须真的对应更低推理开销。关键洞察：验证比生成快。
给定一个序列，大模型可以一次性评估多个 token（像 prefill）。但生成一个 token 需要一次完整的前向传播。
投机解码：用小模型（draft）快速猜几个 token，然后用大模型一次性验证。
投机解码 (Leviathan et al., 2023; Chen et al., 2023)
1. Draft 模型（小）按 \(q\) 生成 γ 个候选 token：\(d_1,d_2,\ldots,d_\gamma\)
2. Target 模型（大）一次并行算出这些位置的 \(p\)
3. 对每个候选抽一个随机数 \(u\sim U(0,1)\)：若 \(u\le \min(1,p(d_i)/q(d_i))\)，接受 \(d_i\)
4. 首次拒绝时，从归一化后的正差分布 \([p-q]_+\) 采样校正 token，并停止本轮验证
5. 若全部候选都接受，再从 target 的下一个位置采一个 token
为什么不能只写“\(p(d_i)>q(d_i)\) 就接受”？当 \(p<q\) 时仍应有 \(p/q\) 的接受概率；少了这个随机步骤与拒绝校正，最终序列就不再严格服从 target 分布。
Exact：规则正确实现时，输出分布与直接从 target 采样相同。Implementation-specific：是否加速取决于接受率、draft 成本、target 验证 shape 与系统实现。
不使用投机解码:
生成 γ 个 token → γ 次大模型前向 → γ × T_large
使用投机解码:
1 次 draft 前向 + 1 次大模型验证
→ γ × T_small + T_verify
→ 只有接受率足够高、draft 足够便宜且 target 能高效并行验证时才会加速
Benchmark-specific：论文与 serving engine 报告的倍数都绑定 target/draft 模型、硬件、batch、prompt/decode 长度、sampling 参数与实现；这里不把某个速度倍数当作普适结论。
训练: 固定 batch, 固定长度
batch_size × sequence_length × hidden_dim
推理: 请求动态到达，长度不一
请求 1: 已生成 10 tokens
请求 2: 刚到达，prefill 中
请求 3: 已生成 100 tokens
朴素方案: 等所有请求完成再返回
→ 先到的请求被后到的拖慢
Yu et al. (2022) - Orca
不是等整个 batch 完成，而是每步都可以加入/移除请求
每个 decode step:
1. 检查是否有新请求到达 → 加入 batch
2. 检查是否有请求完成 → 移除
3. 对所有活跃请求做一步 decode
Paper-specific result：Orca 在论文的 GPT-3 175B、指定硬件与请求负载下，相对 NVIDIA FasterTransformer 报告了同等 latency 下最高 36.9× throughput。这个数字不能脱离 baseline、模型、arrival trace 与 latency target 当成 continuous batching 的普适倍数。
问题: 不同请求长度不同，不能组成规则 tensor
解决方案:
- Attention: 每个序列单独处理（或 padding）
- 非 attention: 拼接所有序列 [total_tokens, hidden_dim]
这样非 attention 部分仍然高效
朴素 KV cache 分配:
每个请求预留 max_seq_len 的空间
问题 1: 内部碎片
大多数请求远不到 max_seq_len
预留的空间浪费
问题 2: 外部碎片
不同请求的 KV cache 大小不一
内存中出现不连续的空洞
Kwon et al. (2023) - vLLM
核心思想: 像操作系统管理内存页一样管理 KV cache
效果:
- 内存利用率从 ~40% 提升到 ~96%
- 支持更大的 batch size
- 支持共享前缀（system prompt）
多个请求共享 system prompt 的 KV cache
当某个请求需要修改时，才复制
类似操作系统的 fork()
| 系统 | 官方 trace 强调的定位 |
|---|---|
| vLLM | PagedAttention 与通用高吞吐 serving 的常用默认选择 |
| SGLang | RadixAttention / prefix reuse，适合 agentic 与共享前缀工作负载 |
| TensorRT-LLM | NVIDIA GPU 上高度优化的 kernel、量化和多 GPU runtime |
| llama.cpp | CPU / 本地推理与紧凑量化生态 |
引擎名称不是性能保证。最终结果取决于模型结构、量化格式、batch 到达过程、prefix reuse、GPU 型号和调度策略，必须在真实请求分布上测 TTFT、吞吐与尾延迟。
Mamba, Jamba, etc.
用 O(1) 状态替代 O(n) KV cache
推理时不需要存储整个 KV cache
但: 不擅长 associative recall
混合架构（Transformer + SSM）是当前趋势
非自回归生成
所有 token 同时生成，多步 refinement
优势: 并行生成，速度快
劣势: 文本质量目前不如自回归
| 方向 | 技术 | 效果 |
|---|---|---|
| 减少 KV cache | GQA, MLA, CLA, local | 2-32× 内存节省 |
| 减少模型成本 | 量化、结构化剪枝、蒸馏 | 降低权重流量与计算量；收益需实测 |
| 投机解码 | draft model | 2-3× 加速 |
| 系统优化 | continuous batching, paged attention | 减少空槽与 KV 碎片；收益依 engine、请求分布与 latency 目标实测 |
核心原则：
问题 1：为什么长上下文的逐 token decode attention 通常特别 bandwidth-sensitive？
问题 2：投机解码为什么能加速推理？
问题 3：PagedAttention 解决了什么问题？
问题 4：Continuous batching 的核心思想是什么？
问题 5：AWQ 量化的核心思想是什么？
答案：新 token 用自己的 \(Q_{new}\) 查询所有历史 \(K_i\)，再读取对应 \(V_i\)。历史 \(Q_i\) 只在 token \(i\) 产生自身 attention output 时消费过一次，未来没有数据依赖。
若回答“因为 Q 更小/更大”，说明你在用存储大小解释机制；真正原因是张量的未来消费者不同。
答案：不是。它并行计算 5000 个 prompt positions 的层表示并建立 KV cache，最后用末位置 logits 开始采样第一个新 token。Prompt tokens 是输入，不是响应输出。
若答“是”，说明你把并行处理输入 positions 与 autoregressive 输出步骤混在了一起。
答案：在其他条件不变时，query heads 与模型输出宽度可以不变；每 token 每 layer 的 K/V 元素约缩到 1/8。质量、kernel 效率与实际显存收益仍需实测。
若认为整个模型 FLOPs/参数/显存都自动缩到 1/8，说明你把局部 KV 结构变化错误扩展到了所有对象。
Scaling law 若只优化 training loss，可能选出部署时权重/KV 太大、延迟太高的模型。下一课会把 inference cost 加进外推与决策，并区分拟合区间内证据和区间外假设。