大模型推理的KV Cache优化全解
大模型推理的KV Cache优化全解 KV Cache是大模型推理中最重要的优化技术,也是最大的内存瓶颈。理解KV Cache的工作原理和优化方法,是构建高效推理服务的基础。 KV Cache是什么 在Transformer的自回归生成中,每生成一个新token,需要计算它对所有之前token的注意力。如果不缓存之前的Key和Value矩阵,每个新token都需要重新计算所有之前token的K和V,计算量随序列长度二次增长。 KV Cache通过缓存之前计算过的K和V矩阵,将每步的计算复杂度从O(n²)降到O(n)。代价是内存占用线性增长——对于Llama-70B模型,生成4K token的KV Cache约占40GB显存。 内存瓶颈 KV Cache的内存占用可以用以下公式计算: KV Cache Size = 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size 以Llama-70B(80层, 64头, 128维, FP16)为例,单序列4096 token的KV Cache约40GB。这意味着一个80GB显存的A100只能服务两个并发请求——这是制约推理吞吐量的最大瓶颈。 PagedAttention:分页管理 vLLM团队提出的PagedAttention是KV Cache管理的革命性创新。灵感来自操作系统的虚拟内存分页机制。 传统分配的问题 传统方法为每个序列预分配一块连续的KV Cache空间,按最大序列长度分配。这导致严重的内存碎片——大多数序列不会用满预分配的空间,但多余的空间不能被其他序列使用。内存利用率通常只有20-40%。 分页方案 PagedAttention将KV Cache划分为固定大小的"页"(通常16个token),每个序列通过页表映射到物理页。页按需分配——序列增长时才分配新页。 效果是显著的:vLLM的内存利用率提升到90%以上,并发吞吐量提升2-4倍。碎片问题被彻底解决,因为不同序列的页可以散布在物理内存中的任意位置。 页的大小选择 页太小(如1 token)会增加页表开销,页太大(如256 token)则回到预分配的问题。16 token是在大多数场景下的最优选择——页表开销不到1%,内存浪费也不显著。 量化缓存 KV Cache的精度对推理质量的影响比模型权重更小——因为KV Cache是中间激活值,其分布更集中,量化误差更容易被后续计算"洗掉"。 FP8 KV Cache 将KV Cache从FP16量化到FP8,内存减半,几乎无损。现代GPU(H100及以后)原生支持FP8运算,所以推理速度也几乎不受影响。这可能是最简单且性价比最高的KV Cache优化。 INT4 KV Cache 更激进的方案是将KV Cache量化到INT4。内存减少到1/4,但精度损失开始显著——在长序列和需要精确注意力的任务上,INT4 KV Cache可能导致输出质量下降。 ...