Jellow 编辑发布,AI 辅助生成并完成技术核对。 计算示例和实测结果会在文中分别说明。
模型权重明明能放进显卡,为什么并发稍微提高,或者对话变长,显存就不够了?
原因之一是生成过程还需要保存 KV Cache。它随着活跃请求的上下文增长,规模并不由模型参数量单独决定。理解这部分显存,可以帮助我们判断该缩短上下文、减少并发,还是调整缓存管理方式。
本文讨论每层都有标准全注意力、各层 KV 头数和维度相同的模型。滑动窗口、混合注意力、MLA 等结构需要按各自的缓存布局另算。所有容量数字均为公式推导。
1. 缓存的是历史 key 和 value
在因果注意力中,历史位置的表示不会因为后面多了一个 token 就重新变化。因此,推理时可以保存每层已计算的 key 和 value,后续步骤追加新位置的缓存。
历史 query 通常不需要保存给后续步骤使用,因为当前输出依赖的是当前 query 与历史 K、V 的交互。每层保存自己的缓存,不能只计算一层的占用。缓存机制说明
2. 从张量形状推导显存公式
定义几个量:
| 符号 | 含义 |
|---|---|
| L | 需要缓存的注意力层数 |
| Hkv | 每层 key/value 头数 |
| D | 每个头的维度,假设 K 与 V 相同 |
| S | 所有活跃序列当前缓存的 token 数之和 |
| b | 每个缓存元素的字节数 |
一个 token 在一层里,需要保存 Hkv × D 个 key 元素和同样多的 value 元素。乘上层数和 token 总数,就得到:
KV Cache 字节数 = 2 × L × Hkv × D × S × b
前面的 2 表示 K 和 V。如果 batch 中 B 条序列长度都为 T,则 S = B × T;长度不同时,使用各序列实际长度求和。
这里计算的是逻辑缓存数据量,没有包含块表、对齐、预分配空间、量化参数、工作区以及其他运行时开销。
举一个假设模型:32 层、8 个 KV 头、头维度 128,缓存采用 BF16,每个元素 2 字节。
每 token = 2 × 32 × 8 × 128 × 2
= 131,072 字节
= 128 KiB
8,192 token 的单条序列 = 1 GiB
16 条这样的活跃序列 = 16 GiB
这里的上下文长度包括已缓存的输入和生成内容。请求还在生成时,缓存也可能继续增加。GB 和 GiB 的换算不同,做容量规划时要统一单位。
下面的 Python 代码可以直接计算:
def kv_cache_gib(layers, kv_heads, head_dim, lengths,
bytes_per_element=2):
total_tokens = sum(lengths)
total_bytes = (
2 * layers * kv_heads * head_dim
* total_tokens * bytes_per_element
)
return total_bytes / 1024**3
print(kv_cache_gib(32, 8, 128, [8192])) # 1.0
print(kv_cache_gib(32, 8, 128, [8192] * 16)) # 16.0
3. 为什么要看 KV 头数,而不是直接用 attention 头数
MHA 中,每个 query 头通常有对应的 K、V 头。MQA 让多个 query 头共享一组 K、V;GQA 则让若干 query 头共享一组 K、V,处于二者之间。GQA 论文
如果模型有 32 个 query 头、8 个 KV 头,容量公式里应该代入 8。如果误用 32,估算会大四倍。
同样,权重采用 4 bit 量化,并不意味着 KV Cache 也自动变成 4 bit。权重存储格式与缓存格式是两件事,需要分别检查。比较容量之前,应读取模型配置和实际推理设置。
4. PagedAttention 处理的是分配问题
假设为每个请求一次性预留最大上下文长度的连续缓存。短请求会浪费空间,长度变化还可能让分配与复用变得困难。
PagedAttention 把逻辑 KV 序列分成固定 token 数的块,通过映射关系访问物理缓存块。不同逻辑块可以放在不同物理位置,请求增长时再按块分配。它借用了操作系统分页的思路,但不等于每次访问都经过 CPU 虚拟内存换页。PagedAttention 论文
设每块容纳 16 个 token,一条序列当前有 35 个 token,则需要 3 个块,总容量 48 个 token,最后一块剩余 13 个位置。
进一步假设三条序列分别有 35、17、48 个 token:
| 分配方式 | 为三条序列分配的 token 容量 |
|---|---|
| 每条固定预留 64 个位置 | 192 |
| 按 16-token 块分配 | 48 + 32 + 48 = 128 |
| 实际使用 | 100 |
这是为解释分配方式构造的例子,不能用来概括所有分配器的浪费比例。分页仍有尾块浪费和管理开销。
实际 attention kernel 需要通过块映射找到相应 K、V,不能把逻辑上不连续的缓存直接当成普通连续数组处理。vLLM Paged Attention 实现说明
5. 分页和前缀缓存不是同一个功能
分页回答“缓存放在哪里”,前缀缓存回答“哪些计算已经做过”。
如果多次请求拥有完全相同、可复用的 token 前缀,系统可以复用这部分前缀的 KV Cache,减少重复 Prefill。语义相近不等于缓存相同,模型配置、token 前缀及其他影响状态的条件必须匹配。vLLM 的 APC 文档也强调,这项优化主要减少输入处理时间,不会直接消除后续生成 token 的计算。自动前缀缓存
比如同一份长文档被连续提问,文档部分适合放在稳定前缀中。把每次都变化的时间戳放在最前面,可能提前打断可复用前缀。是否值得调整 prompt 结构,应结合真实缓存命中情况判断。
6. 如何判断这些优化对自己的服务有没有帮助
可以先做三类对照:相同前缀反复提问、完全不同的输入、混合长短请求。分别记录 Prefill 耗时、输出速度、最大活跃序列数和显存占用。
如果相同前缀的 TTFT 明显改善,而 Decode 变化不大,这是符合预期的结果。如果分页后能容纳更多并发,但单请求生成速度几乎不变,也不能据此判断优化无效:容量利用率和单请求延迟是不同目标。
下一次看到显存不足,可以先算出每条序列的 KV 成本,再查看活跃 token 总量和缓存分配情况。进一步做容量规划时,还要把权重、工作区、安全余量和分页尾块放进同一个预算。
7. 从显存预算反推服务容量
只计算单条序列还不够。部署时更实用的问题是:扣除模型权重和运行时开销后,还能放多少 token?
可以先建立一个保守预算:
可用于 KV 的字节数
= 显存总量 - 权重 - 固定运行时开销 - 安全余量
逻辑 token 容量
= floor(可用于 KV 的字节数 / 每 token KV 字节数)
假设显存总量为 80 GiB,权重占 14 GiB,算子工作区、CUDA context 等固定预留 6 GiB,再留出 4 GiB 安全余量。上一节的假设模型每 token 使用 128 KiB,那么:
可用于 KV = 80 - 14 - 6 - 4 = 56 GiB
token 容量 = 56 × 1024² KiB / 128 KiB
= 458,752 token
如果活跃请求的平均已缓存长度是 4,000 token,简单相除得到约 114 条序列。这个数只是逻辑上限,不能直接写进并发配置:长度分布有长尾,分页有尾块,调度器可能预留尚未生成的空间,运行时开销也会随 batch 改变。
更稳妥的做法是使用长度分布,而不是只用均值。例如分别计算 P50、P95 的上下文长度,再用请求回放观察容量耗尽时是哪些请求占用了缓存。均值相同的两个负载,若其中一个包含少量超长对话,拒绝率和尾延迟可能完全不同。
8. 分页能把浪费压到什么范围
设块大小为 B 个 token,当前有 A 条活跃序列。每条序列只有最后一块可能没有填满,因此仅由尾块造成的未使用位置满足:
尾块浪费 token < A × B
更严格地说,每条非空序列最多浪费 B-1 个位置,所以总量不超过 A × (B-1)。乘上每 token 的 KV 字节数,就能得到这部分显存的上界。
这不是整个分配器开销的上界。块表、元数据、预留策略和暂时无法复用的块还会增加消耗。块越小,尾部浪费通常越少,映射表和调度管理却会更细;块越大则相反。因此应该测真实长度分布下的已用块、空闲块和尾块,而不是只比较配置里的 block size。
分页也没有改变逻辑 KV 数据量。上面的 458,752 token 容量不会因为使用 PagedAttention 就凭空翻倍;它能做的是让物理容量更接近逻辑使用量,避免为每条请求按最大长度提前留出连续大块。PagedAttention 论文把这种分配与共享机制作为提升服务容量的核心。论文中的内存管理设计
9. 用症状决定下一步,而不是看到 OOM 就开分页
| 观测到的现象 | 更可能的原因 | 优先检查 |
|---|---|---|
| 服务启动就几乎占满显存 | 权重或固定工作区过大 | 权重精度、并行切分、加载时峰值 |
| 请求数不多,但长对话后 OOM | 逻辑 KV 数据过大 | 每 token 成本、总活跃 token、上下文上限 |
| 平均长度很短,仍有大量保留空间 | 预留策略或碎片 | 物理块利用率、尾块和调度器配置 |
| 相同长前缀反复输入,TTFT 仍高 | 没有复用 Prefill | 前缀缓存命中条件与 prompt 布局 |
| 降低 KV 精度后容量提高但质量异常 | 缓存量化误差 | 校准方法、敏感任务和长上下文回归 |
容量方案至少要同时回答三个问题:逻辑上每 token 需要多少字节,物理分配额外浪费多少,以及业务长度分布会同时保留多少 token。把这三层分开,才能知道分页、缓存量化、缩短上下文或降低并发中的哪一项真正对症。