☰
KV Cache、Prefill/Decode、内存墙
2026/9/28 20:52:39 网站建设 项目流程
  • Mini-SGLang GitHub 项目
  • SGLang 官方文档
  • vLLM 官方文档

1. 如果没有 KV Cache,会发生什么

LLM 生成文本是**自回归(Autoregressive)**的:模型一次生成一个 token,再把这个 token 接到输入后面,继续生成下一个 token。

假设用户输入 prompt 后,模型依次生成token1、token2、token3。如果没有 KV Cache,每次前向都可能重新计算历史 token 的 Key 和 Value:

用户 Prompt

第 1 次前向
计算 Prompt

生成 token1

第 2 次前向
重新计算 Prompt + token1

生成 token2

第 3 次前向
重新计算 Prompt + token1 + token2

生成 token3

这会造成大量重复计算。序列越长、生成 token 越多,浪费越明显。

我的理解是:KV Cache 就像模型给历史 token 做的一份笔记。已经算过的 Key/Value 不应该每次都重新算一遍,而应该保存下来,下次直接读取。

2. KV Cache 是什么

在 Transformer Decoder 的每一层 Attention 中,每个 token 都会经过投影得到:

  • Query:当前 token 想查询什么信息。
  • Key:当前 token 可以被怎样匹配。
  • Value:当前 token 真正携带的信息。

KV Cache 保存的就是历史 token 在每一层已经计算好的 Key 和 Value。

生成一个新 token 时:

  1. 只计算新 token 的 Query、Key、Value。
  2. 把新 token 的 Key/Value 追加到 KV Cache。
  3. 用新 token 的 Query 去读取历史 KV Cache。
  4. 经过 Attention 和后续网络得到 logits。
  5. Sampler 选出下一个 token。

新 token

Q/K/V 投影

Q_new

K_new / V_new

KV Cache
历史 K/V

Attention

输出 hidden state

Sampler

next token

KV Cache 的作用是:

  • 避免重复计算历史 token 的 K/V。
  • 加快 Decode 阶段。
  • 支持公共前缀复用。

但它也有代价:

  • KV Cache 会占用 GPU 显存。
  • 序列越长,缓存越大。
  • 并发请求越多,总 KV Cache 越大。

一个粗略的单 token KV Cache 显存公式是:

单 token KV 字节数 ≈ 2 × L × H_kv × D_head × B

其中:

  • 2:Key 和 Value 各一份。
  • L:Transformer 层数。
  • H_kv:KV head 数。
  • D_head:每个 attention head 的维度。
  • B:每个元素占多少字节,例如 FP16/BF16 通常是 2 字节,FP8 是 1 字节。

那么N个 token 的 KV Cache 大约是:

总 KV 字节数 ≈ 2 × L × H_kv × D_head × B × N

GQA/MQA 会减少 KV head 数,因此可以明显减少 KV Cache 的显存和读取量。

3. Prefill 和 Decode

一次 LLM 请求通常可以分成两个阶段:Prefill 和 Decode。

3.1 Prefill

Prefill 是“处理用户输入 prompt”的阶段。

它的特点是:

  • prompt 在请求开始时已经完整给出。
  • 可以并行处理 prompt 中的多个 token。
  • 计算量较大,矩阵乘法形状较大。
  • 会为 prompt 中的所有 token 计算并写入 K/V。
  • 主要影响TTFT(Time To First Token,首 token 延迟)。
  • 通常更偏Compute-bound(计算受限)。

3.2 Decode

Decode 是“逐个生成输出 token”的阶段。

它的特点是:

  • 不使用推测解码时,每一步通常只为每个请求生成一个新 token。
  • token 之间有依赖关系,不能像 Prefill 那样一次性并行生成整段答案。
  • 每一步都要把新 token 的 K/V 追加到缓存。
  • 每一步都要读取该请求已经积累的 KV Cache。
  • 主要影响TPOT/TBT(每个输出 token 的时间)。
  • 通常更偏Memory-bound(显存带宽受限)。

未结束

EOS 或达到 max_tokens

用户 Prompt

Tokenizer

Prefill
并行处理 prompt
计算并写入 K/V

Sampler
得到第一个 next token

Decode Loop

计算新 token 的 Q/K/V
追加 K/V

读取全量 KV Cache
生成 next token

返回完整输出

3.3 Prefill 与 Decode 对比

对比项PrefillDecode
处理对象用户输入的完整 prompt每一步新生成的 token
token 是否已知输入 token 已知输出 token 依赖上一步采样结果
并行性可以并行处理多个 prompt token不使用推测解码时,每步每个请求通常只处理一个新 token
KV Cache 行为批量写入 prompt 的 K/V追加一个新 K/V,并读取历史缓存
主要瓶颈通常是计算瓶颈通常是显存带宽瓶颈
关键指标TTFTTPOT、TBT、流式输出速度

3.4 Chunked Prefill

如果 prompt 很长,一次性 Prefill 可能占用大量 GPU 资源,并让正在 Decode 的请求等待很久。

Chunked Prefill 会把长 prompt 切成多个 chunk,分多次调度。这样可以:

  • 避免长 Prefill 长时间阻塞 Decode。
  • 更平滑地调度多个请求。
  • 控制在线服务的 token 间延迟。

但它也可能让单个请求的 TTFT 变长,需要调度器在延迟和吞吐之间做平衡。

4. 内存墙是什么

GPU 中既有计算单元,也有显存:

  • SM / Tensor Core 负责计算。
  • HBM(High Bandwidth Memory)保存权重、KV Cache、激活等数据。

计算单元要工作,必须先把数据从 HBM 搬到芯片上的寄存器或 SRAM。问题是:GPU 计算能力的增长速度长期快于 HBM 带宽的增长速度。

当计算单元很快,但数据供应不上时,GPU 就会等待。这个瓶颈就是内存墙(Memory Wall)。

有限带宽 Bytes/s

低:FLOPs 少但读取数据多

高:计算多且数据复用充分

HBM / 显存
Weights + KV Cache + Activations

SM / Tensor Cores
计算单元

算术强度高吗?

Memory-bound
遇到内存墙

Compute-bound
主要受计算能力限制

判断一个操作更偏计算受限还是内存受限,可以看算术强度:

算术强度 = FLOPs / 访问字节数
  • 算术强度高:单位数据可以支撑很多计算,更容易受计算能力限制。
  • 算术强度低:没算多少东西就要读写大量数据,更容易受显存带宽限制。

一个简化的性能判断方式是:

数据搬运时间 ≈ 访问字节数 / 显存带宽 计算时间 ≈ FLOPs / 计算峰值 实际耗时主要取决于两者中更长的那个

Decode 阶段每一步新 token 的计算量较小,但要读取模型权重和大量 KV Cache,因此很容易变成 Memory-bound。上下文越长,KV Cache 越大,Decode 每一步要读取的数据越多,内存墙越明显。

我对内存墙的理解是:

不是 GPU 不会算,而是它需要的数据来不及从 HBM 送过来。

5. KV Cache、Prefill/Decode 和内存墙的关系

自回归生成

需要历史 token 信息

使用 KV Cache

避免重复计算历史 K/V

占用更多 GPU 显存

显存容量压力
能同时服务多少请求

Decode 每步读取 KV Cache

显存带宽压力
每个 token 生成多快

内存墙

这里其实有两个问题:

  1. 显存容量问题:KV Cache 能不能放下,能支持多少并发和多长上下文。
  2. 显存带宽问题:Decode 时读取 KV Cache 够不够快,token 生成速度是多少。

PagedAttention 更偏向解决多请求下 KV Cache 如何分页、分配和减少碎片;RadixAttention 更偏向解决公共前缀如何自动复用;FlashAttention/FlashInfer 更偏向通过分块、融合和更好的内存访问减少数据搬运。

6. 常见优化手段分别解决什么问题

优化手段主要解决的问题我的理解
KV Cache避免重复计算历史 K/V把历史 token 的 K/V 保存下来
GQA / MQA减少 KV head 数用更少的 K/V 服务更多 Query head
PagedAttentionKV Cache 分页管理和碎片问题像操作系统管理内存页一样管理 KV
RadixAttention公共前缀复用把可复用前缀组织成树,避免重复 Prefill
FlashAttention / FlashInferAttention 的 HBM 数据搬运和 kernel 融合分块读取、片上计算,避免落地巨大 attention 矩阵
Continuous Batching动态拼批,提高 GPU 利用率请求可以动态加入或离开批次
Chunked Prefill长 prompt 阻塞调度把长 Prefill 切成小块,与 Decode 一起调度
KV Cache 量化减少 KV 显存容量和读取字节用 FP8/INT4 等方式保存 K/V
推测解码提高 Decode 每步计算量和并行度一次验证多个候选 token
CUDA Graph减少 CPU 启动 kernel 的开销把多个 kernel 启动流程录下来重放
Prefill/Decode 分离两类阶段资源需求不同让 Prefill 和 Decode 使用不同 GPU 池
Tensor Parallel单卡放不下权重或 KV把权重和 KV 分片到多 GPU,但会引入通信

7. 用自己的话总结

  1. KV Cache 是一份历史 token 的 Key/Value 笔记。它让模型不用在每次生成新 token 时都重新计算整个历史。
  2. Prefill 是批量处理 prompt、批量写笔记的阶段。它通常计算密集,决定首 token 多久出现。
  3. Decode 是一边追加新笔记、一边翻看全部旧笔记的阶段。它通常受显存带宽限制,决定 token 流式输出的速度。
  4. 内存墙是计算单元和显存带宽之间的速度差。GPU 算得很快,但数据来不及从 HBM 搬过来,计算单元就只能等待。
  5. 现代推理系统不是只靠某一个优化,而是同时管理显存、批次、前缀、Attention kernel、通信和调度。

8. Mini-SGLang 中对应源码

阅读 Mini-SGLang 时,可以把笔记和下面文件对应起来:

笔记概念Mini-SGLang 对应文件
KV Cache 抽象python/minisgl/kvcache/base.py
MHA KV 池python/minisgl/kvcache/mha_pool.py
Radix Cachepython/minisgl/kvcache/radix_cache.py
Prefill 调度python/minisgl/scheduler/prefill.py
Decode 调度python/minisgl/scheduler/decode.py
KV 页和 page tablepython/minisgl/scheduler/cache.py
调度主循环python/minisgl/scheduler/scheduler.py
模型执行和采样python/minisgl/engine/engine.py
Attention 层python/minisgl/layers/attention.py
FlashAttention Backendpython/minisgl/attention/fa.py
FlashInfer Backendpython/minisgl/attention/fi.py

9. 参考资料

  1. Mini-SGLang GitHub:https://github.com/sgl-project/mini-sglang
  2. SGLang 官方文档:https://docs.sglang.io/
  3. SGLang 论文:https://arxiv.org/abs/2312.07104
  4. vLLM 官方架构文档:https://docs.vllm.ai/en/stable/design/arch_overview/
  5. vLLM PagedAttention 论文:https://arxiv.org/abs/2309.06180
  6. FlashAttention 论文:https://arxiv.org/abs/2205.14135
  7. NVIDIA GPU Performance Background:https://docs.nvidia.com/deeplearning/performance/pdf/GPU-Performance-Background-User-Guide.pdf
  8. NVIDIA Long-Context Attention 博客:https://developer.nvidia.com/blog/co-designing-ai-model-attention-for-fast-interactive-long-context-inference

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询