文章目录一、自回归生成为什么会重复计算二、Attention 里哪些东西可以缓存三、Prefill 和 Decode两个阶段成本不同四、用代码看 KV Cache 形状增长五、一个可选 PyTorch 形状示例六、KV Cache 到底省了什么七、KV Cache 显存怎么估算八、GQA/MQA 为什么能省 KV Cache九、PagedAttention 的直觉十、多轮对话中的 KV Cache 复用边界十一、工程取舍十二、常见误区总结大模型视角下一篇摘要大语言模型生成回答时是一个 token 一个 token 往外吐。如果每生成一个新 token 都重新计算全部历史上下文推理会非常慢。KV Cache 的核心思想是历史 token 的 Key 和 Value 已经算过就缓存起来下一步生成时只为新 token 计算新的 Q/K/V并让新 Query 关注历史 K/V。本文从自回归生成的重复计算讲起解释 prefill、decode、为什么缓存 K/V 而不是 Q、KV Cache 的张量形状、显存估算、PagedAttention 和工程取舍。前置知识多头注意力 Tokenization混合精度阅读时间约 55-70 分钟代码环境Python 3.10示例只依赖标准库可选 PyTorch 版本用于形状演示入门导读先抓住主线如果你第一次读 KV Cache先抓住一个工程问题生成第 100 个 token 时前 99 个 token 的 Key/Value 其实已经算过了。如果每一步都重新算一遍历史推理会越来越慢。KV Cache 用显存换速度Prefill一次性处理 prompt填好历史 K/V Decode每次只处理新 token复用历史 K/V 代价上下文越长、并发越高KV Cache 显存越大。读完本文你应该能解释KV Cache 省了哪些重复计算为什么它会成为长上下文和高并发推理的核心瓶颈。一、自回归生成为什么会重复计算Decoder-Only 大模型按从左到右生成。给定 prompt请解释 KV Cache 的作用模型先生成第一个 token再把这个 token 接到上下文后面继续生成第二个 token。每生成一步模型都要运行 Transformer。如果没有缓存生成过程会像这样历史 token 被反复计算。上下文越长、生成越长浪费越明显。KV Cache 的核心就是历史 token 的某些中间结果不会变算过后存起来下一步复用。二、Attention 里哪些东西可以缓存Self-Attention 中每层会从 hidden states 计算 Q、K、VQ X W Q , K X W K , V X W V QXW_Q,\quad KXW_K,\quad VXW_VQXWQ​,KXWK​,VXWV​这里X XX当前层输入 hidden statesW Q , W K , W V W_Q,W_K,W_VWQ​,WK​,WV​三组投影矩阵Q QQQuery用来发起查询K KKKey用来和 Query 计算相似度V VVValue被注意力权重加权汇总。生成新 token 时历史 token 的 hidden states 已经确定它们在每层对应的 K/V 也已经确定。因此历史 K/V 可以缓存。新 token 需要做的事是1. 计算新 token 当前层的 Q/K/V 2. 把新 K/V 追加到 cache 3. 用新 Q 和所有历史 K 计算 attention 4. 用 attention 权重加权所有历史 V 5. 进入下一层。注意新 token 的 Query 每一步都不同不能复用历史 Query。历史 Query 对下一步生成通常也没用。三、Prefill 和 Decode两个阶段成本不同推理通常分为两个阶段。Prefill 阶段处理完整 prompt计算所有 prompt token 的 K/V填充 KV Cache。这个阶段可以并行处理整段输入所以吞吐相对高但 prompt 越长prefill 越贵。Decode 阶段每次生成一个新 token。这个阶段会读取已有 KV Cache只计算新 token 的相关内容。单步计算比重新处理全上下文便宜但 decode 必须逐 token 串行因此对延迟很敏感。这就是为什么大模型服务常分开统计prefill latency处理输入需要多久 decode latency每生成一个 token 需要多久 time to first token用户等第一个输出的时间 tokens/s持续生成速度。长 prompt 主要影响 prefill长回答主要影响 decode长上下文和高并发都会增加 KV Cache 显存。四、用代码看 KV Cache 形状增长下面用纯 Python 模拟 cache 长度增长。真实模型里 K/V 是多维张量这里先看序列维度。# Python 3.10k_cache[]v_cache[]forstepinrange(5):new_kfK{step}new_vfV{step}k_cache.append(new_k)v_cache.append(new_v)print(step,step,K cache:,k_cache,V cache:,v_cache)每生成一个 tokencache 里就多一份 K/V。真实张量常见形状可以理解为batch_size × num_kv_heads × seq_len × head_dim每一层都有自己的 K cache 和 V cache所以总显存还要乘以层数。五、一个可选 PyTorch 形状示例如果你的环境有 PyTorch可以运行下面代码看真实张量形状。没有 PyTorch 也不影响理解。# Python 3.10 / PyTorch 2.ximporttorch batch1kv_heads2head_dim4k_cache[]v_cache[]forstepinrange(3):new_ktorch.randn(batch,kv_heads,1,head_dim)new_vtorch.randn(batch,kv_heads,1,head_dim)k_cache.append(new_k)v_cache.append(new_v)all_ktorch.cat(k_cache,dim2)all_vtorch.cat(v_cache,dim2)print(fstep{step}, K{tuple(all_k.shape)}, V{tuple(all_v.shape)})输出里第三个维度会从 1 增长到 3对应已经缓存的 token 数。真实大模型会对每一层都维护这样的缓存。六、KV Cache 到底省了什么没有 KV Cachedecode 第t tt步可能需要重新计算前面t tt个 token 的 K/V。生成越长重复越多。有 KV Cache 后历史 K/V 只算一次。每一步主要计算新 token 的 Q/K/V然后让新 Q 关注历史 K/V。可以用一个粗略计数理解# Python 3.10defwork_without_cache(prompt_len,gen_len):# 每步重新处理 prompt 已生成 token粗略计 token 次数returnsum(prompt_leniforiinrange(gen_len))defwork_with_cache(prompt_len,gen_len):# prefill 处理 promptdecode 每步处理 1 个新 tokenreturnprompt_lengen_lenforgen_lenin[10,100,1000]:no_cachework_without_cache(prompt_len1000,gen_lengen_len)cachework_with_cache(prompt_len1000,gen_lengen_len)print(gen_len,without_cache,no_cache,with_cache,cache,ratio,round(no_cache/cache,1))这不是精确 FLOPs但能说明重复计算差距会随着生成长度扩大。七、KV Cache 显存怎么估算KV Cache 显存大致和这些因素成正比。这里刻意按主流注意力张量[batch, num_kv_heads, seq_len, head_dim]的维度顺序来写公式和实际张量形状对得上B × H_kv × S × D × 2(K和V) × L × dtype_size写成公式M ≈ B × H k v × S × D ⏟ 单层单份的元素数 × 2 × L × bytes M \approx \underbrace{B \times H_{kv} \times S \times D}_{\text{单层单份的元素数}} \times 2 \times L \times \text{bytes}M≈单层单份的元素数B×Hkv​×S×D​​×2×L×bytes其中B BBbatch sizeH k v H_{kv}Hkv​每层的 KV heads 数S SS当前上下文序列长度D DD每个 head 的维度head_dim d_model / n_heads2 22K 和 V 两份L LLTransformer 层数bytes \text{bytes}bytes每个元素字节数FP16/BF16 2FP32 4INT8 1。注意分组顺序先把单层单份 KV 的元素数B ⋅ H k v ⋅ S ⋅ D B\cdot H_{kv}\cdot S\cdot DB⋅Hkv​⋅S⋅D算出来再乘 2K 和 V再乘层数L LL这样和每层里 K/V 的张量形状是[B, H_kv, S, D]的记忆一致也就是下面代码里batch * seq_len * kv_heads * head_dim这一段的语义。估算代码# Python 3.10defkv_cache_gb(batch,layers,seq_len,kv_heads,head_dim,bytes_per_value2):totalbatch*layers*seq_len*kv_heads*head_dim*2*bytes_per_valuereturntotal/1024**3configs[{batch:1,seq:4096},{batch:8,seq:4096},{batch:8,seq:32768},]forcinconfigs:gbkv_cache_gb(c[batch],layers32,seq_lenc[seq],kv_heads8,head_dim128)print(c,fKV Cache≈{gb:.2f}GB)这个例子说明KV Cache 对 batch 和 seq_len 都很敏感。高并发长上下文服务KV Cache 往往比权重之外的其他运行时开销更关键。八、GQA/MQA 为什么能省 KV Cache多头注意力里Query head 可以很多但 Key/Value head 不一定同样多。GQAGrouped Query Attention和 MQAMulti-Query Attention通过减少 KV heads 来降低缓存。在估算公式里KV Cache 和kv_heads成正比。看代码# Python 3.10defkv_cache_gb(layers,seq_len,kv_heads,head_dim,bytes_per_value2):returnlayers*seq_len*kv_heads*head_dim*2*bytes_per_value/1024**3forkv_headsin[32,16,8,1]:print(kv_heads,f{kv_cache_gb(32,32768,kv_heads,128):.2f}GB)这就是为什么很多现代大模型使用 GQA/MQA它们不只是架构细节而是直接影响推理显存和长上下文并发。九、PagedAttention 的直觉高并发推理时每个请求长度不同、生成进度不同。如果给每个请求预留连续大块 KV Cache很容易浪费显存并产生碎片。PagedAttention 的思想类似操作系统分页把 KV Cache 切成小块页按需分配和管理。简化类比普通方式每个请求预留一整块连续缓存 分页方式请求需要多少页就分配多少页生成增长时再追加页。这能提高显存利用率让推理服务支持更多并发。vLLM 的高吞吐能力很大程度就来自这类 KV Cache 管理和调度优化。十、多轮对话中的 KV Cache 复用边界多轮对话看起来适合复用 cache但要注意边界1. 用户新输入会改变上下文不能随便复用不匹配的 cache 2. system prompt、历史消息、工具结果都属于上下文 3. 如果中间插入或删除历史cache 需要重建或做复杂管理 4. 会话太长时KV Cache 显存会持续增长 5. 生产系统通常需要截断、摘要或重建上下文。KV Cache 是对“已有上下文前缀”的缓存。只要前缀完全相同就能复用前缀变了就要谨慎处理。十一、工程取舍目标KV Cache 相关取舍降低生成延迟必须使用 KV Cache支持长上下文KV Cache 显存线性增长提高并发需要分页、调度和内存管理降低显存GQA/MQA、KV 量化、滑动窗口、分页管理降低首 token 延迟优化 prefill 和 prompt 长度降低持续生成延迟优化 decode、batching、cache 读取KV Cache 是推理加速的核心但它不是免费的。它把重复计算问题转成显存管理问题。十二、常见误区误区一KV Cache 会减少模型参数显存。不会。它减少重复计算但增加额外运行时缓存。误区二上下文越长只是计算慢一点。上下文越长KV Cache 显存也会线性增长。误区三有 KV Cache 后生成可以完全并行。单个请求的 decode 仍然逐 token 串行只是每一步更快。误区四所有框架 KV Cache 管理都一样。不同推理框架在分页、调度、continuous batching、cache 复用上差异很大。误区五历史 Query 也应该缓存。不需要。原因是解码时每一步都会为当前新 token生成一个全新的 Query它才是用来发起查询的那一方而历史位置的 Query 只在训练/prefill 时用过一次之后再也不参与后续步骤的 attention 计算——它既不是被查询对象那是 K/V也不再作为 Query 出现所以缓存它没有任何复用价值。缓存 K/V 是因为 attention 里所有历史位置的 K/V 都会被后续每一步的新 Query 用到而 Query 只是当前这一步用完就丢的一次性输入。总结KV Cache 缓存历史 token 的 Key 和 Value避免自回归生成时重复计算历史上下文。Prefill 阶段填充 cacheDecode 阶段逐 token 读取和追加 cache。它显著提升推理速度但显存开销随 batch、层数、KV heads、上下文长度和 dtype 增长是推理系统优化的核心。第一遍记住一句话KV Cache 用显存换速度把历史 token 的 K/V 存起来生成新 token 时不再重复算。大模型视角后面讲 vLLM、PagedAttention、长上下文推理、GQA/MQA、推理部署时KV Cache 都是核心概念。理解它后你就能看懂为什么推理框架如此关注显存管理、调度和连续批处理。下一篇Scaling Law大模型涌现能力的规律—— KV Cache 解释推理成本下一篇看模型规模、数据和算力如何影响能力增长以及为什么大模型训练不能只看参数量。