- Mini-SGLang GitHub 项目
- SGLang 官方文档
- vLLM 官方文档
1. 如果没有 KV Cache,会发生什么
LLM 生成文本是**自回归(Autoregressive)**的:模型一次生成一个 token,再把这个 token 接到输入后面,继续生成下一个 token。
假设用户输入 prompt 后,模型依次生成token1、token2、token3。如果没有 KV Cache,每次前向都可能重新计算历史 token 的 Key 和 Value:
这会造成大量重复计算。序列越长、生成 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 时:
- 只计算新 token 的 Query、Key、Value。
- 把新 token 的 Key/Value 追加到 KV Cache。
- 用新 token 的 Query 去读取历史 KV Cache。
- 经过 Attention 和后续网络得到 logits。
- Sampler 选出下一个 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 × NGQA/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(显存带宽受限)。
3.3 Prefill 与 Decode 对比
| 对比项 | Prefill | Decode |
|---|---|---|
| 处理对象 | 用户输入的完整 prompt | 每一步新生成的 token |
| token 是否已知 | 输入 token 已知 | 输出 token 依赖上一步采样结果 |
| 并行性 | 可以并行处理多个 prompt token | 不使用推测解码时,每步每个请求通常只处理一个新 token |
| KV Cache 行为 | 批量写入 prompt 的 K/V | 追加一个新 K/V,并读取历史缓存 |
| 主要瓶颈 | 通常是计算瓶颈 | 通常是显存带宽瓶颈 |
| 关键指标 | TTFT | TPOT、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)。
判断一个操作更偏计算受限还是内存受限,可以看算术强度:
算术强度 = FLOPs / 访问字节数- 算术强度高:单位数据可以支撑很多计算,更容易受计算能力限制。
- 算术强度低:没算多少东西就要读写大量数据,更容易受显存带宽限制。
一个简化的性能判断方式是:
数据搬运时间 ≈ 访问字节数 / 显存带宽 计算时间 ≈ FLOPs / 计算峰值 实际耗时主要取决于两者中更长的那个Decode 阶段每一步新 token 的计算量较小,但要读取模型权重和大量 KV Cache,因此很容易变成 Memory-bound。上下文越长,KV Cache 越大,Decode 每一步要读取的数据越多,内存墙越明显。
我对内存墙的理解是:
不是 GPU 不会算,而是它需要的数据来不及从 HBM 送过来。5. KV Cache、Prefill/Decode 和内存墙的关系
这里其实有两个问题:
- 显存容量问题:KV Cache 能不能放下,能支持多少并发和多长上下文。
- 显存带宽问题: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 |
| PagedAttention | KV Cache 分页管理和碎片问题 | 像操作系统管理内存页一样管理 KV |
| RadixAttention | 公共前缀复用 | 把可复用前缀组织成树,避免重复 Prefill |
| FlashAttention / FlashInfer | Attention 的 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. 用自己的话总结
- KV Cache 是一份历史 token 的 Key/Value 笔记。它让模型不用在每次生成新 token 时都重新计算整个历史。
- Prefill 是批量处理 prompt、批量写笔记的阶段。它通常计算密集,决定首 token 多久出现。
- Decode 是一边追加新笔记、一边翻看全部旧笔记的阶段。它通常受显存带宽限制,决定 token 流式输出的速度。
- 内存墙是计算单元和显存带宽之间的速度差。GPU 算得很快,但数据来不及从 HBM 搬过来,计算单元就只能等待。
- 现代推理系统不是只靠某一个优化,而是同时管理显存、批次、前缀、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 Cache | python/minisgl/kvcache/radix_cache.py |
| Prefill 调度 | python/minisgl/scheduler/prefill.py |
| Decode 调度 | python/minisgl/scheduler/decode.py |
| KV 页和 page table | python/minisgl/scheduler/cache.py |
| 调度主循环 | python/minisgl/scheduler/scheduler.py |
| 模型执行和采样 | python/minisgl/engine/engine.py |
| Attention 层 | python/minisgl/layers/attention.py |
| FlashAttention Backend | python/minisgl/attention/fa.py |
| FlashInfer Backend | python/minisgl/attention/fi.py |
9. 参考资料
- Mini-SGLang GitHub:https://github.com/sgl-project/mini-sglang
- SGLang 官方文档:https://docs.sglang.io/
- SGLang 论文:https://arxiv.org/abs/2312.07104
- vLLM 官方架构文档:https://docs.vllm.ai/en/stable/design/arch_overview/
- vLLM PagedAttention 论文:https://arxiv.org/abs/2309.06180
- FlashAttention 论文:https://arxiv.org/abs/2205.14135
- NVIDIA GPU Performance Background:https://docs.nvidia.com/deeplearning/performance/pdf/GPU-Performance-Background-User-Guide.pdf
- NVIDIA Long-Context Attention 博客:https://developer.nvidia.com/blog/co-designing-ai-model-attention-for-fast-interactive-long-context-inference