论文
HeadWiseKV:混合长上下文语言模型中按注意力头分配预算的缓存驻留
HeadWiseKV: Budgeted Per-Head Cache Residency for Hybrid Long-Context Language Models
摘要
长上下文推理在解码过程中保留不断增长的键值 (KV) 缓存,这会消耗大量 GPU 内存并降低生成吞吐量。这个瓶颈仍然存在于混合语言模型中,因为它们的剩余全局注意力层可以主导上下文相关的缓存需求。我们研究如何在 KV 驻留总预算下分配此状态。我们引入了 HeadWiseKV,这是一个免训练框架,可以压缩混合语言模型的剩余全局 KV 缓存,同时保留其本地、循环和线性路径。它为每个物理 KV 头分配一个静态的多级历史窗口,使缓存需求在服务之前可预测。我们将这种分配表述为受限操作率失真问题,并提出 SeqCalib 作为 HeadWiseKV 中的核心策略生成算法。 SeqCalib 按执行顺序处理各层,并根据部署时使用的下层策略调整每个决策,从而考虑跨深度的交互。分组缓存运行时将所选策略具体化为实际的每头 KV 驻留而不是完整缓存上的掩码。我们评估了四种混合长上下文模型的下游质量,并研究了 Qwen3.6-27B 上的实际居住和服务行为。 HeadWiseKV 在评估的模型中保留了接近 Full-KV RULER 和 LoCoMo 的质量。在固定模型系统研究中,它在 112K 上下文长度下将采样峰值设备内存减少了 8.59%,并将最大的已验证成功上下文从 114K 扩展到 161K。