博客
文章系列日历
归档关于搜索

鄂ICP备19019526号

© 2026 博客

  1. 文章
  2. KV cache 优化工程实战 2026:从 PagedAttention 到 FlashDecoding 的生产级推理内核

KV cache 优化工程实战 2026:从 PagedAttention 到 FlashDecoding 的生产级推理内核

2026年6月16日·约 20 分钟·5827 字·4 次阅读
AI 原生架构
KV cache 优化工程实战 2026:从 PagedAttention 到 FlashDecoding 的生产级推理内核

目录

  • 引言
  • 1. KV cache 的物理代价:一个 7B 模型在 32K 上下文里要吃多少显存
  • 2. PagedAttention:把 KV cache 从连续张量变成页表
  • 2.1 写时复制(Copy-on-Write)与 beam search
  • 2.2 vLLM 0.6+ 的 prefix caching:跨请求共享 system prompt
  • 3. FlashAttention:从显存访问模式重写 attention
  • 3.1 与 PagedAttention 的关系
  • 4. FlashDecoding:把长上下文推理再压 4-8 倍
  • 5. Speculative Decoding:用小模型为大模型「猜」token
  • 5.1 工程陷阱:哪些场景不适用
  • 6. 端到端架构:把四个优化叠起来
  • 7. Continuous Batching 与 Chunked Prefill:把 GPU 调度到 85%+ 利用率
  • 8. 展望:2026 下半年的推理工程前线
  • 参考文献
  • 一句话摘要

引言

当一个 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 的字节数近似为:

Slayer=2×nkv_head×dhead×sseq×selemS_{\text{layer}} = 2 \times n_{\text{kv\_head}} \times d_{\text{head}} \times s_{\text{seq}} \times s_{\text{elem}}Slayer​=2×nkv_head​×dhead​×sseq​×selem​

其中 nkv_headn_{\text{kv\_head}}nkv_head​ 是 KV 头数(GQA/MQA 下通常远小于 nheadn_{\text{head}}nhead​),dheadd_{\text{head}}dhead​ 是头维度(Llama 系列为 128),sseqs_{\text{seq}}sseq​ 是当前序列长度,selems_{\text{elem}}selem​ 是元素字节数(FP16 为 2)。对 LLaMA-3-70B(nkv_head=8n_{\text{kv\_head}}=8nkv_head​=8, L=80L=80L=80)在 sseq=32768s_{\text{seq}}=32768sseq​=32768 时,单请求 KV cache 就已超过 5 GB——这已经是一张 A100 显存的 60%。

当并发请求数从 1 升到 32,朴素分配策略下碎片化会让实际可用显存跌到理论值的 50%-70%。这是 vLLM 提出 PagedAttention 的根本动因:把操作系统虚拟内存的分页思想搬进 GPU 显存。

2. PagedAttention:把 KV cache 从连续张量变成页表

朴素实现里,每条请求的 KV cache 是一段连续显存,长度等于 max_seq_len。这造成两个问题:

  1. 预分配浪费:为支持 32K 上下文预分配 32K 长度,但用户实际只用了 200 token。
  2. 碎片化:不同长度请求交错分配,碎片让单卡可服务请求数下降。

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 计算本身仍受显存带宽制约。对一个 Lq×LkvL_q \times L_{kv}Lq​×Lkv​ 的 attention 矩阵,朴素实现要 O(Lq⋅Lkv)O(L_q \cdot L_{kv})O(Lq​⋅Lkv​) 的显存读写——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]

关键不变量是 分块计算 + 数值稳定的在线归一化,显存复杂度从 O(Lq⋅Lkv)O(L_q \cdot L_{kv})O(Lq​⋅Lkv​) 降到 O(Lq+Lkv)O(L_q + L_{kv})O(Lq​+Lkv​)。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)阶段 有一个特殊场景没覆盖:Lq=1L_q = 1Lq​=1,即每个 decode 步只生成 1 个 token,但 LkvL_{kv}Lkv​ 已经涨到 32K。这时 Q 是一个 1×d1 \times d1×d 的小矩阵,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 加权合并
    → 最终输出
out=∑g=1Gexp⁡(mg−m∗)⋅lg⋅outg∑g=1Gexp⁡(mg−m∗)⋅lg\text{out} = \frac{\sum_{g=1}^{G} \exp(m_g - m^*) \cdot l_g \cdot \text{out}_g}{\sum_{g=1}^{G} \exp(m_g - m^*) \cdot l_g}out=∑g=1G​exp(mg​−m∗)⋅lg​∑g=1G​exp(mg​−m∗)⋅lg​⋅outg​​

其中 m∗=max⁡gmgm^* = \max_g m_gm∗=maxg​mg​ 是全局行最大值。第一级并行让 SM 满载,第二级归约只是 GGG 个向量的简单合并,开销可忽略。

实测 FlashDecoding 在 32K 上下文 + batch=32 时相对朴素 decode 提速 4-8×,且对长上下文越友好——128K 上下文下提速能到 10×。

5. Speculative Decoding:用小模型为大模型「猜」token

前述优化都聚焦 单次 token 生成 的延迟。Speculative Decoding(Leviathan et al., 2023)从另一个角度入手:接受一定拒绝率的前提下,并行生成多个候选 token。

设大模型为 MqM_qMq​(target),小模型为 MpM_pMp​(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
FlashAttentionprefill 显存带宽3-5×低:CUDA kernel 黑盒
FlashDecodingdecode SM 占用率4-8×(长上下文)中:两级分块 + reduction
Speculative Decodingdecode 串行采样步数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,四个方向正在成为新一轮工程热点:

  1. Multi-Token Prediction(Meta 2024 提出的并行预测训练) 与 speculative decoding 双向融合:训练阶段就让模型学会一次预测 4 个 token,decode 时这些 token 可直接并行验证,省去 draft model 的额外显存与一致性开销。
  2. PagedAttention 的 CPU offload 化:当 KV cache 突破单卡 80 GB(如 128K × 32 batch),H100 单卡已放不下。vLLM 0.7 起的 CPU offload 把冷 block 搬到 host memory + NVMe,代价是 PCIe 往返延迟,对延迟极敏感的场景仍未成熟。
  3. Chunked prefill 的递归深化:从 512-token chunk 推到 128-token chunk + 异步调度,让长 prompt 用户的首 token 延迟降到 100ms 以内。
  4. 跨机张量并行 + 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)或者模型架构本身的根本性变化。

参考文献

  1. Kwon, W., Li, Z., Zhuang, S., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP '23.
  2. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
  3. Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with asynchrony and low-precision. arXiv:2407.08608.
  4. Pope, R., et al. (2023). Efficiently Scaling Transformer Inference. MLSys 2023.
  5. Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023.
  6. Kwon, W., et al. (2024). LLM Inference Unveiled: Survey and Roofline Model Insights. vLLM Technical Report.
  7. 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。

相关文章

  • LLM 可观测性工程实战 2026:从 OpenTelemetry GenAI 语义约定到生产级 trace 架构6月15日
  • RAG 工程实战 2026:从 Naive RAG 到 Agentic RAG 的四层架构跃迁6月14日
  • 推理时计算的范式革命:当大模型学会“多花点时间想”之后,AI 架构发生了什么6月13日

评论

加载评论中…

发表评论

返回文章列表