长序列推理不再卡顿!北大华为KV缓存管理框架实现4.7倍推理加速
LouisKV团队 投稿
量子位 | 公众号 QbitAI
北大华为联手推出KV cache管理新方式,推理速度比前SOTA提升4.7倍!
大模型处理长序列时,KV cache的内存占用随序列长度线性增长,已成为制约模型部署的严峻瓶颈。
为此,来自北京大学与华为的研究团队联合提出了
LouisKV
——一个专为长输入、长输出等各类长序列场景设计的高效KV cache 检索框架。
关键洞察
传统上,学术界与工业界提出了多种KV cache优化方案,其中 KV Cache Retrieval 是极具前景的方向之一。 该类方法将完整的KV cache卸载至容量更大的CPU内存中,并在推理时仅将最关键的KV子集检索回GPU进行计算,从而有效缓解GPU 显存压力。 然而,现有的KV retrieval方法仍面临着 效率 和 精度 的双重瓶颈:- 现有方法通常在生成每个 token 时都触发一次检索操作,这引入了重要性评估的计算开销与 CPU-GPU 间的数据传输开销。在需要生成数千甚至数万 token 的长输出任务中,检索操作带来的累积开销尤为突出,导致模型推理效率不高;
- 现有方法普遍采用固定大小的页(page)作为检索的基本单元。这种粗粒度的划分方式,常常导致被检索的页中仅包含少量真正关键的 KV 条目,而大量无关条目占用了宝贵的 CPU-GPU 带宽和 GPU 上的缓存预算。这不仅造成了数据传输的浪费,更重要的是,在有限的预算下,它挤占了本可以留给其他更关键信息的位置,导致模型推理精度的下降。
- 长输入序列中的稀疏分布:在长文档问答(图a)等任务中,生成答案所需的关键信息在长篇输入中呈稀疏、离散的分布状态;
- 长输出序列中的密集分布:在数学推理(图b)等任务中,模型的注意力会高度集中于先前生成的中间步骤,使得关键 KV 在局部区域内呈现密集分布。
核心设计
基于上述洞察,研究团队提出了一个高效的KV cache检索框架 LouisKV。该框架通过算法与系统的协同设计,解决了现有方法的瓶颈。 其核心包含三大创新。 首先是 语义感知的KV检索策略 (Semantic-Aware KV Retrieval),为利用时序局部性,LouisKV摒弃了“逐token检索”的低效模式,引入了一种自适应的检索策略。 如下图(a)所示,该策略通过轻量级机制监控语义变化。在每个解码步,它会计算当前token与前一token的query向量之间的余弦相似度r。- 若r高于阈值τ,表明模型关注点未发生显著偏移,此时不触发检索,直接复用上一个token检索得到的关键KV cache;
- 仅当r低于阈值τ,表明出现语义边界,才触发一次检索操作,从CPU的KV cache pool中加载新的关键KV cache。
- 输入序列(Prefill Stage):针对关键KV 稀疏分布 的特点,LouisKV采用K-Means聚类。如图(b)所示,它将语义上相似但物理位置上分散的KV聚合为语义簇(Semantic Clusters);
- 输出序列(Decode Stage):针对关键KV 局部密集 的特点,LouisKV将连续生成的token组织成时序(Temporal Segments)。这与模型生成连贯推理步骤的行为天然对齐。
实验结果
为了全面验证LouisKV的高效性,研究团队在多个主流的长序列任务上进行了详尽测试。 这些任务涵盖了 长输入-短输出 (如文档问答)、 短输入-长输出 (如数学推理)和 长输入-长输出 (如长文推理)等多种应用场景。 实验结果表明,LouisKV成功地在推理精度和推理效率之间取得了当前最佳的平衡。- 大幅降低延迟:与先进的KV检索方法Arkvale相比,LouisKV实现了高达1.4倍至4.7倍的端到端推理加速;
- 支持更大批量:当处理大批量任务时,FullCache会因显存不足而失效。相比之下,LouisKV能够在此类高负载场景下稳定运行,从而显著提升了系统的有效吞吐量。
我们正在招聘一名眼疾手快、关注AI的 学术编辑实习生 🎓 感兴趣的小伙伴欢迎关注 👉 了解详情