硬件对齐+原生训练!DeepSeek NSA打造高效稀疏Attention-菜鸟下载

一、研究背景与动机

在自然语言处理领域，长上下文建模对下一代大语言模型至关重要，其应用场景广泛，如深度推理、代码生成、多轮对话等。然而，标准注意力机制计算复杂度高，当处理长序列时，计算成本剧增，成为模型发展的瓶颈。以解码64k长度上下文为例，softmax注意力计算的延迟占总延迟的70 - 80%，这凸显了寻求高效注意力机制的紧迫性。

为提升效率，利用softmax注意力的固有稀疏性是一种可行途径，即选择性计算关键查询 - 键对，在保持性能的同时降低计算开销。现有方法虽各有探索，但在实际应用中存在诸多局限：

推理效率假象：许多稀疏注意力方法在推理时未能实现预期的加速效果。一方面，部分方法存在阶段受限的稀疏性，如H2O在解码阶段应用稀疏性，但预填充阶段计算量大；MInference则只关注预填充阶段稀疏性，导致至少一个阶段计算成本与全注意力相当，无法在不同推理负载下有效加速。另一方面，一些方法与先进注意力架构不兼容，如Quest在基于GQA的模型中，虽能减少计算操作，但KV缓存内存访问量仍较高，无法充分利用先进架构的优势。可训练稀疏性的误区：仅在推理阶段应用稀疏性会导致模型性能下降，且现有稀疏注意力方法大多未有效解决训练阶段的计算挑战。例如，基于聚类的方法（如ClusterKV）存在动态聚类计算开销大、算子优化困难、实现受限等问题；一些方法的离散操作（如MagicPIG中的SimHash选择）使计算图不连续，阻碍梯度传播；HashAttention等方法的非连续内存访问模式，无法有效利用快速注意力技术（如FlashAttention），降低了训练效率。

针对这些问题，本文提出了原生可训练的稀疏注意力机制（Native Sparse Attention，NSA），旨在通过算法创新与硬件对齐优化，实现高效的长上下文建模，平衡模型性能与计算效率。

二、NSA核心工作

（一）方法概述

NSA的技术方法涵盖算法设计与内核优化。其整体框架基于对注意力机制的重新定义，通过设计不同的映射策略构建更紧凑、信息更密集的键值对表示，以减少计算量。同时，针对硬件特性进行内核优化，提升实际运行效率。

背景知识注意力机制：在语言建模中，注意力机制广泛应用。对于输入序列长度为

的情况，注意力操作定义为：

\ o_{t}=Attn\left(q_{t}, k_{: t}, v_{: t}\right)

其中

Attn

表示注意力函数：

\ Attn\left(q_{t}, k_{: t}, v_{: t}\right)=\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}}}}

这里

\alpha_{t, i}

是

与

k_{i}

之间的注意力权重，

d_{k}

是键的特征维度。随着序列长度增加，注意力计算在总计算成本中占比越来越大，给长上下文处理带来挑战。

算术强度：算术强度是计算操作与内存访问的比率，对硬件上的算法优化有重要影响。每个GPU都有由峰值计算能力和内存带宽决定的临界算术强度。对于计算任务，算术强度高于此临界阈值时受GPU浮点运算能力（FLOPS）限制，低于此阈值时受内存带宽限制。在因果自注意力机制中，训练和预填充阶段，批矩阵乘法和注意力计算算术强度高，属于计算受限阶段；而自回归解码时，每次前向传递仅生成一个令牌，但需加载整个键值缓存，算术强度低，受内存带宽限制。这导致不同阶段的优化目标不同：训练和预填充阶段需降低计算成本，解码阶段需减少内存访问。整体框架：为利用注意力的自然稀疏模式，NSA提出用更紧凑的键值对

\tilde{K}_{t}

、

\tilde{V}_{t}

替代原始键值对

k_{: t}

、

v_{: t}

。优化后的注意力输出定义为：

\tilde{K}_{t}=f_{K}\left(q_{t}, k_{: t}, v_{: t}\right), \tilde{V}_{t}=f_{V}\left(q_{t}, k_{: t}, v_{: t}\right)

\ o_{t}^{*}=Attn\left(q_{t}, \tilde{K}_{t}, \tilde{V}_{t}\right)

其中

\tilde{K}_{t}

、

\tilde{V}_{t}

根据当前查询

q_{t}

和上下文内存

k_{: t}

、

v_{: t}

动态构建。通过设计多种映射策略可得到不同类别的

\tilde{K}_{t}^{c}

、

\tilde{V}_{t}^{c}

，并将它们组合起来：

\ o_{t}^{*}=\sum_{c \in C}g_{t}^{c}\cdot Attn\left(q_{t}, \tilde{K}_{t}^{c}, \tilde{V}_{t}^{c}\right)

NSA有三种映射策略

C = \{cmp, slc, win\}

，分别代表压缩、选择和滑动窗口策略，用于处理键值对。

g_{t}^{c} \in [0, 1]

是对应策略

的门控分数，由输入特征经MLP和sigmoid激活得到。令

N_{t}

表示重新映射后的键/值总数：

\ N_{t}=\sum_{c \in C}size[\tilde{K}_{t}^{c}]

通过确保

N_{t} \ll t

，NSA保持较高的稀疏率。

算法设计令牌压缩（Token Compression）：通过聚合连续的键或值块为块级表示，得到压缩后的键值对，以捕获整个块的信息。压缩键表示定义为：

\tilde{K}_{t}^{cmp}=f_{K}^{cmp}\left(k_{: t}\right)=\left\{\varphi\left(k_{id + 1:id + l}\right) | 1 \leq i \leq\left\lfloor\frac{t - l}{d}\right\rfloor\right\}

其中

是块长度，

是相邻块之间的滑动步长，

\varphi

是带有块内位置编码的可学习MLP，用于将块中的键映射为单个压缩键。

\tilde{K}_{t}^{cmp} \in \mathbb{R}^{d_{k} \times\left\lfloor\frac{t - l}{d}\right\rfloor}

是由压缩键组成的张量。通常采用