大模型推理一旦进入解码阶段,算力往往不是瓶颈,显存带宽才是。每个新 token 生成时,模型都要把之前所有 token 的 Key 和 Value 重新读一遍,这份数据就是 KV Cache。它不参与权重更新,却随着上下文长度和并发数线性膨胀,最终把 HBM 带宽吃满,首 token 延迟和单卡并发都卡在这里。

先把账算清楚。单个序列的 KV Cache 大小约等于 2 乘以层数、注意力头数、头维度和单元素字节数,再乘以序列长度。这里的 2 代表 Key 一份、Value 一份。以一个 32 层、每层 32 头、头维 128、FP16 存储的模型为例,单个 token 的 KV 约为 0.5MB;上下文 4096 个 token 时,单个序列就要约 2GB;并发 32 路就是约 64GB。这还没算权重和激活占用,显存先被 KV 占满是常态。

存多少,首先看注意力结构能不能瘦身。分组查询注意力把多个 Query 头共享同一组 Key、Value 头,直接按分组数成比例压缩 KV;多头潜在注意力更进一步,把 KV 压缩成低维隐向量,推理时再还原。两者都是在模型结构层面减少存量,效果最彻底,但需要模型本身支持,事后改结构要重新训练或微调,不能只靠推理框架硬改。

放哪里,核心是分层。热数据留在 HBM,温数据可以放到主机内存甚至 NVMe,通过异步预取在需要前搬回。分层的前提是算好换入换出开销:如果一个序列很久才轮到一次,搬运成本摊薄后可以接受;如果轮询频繁,反复搬运反而比直接放 HBM 更慢。PagedAttention 的思路值得借鉴:把 KV 按块管理,消除按最大长度预分配带来的碎片,让显存利用率从按峰值预留变成按实际占用分配。

怎么复用,是最容易被低估的一环。很多线上请求的系统提示词和开场白高度重合,前缀缓存可以直接复用同一段 KV,避免重复计算。实现时要注意两点:一是复用粒度要按 token 块对齐,不能按整条请求一刀切,否则命中率上不去;二是带状态的对话不能把上一轮的 KV 无条件复用到下一轮,用户切换话题或修改系统提示后必须失效重算。

量化也是常用手段。把 KV 从 FP16 压到 FP8 或 INT8,容量和带宽占用直接减半,但要在代表性长上下文上实测精度回落,尤其注意首 token 附近和长尾位置的敏感性。不同层对量化的容忍度不同,一刀切全量 INT8 容易在个别任务上翻车,分层、分头做混合精度更稳妥。

落地时建议按这个顺序排查:先算单序列和总并发的 KV 上限,确认是容量不够还是带宽不够;再看框架是否已开启分块管理和前缀缓存;最后才考虑跨层 offload 和量化。很多时候打开已有的前缀缓存和分块分配,就能把并发翻倍,而不需要先动模型结构。带宽瓶颈的优化没有银弹,本质是把每一份 KV 的存储成本、搬运成本和复用收益都算到同一个账本里。