原生稀疏注意力：硬件对齐且原生可训练的稀疏注意力

Jingyang Yuan∗1,2, Huazuo Gao1, Damai Dai1, Junyu Luo2, Liang Zhao1, Zhengyan Zhang1, Zhenda Xie1,
Y. X. Wei1, Lean Wang1, Zhiping Xiao3, Yuqing Wang1, Chong Ruan1, Ming Zhang2, Wenfeng Liang1,
Wangding Zeng1
1DeepSeek-AI
2Key Laboratory for Multimedia Information Processing, Peking University, PKU-Anker LLM Lab
3University of Washington
{yuanjy, mzhang_cs}@pku.edu.cn, {zengwangding, wenfeng.liang}@deepseek.com

Paper

Sourcecode

摘要

长上下文建模对于下一代语言模型至关重要，但标准注意力机制的高计算成本带来了巨大的计算挑战。稀疏注意力为提高效率同时保持模型能力提供了一个有希望的方向。我们提出了 NSA，一种本机可训练的稀疏注意力机制，它将算法创新与硬件对齐的优化相结合，以实现高效的长上下文建模。NSA 采用动态分层稀疏策略，将粗粒度标记压缩与细粒度标记选择相结合，以保持全局上下文感知和局部精度。我们的方法通过两项关键创新推进了稀疏注意力设计：（1）我们通过算术强度平衡的算法设计实现了显着的加速，并针对现代硬件进行了实现优化。（2）我们实现端到端训练，在不牺牲模型性能的情况下减少了预训练计算。如图 1 所示，实验表明，使用 NSA 预训练的模型在一般基准、长上下文任务和基于指令的推理中保持或超过了全注意力模型。同时，NSA 在 64k 长度序列的解码、前向传播和后向传播过程中实现了比全注意力机制显著的加速，验证了其在整个模型生命周期中的效率。

1.引言

研究界越来越认识到长上下文建模是下一代大型语言模型的关键功能，其驱动力来自各种现实世界的应用，包括深度推理（DeepSeek-AI，2025 年；Zelikman 等人，2022 年）、存储库级代码生成（Zhang 等人，2023a；Zhang 等人）和多轮自主代理系统（Park 等人，2023 年）。最近的突破包括 OpenAI 的 o 系列模型、DeepSeek-R1（DeepSeek-AI，2025 年）和 Gemini 1.5 Pro（Google 等人，2024 年），使模型能够处理整个代码库、冗长的文档，在数千个 token 上保持连贯的多轮对话，并跨远程依赖关系执行复杂的推理。然而，随着序列长度的增加，普通注意力机制（Vaswani 等人，2017）的高复杂性（Za-heer 等人，2020）成为关键的延迟瓶颈。理论估计表明，在解码 64k 长度的上下文时，使用 softmax 架构的注意力计算占总延迟的 70-80%，这凸显了对更高效注意力机制的迫切需求。

图| 全注意力模型和我们的 NSA 之间的性能和效率比较。左图：尽管 NSA 很稀疏，但在一般基准、长上下文任务和推理评估中，NSA 的平均表现优于全注意力基线。右图：对于 64k 长度的序列处理，NSA 在所有阶段（解码、前向传播和后向传播）都比全注意力模型实现了显着的计算加速。

高效的长上下文建模的一般方法是利用 softmax 注意力机制固有的稀疏性 (Ge et al., 2023; Jiang et al., 2023)，其中选择性地计算关键查询-密钥对可以显着减少计算开销，同时保持性能。最近的进展通过多种策略证明了这种潜力：KV 缓存逐出方法 (Li et al., 2024; Zhang et al., 2023b; Zhou et al., 2024)、块式 KV 缓存选择方法 (Tang et al., 2024; Xiao et al., 2024) 以及基于采样、聚类或散列的选择方法 (Chen et al., 2024; Desai et al., 2024; Liu et al., 2024)。尽管这些策略很有前途，但现有的稀疏注意力方法在实际部署中往往达不到要求。许多方法都无法实现与理论收益相当的加速；此外，大多数方法主要侧重于推理阶段，缺乏有效的训练时间支持来充分利用注意力的稀疏模式。

为了解决这些限制，部署有效的稀疏注意力必须解决两个关键挑战：（1）与硬件对齐的推理加速：将理论上的计算减少转化为实际的速度提升需要在预填充和解码阶段进行硬件友好的算法设计，以缓解内存访问和硬件调度瓶颈；（2）训练感知算法设计：使用可训练的运算符实现端到端计算，以降低训练成本，同时保持模型性能。这些要求对于实际应用实现快速的长上下文推理或训练至关重要。从这两个方面考虑，现有方法仍然存在明显差距。

为了实现更有效、更高效的稀疏注意力，我们提出了 NSA，这是一种集成了分层 token 建模的本机可训练稀疏注意力架构。如图 2 所示，NSA 通过将键和值组织成时间块并通过三条注意路径处理它们来减少每个查询的计算：压缩的粗粒度 token、选择性保留的细粒度 token 和用于本地上下文信息的滑动窗口。然后我们实现专门的内核以最大限度地提高其实际效率。NSA 引入了与上述关键要求相对应的两项核心创新：（1）硬件对齐系统：优化块状稀疏注意力以提高 Tensor Core 利用率和内存访问，确保平衡的算术强度。（2）训练感知设计：通过高效的算法和后向运算符实现稳定的端到端训练。这种优化使 NSA 能够同时支持高效部署和端到端训练。

图2| NSA 架构概览。左图：该框架通过三个并行的注意力分支处理输入序列：对于给定的查询，前面的键和值被处理为针对粗粒度模式的压缩注意力、针对重要标记块的选定注意力和针对局部上下文的滑动注意力。右图：每个分支产生的不同注意力模式的可视化。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

我们通过对真实语言语料库进行全面实验来评估 NSA。在具有 260B 个 token 的 27B 参数 Transformer 主干上进行预训练，我们评估了 NSA 在一般语言评估、长上下文评估和思路推理评估中的表现。我们进一步比较了 A100 GPU 上的内核速度和优化的 Triton（Tillet 等人，2019）实现。实验结果表明，NSA 实现了与全注意力基线相当或更好的性能，同时优于现有的稀疏注意力方法。此外，与全注意力相比，NSA 在解码、前向和后向阶段都实现了显着的加速，并且对于较长的序列，加速比会增加。这些结果验证了我们的分层稀疏注意力设计有效地平衡了模型能力和计算效率。

2. 重新思考稀疏注意力方法

现代稀疏注意力方法在降低 Transformer 模型的理论计算复杂度方面取得了重大进展。然而，大多数方法主要在推理过程中应用稀疏性，同时保留预训练的全注意力主干，这可能会引入架构偏差，从而限制其充分利用稀疏注意力优势的能力。在介绍我们原生的稀疏架构之前，我们通过两个关键视角系统地分析了这些局限性。

2.1 有效推理的错觉

尽管实现了注意力计算的稀疏性，但许多方法未能实现推理延迟的相应减少，这主要是由于两个挑战：
阶段限制稀疏性
诸如 H2O（Zhang 等人，2023b）之类的方法在自回归解码期间应用稀疏性，而在预填充期间需要计算密集型预处理（例如注意图计算、索引构建）。相比之下，像 MInference（Jiang 等人，2024）这样的方法仅专注于预填充稀疏性。这些方法无法在所有推理阶段实现加速，因为至少一个阶段的计算成本与全注意力相当。阶段专业化降低了这些方法在以预填充为主的工作负载（如书籍摘要和代码完成）或以解码为主的工作负载（如长链思维（Wei 等人，2022）推理）中的加速能力。
与高级注意力架构不兼容
一些稀疏注意力方法无法适应现代解码高效架构，如多查询注意力 (MQA) (Shazeer，2019) 和分组查询注意力 (GQA) (Ainslie 等人，2023)，它们通过在多个查询头之间共享 KV 显著减少了解码过程中的内存访问瓶颈。例如，在 Quest (Tang 等人，2024) 等方法中，每个注意力头独立选择其 KV 缓存子集。虽然它在多头注意力 (MHA) 模型中表现出一致的计算稀疏性和内存访问稀疏性，但它在基于 GQA 等架构的模型中呈现出不同的场景，其中 KV 缓存的内存访问量对应于同一 GQA 组内所有查询头的选择的并集。这种架构特性意味着虽然这些方法可以减少计算操作，但所需的 KV 缓存内存访问仍然相对较高。这种限制迫使人们做出一个关键的选择：虽然一些稀疏注意力方法减少了计算量，但它们分散的内存访问模式与高级架构的高效内存访问设计相冲突。

这些限制的出现是因为许多现有的稀疏注意力方法专注于减少键值缓存或理论计算量，但难以在高级框架或后端实现显著的延迟减少。这促使我们开发结合高级架构和硬件高效实现的算法，以充分利用稀疏性来提高模型效率。

2.2. 可训练稀疏性的神话

我们对原生可训练稀疏注意力的追求源于对仅推理方法分析的两个关键见解：（1）性能下降：事后应用稀疏性会迫使模型偏离其预训练的优化轨迹。正如 Chen 等人（2024）所证明的那样，前 20% 的注意力只能覆盖总注意力得分的 70%，导致预训练模型中的检索头等结构在推理过程中容易受到修剪。（2）训练效率要求：高效处理长序列训练对于现代 LLM 开发至关重要。这包括对较长文档进行预训练以增强模型容量，以及后续的适应阶段，例如长上下文微调和强化学习。然而，现有的稀疏注意力方法主要针对推理，而训练中的计算挑战在很大程度上未得到解决。这种限制阻碍了通过高效训练开发更强大的长上下文模型。此外，调整现有稀疏注意力进行训练的努力也暴露出挑战：
不可训练组件
ClusterKV（Liu 等人，2024 年）（包括 k 均值聚类）和 MagicPIG（Chen 等人，2024 年）（包括基于 SimHash 的选择）等方法中的离散操作会在计算图中产生不连续性。这些不可训练的组件会阻止梯度流通过 token 选择过程，从而限制模型学习最佳稀疏模式的能力。
低效的反向传播
一些理论上可训练的稀疏注意力方法在实际训练中效率低下。HashAttention（Desai 等人，2024）等方法中使用的标记粒度选择策略导致在注意力计算期间需要从 KV 缓存中加载大量单个标记。这种不连续的内存访问阻碍了 FlashAttention 等快速注意力技术的有效适应，这些技术依赖于连续的内存访问和分块计算来实现高吞吐量。因此，实现被迫回退到低硬件利用率，从而显著降低训练效率。

2.3. 原生稀疏性是当务之急

这些推理效率和训练可行性方面的限制促使我们从根本上重新设计稀疏注意力机制。我们提出了 NSA，这是一个原生稀疏注意力框架，可同时满足计算效率和训练要求。在以下部分中，我们将详细介绍 NSA 的算法设计和运算符实现。

3.方法

我们的技术方法涵盖算法设计和内核优化。在以下小节中，我们首先介绍我们的方法背景。然后介绍 NSA 的总体框架，然后介绍其关键算法组件。最后，我们详细介绍了硬件优化的内核设计，以最大限度地提高实际效率。

3.1.背景

注意力机制在语言建模中被广泛使用，其中每个查询标记 $q_t$ 计算与所有前面的键 $k_{:t}$ 的相关性分数，以生成值的加权和 $v_{:t}$ 。正式来说，对于长度为 𝑡 的输入序列，注意力操作定义为：
$$o_t=Attn(q_t,k_{:t},v_{:t})(1)$$
其中 Attn 表示注意力函数：
$$Attn(q_t,k_{:t},v_{:t})=\sum _{i=1}^t\frac{{\alpha }_{t,i}{v}_i}{{\sum }_{j=1}^t{\alpha }_{t,j}},{\alpha }_{t,i}=e^{\frac{q_t^{\top }{k}_i}{\sqrt{d_k}}}(2)$$

其中，${\alpha }_{t,i}$ 表示 $q_t$ 和 $k_i$ 之间的注意力权重，$d_k$ 是键的特征维度。随着序列长度的增加，注意力计算在整体计算成本中变得越来越重要，这对长上下文处理提出了重大挑战。

算术强度是计算操作与内存访问的比率。它从本质上决定了硬件上的算法优化。每个 GPU 都有一个关键算术强度，由其峰值计算能力和内存带宽决定，以这两个硬件限制的比率计算。对于计算任务，高于此临界阈值的算术强度将成为计算限制（受 GPU FLOPS 限制），而低于此临界阈值的算术强度将成为内存限制（受内存带宽限制）。

具体来说，对于因果自注意力机制，在训练和预填充阶段，批量矩阵乘法和注意力计算表现出高算术强度，使得这些阶段在现代加速器上计算受限。相反，自回归解码会受到内存带宽的限制，因为它每次前向传递都会生成一个标记，同时需要加载整个键值缓存，从而导致算术强度较低。​​这导致了不同的优化目标——降低训练和预填充期间的计算成本，同时减少解码期间的内存访问。

3.2. 总体框架

为了充分利用具有自然稀疏模式的注意力机制的潜力，我们建议将公式 (1) 中的原始键值对 $k_{:t}$ , $v_{:t}$ 替换为给定每个查询 $q_t$ 的更紧凑、信息密集的表示键值对 𝐾 ̃𝑡,𝑉 ̃𝑡 集。具体来说，我们正式定义优化的注意力输出如下：
$${\tilde{K}}_t={\int }_K(q_t,k_{:t},v_{:t}),{\tilde{V}}_t={\int }_V(q_t,k_{:t},v_{:t})(3)$$
$$o_t^{\ast }=Attn(q_t,{\tilde{K}}_t,{\tilde{V}}_t)(4)$$
其中 ${\tilde{K}}_t$,${\tilde{V}}_t$ 是根据当前查询 $q_t$ 和上下文记忆 k ,v 动态构建的。我们可以设计各种映射策略来获得不同类别的${\tilde{K}}_t^c$,${\tilde{V}}_t^c$并将它们组合如下：
$$o_t^{\ast }=\sum _{c\in C}{g}_t^c\cdot Attn(q_t,{\tilde{K}}_t^c,{\tilde{V}}_t^c).(5)$$

如图 2 所示，NSA 有三种映射策略 C = {cmp, slc, win}，分别表示键和值的压缩、选择和滑动窗口。$g_t^c\in [0,1]$ 是对应策略 𝑐 的门控分数，通过 MLP 和 S 型激活从输入特征中得出。让 $N_t$ 表示重新映射的键/值的总数：
$$N_t=\sum _{c\in C}size[{\tilde{K}}_t^c].(6)$$
我们通过确保$N_t\ll t$ 来保持高稀疏率。

3.3.算法设计

在本小节中，我们介绍了重映射策略${\int }_K$和${\int }_V$的设计：标记压缩、标记选择和滑动窗口。

3.3.1.标记压缩

通过将连续的键或值块聚合到块级表示中，我们获得了捕获整个块信息的压缩键和值。正式地，压缩键表示定义为：
K ~ t c m p = ∫ K c m p ( K : t ) = { φ ( K i d + 1 : i d + l ) ∣ 1 ≤ i ≤ [ t − l d ] } . ( 7 ) \tilde{K}_t^{cmp}=\int_K^{cmp}(K_{:t})=\{ \varphi(K_{id+1}:id+l)|1 \leq i \leq [\frac{t-l}{d}] \}.(7) K~tcmp​=∫Kcmp​(K:t​)={φ(Kid+1​:id+l)∣1≤i≤[dt−l​]}.(7)
其中 𝑙 是块长度，𝑑 是相邻块之间的滑动步幅，$\phi$ 是具有块内位置编码的可学习 MLP，用于将块中的键映射到单个压缩键。
${\tilde{K}}_t^{cmp}\in R^{d_{k}x[\frac{t-l}{d}]}$是由压缩键组成的张量。通常，我们采用 𝑑 < 𝑙 来缓解信息碎片化。压缩值表示 ${\tilde{K}}_t^{cmp}$ 具有类似的公式。压缩表示可捕获更粗粒度的高级语义信息并减少注意力的计算负担。

3.3.2.标记选择

仅使用压缩的键和值可能会丢失重要的细粒度信息，这促使我们有选择地保留单个键和值。下面我们将描述我们高效的标记选择机制，该机制可以以较低的计算开销识别和保留最相关的标记。

块式选择
我们的选择策略在空间连续块中处理键和值序列，其动机有两个关键因素：硬件效率考虑和注意力分数的固有分布模式。块式选择对于在现代 GPU 上实现高效计算至关重要。这是因为与基于索引的随机读取相比，现代 GPU 架构在连续块访问方面表现出更高的吞吐量。此外，块式计算可以最佳地利用 Tensor Cores。这种架构特性已将块式内存访问和计算确立为高性能注意力实现的基本原则，例如 FlashAttention 的基于块的设计。块式选择遵循注意力分数的固有分布模式。 先前的研究（Jiang 等人，2024）表明注意力分数通常表现出空间连续性，这表明相邻的键往往具有相似的重要性水平。我们在第 6.2 节中的可视化也显示了这种空间连续模式。

为了实现逐块选择，我们首先将键、值序列划分为选择块。为了确定对注意力计算最重要的块，我们需要为每个块分配重要性分数。下面我们介绍计算这些块级重要性分数的方法。

重要性分数计算。计算块重要性分数可能会带来很大的开销。幸运的是，压缩标记的注意力计算会产生中间注意力分数，我们可以利用这些分数来得出选择块重要性分数，公式如下：
$$p_t^{cmp}=Softmax(q_t^T{\tilde{K}}_t^{cmp}),(8)$$
其中 $p_t^{cmp}\in R^{[\frac{t-l}{d}]}$ 是$q_t$和压缩密钥之间的注意力得分${\tilde{K}}_t^{cmp}$表示选择块大小。当压缩块和选择块共享相同的分块方案时，即 𝑙′ = 𝑙 = 𝑑，我们可以直接通过 $p_t^{slc}=p_t^{cmp}$ 获得选择块重要性得分 $p_t^{slc}$。对于分块方案不同的情况，我们根据选择块的空间关系推导出它们的 𝑡重要性得分。给定 𝑑 | 𝑙 和𝑑 | 𝑙′，我们有：
$$p_t^{slc}[j]=\sum _{m=0}^{\frac{l^{\prime }}{d}-1}\sum _{n=0}^{\frac{l}{d}-1}{p}_t^{cmp}[\frac{l^{\prime }}{d}j+m+n],(9)$$

其中 [·] 表示用于访问向量元素的索引运算符。对于采用 GQA 或 MQA 的模型，其中键值缓存在查询头之间共享，必须确保这些头之间的块选择一致，以最大限度地减少解码期间的 KV 缓存加载。组中各头之间的共享重要性分数正式定义为：
$$p_t^{sl{c}^{\prime }}=\sum _{h=1}^H{p}_t^{slc,(h)},(10)$$
其中上标中的 (h) 表示头部索引，𝐻 表示每组中查询头部的数量。此聚合可确保同一组内各头部之间的块选择一致。
Top-𝑛 块选择。获得选择块重要性分数后，我们保留按块重要性分数排序的 top-𝑛 稀疏块内的标记，公式如下：
$$I_t=i|rank(p_t^{sl{c}^{\prime }}[i]\le n)(11)$$
K ~ t s l c = C a t [ { K i l ′ + 1 : ( I + 1 ) l ′ ∣ i ∈ I t } ] , ( 12 ) \tilde{K}_t^{slc} = Cat[\{K_{il'+1:(I+1)l'}|i \in I_t\}],(12) K~tslc​=Cat[{Kil′+1:(I+1)l′​∣i∈It​}],(12)
其中 rank(·) 表示按降序排列的位置，rank = 1 对应最高分数，I𝑡 是所选块的索引集，Cat 表示连接操作。${\tilde{K}}_t^{slc}\in R^{d_k\times n{l}^{\prime }}$是由压缩键组成的张量。类似的公式 𝑡 适用于细粒度值${\tilde{V}}_t^{slc}$。然后，所选的键和值参与 𝑡 注意力计算，$q_t$定义如公式 (5)。

3.3.3. 滑动窗口

在注意力机制中，局部模式通常适应得更快，并且可以主导学习过程，从而可能阻止模型有效地从压缩和选择标记中学习。为了解决这个问题，我们引入了一个专用的滑动窗口分支，该分支明确处理局部上下文，允许其他分支（压缩和选择）专注于学习各自的特征，而不会被局部模式所捷径。具体来说，我们在窗口𝑤中维护最近的标记${\tilde{K}}_t^{win}=K_{t-w:t},{\tilde{V}}_t^{win}=V_{t-w:t}$，并将注意力计算不同的信息源（压缩标记和选定标记，滑动窗口）隔离到单独的分支中。然后通过学习门控机制聚合这些分支输出。为了进一步防止在边际计算开销下跨注意力分支进行捷径学习，我们为三个分支提供了独立的键和值。这种架构设计通过防止局部和长距离模式识别之间的梯度干扰实现了稳定的学习，同时引入了最小的开销。
在获得所有三类键和值（${\tilde{K}}_t^{cmp},{\tilde{V}}_t^{cmp};{\tilde{K}}_t^{slc},{\tilde{V}}_t^{slc};和{\tilde{K}}_t^{win},{\tilde{V}}_t^{win}$）后，我们按照公式 (5) 计算最终的注意力输出。结合上面描述的压缩、选择和滑动窗口机制，这构成了 NSA 的完整算法框架。

3.4.核心设计

为了在训练和预填充期间实现 FlashAttention 级别的加速，我们在 Triton 上实现了硬件对齐的稀疏注意内核。鉴于 MHA 占用大量内存且解码效率低下，我们专注于遵循当前最先进的 LLM 的共享 KV 缓存架构，如 GQA 和 MQA。虽然压缩和滑动窗口注意计算与现有的 FlashAttention-2 内核很容易兼容，但我们引入了专门用于稀疏选择注意的内核设计。如果我们遵循 FlashAttention 将时间连续的查询块加载到 SRAM 中的策略，则会导致内存访问效率低下，因为块内的查询可能需要不相交的 KV 块。为了解决这个问题，我们的关键优化在于不同的查询分组策略：对于查询序列上的每个位置，我们将 GQA 组内的所有查询头（它们共享相同的稀疏 KV 块）加载到 SRAM 中。图 3 说明了我们的前向传递实现。提出的内核架构具有以下主要特点：

1. 以群组为中心的数据加载。 对于每个内循环，加载位于位置 𝑡 的组中所有头的查询 $Q\in R^{[h,d_k]}$ 及其共享的稀疏键/值块索引$I_t$ 。
2. 共享键值对获取。 在内层循环中，将索引为$I_t$的连续键值块按照$K\in R^{[B_k,d_k]},V\in R^{[B_k,d_v]}$的顺序加载到 SRAM 中，以最小化内存加载，其中$B_k$是满足$B_k|l^{\prime }$的内核块大小。
3. **网格上的外循环。 **由于内循环长度（与所选块数𝑛成比例）对于不同的查询块几乎保持相同，我们将查询/输出循环放在 Triton 的网格调度程序中，以简化和优化内核。

该设计通过（1）通过分组共享消除冗余的键值传输，以及（2）在 GPU 流多处理器之间平衡计算工作负载，实现了近乎最佳的算术强度。

图 3 | NSA 的内核设计。内核按 GQA 组加载查询（网格循环），获取相应的稀疏 KV 块（内部循环），并在 SRAM 上执行注意力计算。绿色块表示 SRAM 上的数据，蓝色表示 HBM 上的数据。

4.实验

我们从三个角度评估 NSA：（1）一般基准性能、（2）长上下文基准性能和（3）思路链推理性能，并与全注意力基线和最先进的稀疏注意力方法进行比较。我们将稀疏计算范式的效率分析推迟到第 5 节，其中我们将详细讨论训练和推理速度。

4.1. 预训练设置

按照最先进的 LLM 中的常见做法，我们的实验采用了结合分组查询注意力 (GQA) 和混合专家 (MoE) 的主干，具有 27B 总参数和 3B 活动参数。该模型由 30 层组成，隐藏维度为 2560。对于 GQA，我们将组数设置为 4，总共有 64 个注意力头。对于每个头，查询、键和值的隐藏维度分别配置为 𝑑𝑞 = 𝑑𝑘 = 192 和 𝑑𝑣 = 128。对于 MoE，我们采用 DeepSeekMoE（Dai et al.，2024；DeepSeek-AI，2024）结构，其中有 72 个路由专家和 2 个共享专家，并将 top-k 专家设置为 6。为了确保训练稳定性，第一层的 MoE 被 SwiGLU 形式的 MLP 取代。

图 4 | 27B 参数模型上 Full Attention 与我们的 NSA 之间的预训练损失比较。两种模型都表现出稳定的收敛性，NSA 实现了较低的损失值。

表 1 | 在一般基准、知识（MMLU、MMLU-PRO、CMMLU）、推理（BBH、GSM8K、MATH、DROP）和编码（MBPP、HumanEval）任务上，全注意力基线和 NSA 之间的预训练性能比较。尽管稀疏性很高，但 NSA 在大多数基准上都实现了优异的平均性能。

所提出的架构在计算成本和模型性能之间实现了有效的权衡。对于 NSA，我们设置压缩块大小 𝑙 = 32、滑动步幅 𝑑 = 16、选定块大小 𝑙′ = 64、选定块数 𝑛 = 16（包括固定激活 1 个初始块和 2 个局部块）和滑动窗口大小 𝑤 = 512。全注意力模型和稀疏注意力模型均在 8k 长度文本的 270B 个标记上进行预训练，然后使用 YaRN（Peng et al.，2024）在 32k 长度文本上进行持续训练和监督微调，以实现长上下文自适应。两种模型都经过训练以完全收敛，以确保公平比较。如图 4 所示，我们的 NSA 和全注意力基线的预训练损失曲线显示出稳定而平滑的下降，NSA 的表现始终优于全注意力模型。

4.2. 基线方法

除了与 Full Attention 进行比较外，我们还评估了几种最先进的推理阶段稀疏注意方法：H2O（Zhang 等人，2023b）、infLLM（Xiao 等人，2024）、Quest（Tang 等人，2024）和 Exact-Top，后者首先计算完整注意分数并选择与每个查询相对应的 top-𝑛 分数键，然后计算这些位置的注意力。这些方法涵盖了各种稀疏注意范式，包括 KV 缓存驱逐、查询感知选择和精确 top-𝑛 稀疏选择。

表 2 | 我们的 NSA 与 LongBench 上的基线之间的性能比较，包括单文档 QA、多文档 QA、合成和代码任务类别中的子集。NSA 的表现优于包括全注意力在内的大多数基线。

对于一般评估，大多数样本的长度都在稀疏注意力基线的局部上下文窗口内，这些方法实际上等同于全注意力。因此，我们仅展示此设置下 NSA 和全注意力基线之间的比较结果。在长上下文评估中，我们对所有基线方法进行比较，所有稀疏注意力方法的稀疏度设置为相同，以确保公平比较。对于需要长文本监督微调的思路链推理评估，我们将比较限制在全注意力上，因为稀疏注意力基线不支持训练。

4.3. 性能比较

总体评估。 我们在涵盖知识、推理和编码能力的一系列基准上评估了预训练的 NSA 和 Full Attention 基线，包括 MMLU（Hendrycks 等人，2020 年）、MMLU-PRO（Wang 等人，2024 年）、CMMLU（Li 等人，2023 年）、BBH（Suzgun 等人，2022 年）、GSM8K（Cobbe 等人，2021 年）、MATH（Hendrycks 等人，2020 年）、DROP（Dua 等人，2019 年）、MBPP（Austin 等人，2021 年）和 HumanEval（Chen 等人，2021 年）。结果如表 1 所示。尽管 NSA 比较稀疏，但它实现了卓越的整体性能，在 9 个指标中的 7 个上都优于包括 Full Attention 在内的所有基线。这表明，尽管 NSA 可能无法充分利用其在较短序列上的效率优势，但它表现出了强劲的性能。值得注意的是，NSA 在推理相关基准测试中表现出显著的提升（DROP：+0.042，GSM8K：+0.034），这表明我们的预训练有助于模型开发专门的注意力机制。这种稀疏注意力预训练机制迫使模型专注于最重要的信息，通过过滤掉不相关的注意力路径中的噪音，有可能提高性能。在不同评估中的一致表现也证明了 NSA 作为通用架构的稳健性。

长上下文评估。 如图 5 所示，NSA 在 64k 上下文大海捞针 (Kamradt, 2023) 测试中在所有位置上都实现了完美的检索准确率。这种性能源于我们的分层稀疏注意力设计，它结合了压缩标记以实现高效的全局上下文扫描，以及选择标记以实现精确的本地信息检索。粗粒度压缩以较低的计算成本识别相关的上下文块，而对选定标记的标记级注意力可确保关键细粒度信息的保存。这种设计使 NSA 能够同时保持全局意识和局部精度。

我们还在 LongBench (Bai et al., 2023) 上对 NSA 进行了评估，并与最先进的稀疏注意力方法和全注意力基线进行了比较。为了确保一致的稀疏性，我们将所有稀疏注意力基线中每个查询激活的 token 设置为 2560 个 token，这对应于 NSA 在处理 32k 序列长度时激活的 token 的平均数量。按照 Stream-LLM (Xiao et al., 2023)，此 token 预算包括前 128 个 token 和 512 个本地 token。我们从 LongBench 中排除了某些子集，因为它们在所有模型中的得分都较低，可能无法提供有意义的比较。如表 2 所示，NSA 获得了最高平均分数 0.469，优于所有基线（比全注意力高出 0.032，比 Exact-Top 高出 0.046）。这一改进源于两个关键创新：（1）我们原生的稀疏注意力设计，能够在预训练期间对稀疏模式进行端到端优化，促进稀疏注意力模块与其他模型组件之间的同步适应；（2）分层稀疏注意力机制实现了局部和全局信息处理之间的平衡。

图 5 | 64k 上下文长度的上下文位置上的大海捞针检索准确率。NSA 通过其分层稀疏注意力设计实现了完美的准确率。

值得注意的是，NSA 在需要对长上下文进行复杂推理的任务上表现出色，在多跳 QA 任务（HPQ 和 2Wiki）上比全注意力实现了 +0.087 和 +0.051 的改进，在代码理解方面超过了基线的性能（LCC：+0.069），并且在段落检索方面优于其他方法（PassR-en：+0.075）。这些结果验证了 NSA 处理各种长上下文挑战的能力，其原生预训练的稀疏注意力在学习任务最优模式方面提供了额外的好处。

思路链推理评估。为了评估 NSA 与先进的下游训练范式的兼容性，我们研究了其通过后训练获得思路链数学推理能力的能力。鉴于强化学习对小规模模型的有效性有限，我们采用 DeepSeek-R1 的知识提炼，使用 10B 个 32k 长度的数学推理轨迹标记进行监督微调 (SFT)。这产生了两个可比较的模型：Full Attention-R（Full Attention 基线）和 NSA-R（我们的稀疏变体）。我们在具有挑战性的美国数学邀请考试 (AIME 24) 基准上评估了这两个模型。我们使用 0.7 的采样温度和 0.95 的 top-𝑝 值为每个问题生成 16 个答案并获得平均分数。为了验证推理深度的影响，我们针对两个生成上下文限制（8k 和 16k 个标记）进行了实验，以测量扩展推理链是否能提高准确性。附录 A 中提供了模型预测的示例比较。

表 3 | 监督微调后的基于 AIME 指令的评估。我们的 NSA-R 在 8k 和 16k 序列长度上都表现出比 Full Attention-R 更好的性能

图 6 | 基于 Triton 的 NSA 内核与基于 Triton 的 FlashAttention-2 内核的比较。我们的实现显著降低了所有上下文长度的延迟，并且随着输入长度的增加，改进变得更加明显。

如表 3 所示，NSA-R 在 8k 上下文设置下 (+0.075) 的准确率明显高于 Full Attention-R，这一优势在 16k 上下文设置下 (+0.054) 依然存在。这些结果验证了原生稀疏注意力的两个关键优势：(1) 预训练的稀疏注意力模式能够有效捕获对复杂数学推导至关重要的长距离逻辑依赖关系；(2) 我们架构的硬件对齐设计保持了足够的上下文密度，以支持不断增长的推理深度，而不会发生灾难性遗忘。在上下文长度上一致的优异表现证实了稀疏注意力在原生集成到训练管道中时可用于高级推理任务的可行性。

5.效率分析

我们在 8-GPU A100 系统上评估了 NSA 相对于全注意力的计算效率。在效率分析中，我们还将模型配置为 GQA 组 𝑔 = 4、每组头 h = 16、查询/键维度 $d_k=192$和值维度$d_v=128$。按照第 4 节中的相同设置，我们设置 NSA 压缩块大小 𝑙 = 32、滑动步幅 𝑑 = 16、选定块大小 𝑙′ = 64、选定块数 𝑛 = 16 和滑动窗口大小 𝑤 = 512。

5.1. 训练速度

我们将基于 Triton 的 NSA 注意力和全注意力实现与基于 Triton 的 FlashAttention-2 进行比较，以确保在同一后端进行公平的速度比较。如图 6 所示，随着上下文长度的增加，我们的 NSA 实现了逐渐增大的加速，在 64k 上下文长度下，前向加速高达 9.0 倍，后向加速高达 6.0 倍。值得注意的是，序列越长，速度优势就越明显。这种加速源于我们与硬件对齐的算法设计，旨在最大限度地提高稀疏注意力架构的效率：（1）块式内存访问模式通过合并负载最大限度地提高 Tensor Core 利用率，（2）内核中精细的循环调度消除了冗余的 KV 传输。

表 4 | 解码期间每个注意操作的内存访问量（以等效令牌数表示）。由于解码的算术强度低且受内存限制，预期加速比与内存访问量大致呈线性关系。

5.2. 解码速度

Attention 的解码速度主要由内存访问瓶颈决定，这与 KV 缓存加载量密切相关。在每个解码步骤中，我们的 NSA 只需要加载最多$[\frac{s-l}{d}]$个压​​缩标记、𝑛𝑙′ 个选定标记和 𝑤 个邻居标记，𝑑 其中 𝑠 是缓存序列长度。如表 4 所示，随着解码长度的增加，我们的方法延迟显著减少，在 64k 上下文长度下实现了高达 11.6 倍的加速。内存访问效率的这种优势也会随着序列的延长而放大。

6.讨论

在本节中，我们回顾了 NSA 的发展过程，并讨论了从探索不同的稀疏注意力策略中获得的关键见解。虽然我们的方法显示出了有希望的结果，但了解替代策略所遇到的挑战并分析注意力模式为未来的研究方向提供了宝贵的背景。我们首先研究了激发我们设计选择的替代 token 选择策略所面临的挑战，然后进行可视化，以深入了解注意力分布模式。

6.1. 替代标记选择策略的挑战

在设计 NSA 之前，我们探索了将现有的稀疏注意力方法应用于训练阶段。然而，这些尝试遇到了各种挑战，促使我们设计了一种不同的稀疏注意力架构：
基于键聚类的策略。我们研究了基于聚类的策略，如 ClusterKV（Liu 等人，2024）。这些方法将来自同一集群的键和值存储在连续的内存区域中。虽然理论上可用于训练和推理，但它们面临三个重大挑战：（1）动态聚类机制引入的非平凡计算开销；（2）集群间不平衡加剧了运算符优化困难，尤其是在混合专家 (MoE) 系统中，其中倾斜的专家并行 (EP) 组执行时间导致持续的负载不平衡；（3）由于需要强制定期重新聚类和块顺序训练协议而产生的实施限制。这些综合因素造成了巨大的瓶颈，严重限制了它们在实际部署中的有效性。

图 7 | 比较具有全注意力机制和不同 token 选择策略的 3B 参数模型的训练损失。我们的 NSA 实现了更好的性能。

图 8 | 全注意力转换器上的注意力图可视化。浅色区域表示注意力值较高。如图所示，注意力得分呈现块状聚类分布。

其他分块选择策略。我们还考虑了不同于 NSA 的分块键、值选择策略，例如 Quest（Tang 等人，2024 年）和 InfLLM（Xiao 等人，2024 年）。这些方法依赖于计算每个块的重要性分数，并根据它们与$q_t$ 的相似性选择前𝑛 块。然而，现有的方法面临两个关键问题：（1）由于选择操作不可微，基于神经网络的重要性分数计算依赖于辅助损失，这会增加操作员开销并经常降低模型性能；（2）启发式无参数重要性分数计算策略的召回率较低，导致性能不佳。我们在具有相似架构的 3B 参数模型上评估这两种方法，并将它们的损失曲线与 NSA 和全注意力进行比较。对于基于辅助损失的选择方法，我们为每个块引入额外的查询和代表性键来估计块重要性分数。这些分数由每个块内原始查询和键之间的平均注意力分数监督。对于启发式无参数选择方法，我们遵循 Quest 的策略，使用查询之间的乘积和关键块的坐标最小最大值来实现直接选择，而无需引入其他参数。我们还探索了一种冷启动训练方法，其中在过渡到启发式块选择之前，在前 1000 步中应用全注意力。如图 7 所示，两种方法都表现出较低的损失。

6.2.可视化

为了探索 Transformer 注意力分布中的潜在模式并寻求设计灵感，我们在图 8 中可视化了预训练的 27B 全注意力模型的注意力图。可视化揭示了有趣的模式，其中注意力得分往往表现出块状聚类特征，附近的键通常显示出相似的注意力得分。这一观察启发了我们设计 NSA，表明基于空间连续性选择关键块可能是一种有前途的方法。块状聚类现象表明序列中相邻的标记可能与查询标记共享某些语义关系，尽管这些关系的确切性质需要进一步研究。这一观察促使我们探索一种对连续标记块而不是单个标记进行操作的稀疏注意力机制，旨在提高计算效率并保留高注意力模式。

7.相关工作

我们回顾了现有的通过稀疏注意力提高注意力计算效率的方法。这些方法可以根据其核心策略大致分为三类：（1）固定稀疏模式，（2）动态标记修剪，以及（3）查询感知选择。我们介绍了每个类别中的几项代表性作品。

7.1.固定稀疏模式

滑动窗口是一种常用的方法，它允许查询仅在固定窗口内计算注意力。StreamingLLM（Xiao 等人，2023 年）通过维护上下文的两个关键部分来解决处理长文本流的挑战：注意力接收器（早期标记）和本地上下文窗口。虽然这些方法有效地降低了内存和计算成本，但它们忽略上下文的僵化模式限制了它们在需要完全上下文理解的任务上的性能。

7.2.动态标记修剪

H2O（Zhang 等人，2023b）实现了一种自适应方法来减少解码过程中的 KV 缓存内存使用量。该方法根据注意力得分，根据最近的效用动态地逐出被认为对未来预测不太重要的标记。SnapKV（Li 等人，2024）还引入了一种标记修剪策略，通过有选择地仅保留最重要的特征来减少 KV 缓存，从而实现高效的内存使用。SnapKV 在预填充期间通过注意力权重分析和投票来识别重要特征，然后通过将选定的压缩特征与最近的上下文相结合来更新 KV 缓存，以保持及时一致性。

7.3.查询感知选择

Quest (Tang et al., 2024) 采用分块选择策略，其中每个块的重要性通过查询和关键块的坐标最小值-最大值之间的乘积来估计。结果分数有助于选择 top-𝑛 重要键值块以引起注意。InfLLM (Xiao et al., 2024) 通过维护注意力接收器、局部上下文和可检索块将固定模式与检索相结合。该方法从每个块中选择代表性键来估计块重要性。HashAttention (Desai et al., 2024) 通过使用学习函数将查询和键映射到汉明空间，将关键标记识别公式化为推荐问题。ClusterKV (Liu et al., 2024) 通过首先对键进行聚类，然后根据查询-聚类相似性选择最相关的聚类进行注意力计算来实现稀疏性。

8. 结论

我们提出了 NSA，这是一种与硬件对齐的稀疏注意力架构，用于高效的长上下文建模。通过在可训练架构中集成分层标记压缩和逐块标记选择，我们的架构实现了加速训练和推理，同时保持了全注意力性能。NSA 通过展示与全注意力基线相匹配的一般基准性能、在长上下文评估中超越建模能力和增强的推理能力，推动了最先进的技术，所有这些都伴随着计算延迟的可衡量减少和实现显着的加速。

参考文献

J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint
arXiv:2305.13245, 2023.
J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732,
2021.
Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint
arXiv:2308.14508, 2023.
M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. D. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021.
Z. Chen, R. Sadhukhan, Z. Ye, Y. Zhou, J. Zhang, N. Nolte, Y. Tian, M. Douze, L. Bottou, Z. Jia, et al. Magicpig: Lsh sampling for efficient llm generation. arXiv preprint arXiv:2410.16179,
2024.
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. Training verifiers to solve math word problems, 2021. URL https://arxiv.
org/abs/2110.14168, 2021.
D. Dai, C. Deng, C. Zhao, R. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y. Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint
arXiv:2401.06066, 2024.
DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language
model. 2024. URL https://arxiv.org/abs/2405.04434.
DeepSeek-AI. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,
2025. URL https://arxiv.org/abs/2501.12948.
A. Desai, S. Yang, A. Cuadron, A. Klimovic, M. Zaharia, J. E. Gonzalez, and I. Stoica. Hashatten-
tion: Semantic sparsity for faster inference. arXiv preprint arXiv:2412.14468, 2024.
D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint
arXiv:1903.00161, 2019.
S. Ge, Y. Zhang, L. Liu, M. Zhang, J. Han, and J. Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801, 2023.
G. T. Google, P. Georgiev, V. I. Lei, R. Burnell, L. Bai, A. Gulati, G. Tanzer, D. Vincent, Z. Pan, S. Wang, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.
D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
H. Jiang, Q. Wu, C.-Y. Lin, Y. Yang, and L. Qiu. Llmlingua: Compressing prompts for accelerated inference of large language models. arXiv preprint arXiv:2310.05736, 2023.
H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C.-Y. Lin, et al. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint arXiv:2407.02490, 2024.
G. Kamradt. LLMTest NeedleInAHaystack. GitHub repository, 2023. URL https://github.c om/gkamradt/LLMTest_NeedleInAHaystack. Accessed: [Insert Access Date Here].
H. Li, Y. Zhang, F. Koto, Y. Yang, H. Zhao, Y. Gong, N. Duan, and T. Baldwin. Cmmlu: Measuring massive multitask language understanding in chinese. arXiv preprint arXiv:2306.09212, 2023.
Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469, 2024.
G. Liu, C. Li, J. Zhao, C. Zhang, and M. Guo. Clusterkv: Manipulating llm kv cache in semantic space for recallable compression. arXiv preprint arXiv:2412.03213, 2024.
J. S. Park, J. C. O’Brien, C. J. Cai, M. R. Morris, P. Liang, and M. S. Bernstein. Generative agents: Interactive simulacra of human behavior. In S. Follmer, J. Han, J. Steimle, and N. H. Riche, editors, Proceedings of the 36th Annual ACM Symposium on User Interface Software and
Technology, UIST 2023, San Francisco, CA, USA, 29 October 2023– 1 November 2023, pages
2:1–2:22. ACM, 2023.
B. Peng, J. Quesnelle, H. Fan, and E. Shippole. Yarn: Efficient context window extension of large language models. In ICLR. OpenReview.net, 2024.
N. Shazeer. Fast transformer decoding: One write-head is all you need. CoRR, abs/1911.02150, 2019.
M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. V. Le, E. H. Chi, D. Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022.
J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han. Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774, 2024.
P. Tillet, H.-T. Kung, and D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop
on Machine Learning and Programming Languages, pages 10–19, 2019.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems,
2017.
Y. Wang, X. Ma, G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj, X. He, Z. Jiang, et al. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. arXiv preprint arXiv:2406.01574, 2024.
J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. Chain- of-thought prompting elicits reasoning in large language models. Advances in neural
information processing systems, 35:24824–24837, 2022.
C. Xiao, P. Zhang, X. Han, G. Xiao, Y. Lin, Z. Zhang, Z. Liu, and M. Sun. Infllm: Training-free long-context extrapolation for llms with an efficient context memory. In The Thirty-eighth
Annual Conference on Neural Information Processing Systems, 2024.
G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023.
M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. Big bird: Transformers for longer sequences. Advances in neural
information processing systems, 33:17283–17297, 2020.
E. Zelikman, Y. Wu, J. Mu, and N. D. Goodman. Star: Bootstrapping reasoning with reasoning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information
Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 – December
9, 2022, 2022.
F. Zhang, B. Chen, Y. Zhang, J. Keung, J. Liu, D. Zan, Y. Mao, J. Lou, and W. Chen. Repocoder: Repository-level code completion through iterative retrieval and generation. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in
Natural Language Processing, EMNLP 2023, Singapore, December 6–10, 2023, pages 2471–
2484. Association for Computational Linguistics, 2023a.
K. Zhang, J. Li, G. Li, X. Shi, and Z. Jin. Codeagent: Enhancing code generation with tool- integrated agent systems for real-world repo-level coding challenges. In L. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for
Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August
11–16, 2024, pages 13643–13658.
Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710, 2023b.
Z. Zhou, C. Li, X. Chen, S. Wang, Y. Chao, Z. Li, H. Wang, R. An, Q. Shi, Z. Tan, et al. Llm × mapreduce: Simplified long-sequence processing using large language models. arXiv preprint
arXiv:2410.09342, 2024.

A.AIME 结果示例