KV cache 优化工程实战 2026:从 PagedAttention 到 FlashDecoding 的生产级推理内核
约 20 分钟5827 字4 次阅读
引言
当一个 7B 模型在单卡 A100 上推理时,KV cache 的内存占用常常把显存吃满,而 throughput 却被显存的碎片化拖到理论的 30% 以下。2026 年的 LLM 推理工程,已经从「拼 GPU 数量」走向「拼每字节显存的访问效率」。本文从 KV cache 的物理布局出发,沿 PagedAttention → FlashAttention → FlashDecoding → Speculative Decoding 的顺序,剖析生产级推理内核如何把单卡吞吐从 200 tok/s 推到 2000 tok/s 以上的工程路径。
1. KV cache 的物理代价:一个 7B 模型在 32K 上下文里要吃多少显存
对标准的 decoder-only transformer,每生成一个 token 都要为历史所有 token 缓存 K 和 V 两个张量。单层 KV cache 的字节数近似为:
其中 是 KV 头数(GQA/MQA 下通常远小于 ), 是头维度(Llama 系列为 128), 是当前序列长度, 是元素字节数(FP16 为 2)。对 LLaMA-3-70B(, )在 时,单请求 KV cache 就已超过 5 GB——这已经是一张 A100 显存的 60%。
当并发请求数从 1 升到 32,朴素分配策略下碎片化会让实际可用显存跌到理论值的 50%-70%。这是 vLLM 提出 PagedAttention 的根本动因:把操作系统虚拟内存的分页思想搬进 GPU 显存。
2. PagedAttention:把 KV cache 从连续张量变成页表
朴素实现里,每条请求的 KV cache 是一段连续显存,长度等于 max_seq_len。这造成两个问题:
- 预分配浪费:为支持 32K 上下文预分配 32K 长度,但用户实际只用了 200 token。
- 碎片化:不同长度请求交错分配,碎片让单卡可服务请求数下降。
PagedAttention(Kwon et al., SOSP 2023)的核心抽象是 block table——每条请求维护一张 block_id → physical_block 的映射,每个 physical_block 固定存放 16 或 32 个 token 的 K/V。逻辑上 KV cache 仍是连续的 seq_len × head_dim 张量,但物理上散落在若干固定大小的页里。
请求 A 的 block table: [3, 7, 12, 14] → 共 4 页 × 16 token = 64 token
请求 B 的 block table: [5, 9] → 共 2 页 × 16 token = 32 token
请求 C 的 block table: [2, 3, 8] → 共 3 页 × 16 token = 48 token
2.1 写时复制(Copy-on-Write)与 beam search
PagedAttention 一个容易低估的设计是 beam search 下的零拷贝共享。在 beam_width=4 的搜索中,4 个 beam 在第 5 步共享前 4 步的 KV cache。vLLM 通过让 4 个 beam 的 block table 前缀指向同一组 physical_block 实现共享——任何 beam 要写新 token 时才触发 copy-on-write,开辟新页。
这意味着 beam search 的内存成本近似 beam_width=1 而不是 beam_width=4。生产里 beam_width=4 的开销只有非分页实现的 30%。
2.2 vLLM 0.6+ 的 prefix caching:跨请求共享 system prompt
当 100 个并发请求都带同一个 2K token 的 system prompt,朴素实现要分配 100 份 KV。vLLM 0.6 起引入 Automatic Prefix Caching:每个请求结束后不立刻释放 block,而是标记为「可复用」;新请求匹配到相同前缀时直接复用对应 block。
工程上需要权衡三点:
- 哈希开销:每个 block 用其内容 + 父 block hash 做 sha256 标识,纯 C++ 实现下 100 万 token 库扫描约 50ms。
- 淘汰策略:LRU 按请求结束时间排序,最少复用的 block 优先释放。
- 失效传播:当某个 block 命中后被改写(beam 分叉、stop token 命中),其后裔 block 必须 cascade invalidate。
3. FlashAttention:从显存访问模式重写 attention
PagedAttention 解决了 分配与共享,但单次 attention 计算本身仍受显存带宽制约。对一个 的 attention 矩阵,朴素实现要 的显存读写——32K 上下文下这是 10 亿元素、4 GB HBM 流量,单卡延迟被 HBM 带宽(~3 TB/s on H100)卡到毫秒级。
FlashAttention(Dao et al., NeurIPS 2022;FlashAttention-2 2023;FlashAttention-3 2024)通过 tiling + online softmax 重写 kernel:
# 概念伪代码(实际是 CUDA kernel)
def flash_attention(Q, K, V, block_size):
O = zeros_like(Q)
m = -inf * ones(Lq) # 行最大值
l = zeros(Lq) # 行归一化分母
for kv_block in split(K, V, block_size):
S = Q @ kv_block.K.T * scale # 当前块的分数
m_new = max(m, rowmax(S))
P = exp(S - m_new[:, None])
l_new = exp(m - m_new) * l + rowsum(P)
O = exp(m - m_new)[:, None] * O + P @ kv_block.V
m, l = m_new, l_new
return O / l[:, None]
关键不变量是 分块计算 + 数值稳定的在线归一化,显存复杂度从 降到 。FlashAttention-3 进一步利用 H100 的 TMA(Tensor Memory Accelerator)异步加载 + WGMMA 矩阵单元,单卡吞吐相对 FA-2 提升 1.5-2×。
3.1 与 PagedAttention 的关系
工程上两者是 正交 的:PagedAttention 管分配与共享,FlashAttention 管单次计算。vLLM 0.7+ 的实现里两者深度集成——FlashAttention kernel 的 K/V 读取直接走 block table 索引,跳过中间张量拷贝,进一步省去一次 HBM 往返。
图表加载中…
4. FlashDecoding:把长上下文推理再压 4-8 倍
FlashAttention 解决了 前向(prefill)阶段 的 attention,但 自回归解码(decode)阶段 有一个特殊场景没覆盖:,即每个 decode 步只生成 1 个 token,但 已经涨到 32K。这时 Q 是一个 的小矩阵,FlashAttention 的「Q tile × K/V tile」分块因为 Q 维度太小,GPU SM 占用率不足——计算被严重低估。
FlashDecoding(2023 年提出,vLLM 0.4+ 集成)的解法是 两级分块:
第一级(parallel):把 L_kv 切成 G 段(如 G=8),每段独立做 flash attention
→ 得到 G 个局部输出 + G 个 softmax 分母
第二级(reduction):把 G 个局部结果按 softmax 加权合并
→ 最终输出
其中 是全局行最大值。第一级并行让 SM 满载,第二级归约只是 个向量的简单合并,开销可忽略。
实测 FlashDecoding 在 32K 上下文 + batch=32 时相对朴素 decode 提速 4-8×,且对长上下文越友好——128K 上下文下提速能到 10×。
5. Speculative Decoding:用小模型为大模型「猜」token
前述优化都聚焦 单次 token 生成 的延迟。Speculative Decoding(Leviathan et al., 2023)从另一个角度入手:接受一定拒绝率的前提下,并行生成多个候选 token。
设大模型为 (target),小模型为 (draft),speculative decoding 的流程是:
def speculative_step(M_q, M_p, gamma=5):
# 1. 用小模型生成 gamma 个候选 token
draft_tokens = []
for _ in range(gamma):
x = M_p.sample(latest_token)
draft_tokens.append(x)
# 2. 用大模型一次性验证所有候选
q_logits = M_q.forward(draft_tokens) # 一次前向拿到所有位置的分布
# 3. 逐位置按 rejection sampling 决定接受或截断
accepted = 0
for i, t in enumerate(draft_tokens):
p = q_logits[i].softmax()
q = M_p(draft_tokens[:i+1]).softmax()
if random() < min(1, p[t] / q[t]):
accepted += 1
else:
# 用大模型的修正分布替换下一个 token
replace_token = sample_from(p - min(p, q).clamp_min(0))
break
# 4. 把接受的 token + 替换 token 一起输出
return draft_tokens[:accepted] + [replace_token]
关键性质:只要大模型的接受率足够高,单次大步能输出 gamma 个 token(如 5 个),相当于把 decode 速度提升 γ × acceptance_rate 倍。LLama-70B + LLama-7B draft 的典型配置是 gamma=5、acceptance=0.6,单步平均输出 3 个 token,端到端提速约 2-3×。
5.1 工程陷阱:哪些场景不适用
Speculative decoding 并非万能。两类场景加速比会显著下降:
- 小模型接受率 < 0.4:当大模型与小模型能力差距大、或任务领域偏移(draft 是通用 7B、target 是 70B 法律微调),接受率跌到 0.3 以下,verification 的开销大于节省的采样步数。
- batch > 8 的高并发:verification 阶段是
gamma × batch的矩阵乘法,batch 已经让大模型 SM 满载时,verification 的「免费午餐」消失。生产建议:batch ≤ 4 时开 speculative;batch ≥ 8 关闭。
6. 端到端架构:把四个优化叠起来
最后看一段生产级推理服务器的典型请求路径:
图表加载中…
四个优化在生产里的协同效应:
| 优化 | 解决什么问题 | 单卡加速比 | 引入的工程复杂度 |
|---|---|---|---|
| PagedAttention | 显存碎片 + 跨请求共享 | 2-4×(vs 连续分配) | 中:block table + COW + prefix cache |
| FlashAttention | prefill 显存带宽 | 3-5× | 低:CUDA kernel 黑盒 |
| FlashDecoding | decode SM 占用率 | 4-8×(长上下文) | 中:两级分块 + reduction |
| Speculative Decoding | decode 串行采样步数 | 2-3×(接受率 0.6) | 高:draft 模型 + rejection sampling |
把它们按 32K 上下文 + batch=4 的常见生产参数叠加,单卡 decode 吞吐从朴素实现的约 200 tok/s 提升到 2000+ tok/s,约 10× 的工程红利。
7. Continuous Batching 与 Chunked Prefill:把 GPU 调度到 85%+ 利用率
前述四个优化都聚焦「单次前向 / decode 的内核效率」,但生产推理服务器还要面对 请求级别的调度。朴素 static batching 把 N 个请求打包成一个 batch 一起跑,任何一条请求没结束,整个 batch 就不能换人——一旦某条请求的 prompt 特别长(比如 16K),prefill 阶段会阻塞所有其他请求数秒。
Continuous Batching(vLLM 0.2+ 的核心抽象)改写了调度语义:每个 decode 步结束后,已完成的请求立刻让出位置给排队请求。这把单卡 GPU 利用率从 static batching 的 30-50% 提到 70%+。
但 continuous batching 没解决 prefill 阻塞:单个 16K prompt 的 prefill 仍占满 SM 数十毫秒,期间所有 decode 请求排队等待。Chunked Prefill(vLLM 0.5+ 引入,Anthropic / Fireworks 同年跟进)的解法是 把 prefill 切成小块(如 512 token)穿插进 decode 步:
时间步 t=1: [decode_batch=32] ← 利用率 65%
时间步 t=2: [prefill_chunk_1 + decode_batch=32] ← 利用率 82%
时间步 t=3: [prefill_chunk_2 + decode_batch=32] ← 利用率 84%
时间步 t=4: [prefill_chunk_3 + decode_batch=32] ← 利用率 85%
时间步 t=5: [decode_batch=33] ← 新请求 prefill 完成后加入
每个 prefill chunk 的 SM 占用在 30-50%,剩余 50% 容量可同时跑 decode 步。实测在 16K 上下文 + batch=32 的混合负载下,chunked prefill 把 p99 延迟从 800ms 降到 220ms,而总 throughput 提升约 1.4×。
工程上有两个细节容易被忽视:
- chunk 大小权衡:chunk 越小,prefill 摊销越均匀,但 kernel launch overhead 占比升高;chunk 越大,单步阻塞越明显。生产经验值是 512-2048 token,视模型 head_dim 与 GPU 架构微调。
- 调度公平性:长 prompt 用户不应该被「切成 100 个 chunk」拖到 30 秒才出第一个 token。vLLM 的策略是为每个请求保留一个最小 decode 配额——任何等待 prefill 超过 100ms 的请求获得优先权。
8. 展望:2026 下半年的推理工程前线
截至 2026-06,四个方向正在成为新一轮工程热点:
- Multi-Token Prediction(Meta 2024 提出的并行预测训练) 与 speculative decoding 双向融合:训练阶段就让模型学会一次预测 4 个 token,decode 时这些 token 可直接并行验证,省去 draft model 的额外显存与一致性开销。
- PagedAttention 的 CPU offload 化:当 KV cache 突破单卡 80 GB(如 128K × 32 batch),H100 单卡已放不下。vLLM 0.7 起的 CPU offload 把冷 block 搬到 host memory + NVMe,代价是 PCIe 往返延迟,对延迟极敏感的场景仍未成熟。
- Chunked prefill 的递归深化:从 512-token chunk 推到 128-token chunk + 异步调度,让长 prompt 用户的首 token 延迟降到 100ms 以内。
- 跨机张量并行 + PagedAttention:当单卡放不下 70B 模型权重 + KV cache 时,TP=4 把权重切到 4 卡,但 KV cache 仍要按请求分页。vLLM 的 TP+Paged 集成让 70B 模型在 4×H100 上跑到 1500+ tok/s 的 decode 吞吐。
LLM 推理工程的「摩尔定律」不再是制程提升,而是 每一层显存访问与调度粒度的重写。PagedAttention / FlashAttention / FlashDecoding / Speculative Decoding / Continuous Batching / Chunked Prefill 这条主线,本质上都是在同一个瓶颈(显存带宽、SM 占用率、调度粒度)下对内核与调度器的持续打磨。
一个反直觉的结论是:2026 年的推理工程红利已经基本被吃透。从 2024 年的 PagedAttention 起,每一轮新优化带来的边际增益在 1.2-1.5× 之间。下一次 5× 以上的飞跃,可能要等到 2027 年新的硬件架构(如 NVIDIA Rubin 的 HBM4、更大 L2 cache)或者模型架构本身的根本性变化。
参考文献
- Kwon, W., Li, Z., Zhuang, S., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP '23.
- Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
- Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with asynchrony and low-precision. arXiv:2407.08608.
- Pope, R., et al. (2023). Efficiently Scaling Transformer Inference. MLSys 2023.
- Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023.
- Kwon, W., et al. (2024). LLM Inference Unveiled: Survey and Roofline Model Insights. vLLM Technical Report.
- Meta AI (2024). Multi-Token Prediction: Enabling Sample-Efficient Language Modeling. arXiv:2404.19737.
一句话摘要
LLM 推理工程的十年红利,正从「堆 GPU」转向「重写每一层显存访问」——PagedAttention、FlashAttention、FlashDecoding、Speculative Decoding 四件套在生产里的协同效应,已能把单卡 decode 吞吐从朴素实现的 200 tok/s 推到 2000+ tok/s。