ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

KV Cache显存优化指南:从原理到Ollama/vLLM部署踩坑

KV Cache显存优化指南:从原理到Ollama/vLLM部署踩坑 内存不够跑不动大模型瓶颈很可能就藏在 KV Cache 里。这两年大模型推理相关的讨论十句里有八句绕不开 KV Cache。本地部署过 Llama、Qwen 这类模型的朋友应该都有体会模型权重明明不大推理起来显存却飙得飞快上下文一长甚至直接 OOM。我自己第一次用 Ollama 跑 7B 模型、把上下文开满时也盯着监控面板里的显存曲线愣了半天—— weights 只占了几个 G怎么算着算着就爆了答案在于生成每个 token 时模型都会把历史上下文中的 Key 和 Value 矩阵存在显存里供后续注意力计算复用。这堆缓存就是 KV Cache。这篇文章我从原理讲到落地结合我做过的显存估算、Ollama 和 vLLM 部署踩坑经历把 KV Cache 的前因后果、内存计算、优化手段和排障思路全部拆开说透。不管是刚接触大模型推理的新手还是想调优服务端性能的工程师都能从中捞到点能直接用的东西。1. KV Cache 是什么先回答三个最朴素的疑问1.1 为什么生成 token 时要去“翻旧账”先回忆一下大模型生成文本的基本动作每轮预测都是“把当前序列喂进模型输出下一个 token 的概率分布然后取一个结果接到序列末尾再重复”。问题在于自注意力机制每层都要计算当前 token 和序列中所有 token 的相关性。序列越长计算量越大。更关键的是当你已经在生成第 100 个 token 时前 99 个 token 的隐藏状态已经变过头了——它们的 Key 和 Value 早就被各层注意力的投影矩阵变换过这些张量本质上已经算好放在那里了。如果每次预测都要重新对整段历史做一次完整前向传播那种“把所有历史重新跑一遍”的做法计算复杂度会随序列长度平方级上涨300 token 的上下文可能还算凑合到 4k、32k 就是灾难。那么 KV Cache 的思路极为朴实把计算过的 Key 和 Value 张量缓存下来下一个 token 只需要算它自己的 Key、Value 和 Query然后拿 Query 去跟缓存的 Key 做点积得到注意力分数再对缓存的 Value 做加权求和。历史部分彻底免于重复计算新序列的前向代价变成 O(1)单个 token 的新计算量不再随历史长度增长。我画过一条曲线对比两种模式的推理耗时用 2048 长度的上下文无缓存模式下解码 100 个 token内存翻倍不说耗时几乎每隔一段时间就跳一个台阶有缓存后解码阶段每个 token 耗时基本是一条平线。这个差距就是 KV Cache 的价值——它把“读一遍历史”变成了“读一张历史摘要表”。1.2 K 和 V 到底长什么样存到哪里去了以经典的 MHAMulti-Head Attention结构为例输入序列有 L 个 token、隐藏维度为 H注意力头数为 A每个头的维度 D H / A。对于每一层、每一个头K 和 V 的形状通常就是 [L, D]。但是别忽略了多头投影时的维度变换。在 PyTorch 实现中输入 [B, L, H] 经多头注意力权重矩阵 W_k、W_v 投影后会 reshape 成 [B, A, L, D]。所以对单个序列来说每一层缓存一对 K/V 张量形状分别为 [A, L, D]。有 N 层就有 N 对。不同框架存储的 layout 不太一样但这不影响本质——缓存对象就是每层注意力需要复用的一组矩阵。“存到哪里”是很多初学者关心的问题。KV Cache 完全在 GPU 显存里不走内存。因为解码阶段每一轮迭代都要把它读回 SRAM 做矩阵乘如果落到 CPU 内存PCIe 带宽那点速度根本喂不饱 GPU整个推理会被拖成乌龟爬。除非做诸如 Paged KV Cache 和 CPU offload 之类的特殊优化。1.3 预填充阶段和解码阶段KV Cache 的命运截然不同大模型推理有两大阶段Prefill预填充和 Decode解码。Prefill 阶段你丢进去的 prompt 首次过模型此时没有任何缓存存在需要为一个 token 序列中的每个位置都计算一次 K/V然后写入 KV Cache。这个阶段算力密集matrix-matrixGPU 利用率挺高也是首次延迟的主要来源。Decode 阶段每生成一个新 token只需要为新 token 计算自己的 K/V追加到缓存末尾然后用新 Query 与当前全部 Cache 交互得到输出。这个阶段是存储密集、带宽密集的典型matrix-vector 为主GPU 利用率低因为大把带宽都耗在搬运历史 KV 上。我自己实测过一个直观数字7B 模型在 4090 上Prefill 处理 2000 token 的 prompt耗时尚可但是到 Decode 阶段如果序列长度已经跑道 8kCache 一满每秒生成速度立马掉一个档次。这两个阶段对显存的需求、时间占比完全不同理解它们是排障的第一步。2. 显存杀手KV Cache 的大小到底怎么算2.1 一张表算清显存占用从此心里有底很多人拿到一个模型光看权重文件大小就以为掌握了显存需求实际部署的时候经常被 KV Cache 打个措手不及。这里我直接给出一个通用的显存估算公式。对于采用 MHA 的 Transformer 模型当 batch size 1 时KV Cache 所占显存的公式是KV Cache 字节数 ≈ 2K 和 V 两套 × 层数 (N_layers) × 头数 (A) × 维度 (D) × 序列长度 (L) × 字节数 (bytes_per_elem)其中 D 隐藏维度 / 头数。如果是 FP16/BF16bytes_per_elem 2FP8 是 1INT8 也是 1。举例说明模型Llama 3 8B。隐藏层有 32 层32 个头隐藏维度 4096每个头维度 128。假设上下文长度 L 4096单层每个头的 K 形状 [4096, 128]FP16 存储需要 4096 × 128 × 2 1 MB单层 KV 两套 × 32 个头 64 MB32 层全部缓存 2048 MB也就是约 1.9 GiB。所以 8B 模型开满 4096 上下文KV Cache 大约吃掉 2GB 显存权重大约 16GBBF16。7B 系列差距不大。如果你用的是 72B。假设 80 层隐藏维度 819264 个头上下文 L 8192单层 K 缓存量 64 × 128 × 8192 × 2B 128 MB双层 KV 是 256 MB80 层直接 20 GB 往上走。这还只是单请求。多并发时KV Cache 会按 batch size 线性翻倍甚至更爆炸。我把常见配置的估算值整理成表供大家直接对照FP16/BF16 精度模型规模示例层数每层 KV 大小L4kKV Cache 总量L4k, bsz1权重显存约7B~8B如 Llama 3 8B、Qwen2-7B32约 32MB约 1~2GB约 14~16GB13B如 Llama 2 13B40约 40MB约 3~4GB约 26GB70B如 Llama 3 70B80约 128MB约 10GBL4k约 140GB上表只是单序列单 batch 的情况。一旦服务端做并发推理比如 batch size 872B 模型在 8k 上下文长度下光 KV Cache 就能吃掉 320GB 显存。这样你就明白为什么很多私有化部署要 8 张 A100/H100 了——权重是一块缓存才是大头。2.2 计算过程中的额外开销别忽略除了 KV Cache 本身解码阶段每一轮新 token 的前向过程还会产生临时激活张量activation。虽然新 token 的计算量本身不大但注意在并行实现中某些框架为了避免反复分配内存会一次性预留一个最大上下文长度的激活空间——这在长上下文场景会造成不小的浪费。选服务化推理框架时记得关注它的内存分配策略vLLM 的 PagedAttention 就专门缓解了这类浪费。2.3 并发数翻倍Cache 指数飙升在线服务场景中KV Cache 最难受的一点是它和 batch size 直接相乘。你并发来 8 个请求每条上下文 4000那 8B 模型的 KV Cache 占用量直接跑到 16GB 甚至更多。这也是为什么在线推理框架都在做连续批处理、动态批大小以及想办法对 KV 做量化——因为每少 1 bit同等显存可能就能多服务一批请求。3. 优化与部署实操把 KV Cache 压到可控范围3.1 量化 KV Cache牺牲一点精度换一倍显存最早一批 KV Cache 优化方案就是量化。把 FP16 的 K/V 矩阵量化为 INT8 甚至 INT4理论上显存占用直接对半砍甚至再砍半。量化方案一般分两种Per-token / Per-channel 量化对 K 和 V 的每个 token 或每个通道独立算 scale 和 zero-point。这种方式实现简单推理时反量化开销小是目前主流做法。Per-group 量化类似权重分组量化精度更好但计算复杂度略高。我在 vLLM 里常用--kv-cache-dtype fp8_e4m3或fp8_e5m2在 4090 上跑 8B 模型显存峰值明显下降输出质量在短文本上几乎无感。对于线上长上下文场景我建议先做离线评测集效果验证如果质量波动在可接受范围内就把 KV 量化开启能显著提升吞吐。这里特别提醒一句如果发现量化后模型出现了“答非所问”“重复内容变多”之类的退化优先检查是不是 KV Cache 量化精度设置太高步长太大或者某些头对精度极其敏感学界发现部分 attention head 对量化很脆弱可以尝试混合精度策略——给敏感层保留 FP16其余层用较低精度。3.2 GQA / MQA从源头砍缓存头数质量和显存两手抓的路线一定绕不开注意力头结构的改动。GQAGrouped Query Attention目前是绝大多数开源模型Llama 3、Qwen2、Mistral用的方案它的核心思路是查询Query仍然每个头独立投影保持模型表达能力但 Key 和 Value 只有少数几个头共享比如每 4 个 Query 头共享一组 K/V 头。这样 KV Cache 大小直接缩减为原来的 1/GG 是分组大小GQA 里类似于 compression ratio。比如 Llama 3 8B 用 8 个 KV 头对应 32 个 Q 头压缩比就是 4KV Cache 从近 2GB 缩到 500MB 左右。如果模型结构不支持 GQA但你想省显存还有个硬办法MQAMulti-Query Attention所有 Query 头共享同一组 K/V缓存压缩比等于总头数。但 MQA 会对长文本质量有明显影响现在一般模型少用纯 MQA 了除非是小模型极速场景。选型时优先选 GQA 结构的模型这是性价比最高的捷径。3.3 长上下文跑不起来试试 PagedAttention 和 Cache Offloading前面提到解码阶段 KV Cache 是离散增长的传统推理框架HuggingFace 原生 generate会提前预分配 max_seq_len 空间的显存缓冲区导致长上下文场景下显存空洞、浪费严重。vLLM 提出的 PagedAttention 借用了 OS 虚拟内存的分页思想把 KV Cache 切成固定大小的块block每个 block 能装一定数量 token 的 K/V块不要求物理地址连续通过块表block table把逻辑连续的 token 映射到物理分散的块上新 token 增长时只需要新分配一个 block 即可内存复用率大幅提升。这个机制实际效果很猛。用 vLLM 部署时并发请求的显存利用率比原生 HF 实现提升 20% 甚至更多。如果你用 vLLM 跑服务留意--max-num-seqs、--max-model-len和--gpu-memory-utilization这三个参数它们直接控制 KV Cache 预留策略。另一个思路是 Cache Offloading缓存卸载。当显存实在不够时把部分的 KV Cache 挪到 CPU 内存GPU 只保留近期活跃的 token 块。这样代价是增加跨设备传输延迟一般用于离线批处理或长文本一次性生成不适合对首 token 时延要求高的实时场景。如果你只是在自己单卡机器上跑长文档总结这招能救急但别拿来做高并发在线服务。3.4 Ollama 部署中的显存参数怎么调Ollama 可以说是本地跑模型碰 KV Cache 最多的地方。新手最容易踩的坑是模型加载进去了但一生成长文本就 BrokenPipe 或者刚开始很顺畅、越跑越慢最后内存爆掉。这种问题大多数情况下就是 KV Cache 占满了。Ollama 里面对应控制参数主要有num_ctx控制上下文窗口长度默认往往只有 2048当你把上下文拉到 8192 或更高时KV Cache 占用的显存是成倍增加的。num_gpu控制 GPU 层数如果显存不够模型部分会卸载到 CPU但这种混合模式会严重拖垮推理速度。num_gpu_layers不同版本叫法可能不一样但逻辑相同。我给一个实际调优例子单张 309024GB跑 Qwen2.5-7B-InstructBF16权重约 15GB默认 2048 上下文时 KV Cache 占用不到 600MB整体很稳。把num_ctx拉到 16384KV Cache 直接超过 7GB模型 缓存 激活总量逼近 24GB 红线如果再开并发聊天随时 OOM。所以此时我建议开启 KV 量化Ollama 较新版本支持类似环境变量或者把上下文降到 8192。另外说个细节OLLAMA_KV_CACHE_TYPE这类环境变量可以指定 cache 类型q8_0 之类的量化格式实测能在长上下文下救回不少显存。你如果发现模型生成变慢到不可用先看是不是上下文开太高别上来就怪显卡不行。4. 排查实录KV Cache 相关的故障就这么治4.1 现象一显存明明够推理速度却急剧下降这个坑我踩过不止一次。模型权重显存 14GB显卡 24GB理论上余量充足但上下文一长生成速度呈断崖式下跌。排查时我先用nvidia-smi观察显存变化再开ollama ps看缓存分配情况发现 KV Cache 虽没爆显存但预留空间过大导致 GPU 显存带宽在搬运缓存时被占满计算单元饿肚子。解决方案是收小num_ctx或切换到更高效的推理框架vLLM 的 block 管理比 HF 原生实现好得多再不行就上 KV 量化把带宽压力降下来。4.2 现象二多用户并发第二天服务突然 OOM我做线上服务时碰到的情况是头天压测正常第二天早上一看容器重启了。日志里 OOM 发生在 Decode 阶段。原因是并发请求的 KV 总量随 batchsize 动态变化而我当时给的是静态显存上限。后来把 KV Cache 的量化打开、并给 vLLM 设了合理的max-num-seqs和内存池上限解决了问题。核心经验KV Cache 是跟并发线性增长的预估显存时不能只按单条请求算要单条 KV × 最大并发数。另外服务端不要贪心把 GPU 可用显存全留给模型OOM 后重启的代价远比少几个并发位大。4.3 现象三量化 KV 后出现明显“胡说八道”这不代表 KV Cache 量化方案不行很可能是策略太粗糙。常见修正从 FP8 降到更稳妥的量化粒度如 per-head/per-channel 而非 per-tensor把关键层比如第一层和最后几层留成高精度减少量化对 attention softmax 影响有的框架支持把 softmax 计算留在 FP32 中这个也能缓解。我之前对 70B 模型做 KV Cache INT8 量化部分长文档摘要任务中事实性内容出现紊乱把输出 logits 层和部分中间层恢复 FP16 后基本回到正常水平。这类调优没有银弹一定要结合自己业务数据集验证。4.4 一份速查表问题和对应解法问题原因快速解法上下文一长就 OOMKV Cache 占用超显存量化 KV Cache / 开启 GQA / 降低上下文长度越生成越慢Cache 增长挤占带宽用 vLLM 分块 / 换更高带宽卡 / 分布式注意力并发场景 OOMKV 总量按 batch 线性膨胀限制最大并发数 / 调小 max-model-len / 开启 PagedAttention量化后质量下降精度策略粗糙关键层留 FP16 / 调细量化粒度 / 评估集回归测试首 token 延迟很高Prefill 阶段无缓存可复用合理控制输入 prompt 长度 / 升级计算能力5. 工程之外几个值得深入的知识点KV Cache 之所以能成为研究热点和工程重地是因为它直接牵扯到 Transformer 推理的三个核心维度内存、带宽、延迟。顺着这个方向再往深处走你会看到MLAMulti-head Latent AttentionDeepSeek-V2 提出的一种思路把 K/V 压缩成一个隐向量再在计算时恢复KV Cache 占用能压一个量级还多。这是目前极省缓存的结构创新之一。StreamingLLM 和 Attention Sink围绕 KV Cache 的丢弃策略意图用极小的缓存窗口保留大量可用的历史信息适合超长文档场景。Cache 逐出策略在不同请求共享 prefix 时如 RAG 场景的共享文档命中前缀缓存可以减少不少 Prefill 开销。如果你在做 RAG 类应用前缀广告命中带来的收益就很明显——公共知识部分的 K/V 如果被缓存复用了每个用户就省掉大段 Prefill 时间。这算是 KV Cache 在应用侧一个不算新但很实用的玩法。个人经验是做 KV Cache 相关调优时别只看显存数字要结合你的场景确定是“吞吐瓶颈”还是“延迟瓶颈”。离线批处理想方设法缓存量化能省多少省多少在线小并发应用优先级反而是降低缓存访问的延迟这时候花里胡哨的量化未必有直接换成 GQA 模型/更高带宽卡效果好。先把内存模型算清楚再做取舍才不会被动。如果你近期正在部署本地大模型建议第一步先跑一次上文那个显存公式把你目标模型在目标上下文和并发下的 KV Cache 量算出来再决定用哪种优化手段。这个习惯能帮你省下不少瞎调试的时间。
返回列表