大模型推理优化:KV Cache原理、PagedAttention与GQA实战 1. 为什么每个搞大模型推理的人都绕不开 KV Cache如果你最近在折腾大模型本地部署或者读过 vLLM、TensorRT-LLM、SGLang 这些推理框架的源码一定见过KV Cache这个词。它不是什么新鲜概念但绝对是决定推理速度和显存占用的头号选手。我第一次在 24G 显存的卡上跑 7B 模型时发现 batch size 死活上不去显存监控里一大块都被一个叫 “kv_cache” 的东西吃掉了当时就意识到不把 KV Cache 搞明白推理优化根本无从谈起。简单说KV Cache 是大模型自回归生成过程中把历史 token 的 Key 和 Value 张量缓存下来避免每生成一个新 token 都重新计算整段序列的注意力。没有它生成第 N 个 token 的计算量会随序列长度线性增长推理速度会慢到无法接受。有了它每步只需要计算当前 token 的 Q、K、V然后和缓存拼接即可。这个优化带来的加速比在长上下文场景下可以达到几十倍甚至上百倍。这篇文章适合谁看如果你正在做大模型推理优化、私有化部署、显存调优或者单纯想搞懂 Transformer 推理时显存到底花在哪了那这篇内容就是为你准备的。我会从原理讲到实现从朴素 KV Cache 讲到PagedAttention和分组查询注意力GQA再配上实际部署中踩过的坑和调参经验。不堆公式尽量用生活化的类比把事说清楚。2. KV Cache 到底缓存了什么从注意力计算说起2.1 自回归生成的核心矛盾大模型的文本生成是自回归的给定 prompt模型预测下一个 token然后把新 token 拼到输入后面再预测下一个如此循环。问题在于标准 Transformer 的注意力机制每一步都要对整段序列计算 Q、K、V。假设序列长度是 L那每生成一个 token就要做一次 O(L²) 的注意力计算。生成 1000 个 token就是 1000 次 O(L²)其中大量计算是重复的——因为历史 token 的 K 和 V 根本没变。这就像你每次翻书查资料都要把前面所有页重新读一遍。正常人不会这么干而是记住之前读过的内容。KV Cache 做的就是这件事把每一层、每个注意力头在历史位置上算出来的 K 和 V 存起来下一步直接拿来用。2.2 缓存的具体形状与显存占用以 LLaMA 架构为例KV Cache 的形状通常是[batch_size, num_layers, num_kv_heads, seq_len, head_dim]。注意这里有个关键点num_kv_heads 不一定等于 num_heads这就是 GQA 的由来后面会细讲。显存占用怎么算假设一个 7B 模型32 层hidden size 409632 个注意力头head_dim 128FP16 精度2 字节。如果不用 GQAnum_kv_heads 32那么每个 token 的 KV Cache 大小是2 (K和V) × 32 (层) × 32 (头) × 128 (head_dim) × 2 (字节) 524288 字节 ≈ 0.5 MB看起来不大但序列长度到 4096 时单个请求就是 2GB。如果 batch size 是 16直接 32GB 显存没了。这就是为什么长上下文和高并发不能兼得——KV Cache 是线性增长的而显存是固定的。提示很多人以为模型权重占显存大头实际上在长序列、大 batch 场景下KV Cache 往往才是显存杀手。部署时一定要把这块算清楚。2.3 为什么不能简单丢掉缓存有人会想能不能只保留最近几个 token 的缓存可以但会损失长距离依赖。大模型之所以能处理长文档靠的就是注意力能看到很远的位置。如果粗暴截断 KV Cache模型就会“失忆”回答质量断崖式下跌。所以工业界的方向不是丢弃而是更高效地管理——这就引出了 PagedAttention。3. 从朴素实现到 PagedAttentionKV Cache 的管理进化3.1 朴素 KV Cache 的三大浪费最直接的实现是给每个请求预分配一块连续显存大小按最大序列长度算。这种做法有三个明显问题第一内部碎片。你按 4096 长度预分配但实际请求可能只生成 200 个 token剩下 3896 个位置全浪费了。第二外部碎片。不同请求的缓存块大小不一频繁分配释放后显存里会出现很多零散空洞最后明明总空闲够却找不到一块连续空间。第三无法共享。多个请求如果有相同的前缀比如相同的 system prompt它们的 KV Cache 其实可以复用但连续分配模式下很难做到。我早期用 HuggingFace 的generate接口时就吃过这个亏并发一高就 OOM但nvidia-smi看显存占用率又没满典型碎片问题。3.2 PagedAttention 的核心思想把显存当操作系统管vLLM 提出的PagedAttention借鉴了操作系统的虚拟内存分页机制。它把 KV Cache 切成固定大小的 block比如 16 个 token 一块每个请求维护一张 block table记录自己的逻辑块到物理块的映射。物理块可以不连续按需分配。这样做的好处非常直接消除内部碎片最后一块没填满也没关系浪费最多一个 block。消除外部碎片所有 block 大小一致显存池化管理随用随取。支持前缀共享相同前缀的请求可以指向同一批物理块用引用计数管理写时复制。实测下来vLLM 在相同显存下能把吞吐量提升 2 到 4 倍长上下文场景提升更明显。这也是为什么现在几乎所有主流推理框架都在抄这套思路。3.3 Block 大小怎么选一个容易被忽略的调参点PagedAttention 的 block size 默认是 16但这个值不是随便定的。block 太小block table 会很大管理开销上升block 太大内部碎片又会增加。经验上Block Size优点缺点适用场景8碎片小共享粒度细表大调度开销高短序列、高并发16平衡平衡通用默认32表小调度快碎片略大长序列、低并发我一般先用默认 16 跑 benchmark如果发现调度成为瓶颈GPU 利用率上不去但 CPU 很忙再尝试调到 32。反过来如果显存浪费严重就降到 8 试试。4. 分组查询注意力从源头减少 KV Cache4.1 MHA、MQA、GQA 的关系PagedAttention 解决的是“怎么管”而GQAGrouped Query Attention解决的是“怎么少”。标准多头注意力MHA里每个 Query 头都有自己独立的 Key 和 Value 头KV Cache 最大。多查询注意力MQA让所有 Query 头共享一组 K、V缓存直接缩小到 1/num_heads但效果会掉。GQA 是折中把 Query 头分组每组共享一组 K、V。举个例子32 个 Query 头如果分成 8 组那 num_kv_heads 8KV Cache 直接降到原来的 1/4。LLaMA 2 70B、LLaMA 3、Qwen 系列很多模型都用了 GQA这不是偶然——在长上下文和高并发场景下GQA 带来的显存节省是实打实的。4.2 GQA 对推理性能的实际影响有人担心 GQA 会损失模型质量。从公开的评测看在同等参数量下GQA 相比 MHA 的质量下降很小但推理吞吐提升显著。原因在于注意力机制里 K、V 的冗余度本来就高共享并不会丢失太多信息。实际部署时如果你的模型是 MHA 架构又想在长上下文场景下省显存可以考虑用一些工具做注意力头的剪枝或合并但这属于模型改造风险较高。更稳妥的做法是直接选用原生支持 GQA 的模型。4.3 代码层面怎么看 num_kv_heads在 HuggingFace 的 config 里这个参数叫num_key_value_heads。如果它等于num_attention_heads就是 MHA如果等于 1就是 MQA介于两者之间就是 GQA。加载模型时KV Cache 的维度要按num_key_value_heads来算千万别按num_attention_heads算否则显存预估会差好几倍。from transformers import AutoConfig config AutoConfig.from_pretrained(meta-llama/Llama-2-7b-hf) print(config.num_attention_heads) # 32 print(config.num_key_value_heads) # 32 (LLaMA 2 7B 是 MHA) # LLaMA 2 70B 则是 64 和 8GQA5. 实操在本地部署中观测和调优 KV Cache5.1 用 vLLM 跑一个对比实验光说理论没感觉直接上实验。我用一台 24G 显存的机器跑 Qwen2.5-7B-InstructGQAnum_kv_heads4分别测试不同gpu_memory_utilization和max_model_len下的吞吐。# 启动 vLLM 服务 python -m vllm.entrypoints.openai.api_server \ --model Qwen/Qwen2.5-7B-Instruct \ --gpu-memory-utilization 0.9 \ --max-model-len 8192 \ --max-num-seqs 64关键参数解释gpu_memory_utilization控制 vLLM 预分配多少显存给模型权重加 KV Cache 池max_model_len决定单个序列最长多少max-num-seqs是并发请求上限。这三个参数互相制约调大任何一个都可能挤压其他空间。实测结果仅供参考不同硬件有差异max_model_lenmax_num_seqs吞吐 (tokens/s)是否 OOM4096641850否8192321420否819264-是1638416980否可以看到序列长度翻倍并发就得砍半否则直接 OOM。这就是 KV Cache 的线性增长在起作用。5.2 监控 KV Cache 使用率vLLM 的日志里会打印GPU KV cache usage这个指标非常关键。如果它长期接近 100%说明缓存池不够请求会被排队甚至抢占。如果长期低于 50%说明显存分配过多可以调低gpu_memory_utilization给其他进程留空间。我一般会压测到 KV cache usage 稳定在 80% 到 90% 之间这样既不浪费显存又有一定余量应对突发流量。5.3 前缀共享的实战价值如果你的应用有大量相同 system prompt 的请求比如客服机器人、代码助手一定要开启前缀共享。vLLM 默认支持只要请求的 prompt 前缀一致对应的 KV Cache block 就会被复用。我测过一个场景system prompt 有 800 个 token100 个并发请求开启前缀共享后首 token 延迟从 320ms 降到 90ms效果立竿见影。注意前缀共享要求 prompt 完全一致包括空格和换行。如果 system prompt 里有动态内容比如时间戳共享就会失效。设计 prompt 时尽量把固定部分放前面。6. 常见问题与排查技巧实录6.1 为什么我的显存明明够还是 OOM这是最常被问到的问题。原因通常有三个一是显存碎片尤其是用非 PagedAttention 的框架时二是 KV Cache 预估错误比如按 MHA 算了但模型其实是 GQA或者反过来三是max_model_len设得太大预分配就占满了。排查顺序先看框架日志里的 KV cache 分配量再对比理论计算值。如果差很多检查num_key_value_heads和精度设置。如果理论值就对不上那就是碎片问题换用支持 PagedAttention 的框架。6.2 长上下文推理特别慢怎么办长上下文慢一部分是注意力计算本身 O(L²) 的锅另一部分是 KV Cache 读写带宽的锅。优化方向一是用 GQA 模型从源头减少缓存量二是开启 chunked prefill把长 prompt 分块处理避免一次性占用过多缓存三是考虑量化 KV Cache比如 FP8 或 INT8能再省一半显存但要注意精度损失。6.3 并发上不去GPU 利用率却很低这种情况往往是调度瓶颈不是计算瓶颈。检查 block size 是否太小导致 block table 管理开销过大或者max-num-seqs设得太保守。另外如果请求的输入输出长度差异很大调度器可能被长请求拖住可以尝试开启抢占式调度让短请求优先完成。现象可能原因解决方向OOM 但显存监控未满碎片或预估错误换 PagedAttention核对 num_kv_heads长上下文慢注意力计算 缓存带宽GQA、chunked prefill、KV 量化并发低、GPU 闲调度瓶颈调 block size、max-num-seqs、抢占调度首 token 延迟高前缀未共享固定 prompt 前缀开启前缀缓存6.4 一个容易被忽略的坑精度与缓存KV Cache 的精度默认跟随模型精度。如果模型是 FP16缓存也是 FP16。有人为了省显存把缓存改成 FP8结果发现输出质量下降明显尤其是需要精确回忆细节的任务。我的建议是先保证质量再谈省显存。如果确实要量化先在小流量下做 A/B 测试确认质量可接受再全量。7. 我对 KV Cache 优化的一点个人体会折腾了这么多轮我最大的感受是KV Cache 不是孤立的技术点它和模型架构、推理框架、硬件显存三者深度耦合。选模型时看 GQA选框架时看 PagedAttention调参时看显存和吞吐的平衡。任何一环没考虑到都会在压测时暴露出来。另外不要迷信默认参数。vLLM 的默认值适合通用场景但你的业务可能有特殊的长短分布。花半天时间做一轮参数扫描往往比换硬件更有效。我见过太多人一 OOM 就想着加卡其实把max_model_len从 16384 降到 8192再配合前缀共享同样的卡就能多扛一倍并发。最后分享一个小技巧压测时用真实业务的 prompt 分布别用随机生成的等长序列。真实场景里长短差异极大调度器的表现和均匀分布完全不同。这个细节很多官方 benchmark 不会告诉你但实际部署时能帮你少走很多弯路。