/ CS336

[CS336-10] 推理系统

整理 CS336 第十讲：prefill、decode、KV cache、GQA、MLA、speculative decoding、PagedAttention 和动态推理 workload。

2025年5月1日4 minCS336 · Inference · KV Cache

这篇是 Stanford CS336 Spring 2025 第 10 讲 Course Materials 的学习笔记。主题是 LLM inference。

1. Prefill 和 Decode 是两种 workload#

推理不是单一过程。一次请求通常分成 prefill 和 decode。

Prefill 阶段处理 prompt 中已有 token，可以并行计算整段上下文，形态接近训练 forward，通常更 compute-bound。

Decode 阶段每次只生成一个新 token，需要不断读取历史 KV cache，batch 和 sequence 动态变化，通常更 memory-bound。

这就是为什么训练优化和推理优化关注点不同。训练更关心大 batch matmul 吞吐，推理更关心延迟、KV cache、动态 batching 和 memory bandwidth。

第一条是 lossy shortcut：允许模型或表示发生变化，只要质量损失可接受。例如 quantization、pruning、distillation、使用更小模型、减少 KV heads、局部 attention 等。

第二条是 lossless shortcut：保持目标模型采样分布不变，但改变执行方式。例如 speculative decoding 使用 draft model 生成候选 token，再由大模型并行验证；如果设计正确，可以保持 exact sampling。

系统工程里经常需要同时使用两类方法：一个负责降低单 token 成本，一个负责提高整体吞吐。

在自回归生成中，每一层 attention 都需要历史 K/V。缓存它们可以避免重复计算 prefill 结果，但也带来显存压力。

KV cache 大小随 batch size、sequence length、layers、KV heads、head dimension、dtype 增长。长上下文和高并发时，KV cache 可能比模型权重更像瓶颈。

因此 GQA、MLA、CLA、local attention 等方法都可以从减少或压缩 KV cache 的角度理解。

Grouped-query attention 让多个 query heads 共享更少的 key/value heads，减少 KV cache。

Multi-head latent attention 进一步把 key/value 投影到更低维 latent 表示，希望在减少缓存的同时尽量保持质量。

Cross-layer attention 让多层共享或复用部分 attention 信息，减少每层独立存储和计算。

Local attention 限制 token 只看局部窗口，降低长上下文成本，但需要搭配少量 global 层或其他机制保留远距离信息。

这些方法的共同目标是：在尽量不伤害质量的前提下降低 decode memory pressure。

在线推理请求长度不同、到达时间不同、生成长度不同。如果直接为每个请求分配连续大块 KV memory，会产生碎片和浪费。

PagedAttention 借鉴操作系统分页思想，把 KV cache 分成 blocks，通过 block table 管理逻辑序列和物理缓存。这样可以更好地支持动态 batching、request arrival/departure 和 memory reuse。

这说明推理系统不是单纯模型问题，也很像操作系统和内存管理问题。

第 10 讲的核心是：LLM inference 的瓶颈和训练不一样。

如果训练系统关心“多久训完”，推理系统关心的是“每个请求多快、多稳、多便宜”。