11. KV Cache and Memory Growth | KV Cache 与显存增长

难度： Medium | 环境： CPU-first | 标签： 推理显存, Attention, KV Cache | 目标人群： 长上下文推理入门者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

(国内推荐：魔搭社区免费实例)

先把 KV cache 为什么会随上下文增长、为什么多头机制会放大显存压力、以及 PagedAttention / MLA 分别在解决什么问题讲清楚，再去看后面的 Attention 优化和推理系统实现，直觉会更稳。

关键词： KV cache, sequence length, memory growth

前置阅读

导语： 先把数据格式和显存账本对齐，再看 KV cache 的增长规律会更顺。

• Group 1A: Numerical Foundations and Scale Estimation | 1A: 数值基础与算力估算
• Group 1B: Single-GPU Hardware and Memory Optimization | 1B: 单卡硬件与访存优化

相关阅读

导语： 把 KV cache 和 FlashAttention、PagedAttention、MLA 放在一起看，能更快理解不同优化点。

• 14. FlashAttention Memory Model | FlashAttention 显存模型
• 20. FlashAttention Sim | FlashAttention 模拟
• 22. vLLM PagedAttention | vLLM 分页注意力

Q1：为什么 KV cache 会随着上下文长度增长？

点击展开查看解析

KV cache 本质上是在解码阶段保存历史 token 的 Key 和 Value。每生成一个新 token，模型都要把这一 token 在每一层、每一个 KV 头上的 K/V 追加进缓存里，所以它的增长不是“偶尔增加一点”，而是随着上下文长度持续线性增长。

更具体地说，标准自回归推理里，缓存大小通常可以写成：

$$\text{KV Cache Bytes}\approx 2\times L\times B\times H_{kv}\times D\times S$$

其中：

• $L$ 是层数
• $B$ 是 batch size
• $H_{kv}$ 是 KV 头数
• $D$ 是 head dim
• $S$ 是上下文长度
• 前面的 $2$ 表示同时存 K 和 V

这个公式告诉我们一个很关键的事实：KV cache 是“历史状态成本”。上下文越长，历史 token 越多，缓存就越大；batch 越大，缓存也会同步放大；层数越多，这个成本还会在每一层重复一次。

所以长上下文推理里，最先爆的往往不是算力，而是显存。只要把“每个 token 都要被长期保留”这件事想清楚，KV cache 的增长规律就很自然了。

### Q1小验证：KV cache 增长直觉

把上下文长度翻倍，再看缓存大小是不是也几乎翻倍。

def kv_cache_bytes(seq_len, num_layers, num_kv_heads, head_dim, batch_size=1, dtype_bytes=2):
    return 2 * seq_len * num_layers * num_kv_heads * head_dim * batch_size * dtype_bytes

examples = [(1024, 32, 32, 128), (2048, 32, 32, 128), (4096, 32, 32, 128)]
for seq_len, layers, kv_heads, head_dim in examples:
    size_gb = kv_cache_bytes(seq_len, layers, kv_heads, head_dim) / 1e9
    print(f"seq_len={seq_len:4d} -> KV cache ≈ {size_gb:5.2f} GB")

数量级速览

变化项	对 KV cache 的影响	直觉
seq_len 翻倍	近似翻倍	历史 token 变多，缓存线性增长
batch_size 翻倍	近似翻倍	多路请求共享同一层结构，但缓存要按样本复制
num_layers 翻倍	近似翻倍	每层都要单独保存一份 K/V
num_kv_heads 增加	近似线性增加	KV 头越多，缓存越大

这一张表的目的，是先把“增长方向”记牢，再去看后面的 MHA / MQA / GQA 和分页管理。

Q2：为什么 MHA、MQA 和 GQA 会影响 KV cache 压力？

点击展开查看解析

它们影响的是 KV cache 的头数维度。

• MHA (Multi-Head Attention)：每个 query head 都有自己对应的一组 K/V，缓存压力最大。
• MQA (Multi-Query Attention)：多个 query head 共享同一组 K/V，KV cache 立刻变小。
• GQA (Grouped-Query Attention)：介于 MHA 和 MQA 之间，把 query heads 分组共享 K/V，在显存和表达能力之间做折中。

从缓存角度看，真正决定显存大小的不是 query heads 有多少，而是 要存多少组 K/V。所以只要 KV 头数下降，缓存就会按比例下降。

这也是为什么很多长上下文模型会采用 MQA 或 GQA：它们不只是“改了 attention 的形式”，而是在直接压低推理时的 KV cache 成本。

### Q2小验证：头数与显存直觉

固定上下文长度和层数，只改变 KV 头数，观察显存变化。

def kv_cache_gb(seq_len, num_layers, num_kv_heads, head_dim, batch_size=1, dtype_bytes=2):
    return kv_cache_bytes(seq_len, num_layers, num_kv_heads, head_dim, batch_size, dtype_bytes) / 1e9

seq_len = 4096
num_layers = 32
head_dim = 128
for name, kv_heads in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    print(f"{name:>3s}: kv_heads={kv_heads:2d}, KV cache ≈ {kv_cache_gb(seq_len, num_layers, kv_heads, head_dim):5.2f} GB")

Q3：PagedAttention 和 MLA 分别在解决什么问题？

点击展开查看解析

这两个方法解决的层面不一样。

• PagedAttention 主要解决的是 缓存分配和访问组织 问题。

  • 它把 KV cache 按页组织，避免长序列和多请求场景里出现连续大块显存分配困难。
  • 它的重点不是把 K/V 表示本身压缩掉，而是让缓存的存储、搬运和复用更稳定。

• MLA (Multi-Head Latent Attention) 主要解决的是 表示压缩 问题。

  • 它把原本需要长期保存的 KV 表示压到更低维的潜变量空间里。
  • 这样做的核心收益是直接降低每个 token 需要保留的缓存体积。

可以把它们理解成两种不同方向的优化：

• PagedAttention 是在优化“怎么管 cache”。
• MLA 是在优化“cache 本身有多大”。

前者偏系统实现，后者偏表示结构。两者都在缓解长上下文下的显存压力，但切入点不同。

### Q3小验证：方案对照

把“缓存管理”和“表示压缩”分开看，再判断它们各自对显存的影响。

def paged_attention_pages(seq_len, page_size):
    return (seq_len + page_size - 1) // page_size

def mla_cache_bytes(seq_len, num_layers, latent_dim, batch_size=1, dtype_bytes=2):
    return seq_len * num_layers * latent_dim * batch_size * dtype_bytes

seq_len = 4096
page_size = 128
print(f"PagedAttention pages: {paged_attention_pages(seq_len, page_size)}")
print(f"MLA cache example: {mla_cache_bytes(seq_len, 32, 64) / 1e9:.2f} GB (latent_dim=64)")

⚠️ 常见误区

• KV cache 不是只和 token 数有关，它还和层数、batch size、KV 头数一起增长。
• MQA / GQA 不是单纯改名字，而是在实打实地压低缓存体积。
• PagedAttention 解决的是缓存管理和碎片化，不等于表示压缩。
• MLA 解决的是表示体积，不等于把调度和分配问题也一并解决。

这一页最重要的是记住三件事：KV cache 为什么会涨、头数为什么会放大或压缩它、以及不同优化到底在解决“管理”还是“表示”问题。