部署大模型前先算清 KV 缓存显存
模型权重能装入显存,并不代表服务能承受长提示词和多用户并发,KV 缓存可能导致显存耗尽或批处理能力下降。文章解释逐层保存历史 token 的键值张量如何减少重复计算,同时带来随序列长度增长的内存开销。
AI 深度解读
模型权重装得进显卡,不代表服务扛得住真实流量。长提示词、多用户并发、持续生成的回答,还会一起消耗 KV cache。原文的核心建议很直接:先算缓存显存,再决定上下文长度和并发上限。
这篇值得读的地方,是把「模型支持多长上下文」变成了一笔能算的容量账。对自部署模型的团队,它能帮助提前发现显存缺口;对调用托管 API 的程序员,它也解释了为什么提示词里静态内容和动态内容的顺序会影响缓存复用。
显存为什么会在请求进来后不够用
自回归生成每次预测一个 token。对于原文讨论的 decoder-only Transformer,因果遮罩保证了前面位置的状态不会被后续 token 改写,因此运行时可以保存每层已经计算好的 Key 和 Value,供后续生成复用。
这就是 KV cache:用显存换取少做重复计算。
但缓存会随上下文长度、活跃序列数量一起增长。这里的上下文包含输入提示词和已经生成的输出,不能只看请求刚进来时的长度。
单层 Key、Value 缓存各自的逻辑形状都是:
[batch, kv_heads, seq_len, head_dim]
把所有层、两份张量和元素字节数乘起来,就得到原文的基础公式:
bytes ≈ 2 × layers × kv_heads × head_dim × seq_len × bytes_per_element × batch
其中,layers 对应 num_hidden_layers,kv_heads 对应 num_key_value_heads。head_dim 通常由 hidden_size / num_attention_heads 得到,但仍要核对目标模型配置。
seq_len 是输入与输出 token 总数;batch 是同时持有上下文的活跃序列数量;FP16、BF16 的 bytes_per_element 为 2,8-bit 格式则为 1。
这笔计算给出的是缓存数据的理论下限,不是整个服务的显存预算。
实际部署还需要支付内存分配填充、块大小取整、页表、注意力内核临时空间等成本。模型权重之外,还得给激活峰值、CUDA 上下文和运行时缓冲留空间,原文没有给出通用的安全余量比例。
下面的脚本解决部署前快速估算的问题,输出单位是 GiB,即 1024^3 字节:
def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
return 2 * layers * kv_heads * head_dim * seq_len * bytes_per_elem * batch
def print_kv_gib(layers, kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
total_bytes = kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per_elem)
gib = total_bytes / (1024 ** 3)
print(f"KV cache floor: {gib:.2f} GiB ({total_bytes:,} bytes)")
if __name__ == "__main__":
# 示例 GQA 配置:8 条并发序列,每条 32,768 个 token
print_kv_gib(layers=32, kv_heads=8, head_dim=128, seq_len=32768, batch=8, bytes_per_elem=2)
重点看最后一行:它假设 8 条序列同时达到长度上限,算的是这个场景的缓存需求,不能把它当成日常平均占用。
同样的上下文,为什么能差出四倍显存
原文用了一组假想架构,方便算数,并不对应某个真实发布的模型:32 层、32 个 query heads、每个 head 的维度为 128,缓存使用每元素 2 字节的 FP16 或 BF16。
| 注意力配置 | 单 token 缓存 | 单序列 32,768 tokens | 8 条满长序列 |
|---|---|---|---|
| MHA,32 个 KV heads | 524,288 bytes,即 0.5 MiB | 16 GiB | 按公式为 128 GiB |
| GQA,8 个 KV heads | 131,072 bytes,即 128 KiB | 4 GiB | 32 GiB |
| GQA,8 个 KV heads,8-bit 缓存 | 65,536 bytes,即 64 KiB | 2 GiB | 16 GiB |
其中,第一行的 8 条序列以及第三行的单 token、单序列数字,是根据原文公式推算的。原文还给出一个参照:MHA 单序列达到 8,192 tokens 时,缓存就占 4 GiB。
差异来自 KV heads 的数量。多头注意力(MHA)让每个 query head 对应自己的 Key、Value;分组查询注意力(GQA)则让一组 query heads 共享一对 KV heads。上述 GQA 配置里,每 4 个 query heads 共享一组,缓存因此缩小到四分之一。
多查询注意力(MQA)更进一步,每层所有 query heads 共用一个 KV head。
这是选模型时就要看的架构差别。已经选定 checkpoint 后,不能把 num_key_value_heads 当成普通部署参数随手改小。
哪些设置能省显存,哪些只是减少浪费
原文列出的五个调节方向,需要分清各自解决什么问题。
限制最大序列长度,直接限制单请求的缓存增长。限制最大并发序列数,则控制同时占用缓存的请求数量,代价是更多请求排队。两者必须一起算:单请求能跑满上下文,不代表多个请求同时跑满也能撑住。
缓存量化减少每个元素的存储成本。8-bit 可以把上述 2 字节缓存减半,部分实验性或专用运行时还支持 4-bit。但量化会引入误差,长上下文精确检索和数学推理都可能受影响,原文没有提供准确率对照数据。
选择 GQA、MQA 模型是在架构层面缩小乘数。分页或分块分配解决的则是另一件事:避免为每条请求提前保留一整段最大长度的连续显存。
vLLM 的 PagedAttention 按需分配固定大小的缓存块,减少空槽和碎片,并为相同提示词前缀共享物理缓存块提供机制。原文引用 Kwon 等人的论文结果:在论文所比较的系统和相近延迟下,吞吐提升为 2 到 4 倍。这是特定实验结果,不能直接写进自己服务的容量承诺。
分页分配能减少浪费,但不会让每个独有 token 的 KV 数据凭空消失。
托管 API 的提示词缓存也受相同的因果依赖约束。位置 t 的 KV 状态依赖从开头到 t 的完整 token 序列,开头插入时间戳、请求 UUID 或变化的用户资料,后面的内容即使相同,也无法沿用原来的完整前缀缓存。
因此,静态系统指令、工具 schema、固定参考文档适合放在前面,每次变化的输入放在后面。匹配要求是 token 序列一致,语义相近并不够。
我的判断
我认同先算容量、再放开上下文的顺序。这篇最可靠的部分,是张量形状推导和明确标成假想配置的算例,读者可以自己复算。
不过,文中的内容不能统称为作者实测。缓存占用主要来自公式计算,2 到 4 倍吞吐来自被引用论文;作者没有展示自己的部署压测,也没有比较量化前后的任务质量。
原文对计算收益的表述也略简化了。复用历史 KV 能避免反复计算历史状态,但每步生成仍需处理对历史上下文的注意力,不能把「线性查找」理解成上下文变长后生成成本不再增长。
相比单纯降低并发,GQA 和缓存量化更可能保留服务容量,但前者受模型选型约束,后者需要质量验证。相比直接缩短上下文,分页分配能保留长请求的空间,却不能解决所有请求同时接近上限时的真实数据占用。
托管 API 部分则应停留在机制解释。原文没有说明具体供应商的命中门槛、缓存保留时间和计费规则,不能只凭调整提示词顺序就推断收益。
能拿来做什么
- 自部署代码助手前,先从
config.json读取层数、query heads、KV heads 和隐藏维度,用计划中的上下文上限与最大并发计算缓存下限,再扣除权重和运行时开销,确定可用预算。 - 为 Node 服务接入自托管推理接口时,显式配置推理引擎的最大模型长度和最大并发序列数,并用接近上限的输入与输出做并发压测,观察排队延迟和显存峰值。
- 开启 8-bit 或 4-bit KV cache 前,准备业务中的长上下文检索、代码定位等代表性任务,对比输出质量;能避免内存溢出,只完成了部署验证的一部分。
- 检查日常 AI 编程的提示词模板,把稳定的系统指令、工具 schema 和固定资料集中到前缀,把时间戳、会话资料、用户问题等动态内容后置,再依据服务提供方暴露的缓存信息验证效果。