HySparse：一种基于先验Token选择与KV缓存共享的混合稀疏注意力架构 HySparse: A Hybrid Sparse Attention Architecture with Oracle Token Selection and KV Cache Sharing
用全注意力层当Oracle指导稀疏层选token并共享KV缓存，降低计算和内存开销
前置知识
自注意力机制（Self-Attention）
Transformer的核心计算模块。对于序列中每个token $t$，通过查询向量 $\mathbf{q}_t$ 与所有键向量 $\mathbf{k}_i$ 的点积计算注意力分数，再对值向量 $\mathbf{v}_i$ 加权求和得到输出。标准公式为 $$\mathbf{o}_t = \sum_{i=1}^{t} \frac{\exp(\mathbf{q}_t^\top \mathbf{k}_i / \sqrt{d})}{\sum_{j=1}^{t} \exp(\mathbf{q}_t^\top \mathbf{k}_j / \sqrt{d})} \mathbf{v}_i$$，其中 $d$ 是头维度。计算复杂度为 $O(n^2)$，随序列长度二次增长，这是长上下文场景的主要瓶颈。
本文的核心动机就是解决自注意力的二次计算和内存开销问题，理解标准注意力公式是理解稀疏注意力如何节省计算的前提。
稀疏注意力（Sparse Attention）
一种降低注意力计算复杂度的方法，不计算所有token对的注意力分数，而是只选择一部分重要token进行注意力计算。现有方法分为两类：训练无关方法（如固定滑动窗口、基于启发式的选择）和可训练方法（通过学习选择模块在训练时学习哪些token重要）。关键挑战在于如何高效准确地识别重要token，以及如何减少KV缓存的内存占用。
HySparse本身就是一种新型稀疏注意力架构，理解稀疏注意力的基本思想和现有局限性是理解本文创新点的基础。
KV缓存（Key-Value Cache）
在自回归生成过程中，每个token的键向量 $\mathbf{k}$ 和值向量 $\mathbf{v}$ 会被缓存起来，供后续token的注意力计算使用。KV缓存的大小与层数、头数、序列长度成正比，是推理时GPU显存的主要消耗者。例如一个7B模型在32K上下文下，KV缓存可能占用数十GB显存，直接限制了批处理大小和吞吐量。
HySparse的一个核心贡献就是通过跨层KV缓存共享大幅减少KV缓存占用，理解KV缓存的内存瓶颈才能理解这一设计的价值。
混合注意力架构（Hybrid Attention）
将不同类型的注意力层交错排列在同一模型中，以平衡计算效率和模型能力。常见设计包括将全局全注意力层与局部滑动窗口注意力（SWA）层交错排列，例如GPT-OSS使用1:N的比例将全注意力层与SWA层交替。SWA层只关注最近的 $w$ 个token（如 $w=128$），计算量极小但能捕获局部信息。全注意力层则提供全局信息聚合能力。
HySparse建立在混合注意力架构的基础上，是在SWA混合架构中引入稀疏注意力的扩展，理解混合架构的设计范式是理解本文定位的关键。
FlashAttention
一种高效的注意力计算内核，通过分块（tiling）和在线归一化技术避免显式存储完整的注意力矩阵，显著降低了注意力计算的内存带宽需求。FlashAttention将查询、键、值矩阵分成小块在SRAM中计算，通过在线更新softmax的归一化因子来保证数值正确性，使得注意力计算可以从 $O(n^2)$ 内存降至 $O(n)$。
HySparse修改了FlashAttention内核来提取块级注意力分数，这是实现Oracle token选择的技术基础，理解FlashAttention的工作原理才能理解如何巧妙地复用其内部中间结果。
分组查询注意力（Grouped-Query Attention, GQA）
一种注意力头共享机制，多个查询头共享同一组键值头。例如7B模型使用32个查询头和8个KV头（即每4个查询头共享1个KV头）。GQA在保持模型质量的同时显著减少KV缓存大小和计算量。在HySparse中，GQA还被用于在稀疏选择时聚合同一组内的注意力分数，使同一组内的所有头共享相同的稀疏索引。
HySparse在GQA设定下进行稀疏索引的组内聚合，这是一个重要的工程细节，影响稀疏注意力内核的效率和索引开销。
混合专家模型（Mixture-of-Experts, MoE）
一种稀疏化的模型架构，每一层包含多个专家子网络，但每次前向传播只激活其中一小部分。例如本文的80B MoE模型总参数80B，但每次只激活约3B参数（80B-A3B，8个激活专家/512个总专家）。MoE能在不显著增加计算量的情况下大幅扩展模型容量，是当前大模型的主流架构之一。
HySparse同时在7B密集模型和80B MoE模型上验证，MoE场景下的实验是本文的重要贡献，展示了方法在不同模型架构下的通用性。
研究动机
稀疏注意力虽然能有效缓解Transformer的二次计算瓶颈，但现有方法存在两个根本性缺陷。第一个是代理式（Proxy-based）token选择问题：现有方法通常依赖额外的轻量代理来预测token重要性，如预定义模式、启发式规则、近似估计或额外的选择模块。这些代理本质上是近似的，在长上下文和动态演化的情境中可能无法准确捕获真正的token重要性，从而限制了稀疏注意力的表达能力。例如SeerAttention和DSA等方法虽通过可学习的选择模块缓解了选择误差，但并没有从根本上消除代理瓶颈，反而引入了额外的训练复杂性。第二个是计算减少但内存未减的问题：动态稀疏注意力方法为了保持模型精度通常保留完整的KV缓存，因为token重要性会随生成过程和上下文演化而变化，完全驱逐KV缓存是不可逆的破坏性操作。这意味着虽然计算量降低了，但内存消耗没有任何缓解，KV缓存仍然占据大量GPU显存，限制了长上下文场景下的服务吞吐量和最大批处理大小。
本文的目标是本文的目标是设计一种新的混合注意力架构，同时解决上述两个问题：（1）利用全注意力层自身作为精确的Oracle来识别重要token，消除对代理模块的依赖；（2）让稀疏注意力层复用全注意力层的KV缓存，在不增加额外KV缓存开销的前提下实现稀疏计算。最终在保持甚至超越全注意力模型性能的同时，大幅减少KV缓存内存占用（目标约10倍），推动混合注意力架构中全注意力层比例的极限。
与已有工作不同的是，本文的独特切入角度在于两个关键观察的结合：（1）跨层显著token稳定性——连续层之间的显著token（注意力分数较高的token）倾向于保持相对稳定，这意味着一个全注意力层识别出的重要token可以安全地被后续稀疏层复用；（2）跨层KV缓存共享——相邻层的KV缓存具有高度相似性，可以在不损害模型精度的前提下跨层共享。HySparse将这两个经验观察提升为一种预训练架构设计原则，通过将每个全注意力层与多个稀疏注意力层交错排列，让稀疏层直接从前序全注意力层获取token选择和KV缓存。这与现有工作（如SeerAttention仅在推理时利用层间相似性）不同，HySparse是从预训练阶段就内建这一机制的端到端架构方案。
核心方法
HySparse的核心思想是将标准Transformer骨架替换为重复的混合块，每个混合块由一个全注意力层后接 $N$ 个连续的稀疏注意力层组成。直觉上，全注意力层像一个侦察兵，精确计算所有token之间的注意力关系并输出重要token的索引，同时其产生的KV缓存已经包含了全局信息；后续的稀疏注意力层像突击队，直接复用侦察兵的发现（重要token索引和KV缓存），只关注最相关的token子集进行高效计算。具体技术路线为：（1）全注意力层通过修改的FlashAttention内核输出块级注意力分数矩阵 $\mathbf{S}$，从中提取TopK块索引 $\mathbf{I}$；（2）稀疏注意力层包含两个分支——块稀疏注意力分支（复用全注意力层的KV缓存和索引 $\mathbf{I}$）和滑动窗口注意力SWA分支（维护自己独立的小型KV缓存）；（3）两个分支的输出通过sigmoid门控融合。这种设计使得在7B密集模型中全注意力与稀疏层比例为1:3，在80B MoE模型中更激进地达到1:11。
HySparse的核心创新在于用全注意力层本身作为Oracle指导稀疏层的token选择，这与现有方法有本质区别。现有稀疏注意力方法要么依赖预定义模式（如滑动窗口、块状模式），要么需要额外的选择模块（如SeerAttention的可学习门控、DSA的自蒸馏索引器），这些代理方案本质上是近似的，可能无法捕获真正的token重要性。HySparse的洞察是：全注意力层的注意力分数本身就是token重要性的最精确度量，因为它就是完整的 $O(n^2)$ 注意力计算的结果。通过从全注意力层的FlashAttention计算过程中提取块级最大注意力分数（仅需存储在线softmax过程中的行最大值并适当缩放），可以零额外开销地获得高质量的稀疏选择。此外，全注意力层已经计算了KV缓存，稀疏层直接复用这些KV缓存（仅用于块稀疏注意力分支），无需为稀疏层额外存储KV，实现了一份KV缓存多层使用的内存优化。
方法步骤详情
HySparse的方法流程可以分为三个阶段。第一阶段是全注意力层计算与分数提取：输入序列经过标准缩放点积自注意力计算 $\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t = \mathbf{W}_{q/k/v}\mathbf{x}_t$，修改的FlashAttention内核在计算过程中额外输出块级最大注意力分数 $\mathbf{S} \in \mathbb{R}^{t \times \lceil t/B \rceil}$，其中 $S_i^t = \max_{i' \in \mathcal{B}_i} \frac{\exp(\mathbf{q}_t^\top \mathbf{k}_{i'} / \sqrt{d})}{\sum_{j=1}^{t} \exp(\mathbf{q}_t^\top \mathbf{k}_j / \sqrt{d})}$，$B$ 为块大小（默认64）。在GQA设定下，同一查询组内的注意力分数通过组内最大值聚合，使同组所有头共享相同索引。随后对 $\mathbf{S}$ 应用TopK算子选择 $k/B$ 个最重要块的索引 $\mathbf{I}$（默认 $k=1024$，$B=64$，即选16个块）。第二阶段是稀疏注意力层计算：每个稀疏层包含两个分支。SWA分支使用独立的QKV投影 $\mathbf{q}'_t, \mathbf{k}'_t, \mathbf{v}'_t = \mathbf{W}_{q'/k'/v'}\mathbf{x}_t$，在大小为 $w=128$ 的滑动窗口内计算标准注意力，维护自己的小型KV缓存。块稀疏注意力分支复用相同的查询 $\mathbf{q}'_t$，但键值来自全注意力层的共享KV缓存，仅索引 $\mathbf{I}$ 对应的块，然后计算稀疏注意力输出。第三阶段是门控融合：两个分支的输出通过可学习的sigmoid门控 $\tilde{g}_t, g'_t = \sigma(\mathbf{W}_{\tilde{g}/g'}\mathbf{x}_t)$ 加权求和得到最终输出 $\mathbf{o}_t = \tilde{g}_t \odot \tilde{\mathbf{o}}_t + g'_t \odot \mathbf{o}'_t$。最后一个层始终使用全注意力以保留全局聚合能力。
技术新颖性
HySparse的技术新颖性体现在多个层面。首先，它将全注意力层作为Oracle这一思想从推理时的启发式利用提升为预训练架构的内建设计原则，这是与SeerAttention、OmniKV等仅在推理阶段利用层间稳定性的方法的根本区别。其次，它巧妙地修改了FlashAttention内核，在计算标准注意力输出的同时零额外开销地提取块级注意力分数，复用了FlashAttention在线softmax过程中已有的行最大值中间结果，仅需存储和适当缩放，算法复杂度几乎没有增加。第三，HySparse实现了token选择和KV缓存共享的双重复用：稀疏层既复用全注意力层的token选择索引（Oracle选择），又复用其KV缓存（跨层共享），而SWA分支保持独立KV缓存以保留局部建模能力。这种选择共享加KV共享但SWA独立的设计是基于消融实验得出的最优配置。第四，HySparse将混合SWA架构推向更激进的全注意力层比例（1:11），在此前被认为可能导致严重性能退化的设置下仍保持强劲性能，证明了稀疏注意力分支在补充局部SWA能力方面的关键价值。
实验结果
HySparse在所有实验设置下均展现出一致且显著的性能优势。在7B密集模型（36层，全注意力:稀疏层=1:3）上，HySparse在大多数基准上超越了全注意力基线和混合SWA基线，例如MMLU达到58.8（vs. 全注意力56.9）、MMLU-Redux达到61.6（vs. 59.6）、MMLU-Pro达到29.0（vs. 26.8）、GSM8K达到37.9（vs. 33.3），数学推理能力提升尤为明显。在更具挑战性的80B MoE模型（49层，全注意力:稀疏层=1:11，仅5层全注意力）上，HySparse的表现更加亮眼：在几乎所有基准上同时超越全注意力和混合SWA，例如MMLU达到62.2（vs. 全注意力61.8，混合SWA仅54.9）、GSM8K达到54.1（vs. 53.8）、MATH达到30.8（vs. 28.6）、HumanEval达到38.4（vs. 35.4）。值得注意的是，在如此激进的混合比例下，混合SWA出现了严重的精度退化（如BBH从56.1降至48.2，MMLU从61.8降至54.9），而HySparse通过稀疏注意力分支恢复了对全局重要token的访问，不仅弥补了退化甚至超越了全注意力基线，同时实现了近10倍的KV缓存减少。在RULER长上下文基准上，HySparse同样表现出色：7B模型在16K和32K上下文长度下总分分别为94.1和89.3（vs. 全注意力93.0和88.2，混合SWA 91.6和84.2）；80B模型在32K下HySparse达到87.4，甚至超越了全注意力的82.1，而混合SWA仅69.5。消融实验进一步验证了两个关键设计：移除SWA分支导致DROP从52.2降至46.4、GSM8K从37.7降至29.7；将KV缓存同时共享给SA和SWA分支则导致MMLU从58.4降至52.8、MMLU-Pro从29.0降至23.2，证明SWA分支需要独立的KV缓存来保持局部建模能力。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 通用语言理解与推理（MMLU） | 5-shot准确率 | 7B: 58.8, 80B: 62.2 | 7B Full-Attn: 56.9, 7B Hybrid SWA: 57.5; 80B Full-Attn: 61.8, 80B Hybrid SWA: 54.9 | 7B: +1.9 vs Full-Attn, +1.3 vs Hybrid SWA; 80B: +0.4 vs Full-Attn, +7.3 vs Hybrid SWA |
| 数学推理（GSM8K） | 8-shot准确率 | 7B: 37.9, 80B: 54.1 | 7B Full-Attn: 33.3, 7B Hybrid SWA: 35.6; 80B Full-Attn: 53.8, 80B Hybrid SWA: 45.3 | 7B: +4.6 vs Full-Attn, +2.3 vs Hybrid SWA; 80B: +0.3 vs Full-Attn, +8.8 vs Hybrid SWA |
| 高级数学推理（MATH） | 4-shot准确率 | 7B: 10.1, 80B: 30.8 | 7B Full-Attn: 9.2, 7B Hybrid SWA: 9.2; 80B Full-Attn: 28.6, 80B Hybrid SWA: 25.8 | 7B: +0.9 vs Full-Attn, +0.9 vs Hybrid SWA; 80B: +2.2 vs Full-Attn, +5.0 vs Hybrid SWA |
| 代码生成（HumanEval） | 0-shot pass@1 | 7B: 23.5, 80B: 38.4 | 7B Full-Attn: 25.0, 7B Hybrid SWA: 22.0; 80B Full-Attn: 35.4, 80B Hybrid SWA: 31.7 | 7B: -1.5 vs Full-Attn, +1.5 vs Hybrid SWA; 80B: +3.0 vs Full-Attn, +6.7 vs Hybrid SWA |
| 中文理解（C-Eval） | 5-shot准确率 | 7B: 52.2, 80B: 65.0 | 7B Full-Attn: 50.6, 7B Hybrid SWA: 50.6; 80B Full-Attn: 64.6, 80B Hybrid SWA: 58.8 | 7B: +1.6 vs Full-Attn, +1.6 vs Hybrid SWA; 80B: +0.4 vs Full-Attn, +6.2 vs Hybrid SWA |
| 长上下文（RULER 32K） | 总分 | 7B: 89.3, 80B: 87.4 | 7B Full-Attn: 88.2, 7B Hybrid SWA: 84.2; 80B Full-Attn: 82.1, 80B Hybrid SWA: 69.5 | 7B: +1.1 vs Full-Attn, +5.1 vs Hybrid SWA; 80B: +5.3 vs Full-Attn, +17.9 vs Hybrid SWA |
局限与改进
尽管HySparse取得了令人印象深刻的结果，但仍存在一些局限性。首先，论文仅在7B和80B两个模型规模上进行了验证，尚缺乏在更大规模模型（如数百B参数）上的实验，以及在更长上下文（如128K、1M tokens）下的系统性评估。其次，HySparse的设计假设跨层显著token稳定性在预训练全程都成立，但这一假设在训练初期或不同训练阶段是否始终有效尚未充分讨论。第三，论文中固定了TopK token数为1024和块大小为64，缺乏对这些超参数的敏感性分析，不清楚在不同任务或上下文长度下是否需要动态调整这些参数。第四，虽然HySparse减少了KV缓存的内存占用，但块级注意力分数的提取和TopK选择操作本身引入了一定的计算和索引开销，论文未提供详细的延迟和吞吐量分析。第五，论文主要聚焦于预训练阶段，未讨论HySparse与现有后训练技术（如RLHF、指令微调）的兼容性，也未涉及推理阶段的进一步优化（如KV缓存驱逐策略）。最后，SWA分支需要维护独立KV缓存这一设计虽然保持了模型质量，但意味着每个稀疏层仍有一定的KV缓存开销，虽然很小（窗口仅128 tokens），但并非完全为零。
独立分析的弱点
HySparse的几个潜在弱点值得关注。第一，块级TopK选择的粒度可能过于粗糙：块大小为64意味着选择的最小单位是64个连续token，如果重要token分散在多个块中，可能会选中大量不重要的token，或者遗漏块内个别关键token。改进方向可以是探索自适应块大小或token级选择与块级选择的混合策略。第二，所有稀疏层共享相同的TopK索引 $\mathbf{I}$，这意味着不同稀疏层无法根据自身的注意力模式做出差异化选择。虽然论文引用了层间稳定性的证据，但在某些任务（如多跳推理）中不同层可能需要关注不同token子集，改进方向可以是允许每个稀疏层在Oracle索引的基础上进行微调或自适应调整。第三，sigmoid门控融合的两个分支权重是通过简单的线性投影+sigmoid计算的，这种静态融合机制可能无法很好地适应不同位置和不同上下文的动态需求，改进方向可以是引入更复杂的融合机制如MoE风格的路由。第四，论文未讨论HySparse与KV缓存量化技术的兼容性，如果两者可以结合，可能进一步压缩内存占用。
未来方向
作者在论文中提出了几个未来方向。首先，计划将HySparse扩展到更大的模型规模和更多的训练token，以充分释放其在高效长上下文建模方面的潜力。其次，提出了HySparse用于高效KV缓存卸载的系统级策略：将全注意力层的KV缓存卸载到外部存储（如CPU内存或NVMe），在计算前预取，而GPU上只保留持久化的稀疏KV缓存，这有望进一步减少GPU显存占用，支持更大的批处理和更高的推理效率。基于已有成果，还可以延伸多个方向：（1）将HySparse与后训练技术结合，探索在指令微调和RLHF阶段保留稀疏注意力结构的可能性；（2）研究自适应稀疏度，根据输入序列的特征动态调整TopK数量和全注意力层比例；（3）将HySparse的Oracle选择思想与其他高效注意力变体（如线性注意力、核注意力）结合；（4）探索HySparse在多模态模型中的应用，利用视觉token的稀疏性进一步优化效率。
复现评估
从复现角度来看，HySparse面临一定的挑战。论文未开源代码或预训练模型，这限制了独立复现和验证的可能性。训练成本较高：7B模型需要在1T tokens上训练（序列长度8192），再用200B tokens进行长上下文训练（序列长度32768）；80B MoE模型需要在500B tokens上训练（序列长度32768），这对算力要求非常高，非大型研究团队难以负担。技术实现上，需要修改FlashAttention内核以输出块级注意力分数，这需要CUDA编程和底层内核优化的专业知识。此外，论文中的训练超参数（如学习率调度、RoPE频率调整等）都是针对特定模型规模精心调优的，在不同规模下可能需要重新调参。积极的方面是，论文提供了详细的模型配置表（Table 1）和训练超参数描述，方法的核心思想（Oracle选择 + KV共享 + SWA分支）概念清晰，如果获得足够的算力和FlashAttention内核修改能力，技术路线是可复现的。
论文图表