大模型并发容量为何受 KV 缓存限制

dev.to · AI#推理优化#显存#KV缓存

Editor brief · 编辑速览

文章以 Llama 2 7B 结构为例，计算上下文与并发请求如何消耗 KV 缓存显存。它解释了缓存复用为何成立，以及缓存能减少历史投影重算、却不能消除注意力读取开销。

文章思维导图

Knowledge map

拖拽缩放

AI deep read

AI 深度解读

一张 80 GB 显卡装下权重约 13.5 GB 的 7B 模型，为什么同时服务几十个长对话就可能吃紧？作者把账算到了 KV cache：权重只需存一份，每条生成序列却要单独保存历史 token 的缓存，上下文越长、并发越多，显存占用就越大。

这篇的价值是把模型大小、上下文长度和服务吞吐连到同一条公式上。不过，文中的约 30 路并发属于特定配置下的容量估算，不能直接当成压测结果。

自回归模型逐个生成 token。每一步，新 token 的 Query 都要读取此前 token 的 Key 和 Value。如果没有缓存，每次都要重新计算整个前缀的 K、V。

因果掩码（causal mask）保证了缓存可以复用：前面的 token 看不到后面的内容，因此追加 token 不会改变已经算出的 K、V。每层保存历史结果，再逐步追加即可。

作者举的例子是：输入 5 个 token，生成 64 个 token。不缓存时，每层需要做 2,336 次 K/V 投影；缓存后只需 68 次，减少约 34 倍。但每一步注意力仍要读取历史缓存，注意力计算的成本没有被一起消除。

显存账按一个 Llama-2-7B 形状的模型计算：32 层、32 个注意力头、每头维度 128，使用 fp16。下面计算的是每个 token 需要多少缓存：

2 (K and V) × layers × kv_heads × head_dim × bytes
= 2 × 32 × 32 × 128 × 2
= 524,288 bytes ≈ 0.5 MiB

要盯住第一行的 kv_heads 和 bytes：缓存大小既取决于模型结构，也取决于缓存精度。

每个 token 约 0.5 MiB，4,096 个 token 就约占 2.1 GB；一条 32K-token 对话的缓存，已经超过文中约 13.5 GB 的模型权重。

权重是服务共享的固定开销，KV cache 是随序列数量和长度增长的开销。

于是，80 GB 显卡剩下约 66 GB，按文中更具体的算式 66.5 GB / 2.1 GB ≈ 30，只能容纳约 30 条达到该长度的序列。这还没扣除激活值和运行时缓冲区。

生成分为两个阶段。

Prefill 一次处理整个输入，读入的权重可以用于多个输入 token。Decode 每一步生成一个 token，低批量时，同样一次权重读取能完成的工作少得多，容易受显存带宽限制。

批处理（batching）让一次读取的权重服务多条序列，提高利用率。但每增加一条序列，就要增加它自己的 KV cache。缓存装不下，batch 就加不上去，吞吐也随之受限。

作者给出的优化，分别减少公式里的不同项：

减少 KV 头数。 Multi-query attention，出自 Shazeer（2019），让所有 Query 头共享一组 K/V；Grouped-query attention，出自 Ainslie 等（2023），则分组共享。同样的 7B 配置，把 KV 头从 32 减到 8，估算容量可从约 30 条升至 123 条。Llama 2 70B 使用 8 个 KV 头，7B 使用 32 个。
减少每个数值的字节数。在 8 个 KV 头的基础上，把缓存存成 int8，容量再减半，估算可容纳约 247 条序列。作者称通常只有较小的精度损失，但没有提供具体评测。
减少保留的 token。滑动窗口只保留最近的上下文，原文以 Mistral 7B 的 4,096-token 窗口为例。
减少空置空间。vLLM 使用的 PagedAttention，出自 Kwon 等（2023），把缓存放进固定大小的块，避免每条序列都按最大长度预留空间。

前两类结构调整中的 KV 头数，以及滑动窗口，属于训练前或训练期间的架构选择；缓存量化和分页则可以在服务阶段应用。它们不是四个随时可以互换的开关。

这套解释适合用来纠正一种容量误判：显卡能装下模型，不代表还能承载足够多的长对话。把缓存按「每 token、每序列」拆开计算，比只看参数量更接近服务端实际约束。

但作者提供的是算术推导，没有给出硬件压测、请求长度分布或吞吐对照。30、123、247 都应视为估算值，不能当成生产承诺。原文也交替使用 GB 和 MiB，精确规划时需要统一单位。

Decode 受带宽限制的判断适合解释低批量生成，实际瓶颈仍会受批量和硬件影响。滑动窗口会缩小可直接关注的历史范围；量化需要验证质量；分页主要减少分配浪费，不会降低每个有效 token 的理论缓存量。

原文最后介绍了作者制作的 Lensa ML 交互讲解器，可以调整上下文、batch、KV 头数、精度和窗口观察变化。它是理解公式的工具，不能替代服务压测。