HeadWiseKV：针对混合长上下文语言模型的预算化每头缓存驻留

# HeadWiseKV：面向混合长上下文语言模型的按头预算缓存驻留策略
来源：https://arxiv.org/html/2609.02029  
杨骏程¹ 胡傲庭² 张明曦³ 吴立瑶² 洪哲衡⁴ 许威⁵  

###### 摘要  
长上下文推理在解码过程中会保留不断增长的键值（KV）缓存，这会消耗大量GPU内存并可能降低生成吞吐量。该瓶颈在混合语言模型中依然存在，因为其残余全局注意力层可能主导上下文相关的缓存需求。本文研究如何在总KV驻留预算下分配这一状态。我们提出 HeadWiseKV，这是一个无需训练的框架，旨在压缩混合语言模型中的残余全局KV缓存，同时保留其原生的局部、循环和线性路径。它为每个物理KV头分配一个静态的多级历史窗口，使缓存需求在服务前即可预测。我们将此分配问题形式化为一个受限的运行率失真问题，并提出 SeqCalib 作为 HeadWiseKV 中的核心策略生成算法。SeqCalib 按执行顺序处理各层，并将每个决策建立在部署时实际使用的低层策略基础之上，从而考虑了深度方向的交互作用。分组缓存运行时将选定策略实例化为实际的每头KV驻留，而非在完整缓存上的掩码。我们在四个混合长上下文模型上评估了下游质量，并在 Qwen3.6-27B 上研究了物理驻留和行为。HeadWiseKV 在所评估的模型中保留了接近完整KV的 RULER 和 LoCoMo 质量。在固定模型的系统研究中，它在112K上下文长度下将采样的峰值设备内存降低了8.59%，并将最大验证成功上下文从114K扩展到161K。  
¹南京邮电大学 ²Tylogi AI 实验室 / TAIL ³安徽工业大学 ⁴上海交通大学 ⁵东南大学信息科学与工程学院  
[email protected], [email protected], [email protected], [email protected], [email protected], [email protected], [email protected]  

## 1 引言  
在长上下文自回归推理期间，每个全注意力层会为每个缓存的令牌存储键向量和值向量。此KV缓存随上下文长度和并发量增长，可能耗尽设备内存（Zhang 等，2023 (https://arxiv.org/html/2609.02029#bib.bib5)）。混合语言模型减轻了但并未消除这一负担。Gemma 2 和 Gemma 3 交错使用全局和局部滑动窗口注意力（Gemma Team 等，2024 (https://arxiv.org/html/2609.02029#bib.bib18)；Gemma Team 等，2025 (https://arxiv.org/html/2609.02029#bib.bib25)），而 Qwen3-Next 和 Qwen3.6 则交错使用全注意力或门控注意力与循环 Gated-DeltaNet 块（Qwen Team 2025 (https://arxiv.org/html/2609.02029#bib.bib29)；Qwen Team 2026 (https://arxiv.org/html/2609.02029#bib.bib30)；Yang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib23)）。它们剩余的全局块仍然保留着随提示词增长的历史记录。因此，混合化是重新定位而非消除KV瓶颈：一小部分残余全局层可能决定一个长请求是否能放入内存。  

KV压缩涉及两个不同的选择：何时确定保留策略，以及该策略如何在内存中实现。依赖于提示的方法从当前请求、观察到的注意力或在线驱逐状态中推导出令牌或头部决策（Ge 等，2024 (https://arxiv.org/html/2609.02029#bib.bib7)；Feng 等，2025 (https://arxiv.org/html/2609.02029#bib.bib16)；Fu 等，2025 (https://arxiv.org/html/2609.02029#bib.bib21)；Rehg 2024 (https://arxiv.org/html/2609.02029#bib.bib19)；Zhang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib33)；Qin 等，2025 (https://arxiv.org/html/2609.02029#bib.bib34)）。这种适应性可以保留对特定提示重要的孤立远距离证据。然而，其系统收益取决于运行时。逻辑稀疏掩码可以减少注意力计算量，同时保持分配的KV缓存不变。压缩或驱逐可以减少解码时的驻留，但首先构建完整缓存的方法不一定能降低预填充的内存峰值。在线评分、索引和KV移动也会增加请求时的开销，因此稀疏注意力并不自动意味着更高的端到端吞吐量。  

静态策略做出了互补的权衡。在预填充前固定容量允许运行时直接分配有限存储，从而产生可预测的内存需求，无需请求时选择。该策略无法检查当前提示词，因此无法识别未来查询可能需要的孤立历史令牌。现有的静态方法通过为检索头分配完整历史，为其他头分配压缩或流式缓存来缓解此限制（Tang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib17)；Xiao 等，2025 (https://arxiv.org/html/2609.02029#bib.bib20)）。这种二元拆分部署起来规律，但对于残余全局头需要不同历史量的混合模型来说可能过于粗糙。表1 (https://arxiv.org/html/2609.02029#S3.T1) 将策略时机与物理缓存实现分离开来，并总结了这种权衡。  

静态分配还会在深度方向上创建一个校准问题。敏感性因头和层而异（Ge 等，2024 (https://arxiv.org/html/2609.02029#bib.bib7)；Tang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib17)；Fu 等，2025 (https://arxiv.org/html/2609.02029#bib.bib21)；Zhang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib33)；Wang 等，2025b (https://arxiv.org/html/2609.02029#bib.bib31)；Cai 等，2025 (https://arxiv.org/html/2609.02029#bib.bib14)；Wan 等，2025 (https://arxiv.org/html/2609.02029#bib.bib32)；Qin 等，2025 (https://arxiv.org/html/2609.02029#bib.bib34)）。通用窗口可能在宽容的头上浪费内存，同时截断敏感的头。此外，缩短早期缓存会改变后续层所消耗的表示。因此，在全完整前缀下校准的高层在部署后可能会看到不同的输入。可部署的静态策略必须在实际将要激活的低层决策下，为每个头选择细粒度的容量，并且运行时必须将这些容量实现为物理存储。  

我们引入 *HeadWiseKV* 来满足这些需求。它为每个残余全局KV头离线分配几个连续历史长度之一，并在服务期间保持生成的策略固定。这种多级分配比二元头特化更具表现力，同时保持规律的、独立于提示的布局。其运行时仅存储选定的历史，因此策略决定的是物理驻留，而非对完整缓存的逻辑访问。HeadWiseKV 因此用请求特定的令牌选择交换了预填充前已知的细粒度内存计划。  

HeadWiseKV 在一个从离线到在线的流程中连接了三个技术组件。一个结构化的驻留模型首先定义了候选历史长度及其存储成本。HeadWiseKV 的核心是 *SeqCalib*，这个离线算法为每个可配置的KV头选择一个历史长度。对于每一层，SeqCalib 保持先前选择的低层窗口活跃，并在相同的低层策略下，将候选的预门控注意力输出与有条件的完整历史参考进行比较。它选择满足平均余弦阈值的最低成本码本条目。对于已实现的低层前缀，这个有限搜索是逐阶段精确的，并非联合全局最优性的保证。在有限阈值集上运行 SeqCalib 会产生跨越不同驻留成本的静态策略矩阵。然后，一个确定性的预算选择器在可行的矩阵中进行选择。分组缓存运行时一次加载选定的矩阵，并为每个头分配相应的物理历史。简而言之，驻留模型定义选项，SeqCalib 决定保留什么，运行时使该决策成为现实。这些步骤都不需要重新训练或架构更改。  

我们将有限分配形式化为一个*受限的运行率失真问题*（Shannon 1959 (https://arxiv.org/html/2609.02029#bib.bib2)；Cover 和 Thomas 2006 (https://arxiv.org/html/2609.02029#bib.bib1)）。其逐阶段保证并不建立全局最优的联合分配或下游质量。  

图1：Qwen3.6-27B 内存、容量和质量总结。HeadWiseKV 在保持长上下文质量的同时，降低了采样的峰值设备内存并扩展了验证的上下文范围。面板(a)结合了匹配的内存测量和独立的相邻网格容量区间；面板(b)报告了记录的 RULER 和 LoCoMo 分数。系统和质量面板使用分别记录的 τ=.995 操作点，并不构成联合质量-系统结果。  
图1 (https://arxiv.org/html/2609.02029#S1.F1) 总结了 HeadWiseKV 的核心优势：它在保持长上下文质量的同时，降低了设备内存需求并支持更长的上下文。在 Qwen3.6-27B（Qwen Team 2026 (https://arxiv.org/html/2609.02029#bib.bib30)）上，HeadWiseKV 在112K令牌时将采样的峰值设备内存降低了8.59%，并将最大验证成功上下文从114K增加到161K，同时在 RULER 和 LoCoMo 上保留了接近完整KV的性能。系统和质量面板分别报告了操作点；完整的协议和结果见第4节 (https://arxiv.org/html/2609.02029#S4)。  

我们的贡献是：  
- • 针对混合模型的残余全局KV头提出了一种预算化的物理驻留公式，以及分析支持非均匀分配和部署条件化校准；  
- • HeadWiseKV 及其前缀条件化的 SeqCalib 算法，该算法在规定的KV驻留预算下构建静态的、多级的每头后缀策略；  
- • 一个分组物理每头缓存运行时，它将选定策略实例化为实际的缓存驻留，并在模型系列、压缩预算和长上下文工作负载上进行了评估。  

## 2 相关工作  
先前的工作沿两个独立轴变化。策略时机决定了保留是离线固定还是从当前请求中选择。缓存实现决定了该选择是改变物理驻留，还是仅改变稀疏注意力访问。依赖于请求的令牌选择使用注意力、位置或提示特定信号（Liu 等，2023 (https://arxiv.org/html/2609.02029#bib.bib4)；Zhang 等，2023 (https://arxiv.org/html/2609.02029#bib.bib5)；Li 等，2024 (https://arxiv.org/html/2609.02029#bib.bib12)；Wang 等，2025a (https://arxiv.org/html/2609.02029#bib.bib24)）。头和层感知方法进一步在模型中改变预算（Ge 等，2024 (https://arxiv.org/html/2609.02029#bib.bib7)；Tang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib17)；Feng 等，2025 (https://arxiv.org/html/2609.02029#bib.bib16)；Fu 等，2025 (https://arxiv.org/html/2609.02029#bib.bib21)；Xiao 等，2025 (https://arxiv.org/html/2609.02029#bib.bib20)；Rehg 2024 (https://arxiv.org/html/2609.02029#bib.bib19)；Zhang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib33)；Wang 等，2025b (https://arxiv.org/html/2609.02029#bib.bib31)；Cai 等，2025 (https://arxiv.org/html/2609.02029#bib.bib14)；Wan 等，2025 (https://arxiv.org/html/2609.02029#bib.bib32)；Qin 等，2025 (https://arxiv.org/html/2609.02029#bib.bib34)）。一个密切相关的方法是减少每个查询所咨询的上下文。QUEST 使用查询相关边界选择KV页面，而 TokenSelect 跨头执行动态令牌级选择（Tang 等，2024 (https://arxiv.org/html/2609.02029#bib.bib15)；Wu 等，2025b (https://arxiv.org/html/2609.02029#bib.bib22)）。InfLLM 从辅助上下文记忆中检索相关块，而 PyramidInfer 根据逐层注意力一致性减少保留的令牌（Xiao 等，2024a (https://arxiv.org/html/2609.02029#bib.bib10)；Yang 等，2024 (https://arxiv.org/html/2609.02029#bib.bib13)）。这些方法共同使令牌或块访问适应输入或查询。它们的物理内存和延迟效应仍然取决于运行时是掩码、检索、压缩还是驱逐选定的状态。  

相比之下，静态方法将保留决策移至请求之前。StreamingLLM 固定了 sink-recent 策略，而 RazorAttention 和 DuoAttention 使用离线的头角色，主要暴露完整和压缩历史之间的二元选择（Xiao 等，2024b (https://arxiv.org/html/2609.02029#bib.bib6)；Tang 等，2025 (https://arxiv.org/html/2609.02029#bib.bib17)；Xiao 等，2025 (https://arxiv.org/html/2609.02029#bib.bib20)）。HeadWiseKV 则相反，为每个物理KV头固定一个多级后缀长度，在选定的低层策略下校准高层，并在预填充前分配产生的容量。表1 (https://arxiv.org/html/2609.02029#S3.T1) 比较了这些机制。“独立于提示的保留”和“预填充前物理”列特意将策略依赖性与存储实现分离。“部署条件化校准”列捕获了高层决策是否在部署时使用的低层策略下做出。勾号表示已记录的属性，并非假设该属性普遍优先。  

这种保留策略与缓存精度和推理时内存管理正交，后者通过表示成本或管理分配来降低，而不为每个头选择历史长度（Liu 等，2024 (https://arxiv.org/html/2609.02029#bib.bib9)；Hooper 等，2024 (https://arxiv.org/html/2609.02029#bib.bib8)；Kwon 等，2023 (https://arxiv.org/html/2609.02029#bib.bib35)；Prabhu 等，2025 (https://arxiv.org/html/2609.02029#bib.bib36)）。HeadWiseKV 专门针对混合模型的残余全局注意力KV历史，保持其有界内存的局部或循环路径不变。  

图2：HeadWiseKV 概览。残余全局注意力层在 KV 头之间可能具有异构历史要求。离线 SeqCalib 在冻结的低层策略下逐层评估有限的后缀码本，并生成静态策略矩阵。运行时一次加载该矩阵，并实例化具有不同物理长度的分组每头缓存。窗口值是示意性的。  

## 3 方法  
图2 (https://arxiv.org/html/2609.02029#S2.F2) 总结了从离线到在线的流程。HeadWiseKV 定义可行的每头历史长度及其存储成本，使用 SeqCalib 选择一个层条件化的策略，并将该策略实现为物理缓存驻留。接下来的小节将形式化这些阶段。  

### 3.1 结构化驻留模型  
我们首先形式化静态的每头驻留问题。在固定上下文长度下，策略决定每个可配置的KV头保留多少近期历史。只有残余的全注意力层参与：ℒ_F 按前向执行顺序列出这些层，ℋ_ℓ 包含层 ℓ 中的物理KV头。在分组查询注意力（GQA）下，多个查询头可能共享一个KV头。我们用 𝒢_ℓh 表示该组。因此，决策单元是物理KV头 (ℓ, h)，而非单个查询头。所有决策单元的集合是 𝒰 = {(ℓ, h) : ℓ ∈ ℒ_F, h ∈ ℋ_ℓ}，而 HeadWiseKV 为每个头分配 a_ℓh ∈ 𝒲 = {w₁, ..., w_K, full}。