HySparse: A Hybrid Sparse Attention Architecture with Oracle Token Selection and KV Cache Sharing

发表时间: 2026-02 · arXiv:2602.03560 (LLM-Core Xiaomi)

原文: https://arxiv.org/abs/2602.03560

作者/机构: Yizhao Gao, Jianyu Wei, Qihao Zhang, Yu Cheng, Shimao Chen, Zhengju Tang, Zihan Jiang, Yifan Song, Hailin Zhang, Liang Zhao, Bo Yang, Gang Wang, Shijie Cao, Fuli Luo⋄ (LLM-Core Xiaomi)

A1 主要贡献

本文介绍了一种名为混合稀疏注意力(HySparse)的新型架构,其核心思想是在每个全注意力层之后交错插入几个稀疏注意力层。该架构旨在解决现有稀疏注意力方法存在的两个基本限制:

  1. 对代理(Proxy)的依赖:传统稀疏注意力方法通常依赖预定义模式、启发式规则或额外的选择模块等轻量级代理来预测重要词元(token),这不仅增加了复杂性,还可能因代理的近似性而导致次优性能。HySparse通过将前面的全注意力层作为精确的“神谕”(oracle),直接识别重要词元,从而消除了对代理的依赖,避免了因代理保真度不足而限制稀疏注意力的表达能力。

  2. 计算减少但内存未节约:许多动态稀疏注意力方法虽然能有效减少计算量,但通常仍需保留完整的键值缓存(KV cache),因为词元的重要性会随上下文演变而变化,完全驱逐缓存是不可逆的。这使得KV缓存在长上下文场景下依然是内存瓶颈。HySparse通过让稀疏注意力层直接重用前一个全注意力层的KV缓存,实现了计算和内存的双重节省。

研究目标与创新点
HySparse的核心设计是利用两个经验性观察:(1) 相邻层之间的显著词元具有稳定性;(2) 跨层KV缓存共享能在不损害性能的情况下减少内存。基于此,HySparse架构让全注意力层精确地识别重要词元,并生成KV缓存,后续的稀疏注意力层可以直接重用这些词元索引和KV缓存。这种设计具有以下创新点:

  • 神谕式词元选择:通过重用全注意力层识别的重要词元索引,稀疏选择变得由“神谕”引导,无需额外的代理模块,保证了端到端训练的稳定性。
  • 零开销KV缓存:稀疏注意力层重用全注意力层的KV缓存,自身不产生额外的KV开销,有效缓解了动态稀疏注意力的内存压力。
  • 混合架构增强:受混合滑动窗口注意力(SWA)架构的启发,HySparse为每个稀疏注意力层增加了一个SWA分支,该分支维护一个小的、独立的本地KV缓存,以增强模型的短距离建模能力。

实验结果表明,无论是在7B密集模型还是80B MoE模型上,HySparse的性能均优于全注意力和混合SWA基线。特别是在一个包含49层的80B MoE模型中,仅使用5个全注意力层(全注意力与稀疏注意力的层级比例为1:11),HySparse不仅取得了显著的性能提升,还将KV缓存存储减少了近10倍。这证明了HySparse能够在不牺牲模型能力的情况下,大幅减少全注意力层的数量,为构建高效且强大的长上下文模型提供了有效的架构方案。

A3 背景知识与动机

2.1 免训练与可训练的稀疏注意力

稀疏注意力方法的分类。稀疏注意力方法可分为免训练(training-free)和可训练(trainable)两类。免训练方法依赖固定的模式或启发式规则来识别重要词元,它们在推理时作为即插即用的修改,以最小的计算成本实现快速的稀疏决策【【53,Quest: Query-aware sparsity for efficient long-context llm inference,2024,arXiv】,【55,Efficient streaming language models with attention sinks,2023,arXiv】,【56,Duoattention: Efficient long-context llm inference with retrieval and streaming heads,2024,arXiv】,【61,H2o: Heavy-hitter oracle for efficient generative inference of large language models,2023,NeurIPS】】。然而,仅在推理时应用稀疏性会造成训练-推理不匹配,可能导致在长解码或多步推理中误差累积【【24,Defeating nondeterminism in llm inference,2025,Thinking Machines Lab】,【28,Lil: Less is less when applying post-training sparse-attention algorithms in long-decode stage,2026,arXiv】】。相比之下,可训练的稀疏注意力方法在训练期间通过轻量级选择模块学习词元的重要性。通过将稀疏性整合到训练过程中,它们改善了训练和推理之间的一致性,能够以更高的召回率和整体准确性选择信息更丰富的词元【【17,Seerattention: Learning intrinsic sparse attention in your llms,2024,arXiv】,【18,Seerattentionr: Sparse attention adaptation for long reasoning,2025,arXiv】,【40,Deepseekv3. 2: Pushing the frontier of open large language models,2025a,arXiv】,【42,Moba: Mixture of block attention for long-context llms,2025,arXiv】,【43,Minicpm4: Ultra-efficient llms on end devices,2025,】,【59,Native sparse attention: Hardware-aligned and natively trainable sparse attention,2025,arXiv】,【62,Infllmv2: Dense-sparse switchable attention for seamless short-to-long adaptation,2025,arXiv】】。然而,训练这些选择模块并非易事。一种方法是使用辅助损失(如自蒸馏)来使门控或索引模块与原始的密集注意力对齐【【17,Seerattention: Learning intrinsic sparse attention in your llms,2024,arXiv】,【18,Seerattentionr: Sparse attention adaptation for long reasoning,2025,arXiv】,【40,Deepseekv3. 2: Pushing the frontier of open large language models,2025a,arXiv】】。这些方法虽然简单,但并非最优。另一种方法,如NSA【59,Native sparse attention: Hardware-aligned and natively trainable sparse attention,2025,arXiv】,通过将压缩注意力(选择模块)的输出注入主注意力中,进行端到端的稀疏预训练。这种设计使得选择模块只能通过最终的注意力输出间接接收学习信号,而没有对其词元选择决策的直接监督。

2.2 混合注意力架构

混合注意力作为扩展上下文长度的方案。为了降低二次方的计算和KV缓存成本,混合注意力已成为扩展上下文长度的一种有前途的解决方案。例如,MiniMax-01【35,Minimax01: Scaling foundation models with lightning attention,2025a,arXiv】以结构化模式集成了线性和softmax注意力机制。类似地,Qwen3-Next【47,Qwen3-next: Towards ultimate training & inference efficiency,2025,Qwen Team Blog】和Kimi Linear【34,Kimi linear: An expressive, efficient attention architecture,2025b,arXiv】融入了Gated DeltaNet【57,Gated delta networks: Improving mamba2 with delta rule,22024b,arXiv】或其变体。Nemotron系列【【5,Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model,2025,arXiv】,【6,Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models,2025,arXiv】】和Jamba【38,Jamba: A hybrid transformer-mamba language model,2024,arXiv】将Mamba模块【【13,Transformers are ssms: Generalized models and efficient algorithms through structured state space duality,2024,arXiv】,【21,Mamba: Linear-time sequence modeling with selective state spaces,2023,arXiv】】与标准自注意力模块集成。GPT-OSS【1,gpt-oss-120b & gpt-oss-20b model card,2025,arXiv】、Gemma3【20,Gemma 3 technical report,2025,arXiv】和MiMo-V2-Flash【56,Mimo-v2-flash technical report,2026,arXiv】等模型采用了滑动窗口注意力和全局全注意力层的异构交错。滑动窗口的大小可以小至128个词元,KV缓存开销可忽略不计。然而,带有动态稀疏注意力的混合模型尚未被充分探索。

2.3 跨层显著词元的稳定性

显著词元在相邻层间的稳定性观察。多项同期的工作观察到,在标准Transformer模型中,显著词元(具有较高注意力分数的稀疏词元)在连续的层之间倾向于保持相对稳定【【15,Kascade: A practical sparse attention method for long-context llm inference,2025,arXiv】,【23,Omnikv: Dynamic context selection for efficient long-context llms,2025,ICLR】,【56a,Tidaldecode: Fast and accurate llm decoding with position persistent sparse attention,2024a,arXiv】,【56b,Less is more: Training-free sparse attention with global locality for efficient reasoning,2025a,arXiv】,【60,Delta: Dynamic layer-aware token attention for efficient long-context reasoning,2025,arXiv】】。这些方法利用这一特性,以免训练的方式加速推理。具体来说,它们使用一个全注意力层来识别重要词元,并在后续层中重用这些显著词元的索引来执行稀疏注意力计算。受这些工作的启发,我们将这一经验性观察提升为一个用于预训练的混合注意力架构,其中全注意力层识别的重要词元被随后的稀疏注意力层重用。

2.4 跨层KV缓存共享

跨层KV缓存共享作为内存优化技术。跨层KV缓存共享是一种内存优化技术,即在一个层中计算的Key和Value张量被后续层重用,而不是为每个层独立地重新计算和存储。这种设计显著减少了KV缓存的内存占用,同时实证研究表明,它几乎不会导致模型准确性的下降。YOCO【51,You only cache once: Decoder-decoder architectures for language models,2024,NeurIPS】、CLA【7,Reducing transformer key-value cache size with cross-layer attention,2024,NeurIPS】、苹果基础模型【36,Apple intelligence foundation language models: Tech report 2025,2025b,arXiv】和Gemma 3n【20,Gemma 3 technical report,2025,arXiv】直接将跨层KV缓存共享机制集成到它们的模型架构中。SwiftKV【45,Swiftkv: Fast prefill-optimized inference with knowledge-preserving model transformation,2025,EMNLP】通过蒸馏使标准的预训练模型能够在一部分层上支持跨层KV缓存共享。MiniCache【39,Minicache: Kv cache compression in depth dimension for large language models,2024,NeurIPS】也观察到,在LLM的中层到深层区域,相邻层之间的KV缓存表现出高度相似性,并提出了利用这种冗余的跨层压缩方法。

A2 方法细节

3.1 HySparse 概述

图1 HySparse架构图。每个全注意力层与多个稀疏注意力层交错排列。稀疏注意力层直接重用来自前一个全注意力层的KV缓存和重要词元索引。
图1 HySparse架构图。每个全注意力层与多个稀疏注意力层交错排列。稀疏注意力层直接重用来自前一个全注意力层的KV缓存和重要词元索引。

HySparse架构的核心组成。如图1所示,HySparse架构用重复的混合块替换了标准的Transformer主干。每个混合块由一个全注意力层和随后的 $n$ 个连续的稀疏注意力层组成。其核心在于,这些稀疏层中使用的稀疏重要词元索引和KV缓存都直接源自同一块内前方的全注意力层。

HySparse的工作流程。全注意力层计算标准的缩放点积自注意力,但额外输出块级(block-wise)的注意力重要性分数 $S$,我们从中导出TopK块索引。这些索引随后被接下来的 $n$ 个稀疏层重用。为了减少KV缓存的内存和带宽消耗,HySparse进一步引入了跨层KV缓存共享。稀疏注意力层为其块稀疏注意力分支重用由混合块内前一个全注意力层产生的KV缓存。相比之下,滑动窗口注意力(SWA)分支则维护其自身的轻量级KV缓存,以增强短距离建模能力。最后,在两个分支的输出上应用sigmoid门【46,Gated attention for large language models: Non-linearity, sparsity, and attention-sinkfree,2025,arXiv】,然后相加作为最终的注意力输出。

3.2 全注意力层

全注意力层的标准与挑战。全注意力层遵循Transformer中使用的标准softmax自注意力公式。为了给后续的稀疏注意力识别显著词元,注意力分数必须被暴露出来以供选择。然而,物化完整的注意力矩阵在内存和带宽方面的开销都极其高昂。因此,现代Transformer依赖于FlashAttention【【12,Flashattention-2: Faster attention with better parallelism and work partitioning,2023,URL】,【14,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,NeurIPS】,【50,Flashattention-3: Fast and accurate attention with asynchrony and low-precision,2024,NeurIPS】】,它通过以分块(tiled)方式结合在线归一化来计算softmax,从而避免了显式存储注意力分数。

HySparse的解决方案:块级最大注意力分数。为了解决这个问题,HySparse不输出完整的注意力分数矩阵,而是只物化块级(tile-level)的最大注意力分数用于TopK选择。根据标准的自注意力公式:
公式1 公式2
其中 $d$ 是头维度大小。设 $B$ 是我们注意力分数输出的块大小,将会有 $\lceil N/B \rceil$ 个块。我们将块索引 $j$ 处的列词元索引集定义为 $B_j = \{(j-1)B+1, \dots, \min(jB, N)\}$。那么,块级最大注意力分数 $S \in \mathbb{R}^{N \times \lceil N/B \rceil}$ 将是:
公式3

实现方式:修改FlashAttention内核。我们发现,$S_{ij}$ 可以通过对FlashAttention内核进行微小修改而轻松获得,这类似于【17,Seerattention: Learning intrinsic sparse attention in your llms,2024,arXiv】和【18,Seerattentionr: Sparse attention adaptation for long reasoning,2025,arXiv】中的方法。具体来说,FlashAttention在其在线softmax过程中已经计算了注意力logits的行最大值,这个中间结果可以通过存储并进行适当的重新缩放来推导出块级注意力分数。

修改后的FlashAttention算法。算法1总结了修改后的FlashAttention过程,为简化起见,我们假设稀疏注意力的块大小 $B$ 与 $B_C$ 相同。除了标准的注意力输出外,该内核还会输出块级注意力分数,这些分数可以被后续的稀疏注意力层直接用于显著块选择,且开销可忽略不计。

算法1:带块级注意力分数输出的FlashAttention
需要: Queries $Q \in \mathbb{R}^{N \times d}$, Keys $K \in \mathbb{R}^{N \times d}$, Values $V \in \mathbb{R}^{N \times d}$, softmax 缩放因子 $s$
确保: 注意力输出 $O \in \mathbb{R}^{N \times d}$, 块注意力分数 $S \in \mathbb{R}^{N \times \lceil N/B \rceil}$
1: $T_R \leftarrow \lceil N/B_R \rceil$, $T_C \leftarrow \lceil N/B_C \rceil$
2: 初始化: $O_i$, $S_{ij}$, $m_i$, $l_i$
3: for $i = 0, \dots, T_R - 1$ do
4: 从HBM加载 $Q_i \in \mathbb{R}^{B_R \times d}$ 到SRAM
5: for $j = 0, \dots, T_C - 1$ do
6: 从HBM加载 $K_j, V_j \in \mathbb{R}^{B_C \times d}$ 到SRAM
7: $A_{ij} \leftarrow Q_i K_j^\top \cdot s$
8: $\tilde{m}_{ij} \leftarrow \text{rowmax}(A_{ij})$, 将 $\tilde{m}_{ij}$ 存储到 $S_{ij}$
9: $m_i^{new} \leftarrow \max(m_i, \tilde{m}_{ij})$
10: 如FLASHATTENTION中更新 $O_i, l_i$
11: end for
12: $O_i \leftarrow O_i / l_i$, 将 $O_i$ 写回HBM
13: for $j = 0, \dots, T_C - 1$ do
14: $S_{ij} \leftarrow (S_{ij} - m_i) / l_i$, 将 $S_{ij}$ 写回HBM
15: end for
16: end for
17: return $O, S$

TopK选择与GQA优化。利用块级注意力分数 $S$,我们应用TopK算子来选择键块(key-block)索引 $I$,这些索引将被后续的稀疏注意力层重用。需要注意的是,在稀疏注意力中关注 $k$ 个词元对应于选择 $k/B$ 个TopK词元块,其中 $B$ 是块大小。在HySparse中,默认的 $k$ 和 $B$ 分别是1024和64。在分组查询注意力(GQA)【2,Gqa: Training generalized multi-query transformer models from multi-head checkpoints,2023,arXiv】下,我们进一步在每个查询组内聚合 $S$(通过组内最大值),这样同一组中的所有头共享相同的稀疏索引,从而提高了稀疏注意力内核的效率并减少了索引开销。

3.3 稀疏注意力层

稀疏层的双分支结构。每个稀疏层包含两个注意力分支,它们操作于相同的查询但使用不同的KV源。块稀疏注意力(Block Sparse Attention) 分支仅关注由索引 $I$ 标识的键值块,其中索引 $I$ 和KV缓存都源自前一个全注意力层。SWA 分支则关注一个大小为 $w$ 的局部滑动窗口,并拥有自己小型的KV缓存,以提高局部性和表达能力。在HySparse的实现中,$w$ 设置为128。然后,两个分支的输出通过一个轻量级的sigmoid门进行融合。详细过程可以建模如下。首先,我们使用以下公式计算标准的SWA分支输出:

公式4 公式5

稀疏注意力分支的计算。接着,在计算稀疏注意力分支时,我们使用块索引 $I$ 从全注意力层共享的 $K, V$ 中拼接选定的键和值块。稀疏注意力分支和SWA分支使用相同的查询 $Q'_{ij}$,其输出可以写为:

公式6 公式7

分支融合。最后,我们对两个分支的输出应用sigmoid门,并将它们相加得到最终的注意力层输出【46,Gated attention for large language models: Non-linearity, sparsity, and attention-sinkfree,2025,arXiv】。

公式8
公式8

SWA分支独立KV缓存的重要性。通过我们的实验发现,为SWA分支维护一个独立的KV缓存对于保持模型表达能力至关重要。一种可能的解释是,SWA主要作为局部信息通路,需要不同的表示来捕捉短程连贯性,而从前一个全注意力层共享的KV缓存是为全局检索而优化的,可能缺乏足够的局部特征。双分支门控融合在保持计算和内存效率的同时,实现了全局和局部信息的动态混合。

A4 实验环境

模型配置

实验中使用了两种模型:一个标准的7B密集Transformer模型(36层)和一个80B-A3B MoE模型(49层)。两种模型均采用分组查询注意力(GQA),其中7B模型使用32个查询头和8个KV头,80B MoE模型使用64个查询头和4个KV头。

评估了三种架构:
1. Full-Attn: 所有层都使用标准的全注意力。
2. Hybrid SWA: 混合滑动窗口注意力模型。7B模型中全注意力层与SWA层的比例为1:3,80B MoE模型中为1:11。
3. HySparse: 混合模型,混合比例与Hybrid SWA相同,但在SWA的基础上增加了本文提出的稀疏注意力(选择Top-1024词元,块大小为64)。

所有混合模型的最后一层都使用全注意力。对于稀疏注意力和滑动窗口注意力,采用了类似gpt-oss【1,gpt-oss-120b & gpt-oss-20b model card,2025,arXiv】的方法,为每个头引入了可学习的sink偏差。在MoE设置中,Full-Attn模型额外使用了门控注意力【46,Gated attention for large language models: Non-linearity, sparsity, and attention-sinkfree,2025,arXiv】以稳定训练。

表1 模型架构配置
表1 模型架构配置

硬件与软件配置

  • 训练精度: BF16。
  • 优化器: AdamW ($β_1 = 0.9, β_2 = 0.95, ε = 10^{-10}$),权重衰减为0.1,梯度裁剪范数为1.0。
  • 学习率调度: WSD调度。

数据集与训练超参数

  • 7B模型:

    • 第一阶段:在1T词元上进行训练,序列长度为8,192,最大学习率为$8.3 \times 10^{-4}$。
    • 第二阶段(长上下文扩展):在200B词元上进一步训练,序列长度为32,768,学习率为$3.0 \times 10^{-5}$。此阶段RoPE基础频率调整为640,000。
  • 80B MoE模型:

    • 在500B词元上进行训练,序列长度为32,768,最大学习率为$1 \times 10^{-3}$,RoPE基础频率也设置为640,000。

评估基准

  • 通用语言理解与推理: BBH, MMLU, MMLU-Redux, MMLU-Pro, DROP, ARC, HellaSwag, WinoGrande, TriviaQA。
  • 数学推理: GSM8K, MATH。
  • 代码能力: HumanEval, MBPP。
  • 中文理解: C-Eval, CMMLU。
  • 长上下文: Ruler。

A4 实验结果

通用基准测试性能

表2 在7B密集模型(1T词元训练)和80B MoE模型(500B词元训练)上,HySparse与Full-Attn和Hybrid SWA的性能比较。每行中的最佳结果以粗体突出显示。
表2 在7B密集模型(1T词元训练)和80B MoE模型(500B词元训练)上,HySparse与Full-Attn和Hybrid SWA的性能比较。每行中的最佳结果以粗体突出显示。

7B密集模型性能
如表2所示,在7B模型上,HySparse在通用基准、数学和中文理解方面表现出色。它在多个知识和推理基准上超越了Full-Attn基线,例如MMLU(58.8 vs 56.9)、MMLU-Redux(61.6 vs 59.6)和MMLU-Pro(29.0 vs 26.8),表明稀疏词元选择在减少计算的同时能够保持甚至增强全局推理和事实回忆能力。在GSM8K和MATH等多步推理任务上,HySparse也取得了稳定的增益。在中文基准C-Eval和CMMLU上,HySparse同样表现出明显优势。Hybrid SWA作为一个计算高效的基线,在BBH和MBPP+等任务上表现具有竞争力。

扩展到80B MoE模型
在80B MoE模型上,采用了更激进的1:11全注意力与稀疏注意力层比例。尽管只有五个全注意力层,HySparse在几乎所有基准上都优于Full-Attn和Hybrid SWA。值得注意的是,性能增益通常比7B密集模型中观察到的更大。在这种设置下,Hybrid SWA在多个基准(如BBH, DROP, MMLU系列, GSM8K)上相比Full-Attn出现明显的准确率下降,表明当混合比例变得激进时,仅依赖局部窗口注意力是不够的。HySparse通过引入稀疏注意力分支,弥补了这一差距,恢复了对全局相关词元的访问,在许多情况下甚至超越了Full-Attn基线,同时所需的KV缓存减少了10倍。这突显了HySparse的核心优势:可以大幅减少全注意力层的数量而不牺牲模型能力。

长上下文基准测试

表3 7B密集模型和80B MoE模型的RULER基准性能。注意,所有7B模型首先在8K序列长度的1T词元上训练,然后在32K序列长度的200B词元上进一步训练。所有80B MoE模型都在32K序列长度的500B词元上训练。
表3 7B密集模型和80B MoE模型的RULER基准性能。注意,所有7B模型首先在8K序列长度的1T词元上训练,然后在32K序列长度的200B词元上进一步训练。所有80B MoE模型都在32K序列长度的500B词元上训练。

如表3所示,HySparse在长上下文任务中始终保持强大的准确性。
* 7B密集模型: 在16k和32k上下文长度下,HySparse的总分均优于两个基线,尤其在较难的多键/值和推理密集型子集(如CWE)上增益最为明显,显示出随上下文增长更好的鲁棒性。
* 80B MoE模型: 在激进的混合比例下,Hybrid SWA的性能急剧下降。相比之下,HySparse在16k上下文长度下与Full-Attn相当(90.6 vs 93.6),在32k时甚至显著超越了Full-Attn(87.4 vs 82.1),这得益于在MK3等困难任务上的巨大性能恢复和更强的稳定性。

总体而言,HySparse在不同设置下提供了与Full-Attn相当的长上下文能力,同时显著降低了计算和KV缓存大小。

消融研究

表4 7B实验中不同架构设计选择的消融研究。 图2 HySparse准确率与训练迭代次数的关系。

所有消融研究均在7B密集模型上进行,结果如表4和图2所示。

研究1:与SWA的层内混合
该研究探讨了在稀疏索引由“神谕”提供的情况下,稀疏注意力层是否仍需要一个额外的SWA分支。结果表明,移除SWA分支会导致明显的准确率下降。如表4所示,增加层内SWA分支在大多数基准上都带来了一致的提升(例如,DROP提升5.8分,GSM8K提升8.0分)。这表明,即使有高质量的稀疏选择,一个专用的滑动窗口路径对于建模短程连贯性和局部计算模式仍然很重要。

研究2:跨层KV缓存共享配置
该研究探讨了当每个稀疏层包含稀疏注意力(SA)和滑动窗口注意力(SWA)两个分支时,应如何应用KV缓存共享。实验比较了两种配置:(1) SA和SWA分支都共享前一个全注意力层的KV缓存;(2) 仅SA分支共享KV缓存,SWA分支维持独立的KV缓存。
结果如表4所示,为SA和SWA两个分支都共享KV缓存会严重降低准确性。相比之下,仅为SA分支共享KV缓存,同时为SWA保留独立的KV缓存,则能恢复并提升所有任务的性能(例如,MMLU提升5.6分,MMLU-Pro提升5.8分)。这表明SA可以安全地重用跨层KV缓存以节省GPU内存,而SWA则应保持自己的KV缓存以保留强大的局部信息。

A7 补充细节

我们最终能避免全注意力吗?

关于完全消除全注意力的讨论。我们的发现与高效注意力的更广泛趋势相符,包括混合注意力和稀疏注意力方法。一个反复出现的主题是,在实践中完全消除 $O(N^2)$ 风格的全注意力组件仍然具有挑战性:混合模型保留了明确的全注意力层,而像SeerAttention【17,Seerattention: Learning intrinsic sparse attention in your llms,2024,arXiv】和DSA【40,Deepseekv3. 2: Pushing the frontier of open large language models,2025a,arXiv】这样的稀疏注意力方法通常依赖于门控或索引机制,这些机制虽然以压缩形式存在,但仍在 $O(N^2)$ 复杂度下运行。在这种背景下,最重要的是昂贵的全局计算与廉价的局部或稀疏计算的比例,以及GPU内存的使用。我们的结果表明,HySparse通过神谕式词元选择和跨层KV共享,提供了一种有前途的方法来降低这一比例,同时保持长上下文建模能力。

HySparse在高效KV缓存卸载方面的潜力

HySparse对系统级优化的启示。HySparse还指向了一种直接的系统级策略,用于长上下文服务:将全注意力层的KV缓存卸载到外部存储器,并在计算前预取,而只在GPU上保留持久的、被选择的/稀疏的KV,供后续的稀疏注意力层使用。之前的工作如OmniKV【23,Omnikv: Dynamic context selection for efficient long-context llms,2025,ICLR】在后训练设置中探索了类似的方法。这种技术有潜力显著减少GPU上的KV缓存占用,从而实现更大的批量大小并提高整体推理效率。

A5 结论

本文介绍了混合稀疏注意力(HySparse),一种简单而有效的混合注意力架构,它在每个全注意力层之后交错插入多个稀疏注意力层。通过策略性地从前面的全注意力层派生重要词元选择和KV缓存,HySparse消除了对基于代理的词元选择的需求,并使稀疏层能够在没有额外内存开销的情况下运行。重要的是,HySparse允许在不损害模型能力的情况下大幅减少混合模型中全注意力层的数量。在未来的工作中,我们计划将HySparse扩展到更大的模型尺寸,并在更多词元上进行训练,以充分发掘其在高效和准确的长上下文建模方面的潜力。

引用的参考文献

编号 引用段落及内容描述 参考文献信息
1 HySparse的SWA分支采用类似gpt-oss的方法,引入了可学习的sink偏差。 Agarwal, S. et al. (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv preprint arXiv:2508.10925.
2 在GQA下,HySparse通过在查询组内聚合注意力分数来共享稀疏索引,以提高效率。 Ainslie, J. et al. (2023). Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245.
7 CLA是实现跨层KV缓存共享的代表性工作之一。 Brandon, W. et al. (2024). Reducing transformer key-value cache size with cross-layer attention. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
9 稀疏注意力通过计算部分重要词元的注意力来缓解二次方瓶颈,本文引用Sparse Transformers作为开创性工作。 Child, R. et al. (2019). Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509.
12 HySparse通过修改FlashAttention内核来高效提取块级注意力分数。 Dao, T. (2023). Flashattention-2: Faster attention with better parallelism and work partitioning. URL: https://arxiv.org/abs/2307.08691.
13 Jamba模型集成了Mamba模块与标准自注意力。 Dao, T. and Gu, A. (2024). Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060.
14 HySparse通过修改FlashAttention内核来高效提取块级注意力分数。 Dao, T. et al. (2022). Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359.
15 Kascade等工作观察到显著词元在相邻层间的稳定性,并利用此特性加速推理。 Deshmukh, D. et al. (2025). Kascade: A practical sparse attention method for long-context llm inference. arXiv preprint arXiv:2512.16391.
17 HySparse提取块级注意力分数的方法类似于SeerAttentionSeerAttention是可训练稀疏注意力的代表,通过自蒸馏对齐门控模块。讨论部分提到其仍依赖$O(N^2)$复杂度的机制。 Gao, Y. et al. (2024). Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410.13276.
18 HySparse提取块级注意力分数的方法类似于SeerAttentionRSeerAttentionR是可训练稀疏注意力的代表,通过自蒸馏对齐门控模块。 Gao, Y. et al. (2025). Seerattentionr: Sparse attention adaptation for long reasoning. arXiv preprint arXiv:2506.08889.
20 Gemma3Gemma 3n是采用混合注意力架构(SWA+全注意力)和跨层KV缓存共享的代表模型。 Gemma Team, et al. (2025). Gemma 3 technical report. arXiv preprint arXiv:2503.19786.
21 Jamba模型集成了Mamba模块与标准自注意力。 Gu, A. and Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752.
23 OmniKV等工作观察到显著词元在相邻层间的稳定性。讨论部分提到HySparse的卸载策略与OmniKV类似。 Hao, J. et al. (2025). Omnikv: Dynamic context selection for efficient long-context llms. In The Thirteenth International Conference on Learning Representations.
24 免训练稀疏注意力方法可能因训练-推理不匹配导致误差累积。 He, H. and Lab, T. M. (2025). Defeating nondeterminism in llm inference. Thinking Machines Lab: Connectionism. doi: 10.64434/tml.20250910.
28 免训练稀疏注意力方法可能因训练-推理不匹配导致误差累积。 Hu, J. et al. (2026). Lil: Less is less when applying post-training sparse-attention algorithms in long-decode stage. arXiv preprint arXiv:2601.03043.
31 Minference 1.0是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。 Jiang, H. et al. (2024). Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint arXiv:2407.02490.
34 Kimi Linear是采用混合注意力架构(集成Gated DeltaNet变体)的代表模型。 Kimi Team, et al. (2025b). Kimi linear: An expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692.
35 MiniMax-01是采用混合注意力架构(集成线性和softmax注意力)的代表模型。 Li, A. et al. (2025a). Minimax01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313.
36 苹果基础模型是实现跨层KV缓存共享的代表性工作之一。 Li, E. et al. (2025b). Apple intelligence foundation language models: Tech report 2025. arXiv preprint arXiv:2507.13575.
38 Jamba是采用混合注意力架构(集成Mamba模块)的代表模型。 Lieber, O. et al. (2024). Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887.
39 MiniCache观察到相邻层KV缓存的高度相似性,并提出跨层压缩方法。 Liu, A. et al. (2024). Minicache: Kv cache compression in depth dimension for large language models. Advances in Neural Information Processing Systems, 37:139997–140031.
40 Deepseek-v3.2 (DSA) 是可训练稀疏注意力的代表,通过自蒸馏对齐门控模块。讨论部分提到其仍依赖$O(N^2)$复杂度的机制。 Liu, A. et al. (2025a). Deepseekv3. 2: Pushing the frontier of open large language models. arXiv preprint arXiv:2512.02556.
42 Moba是可训练稀疏注意力的代表,通过学习选择信息丰富的词元。 Lu, E. et al. (2025). Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189.
43 MiniCPM是可训练稀疏注意力的代表,通过学习选择信息丰富的词元。 MiniCPM Team. (2025). Minicpm4: Ultra-efficient llms on end devices.
45 SwiftKV通过蒸馏使标准模型支持跨层KV缓存共享。 Qiao, A. et al. (2025). Swiftkv: Fast prefill-optimized inference with knowledge-preserving model transformation. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing.
46 HySparse的稀疏层输出融合使用了Gated Attention中的sigmoid门。MoE设置中的Full-Attn模型也使用门控注意力以稳定训练。 Qiu, Z. et al. (2025). Gated attention for large language models: Non-linearity, sparsity, and attention-sinkfree. arXiv preprint arXiv:2505.06708.
47 Qwen3-Next是采用混合注意力架构(集成Gated DeltaNet变体)的代表模型。 Qwen Team. (2025). Qwen3-next: Towards ultimate training & inference efficiency. Qwen Team Blog. URL: https://qwen.ai/blog?id=4074cca80393150c248e508aa62983f9cb7d27cd.
50 HySparse通过修改FlashAttention内核来高效提取块级注意力分数。 Shah, J. et al. (2024). Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658–68685.
51 YOCO是实现跨层KV缓存共享的代表性工作之一。 Sun, Y. et al. (2024). You only cache once: Decoder-decoder architectures for language models. Advances in Neural Information Processing Systems, 37:7339–7361.
53 Quest是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。 Tang, J. et al. (2024). Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774.
55 Attention Sinks是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。 Xiao, G. et al. (2023). Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453.
56 Duoattention是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。MiMo-V2-Flash是采用混合注意力架构(SWA+全注意力)的代表模型。 Xiao, G. et al. (2024). Duoattention: Efficient long-context llm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10819. / Xiaomi Team, B. et al. (2026). Mimo-v2-flash technical report. arXiv preprint arXiv:2601.02780.
57 Qwen3-NextKimi Linear等混合模型集成了Gated DeltaNet或其变体。 Yang, S. et al. (2024b). Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464.
58 Lserve是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。 Yang, S. et al. (2025b). Lserve: Efficient long-sequence llm serving with unified sparse attention. arXiv preprint arXiv:2502.14866.
59 Native Sparse Attention (NSA) 是可训练稀疏注意力的代表,通过端到端稀疏预训练学习选择模块。 Yuan, J. et al. (2025). Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502.11089.
60 Delta等工作观察到显著词元在相邻层间的稳定性,并利用此特性加速推理。 Zarch, H. E. et al. (2025). Delta: Dynamic layer-aware token attention for efficient long-context reasoning. arXiv preprint arXiv:2510.09883.
61 H2O是免训练稀疏注意力方法的例子,依赖启发式规则选择重要词元。 Zhang, Z. et al. (2023). H2o: Heavy-hitter oracle for efficient generative inference of large language models. In Thirty-seventh Conference on Neural Information Processing Systems.
62 InfLLMv2是可训练稀疏注意力的代表,通过学习选择信息丰富的词元。 Zhao, W. et al. (2025). Infllmv2: Dense-sparse switchable attention for seamless short-to-long adaptation. arXiv preprint arXiv:2509.24663.