大模型推理时显存暴涨的根本原因在于KV Cache的持续累积,只要显存被KV Cache占满,首要优化手段是KV Cache量化、前缀复用与动态批处理的组合;如果单卡依然放不下,就需要把不再经过Transformer计算的KV Cache层卸载到CPU共享内存。低效的缓存策略会让相同参数量的模型服务同时只能处理一两个请求,而合理的调度策略可以将并发吞吐提升数倍。理解这组分配关系,才能够在预算固定的前提下,把大模型推理服务的承载量发挥出来。
KV Cache占了自回归大模型推理阶段的大部分显存消耗,模型参数与KV Cache的占比会随着batch size和序列长度迅速倾斜。这是因为每一步解码都要依赖之前所有tokens的Key和Value,这些缓存按层数、隐藏层尺寸、序列长度线性放大,序列越长,KV Cache增长得越陡。
具体来说,单个请求在某一层的KV Cache内存占用可以按如下公式预估:2 × hidden_size × seq_len × num_layers × batch_size。当模型层数为50,hidden_size为8000,序列长度32万时,单一请求的KV Cache就已经远远超过模型本身的权重体积,这就是为什么长上下文任务经常直接把显存打爆的原因。
将KV Cache从FP16压缩到INT8或FP8可以让缓存显存占用直接减半,而模型最终输出质量几乎察觉不到下降。当前主流推理框架比如vLLM已经原生支持FP8与INT8两种KV Cache量化模式,不必额外改写模型权重,只对缓存环节生效。
Hugging Face的KV Cache量化测评显示,使用8bit KV Cache后,Huggingface 的KV Cache量化使显存占用下降超过50%,输出困惑度增加通常不到0.5%。KV Cache量化对缓存而言,也意味着更大的batch容量、更长的sequence受け入れ。联合用途是启用vLLM的--kv-cache-dtype fp8标签就能在项目里直接生效。
# vLLM启动服务时启用FP8 KVCache支持 model=Qwen2.5-7B-Instruct tokenizer-mode auto \ --max-model-len 32768 \ --kv-cache-dtype fp8 \ --gpu-memory-utilization 0.9
不过量化之后不要盲目追求极限,模型层和训练数据都会被丢到INT8上,效果呢?相关的FP8方案要同一批次存在。量化误差的主要来源是值域分布,KV Cache中Key和Value的分布差异明显,因此工程上往往会为KV Cache独立做per-channel量化,而不是直接把权重矩阵的量化方式复用起来。
对于多轮对话和长文档回答,如果系统提示词或上下文前缀一直没有变,就把这些前缀对应的KV Cache缓存起来,避免整个提示词重新参与计算。Prefix Cache复用的效果并不需要算法模型去判断,它直接改写词解码流程中过去没有的片段;社区常用做法是把prompt里较长的system message和上下文片段加一个整段的前缀表,转发时让推理引擎自动命中缓存。
vLLM中开启--enable-prefix-caching即可自动保存前缀并复用,同一个前缀的预热肯定需要第一次消耗一次KV Cache才完成计算,但后续每次请求都能直接节省掉该前缀从头计算的成本。实测单个固定的长系统提示词,启用Prefix Cache后首token时间不再随系统提示词长度线性增长,等于把没意义的大上下文彻底从每轮中间移除。
在KV Cache总量受限的前提下,动态批处理和连续的TPS吞吐才更决定实际感受层面的收益,动态Batch不加限制的时候,多个请求会在注意力计算中重叠,但如果一个输出子段长一个输出子段短,在旧死死的批策略下会严重浪费缓存。
目前推理框架普遍采用continuous batching,让每个序列少了任意地方自己结束。Continuous batching不再是按整批等待最长的子输出结束,而是提前删除已经结束序列占有的KV Cache空间,并立刻把正在等待的新请求拉入当前计算批次。这是让显存闲置指标从30%降到5%以下的最强调度手段。
| 优化手法 | 显存影响程度 | 速度影响 | 落地成本 |
|---|---|---|---|
| KV量化 INT8 | 减少约50% | 略微提升 | 低,框架直接支持 |
| Prefix Cache | 减小重复前缀开销 | 首token显著加快 | 低,需识别前缀 |
| Continuous Batching | 降低闲置缓存 | 吞吐高峰更好 | 中,需评测排查 |
| KV Offload到CPU | 降低单卡峰值占用 | 增加延迟 | 中,适合长序列 |
如果量化、复用和批处理做完之后KV Cache还是把单卡占满,就应该把处理器缓存转移到CPU内存中,而不是直接扩卡。KV Cache Offload把最多绝大多数缓存放进CPU内存,只有当前计算需要的近窗口留在GPU上,每层计算前再异步搬回对应部分。
这种做法在长上下文文本生成时效果很突出,因为HTTP推理场景里,输出不同位置都不再需要历史全部KV Cache同时在显存中的高带宽承载,且推理框架会按tile粒度搞迁移,因此不会每个token都把整条历史cache拖一遍。但有一点:offload之后单个token延迟会提高,这个方案适合喂给长文本问答、离线总结,并不适合延迟极敏感的交互。
接下来另一条规模路线是把Tensor Parallel跨卡展开,把例如Q、K、V路径分散到多张GPU,实际上KV Cache也必须跟着权重数量被分割到各卡上。横过多卡之后不要盲目把压缩舍弃,低精度KV Cache和跨卡分割可以同时共存,DBaba推理的总显存预算反而容易从3节增长到半卡状态。
因此在解决限的大模型推理负载上,首先利用KV Cache量化和Prefix缓存把单卡能容纳的token长度恢复,再配合Continuous Batching提升卡上的吞吐值,最后在超长上下文上调用KV Cache Offload和Tensor Parallelism。只要把需求形态和业务性能对照优先级排好,没有必要一上来就铺多卡。
评论
加载中...