大模型推理百倍加速之KV cache篇

原文：安全验证 - 知乎

目录

收起

为什么要研究KV cache？

KV cache的作用

基于KV cache的加速策略

Window--窗口

Sparse--稀疏化

Quantization--量化

Allocator--显存分配

Share--KV cache共享

总结

参考资料

为什么要研究KV cache？

设输入序列的长度为 s ，输出序列的长度为 n ，模型深度为l，维度为h,以 FP16 来保存KV cache，那么KV cache的峰值显存占用大小为 b(s+n)h∗l∗2∗2=4blh(s+n) 。这里第一个2表示K/V cache，第二个2表示 FP16 占2个bytes。
以 GPT3 (175B) 为例，对比 KV cache 与模型参数占用显存的大小。GPT3 模型weight占用显存大小为350GB (FP16)，层数 l为96，维度h为12888。

batch size	s+n	KV cache(GB)	KV cache/weight
4	4096	75.5	0.22
16	4096	302	0.86
64	4096	1208	3.45

参考上图，随着 batch size 和 长度的增大，KV cache 占用的显存开销快速增大，甚至会超过模型本身。从LLM的趋势上而讲，主要有三个方面来说明kv cache优化的必要性：
1、总体趋势上LLM 的窗口长度在不断增大，因此就出现一组主要矛盾，即：对不断增长的 LLM 的窗口长度的需要与有限的 GPU 显存之间的矛盾。因此优化 KV cache 非常必要。
OpenAI API场景，API最烧钱的是输入而非输出，输入包括prefill prompt 和conversation，长度动辄数十K token。虽说每输入token比每输出token便宜，但能够降低kv重新计算的开销，无论是硬件资源门槛，还是模型推理降本，都有着极为积极的作用。
2、对于消费级显卡这种性价比较高的显卡而言，显存容量相对较小，KV cache从一定程度上降低了模型的batch size，因而KV cache优化在工程落地中更显重要。

框架	模型	input/output	机器配置	tokens/s	最佳batch_size	gpu/cpu负载
TGI	llama2 70B	128/512	4090*8，32c	291	13	cpu25%，gpu95%
TGI	llama2 70B	128/512	A800，32c	1122	43	cpu25%，gpu95%

从上表能够看出，类似4090等消费级显卡，其主要的瓶颈时batch_size，而影响batch_size的主要因素就是显存容量，而KV cache在推理过程中占用大量的显存。如果能通过KV cache降低显存占用，从一定程度上就能提升消费级显卡的性价比，带来非常高的商业收益。
3、sora/sd3等文生视频或者文生图的模型，纷纷放弃u-net架构，转而支持DIF（diffusion transformer）架构。对此类AIGC模型而言， KV cache同样能起到类似LLM上的加速效果。
根据资料，Sora类训练任务的特点是模型本体不大（10B以下），但是由于视频复杂性带来的序列长度特别长（接近1000kpatches的长度），可以对模型推理进行简易测算：

• 按照batch size = 1 进行测算，kv cache和模型权重对显卡占比能达到10:1（例如4090的24G显存，2G分给模型，22G分给kv cache）左右，这个场景的显存分配占比与LLM差异性还是非常的大。
• 按照batch size = 4 进行测算，kv cache和模型权重对显卡占比能达到40:1（batch size越大，kv cache的显存越大）

由此可见，KV cache会成为视频生成领域的一个重要瓶颈，但不排除有替代kv cache的加速方案。

KV cache的作用

解释kv cache之前，先看一组对话：
paki：What is the apples？
llama：Apples are a boring fruit.
假设上述对话中，paki是自然人并输入prompt（What is the apples？），llama是模型并返回result（Apples are a boring fruit）。如果对上述对话的模型推理流程进行分析，实际上需要将推理步骤分成两个阶段，即prefill和decode。
prefill阶段：输入为Q，即‘What is the apples？’，返回了第一个token，即‘Apples’，同时初始化了kv cache。
decode阶段：输入为单个词或者说q，通过自回归的方式，生成‘Apples are a boring fruit’这个句子。需要注意的是，decode计算的过程中，q的长度为1，即当前词，返回下一个词，例如通过‘Apples’生成‘are’，同时更新kv cache。