跳转至

CS336 2026 Lecture 10:Inference、KV Cache 与服务系统

LaTeX 源码 · 观看视频

字段 内容
作者/整理 基于 Stanford CS336 Spring 2026 官方可执行讲义重新整理
来源 Stanford CS336
日期 2026 年春季

CS336 2026 Lecture 10:Inference、KV Cache 与服务系统

Lecture 10: inference 的总图

本讲从一个端到端服务视角组织推理优化:先理解 prefill 与 generation 的算术强度和 KV cache,再区分有损 shortcut、带校验的无损 shortcut,以及处理动态请求的 continuous batching 与 PagedAttention。所有方法最终都要回到 latency、throughput、memory footprint 与质量四项账本。

Lecture 10: inference 的整体结构:理解工作负载,有损 shortcut,无损校验,以及动态服务系统。
Lecture 10: inference 的整体结构:理解工作负载,有损 shortcut,无损校验,以及动态服务系统。 查看原图

读图:inference-schema 如何串起整讲

图里把本讲分成四块。第一块是 Understanding the inference workload:为什么推理和训练不同,为什么 generation memory-bound。第二块是 Taking shortcuts (lossy):GQA/MLA/CLA、local attention、quantization、pruning/distillation 等可能改变模型或精度的做法。第三块是 Use shortcuts but double check (lossless):speculative sampling 利用便宜 draft model 但保持目标模型精确分布。第四块是 Handling dynamic workloads:continuous batching 和 PagedAttention 解决真实请求到达、长度不同、共享前缀和 KV cache 碎片化。

本讲的核心判断

推理优化首先是内存流量优化,其次才是 FLOPs 优化。这里的 HBM 是 High Bandwidth Memory,即 GPU 上的高带宽显存;generation 阶段一次生成一个 token,不断从 HBM 读取模型权重和请求自己的 KV cache,常常 memory-bound。和 ZeRO/FSDP 训练状态分片不同,推理侧通常不维护 optimizer state;若要 shard 模型或 KV cache,目标是降低单卡 HBM 压力或提升吞吐,而不是同步训练梯度。

讲义提醒:不要只报告单请求 latency

推理系统的“快”必须绑定 workload。单请求、固定 batch、离线吞吐和在线到达流会得到不同最优策略;如果不报告 prompt/output 长度、并发度、硬件、精度与服务级目标,任何 tokens/s 数字都很难复现或比较。

Understanding the inference workload

推理为什么重要,以及什么叫快

Inference 出现在实际产品、模型评测、agent 内部轨迹、批量数据处理和 RL 采样中。训练是阶段性成本,推理是每次使用都要付的成本。聊天机器人里,人类等待 time-to-first-token;agent 工作流里,内部 trace 可能生成大量不可见 token;RL 中则要采样许多 completions 再打分。换句话说,tokens generated = compute spent。

课堂提示:训练付一次,推理会重复付费

老师用数量级对比强调 inference efficiency:训练是一次性阶段成本,推理却会在产品、评测和 RL 采样中反复发生;agent 的内部 trace 还可能让生成 token 数无界增长。于是“tokens generated = compute spent”不只是口号,而是要求每个产品和平台都把推理成本当作长期系统预算。

术语消化:推理服务指标

指标 含义 主要优化方向
TTFT Time-to-first-token,用户等到第一个 token 的时间。 prefill latency,小 batch、prefix cache、prompt 处理。
Latency 单个请求 token 出现的速度,常以 seconds/token 计。 generation step time、KV cache 读取、batch size。
Throughput 系统总体 tokens/second。 batching、并发、KV memory 管理、kernel efficiency。
Memory-bound 速度受 HBM 读写限制,而不是算力限制。 减少 KV cache、量化、分页、提高复用。

Transformer notation 与算术强度

为了让后面的优化能做数量级比较,本节先统一 \(B,T,D,H,N,K,G,S,F\) 等维度,并把主要算子写成 FLOPs 与 bytes。核心问题是同一层在 prefill 时更像大矩阵乘,generation 时更像矩阵向量乘;shape 改变会让算术强度和瓶颈从 compute 转向 HBM bandwidth。

Transformer 结构与维度约定:\(B,T,D,H,N,K,G,S,F\) 等符号用于后续 FLOPs 和 IO 账本。
Transformer 结构与维度约定:\(B,T,D,H,N,K,G,S,F\) 等符号用于后续 FLOPs 和 IO 账本。 查看原图

读图:Transformer 维度怎么读

\(B\) 是 batch size,\(T\) 是要计算的 token 数,\(S\) 是已条件化的历史 token 数,\(D\) 是 model dimension,\(H\) 是 head dimension,\(N\) 是 query heads 数,\(K\) 是 key/value heads 数,\(G=N/K\) 是 GQA group 数,\(F=4D\) 是 MLP up-projection width。后面所有推理账本都在追踪这些维度如何进入 FLOPs、HBM 读写和 KV cache。

算术强度 arithmetic intensity 定义为每搬运一个 byte 做多少 FLOPs:

\[ I = \frac{\text{FLOPs}}{\text{HBM bytes read/write}}. \]

\(I\) 高于硬件的 compute/bytes 比,算子 compute-bound;若低于该阈值,算子 memory-bound。矩阵乘 \(X(B,D)W(D,F)\) 的 FLOPs 约为 \(2BDF\),但权重读取约为 \(DF\) 元素。当 \(B=1\) 时,矩阵乘退化成 matrix-vector product,读一整个权重矩阵只服务一个 token,算术强度极低。

为什么 generation 像 batch size 为 1 的矩阵向量乘

训练时 \(B\times T\) 通常很大,同一个权重矩阵服务许多 tokens;generation 时每个请求每步只生成一个新 token,\(T=1\),并发 \(B\) 又由在线流量决定。因此 generation 很难稳定形成大矩阵乘,往往被参数和 KV cache 的 HBM 读取限制。

MLP layers (only looking at the matrix multiplications)

把 gated MLP 只按三个矩阵乘计数:up projection、gate projection 和 down projection。若输入为 \(X\in\mathbb{R}^{B\times T\times D}\),中间宽度 \(F\approx4D\),则主要计算量与算术强度近似为

\[ \mathrm{FLOPs}_{\mathrm{MLP}}\approx 6BTDF, \qquad I_{\mathrm{MLP}}\approx BT. \]

第二个近似假设权重读取 \(O(DF)\) 主导 bytes。Prefill 可用长序列和 batch 把 \(BT\) 做大;generation 的 \(T=1\),只剩在线并发 \(B\) 能提高权重复用,因此更容易落入 memory-bound 区域。

课堂提示:prefill 与 generation 的 MLP 只差 shape,不差权重

老师把同一组 MLP 权重放进两个阶段比较:prefill 的 \(BT\) 往往足够大,权重一次读入服务许多 tokens;generation 每次只处理一个新 token,并发请求数又不可预测。这个差异解释了为什么训练优化经验不能原样搬到在线 decode,也解释了 batching 为什么对 generation 如此关键。

Attention layers (focusing on the matrix multiplications with FlashAttention)

FlashAttention 消除了显式 \(T\times S\) attention matrix 的 HBM 往返,但 QKV projection、attention matmul 与 output projection 仍要读取权重和 KV cache。Prefill 中 \(T=S\),attention 的复用随上下文增大;generation 中 \(T=1\),每个新 query 都要读取该请求已有的 \(S\) 个 keys/values,因此 attention intensity 仍接近常数且通常低于 1。

Attention generation 为什么 batching 帮助有限

不同请求拥有不同 KV cache。增大 batch 可以共同执行 projection,却不能让一个请求的历史 KV 被另一个请求复用;所以 MLP 的权重读取可由 batch 摊薄,attention generation 的 KV 读取则随每条序列单独增长。

naive inference、KV cache、prefill 与 generation

前面的 shape 账本解释了为什么逐 token 重算 prefix 非常浪费,本节比较 naive autoregressive inference 与 KV-cached inference。KV cache 把历史 key/value 从重复计算改成持续读取,显著减少 FLOPs,却让 cache 容量、HBM 带宽和请求生命周期成为新的主瓶颈。

Naive inference:每生成一个 token 都把完整历史重新喂进 Transformer,重复计算大量 prefix。
Naive inference:每生成一个 token 都把完整历史重新喂进 Transformer,重复计算大量 prefix。 查看原图

读图:naive inference 为什么是 \(O(T^3)\)

\(t\) 个 token 的 forward pass 会处理长度约为 \(t\) 的上下文;单次 attention 近似 \(O(t^2)\)。把 \(t=1\)\(T\) 累加,整体生成 \(T\) 个 token 需要 \(O(T^3)\) 级别的重复计算。图里的关键是 prefix 被反复重算。

Cached inference:把历史 token 的 key/value 存进 KV cache,后续生成只计算新 token 与历史 cache 的交互。
Cached inference:把历史 token 的 key/value 存进 KV cache,后续生成只计算新 token 与历史 cache 的交互。 查看原图

读图:KV cache 省了计算,但制造了内存瓶颈

KV cache 为每个 sequence、token、layer、head 保存 \(K,V\) 向量。它避免历史 token 反复投影,复杂度大幅降低;但每个新 token 都要读取该请求历史的 KV cache。历史越长、batch 越大、层数越多,HBM 读写越重。

\[ M_{\text{KV}} \approx 2 \cdot B \cdot S \cdot L \cdot K \cdot H \cdot \text{bytes}. \]

其中 \(2\) 来自 key 和 value,\(B\) 是并发请求数,\(S\) 是历史上下文长度,\(L\) 是层数,\(K\) 是 KV heads 数,\(H\) 是 head dimension。这个公式解释了为什么后面 GQA、MLA、CLA、local attention、PagedAttention 都围绕 KV cache 做文章。

prefill 和 generation 的差别

Prefill:给定 prompt,一次性编码 \(S\) 个 tokens,能像训练一样并行,MLP intensity 约随 \(B S\) 提升,attention intensity 约随 \(S/2\) 提升。Generation:每次只生成一个 token,MLP intensity 约随并发 \(B\) 提升,attention generation intensity 小于 1,几乎无法靠 batching 改善,因为每个请求读取自己的 KV cache。

latency-throughput tradeoff

KV cache 建立后,服务系统仍要决定一次处理多少请求以及是否切分模型。本节的核心 tradeoff 是:更大 batch 和更多并发通常提高 hardware utilization 与总 throughput,却增加排队时间、每请求 cache 占用和 tail latency;因此复制模型、tensor parallel 与 pipeline parallel 要按 SLO 和 workload 选择。

Batch size 增大时,模型权重读取可被更多请求摊薄,throughput 上升;但 KV cache 随 \(B\) 增大,单请求等待和每步读取也变重,latency 变差。一个简单策略是 prefill 用较小 batch 改善 TTFT,generation 用较大 batch 提高吞吐。

复制模型和切分模型的差别

若启动 \(M\) 个完整模型副本,latency 基本不变,throughput 近似乘 \(M\)。若 shard 模型和 KV cache,则能服务更大模型或更大上下文,但每步推理需要跨设备通信,latency 可能上升。

本章小结

理解 inference 的关键是把请求拆成 prefill 和 generation,把算子拆成 FLOPs 与 HBM bytes,把状态拆成 weights 和 KV cache。Generation 的 memory-bound 特性决定了后续所有优化方向。

Taking shortcuts (lossy):减少 KV cache 和模型字节

前面的 workload 分析指出 generation 常被权重与 KV cache 读取限制,因此本部分允许改变模型表示来减少 bytes。GQA/MQA、MLA/CLA、local attention、quantization 与 pruning 都可能换来更高吞吐,但它们会改变可表达函数或数值精度,必须用任务质量和长上下文评测重新确认边界。

Grouped-query attention (GQA):按 head 共享 KV

本节先从最直接的 KV head 共享开始。MHA 为每个 query head 保存独立 K/V,MQA 让所有 query heads 共享一组 K/V,GQA 则在两者之间按 group 共享;\(N/K\) 越大,cache 越小、读取越少,但不同 heads 可使用的独立记忆子空间也越少。

Grouped-query attention:\(N\) 个 query heads 共享 \(K\) 个 key/value heads,\(K=1\) 是 MQA,\(K=N\) 是 MHA。
Grouped-query attention:\(N\) 个 query heads 共享 \(K\) 个 key/value heads,\(K=1\) 是 MQA,\(K=N\) 是 MHA。 查看原图

读图:GQA 图应该看 \(N/K\)

MHA 为每个 query head 保存独立 KV;MQA 所有 query heads 共享一个 KV head;GQA 介于两者之间,每组 query heads 共享一组 KV。KV cache 大小约按 \(N/K\) 缩小,因此 latency 和 throughput 改善来自更少 HBM 读取。

GQA 的 latency/throughput 结果:减少 KV heads 后更容易 fit,throughput 改善。
GQA 的 latency/throughput 结果:减少 KV heads 后更容易 fit,throughput 改善。 查看原图

读图:GQA speed 图说明什么

横向比较时,KV cache 小的配置通常 latency 更低或能容纳更大 batch,从而 throughput 更好。若原 MHA 配置不 fit,GQA 首先解决容量;若能 fit,GQA 仍可能通过减少内存流量改善吞吐。

GQA accuracy 检查:减少 KV cache 后必须验证准确率是否下降。
GQA accuracy 检查:减少 KV cache 后必须验证准确率是否下降。 查看原图

GQA 是 lossy shortcut

GQA 改变 attention 结构,通常会影响模型质量。速度图只说明系统收益,accuracy 图才说明是否值得用。课程主线是:有损 shortcut 必须以质量检查收尾。

课堂提示:GQA 的速度来自 \(N/K\) 倍 KV 缩减

老师先重申 inference 的瓶颈是 memory,再给出 GQA 的直接机制:\(N\) 个 query heads 只配 \(K\) 个 KV heads,cache 约缩小 \(N/K\) 倍,因此 latency/throughput 改善来自更少 HBM bytes。源码随后立即要求检查 accuracy,说明“更快且 fit”仍不等于可以无条件替换 MHA。

讲义提醒:速度图和准确率图必须成对出现

GQA 的价值不能只看 tokens/s。更少 KV heads 可能让更大 batch fit 进显存,所以性能收益包含容量效应;同时质量下降可能只在长上下文、知识密集或特定下游任务暴露。部署决策必须在目标 workload 上同时画出 latency/throughput 与 accuracy frontier。

Multi-head latent attention (MLA):压缩 KV

GQA 沿 head 维共享 K/V,MLA 则进一步改变 cache 的表示:保存低维 latent,并在使用时投影成 K/V。它减少长期存储和读取,却引入重建计算与潜在表达瓶颈,需要在相同 cache budget 下比较质量。

Multi-head latent attention:存压缩 latent \(c=W_c h\),需要时再投影出 \(K,V\)。
Multi-head latent attention:存压缩 latent \(c=W_c h\),需要时再投影出 \(K,V\)。 查看原图

读图:MLA 如何压缩 KV cache

普通 attention 存 \(K=W_Kh,V=W_Vh\),维度约为 \(N H\)。MLA 存低维 latent \(c\),例如 DeepSeek v2 中从 \(16384\) 维压到 \(512\) 维,再额外保留 RoPE 所需维度。它把长期 cache 从高维 KV 变成低维 latent,换来每步投影计算。

MLA/MHA/GQA 准确率比较之一:先确认 MHA 与 GQA 的质量差异。
MLA/MHA/GQA 准确率比较之一:先确认 MHA 与 GQA 的质量差异。 查看原图
MLA 准确率比较之二:MLA 在某些设置下可接近甚至超过 MHA,同时更省 KV cache。
MLA 准确率比较之二:MLA 在某些设置下可接近甚至超过 MHA,同时更省 KV cache。 查看原图

两张 accuracy 表需要作为同一组证据读取:先比较 MHA 与 GQA 的基线差异,再观察 MLA 在相同或更小 cache budget 下是否保持质量。若不同实验的模型大小、训练 token 或 cache 尺寸不一致,就不能把表中绝对分数直接解释成 MLA 机制本身的收益。

读图:MLA accuracy 两张表一起看

第一张表确认不同 attention 结构的质量基线;第二张表展示 MLA 的质量和成本折中。读表时不能只看平均分,还要看任务类别和模型规模,因为 KV 压缩可能对长上下文、检索或数学任务影响不同。

Cross-layer attention (CLA):跨层共享 KV

MLA 压缩每一层保存的表示,CLA 则改变“哪些层必须各自拥有 cache”:相邻或分组 layers 共享 K/V,使层维度上的冗余像 GQA 在 head 维度上的冗余一样被压缩。代价是不同层无法自由保存各自的历史表示,因此应看 accuracy--cache Pareto frontier。

Cross-layer attention:跨层共享 KV,类似 GQA 在 heads 维共享 KV。
Cross-layer attention:跨层共享 KV,类似 GQA 在 heads 维共享 KV。 查看原图
CLA 结果:改善 accuracy 与 KV cache size 的 Pareto frontier。
CLA 结果:改善 accuracy 与 KV cache size 的 Pareto frontier。 查看原图

读图:CLA 的 Pareto frontier

CLA 不是只追求最小 KV cache,而是在相同 cache budget 下尽量提高 accuracy,或在相同 accuracy 下减少 cache。Pareto frontier 向左下移动才说明结构改动有系统价值。

Local (sliding window) attention 与 hybrid/DeepSeek attention

前面的压缩方法仍让每个新 token 访问全部历史 cache,本节转向选择“访问哪些位置”。局部注意力与 sliding-window attention 限制每层感受野,hybrid 模型周期性插入 full 或 global 层恢复远程信息,更新的压缩与选择机制则试图在超长上下文中只读取少量高价值状态。

Local/sliding-window attention:每层只看局部窗口,多层叠加后 effective context 线性扩大。
Local/sliding-window attention:每层只看局部窗口,多层叠加后 effective context 线性扩大。 查看原图

读图:local attention 的收益和代价

局部窗口让每层 KV cache 与可见窗口相关,而不是与完整 sequence length 线性相关。多层堆叠可以传播信息,但远距离精确依赖会变弱。因此实践中常用 hybrid layers:部分层 local,部分层 global。

DeepSeek v4 attention:CSA、DSA、HCA 等多级压缩/选择支持超长上下文。
DeepSeek v4 attention:CSA、DSA、HCA 等多级压缩/选择支持超长上下文。 查看原图

这张结构图应沿数据流阅读:先看历史状态如何被压缩成候选表示,再看 selector 如何减少需要精确 attention 的位置,最后检查被选状态如何进入当前 token 计算。它表达的是“先便宜筛选、再昂贵计算”的分层机制,而不是免费获得无限上下文;selector 误差和额外索引开销仍需计入。

术语消化:DeepSeek v4 attention 三个缩写

术语 机制 推理意义
CSA Compressed Sparse Attention,把每 \(m\) 个 tokens 压成一个表示。 降低长上下文 KV 读取。
DSA DeepSeek Sparse Attention,选择 top \(k\) 相关位置。 把注意力集中到更有用 token。
HCA Heavily Compressed Attention,更激进压缩。 服务超长上下文时进一步控成本。

quantization 与 Activation-aware quantization (AWQ)

Attention 结构优化之后,本节处理另一大类 bytes:模型权重与 activation。Quantization 用更少 bit 表示数值,直接降低 HBM 容量和带宽压力,并可能启用更快低精度矩阵单元;代价是 rounding、clipping 和 outlier 误差,因此训练时量化与训练后量化有不同适用边界。

不同数值精度示意:fp32、bf16、fp8、int8、int4 代表不同 bytes 和误差。
不同数值精度示意:fp32、bf16、fp8、int8、int4 代表不同 bytes 和误差。 查看原图

读这张格式图时要同时看 exponent、mantissa 与总 bytes:FP8 相比 BF16/FP16 降低存储和带宽,但动态范围与精度由具体 E4M3/E5M2 变体决定;INT8/INT4 还需要 scale、zero point 或 group-wise metadata。理论上的 2 倍/4 倍压缩不会自动转化成同等速度,kernel 和 dequantization 也必须匹配。

读图:量化为什么直接影响推理

推理 memory-bound 时,参数和 KV cache 的 bytes 直接进入 latency/throughput。bf16 是常见默认;fp8/int8/int4 减少内存流量,但会引入量化误差。位宽越低,越需要校准、补偿或训练时适配。

术语消化:QAT、PTQ、GPTQ、AWQ

方法 做法 代价/适用性
QAT Quantization-aware training,训练中模拟量化误差。 质量好但昂贵。
PTQ Post-training quantization,训练后用校准数据确定 scale/zero point。 便宜,常用于部署。
GPTQ 用 Hessian 信息补偿量化误差。 更精细但实现复杂。
AWQ Activation-aware quantization,保留少量重要权重高精度。 关注大 activation channel 命中的权重。
AWQ:根据 activation 重要性选择少量权重保留更高精度。
AWQ:根据 activation 重要性选择少量权重保留更高精度。 查看原图

AWQ 的关键不是把所有权重一视同仁地压到低 bit,而是根据 activation 统计识别对输出更敏感的通道,并通过 scaling 或少量高精度保留降低误差。读图时应区分离线校准、权重变换、量化存储和运行时反量化四个阶段,避免把校准成本误算成每 token 开销。

读图:AWQ schema 说明什么

AWQ 的直觉是并非所有权重量化误差同等重要。若某些 activation channels 很大,与其相乘的权重误差会被放大。保留 0.1%-1% 重要权重的高精度,可能用很小内存代价换来明显质量改善。

课堂提示:AWQ 让 activation 决定哪些权重值得精度

老师把 AWQ 的观察压缩成三步:某些 activation channels 很大;与它们相乘的权重更重要;因此只为约 0.1%--1% 的关键权重分配更多精度。课程给出的案例从 fp16 到 int3 可带来约 4 倍更低内存和 3.2 倍 speedup,但这些数字依赖校准数据、kernel 与目标模型,机制比单一 benchmark 更可迁移。

pruning 与 distillation

量化保留稠密结构但降低每个元素的 bit,本节则直接删除权重、通道、层或注意力头。Pruning 会改变模型容量和 kernel shape,若硬件无法利用非结构化稀疏,参数更少也未必更快;distillation 用原模型输出修复裁剪后的行为,因此两者应作为一个闭环而不是两个独立技巧。

Pruning + distillation loop:识别重要结构、裁剪,再用原模型蒸馏修复。
Pruning + distillation loop:识别重要结构、裁剪,再用原模型蒸馏修复。 查看原图

读图:pruning 为什么必须 distill

直接移除 layers、heads 或 hidden dimensions 会破坏模型函数。蒸馏让小模型模仿原模型输出,把结构损伤修复回来。它是有损 shortcut,但通过 teacher signal 降低质量损失。

老师强调:先 rip out,再 fix it up

源码用非常直接的语言概括 pruning:先从昂贵模型中移除不重要的 layer、head 或 hidden dimension,再用原模型蒸馏修复。它同时对应两种 recipe:from scratch 是定义并训练更快架构;distillation 则从原模型初始化不同结构,再通过 teacher signal repair。少参数只是开始,修复后的质量才决定 shortcut 是否成立。

Pruning + KD 结果:裁剪后经蒸馏可恢复部分质量,同时降低推理成本。
Pruning + KD 结果:裁剪后经蒸馏可恢复部分质量,同时降低推理成本。 查看原图

有损 shortcut 的共同风险

GQA、MLA、CLA、local attention、quantization、pruning 都可能改变输出分布。它们的验收不能只看速度,必须看准确率、长上下文能力、校准、下游任务和安全性指标。

Use shortcuts but double check (lossless):speculative sampling

Speculative sampling 利用一个便宜 draft model \(p\) 一次猜多个 token,再让 target model \(q\) 并行检查。直觉是:generation 慢,但 checking 多个候选 token 比逐个生成便宜。关键性质是它仍然精确采样自 \(q\),因此属于 lossless shortcut。

课堂提示:checking 比 generation 更容易并行

老师先回顾 prefill 可并行且 compute-bound、generation 逐 token 且 memory-bound,再得出不对称性:checking 一串候选比逐个生成更快。Draft model 只负责便宜地提案,target model 仍决定分布;修改过的 rejection sampling 保证最终是 target 的 exact sample,因此这里的 speedup 不靠牺牲输出分布。

Speculative sampling 算法:draft model 提案,target model 校验,必要时从 residual distribution 采样。
Speculative sampling 算法:draft model 提案,target model 校验,必要时从 residual distribution 采样。 查看原图

读图:为什么它仍是 exact sample

若 draft \(p\) 对某 token 给出概率高于 target \(q\),该 token 只以 \(q/p\) 的概率接受;被拒绝的概率质量会转移到 residual distribution \(\max(q-p,0)\)。这样最终每个 token 的边际概率仍等于 \(q\),只是把多个候选的校验并行化。

\[ P(A) = p(A)\frac{q(A)}{p(A)} + p(B)\cdot 1\cdot 0 = q(A). \]

这个两 token 例子说明:如果 \(p(A)>q(A)\),过多提出的 \(A\) 会被拒绝一部分;若 \(p(B)<q(B)\),不足的概率通过 residual 补回来。

Speculative sampling 结果:draft/target 组合可显著减少 target model 的 sequential generation 次数。
Speculative sampling 结果:draft/target 组合可显著减少 target model 的 sequential generation 次数。 查看原图

结果图中的 speedup 需要拆成三项:draft 生成候选的成本、target 并行验证候选的成本,以及一次能接受多少 token。若 draft 太弱,接受率低;若 draft 太大,提案本身昂贵;若 batch 或序列太短,target verification 的并行优势也无法摊薄。

Speculative sampling 统计:性能取决于 draft model 接近 target 的程度和候选长度。
Speculative sampling 统计:性能取决于 draft model 接近 target 的程度和候选长度。 查看原图

读图:speculative speedup 的两个控制旋钮

Draft model 越接近 target,接受率越高;一次猜的 tokens 越多,潜在并行度越高,但被拒绝时浪费也越多。最佳点取决于 draft 成本、target 成本、接受率和服务端 batch 调度。

Medusa 和 EAGLE:改进 draft model 的方法,让提案更接近 target。
Medusa 和 EAGLE:改进 draft model 的方法,让提案更接近 target。 查看原图

读图:Medusa/EAGLE 的共同目标

Medusa 让 draft 头并行提出多个未来 token;EAGLE 利用 target model 的高级特征帮助 draft。二者都服务同一个目标:提高接受率,同时保持 draft 成本足够低。

本章小结

Speculative sampling 的美感在于:它利用系统不对称性,checking 比 generation 更可并行;又用 rejection-sampling 逻辑保证目标分布不变。它是本讲中最典型的“系统技巧 + 概率正确性”结合。

Handling dynamic workloads:continuous batching

前面的 shortcut 都假设已经有一个 batch,但在线服务中的请求会持续到达、长度不同、结束时间也不同。本部分把优化对象从单条序列提升到调度器:系统必须在每个 decoding step 重新选择活跃请求,让已完成序列立即退出、新请求及时加入,并控制 cache 与公平性。

真实推理流量不是训练中的矩形 batch。请求到达时间不同、prompt 长度不同、生成长度不同,还可能共享 system prompt 或要求多样本生成。静态 batching 会让早到请求等待,也会因 padding 浪费大量计算。

课堂提示:动态流量先变成 ragged array,再拆算子 batching

老师把训练与推理的形状差异概括为 dense block 对 ragged array。Iteration-level scheduling 每个 decode step 加入新请求、移除完成请求;selective batching 再区分算子:attention 按各自长度分别处理,non-attention 部分把 \([3,H],[9,H],[5,H]\) 拼成 \([17,H]\) 做大矩阵计算。动态 batching 不是把所有张量硬 pad 成同一长度。

静态 batching 的问题:请求长度和到达时间不同,等待和 padding 都会浪费。
静态 batching 的问题:请求长度和到达时间不同,等待和 padding 都会浪费。 查看原图

读图:continuous batching 解决哪个问题

图中不同请求像长度不一的条带。训练可以把 \(B\times S\times H\) 看成稠密矩形;推理则是 ragged array。Iteration-level scheduling 每个 decode step 都可以加入新请求、移除完成请求,让 GPU 更持续地工作。

selective batching 的核心

Attention 需要按 sequence 单独处理,因为每个请求历史长度不同;non-attention MLP 等操作可以把所有 active tokens concatenate 成 \([3+9+5,H]\) 这样的矩阵统一计算。服务系统要在“保持语义正确”和“形成大矩阵乘”之间拆算子。

讲义提醒:动态 batching 的难点是 tail,不是平均值

平均 token 长度会掩盖少量超长请求对显存和排队的影响。调度器需要关注 p95/p99 latency、请求取消、优先级、prefill/generation 混排和 cache eviction;只在固定长度 synthetic benchmark 上测吞吐,很容易高估线上收益。

PagedAttention:用操作系统思想管理 KV cache

前面的 Continuous batching 让请求集合动态变化,传统为每条序列预留连续最大 cache 的方法会产生严重碎片。本部分借用虚拟内存思想,把逻辑 KV 序列映射到固定大小的物理 blocks,使系统能够按需分配、非连续存储、共享 prefix,并在写入时执行 copy-on-write。

课堂提示:PagedAttention 直接借用操作系统 paging

老师先把旧方案类比为硬盘碎片:按最大长度预留会产生 internal fragmentation,请求释放后还会留下 external fragmentation。解决方案是把每条序列的 KV cache 切成非连续 blocks,用逻辑映射隐藏物理位置;共享 prefix 时复用 blocks,分叉写入时 copy-on-write。这里的模型服务创新来自操作系统思想,而不是新的 attention 数学。

fragmentation 与 blocks

本节先区分 internal fragmentation 与 external fragmentation。前者来自为短请求预留未使用空间,后者来自空闲区域不连续而无法容纳新请求;固定大小 blocks 将分配单位标准化,再通过 block table 把逻辑 token 位置映射到任意物理页。

传统 KV cache 预分配会产生 internal/external fragmentation。
传统 KV cache 预分配会产生 internal/external fragmentation。 查看原图

读图:KV cache fragmentation

如果为每个请求按最大长度预留连续空间,实际生成较短时会产生 internal fragmentation;不同请求释放/分配后,空洞之间不连续,又会产生 external fragmentation。显存看似充足,却无法容纳新的长请求。

PagedAttention 把每个序列的 KV cache 切成非连续 blocks。
PagedAttention 把每个序列的 KV cache 切成非连续 blocks。 查看原图

读图:blocks 像虚拟内存页

每个请求看到的是逻辑连续的 token 序列,但物理 KV blocks 可以分散在显存里。调度器维护 block table,把逻辑位置映射到物理块。这样可以按需增长,而不是一开始预留最大长度。

sharing、copy-on-write 与并行读取

有了逻辑到物理 block 映射后,系统可以进一步共享相同 system prompt 或多样本生成的 prefix。多个请求先引用同一只读 blocks,只有某个分支写入新 token 时才复制对应页;attention kernel 则根据 block table 并行读取非连续 cache,避免先搬成连续大缓冲区。

PagedAttention 的逻辑视图:请求可以共享 prefix 的 KV blocks。
PagedAttention 的逻辑视图:请求可以共享 prefix 的 KV blocks。 查看原图
KV cache sharing 场景:共享 system prompt、多样本生成、程序合成等。
KV cache sharing 场景:共享 system prompt、多样本生成、程序合成等。 查看原图

读图:prefix sharing 为什么省显存

多个请求若有相同 system prompt 或同一 prompt 的多个 samples,前缀 KV 完全相同。PagedAttention 可让它们共享同一批 blocks;当某个请求继续生成并发生分叉时,再 copy-on-write 新 blocks。这和操作系统共享内存页的思路一致。

PagedAttention 并行读取 blocks 并执行 attention。
PagedAttention 并行读取 blocks 并执行 attention。 查看原图

读图:PagedAttention 不只是内存分配器

为了让非连续 blocks 不拖慢 attention,系统还需要 kernel 支持:融合 block read 和 attention,减少 launch overhead,配合 FlashAttention/FlashDecoding/CUDA graphs。内存管理和 kernel design 必须一起做。

PagedAttention 的系统意义

PagedAttention 把 KV cache 从“连续大数组”变成“可分页的逻辑地址空间”。它直接服务动态 workload:变长请求、共享前缀、多样本生成、早停、长上下文。vLLM 的成功说明,LLM serving 的关键创新不只在模型结构,也在操作系统式资源管理。

本章小结

Continuous batching 解决时间维度的动态性,PagedAttention 解决空间维度的动态性。前者让 batch 随请求到达和完成不断变化;后者让 KV cache 随序列增长和共享动态分配。

总结与延伸

Summary:从三类 shortcut 到动态 serving

Lecture 10 的主线是:推理不同于训练,因为它反复发生、在线动态、generation sequential 且 memory-bound。优化方向可以分三类:改变模型或精度以减少内存流量,利用概率校验保持 exact sampling,借鉴系统思想管理动态请求和 KV cache。

最终 takeaways

  1. 推理成本长期重复发生,TTFT、latency、throughput 是不同目标。
  2. Prefill 通常较 compute-bound,generation 通常 memory-bound。
  3. KV cache 大小约为 \(2BSLKH\cdot\text{bytes}\),是推理显存和带宽的核心对象。
  4. GQA、MLA、CLA、local/hybrid attention 都在减少 KV cache,但需要质量检查。
  5. Quantization、AWQ、pruning、distillation 减少模型字节或计算,但属于有损 shortcut。
  6. Speculative sampling 用 draft model 提案、target model 校验,能保持 exact target distribution。
  7. Continuous batching 和 PagedAttention 把在线 serving 变成调度和内存管理问题。

拓展阅读

  • vLLM and PagedAttention paper.
  • Speculative decoding / speculative sampling papers.
  • GQA, MLA, CLA and DeepSeek attention reports.
  • AWQ, GPTQ, QAT/PTQ quantization references.
  • Orca continuous batching serving system.

部署验收:五本账必须同时过线

任何 inference 优化都应在目标流量上同时提交以下证据;只优化其中一列,往往会把成本转移到另一列:

账本 必须报告 常见误判
Workload prompt/output 长度分布、arrival rate、并发、prefix sharing 用固定长度 synthetic batch 代替真实 ragged traffic。
Latency TTFT、inter-token latency、p50/p95/p99、排队时间 只报平均 latency,忽略长请求与 prefill 干扰。
Throughput tokens/s、requests/s、active sequences、GPU utilization 用更大 batch 提吞吐,却违反交互式 SLO。
Memory weights、KV cache、fragmentation、临时 workspace、峰值 只按参数 bytes 估算,忽略 cache 生命周期和碎片。
Quality / exactness 下游准确率、长上下文、校准,或 exact-distribution 证明 把 lossy shortcut 当作无损加速,或只在平均 benchmark 上验收。

最后的判断不是“哪个技巧最新”

正确问题是:在给定硬件、流量和质量约束下,哪组 architecture、precision、scheduler 与 KV-memory 策略位于 Pareto frontier。若没有统一 workload 和四类指标,GQA、AWQ、speculative sampling、continuous batching 与 PagedAttention 之间无法公平比较。