BLASST 刚刚在 #MLSys26 会议上荣获最佳论文奖！在这篇论文中，我们提出了一种简单、无需训练的动态稀疏注意力机制，该机制利用在线 softmax 统计量上的单一标量阈值来跳过可忽略的注意力块。遗憾的是我无法亲临现场，但请向我的优秀合著者们问好！ 论文地址: https://arxiv.org/abs/2512.12087

1 引言

来源: https://arxiv.org/html/2512.12087 marginparsep 已被修改。topmargin 已被修改。marginparwidth 已被修改。marginparpush 已被修改。页面布局违反了 ICML 样式。请不要修改页面布局，或引入像 geometry、savetrees、fullpage 这样的包来改变它。我们无法可靠地撤销对样式的任意更改。请移除违规的包或布局改变命令，然后重试。

BLASST: 通过 Softmax 阈值实现动态块级注意力稀疏性

匿名作者1

摘要

大型语言模型 (LLM) 对长上下文推理能力的需求日益增长，加剧了自注意力机制固有的计算和内存瓶颈。为了解决这一挑战，我们引入了 BLASST，一种即插即用的动态稀疏注意力机制，它仅使用一个固定的标量阈值来跳过注意力块，从而加速推理。我们的方法针对实际推理部署，消除了现有工作中存在的采纳障碍。因此，BLASST 消除了训练需求，避免了昂贵的预计算过程，同时加速了所有主要注意力变体 (MHA, GQA, MQA 和 MLA) 的预填充和解码阶段，为现代硬件提供了优化支持，并能轻松集成到现有框架中。这是通过重用在线 softmax 统计量来识别可忽略的注意力分数，从而跳过 softmax、值块加载以及随后的矩阵乘法来实现的。我们通过提供延迟开销可忽略的优化内核来展示 BLASST 算法。我们的自动阈值校准过程揭示了最优阈值与上下文长度之间存在简单的反比关系，这意味着每个模型只需为预填充和解码分别设置一个阈值。在保持基准准确率的同时，我们在现代 GPU 上展示了预填充阶段 71.9% 稀疏度下 1.52 倍的加速，以及解码阶段 73.2% 稀疏度下 1.48 倍的加速。 ††脚注: 1 匿名机构，匿名城市，匿名地区，匿名国家。通讯作者: 匿名作者。初步工作。正在由机器学习与系统 (MLSys) 会议审稿。请勿分发。

大型语言模型 (LLM) 彻底改变了自然语言处理，在各项任务中取得了卓越的性能。然而，它们的实际部署面临一个关键瓶颈：注意力机制的二次计算复杂度。随着应用越来越需要更长的上下文窗口——从处理整个代码库 (Roziere 等人, 2023) 到分析冗长文档 (Zeng 等人, 2025) 以及维持长时间对话 (Achiam 等人, 2023)——这个瓶颈变得越来越严重。像 Deepseek-R1 (Guo 等人, 2025) 和 Qwen3 (Yang 等人, 2025) 这样的最新模型支持高达 128K 令牌的上下文长度，一些模型甚至推到了 1M 令牌 (Comanici 等人, 2025)。然而，处理如此长的序列在计算上仍然代价高昂，注意力计算在延迟和内存消耗方面都占主导地位。对于长度为 n 的序列，注意力机制需要 O(n²) 次操作和内存访问，这使得长上下文模型在即使是最先进的硬件上实现实际部署也充满挑战。虽然 FlashAttention (Dao 等人, 2022; Zadouri 等人, 2026) 及其后继者通过分块和内核融合优化了内存带宽利用率，但它们仍然计算完整的注意力矩阵，未能解决基本的二次复杂度问题。

稀疏注意力方法作为一种有前景的解决方案出现，它只计算完整注意力矩阵的一个子集。虽然这些方法巧妙地决定了跳过哪些注意力分数，但它们增加的复杂性阻碍了实际使用。我们确定了其采用的五个关键障碍：(1) 一些方法需要昂贵的预计算来确定稀疏模式，这往往会抵消其理论上的加速效果 (Jiang 等人, 2024; Xu 等人, 2025)。(2) 其他方法引入了需要模型微调 (Xiao 等人, 2025) 或训练全新架构 (DeepSeek-AI, 2025) 的新层。(3) 大多数现有工作只专注于预填充或解码阶段之一，错过了端到端推理加速的机会。(4) 它们缺乏对新 GPU 的内核支持，使得其加速效果能否转化到现代硬件（如 Blackwell 和 Hopper）的特性上尚不明确。(5) 这些方法常常阻碍框架集成，需要对模型架构或注意力接口进行侵入式修改，并对现有 API 进行重大更改。

为了解决这些障碍，我们提出了 BLASST（通过 Softmax 阈值实现块级注意力稀疏性），一种简单而有效的、无需训练的稀疏注意力方法，它在预填充和解码阶段都能动态地修剪可忽略的注意力块，并且没有预计算开销。我们的关键洞察是：在 FlashAttention 的逐块在线 softmax 过程中，我们可以仅基于已计算的信息，识别并跳过那些对最终输出贡献可忽略的块。具体来说，当顺序处理块时，我们维护一个注意力分数的运行最大值。如图 1 所示，如果一个块的局部最大分数显著小于这个运行最大值（差值超过阈值 λ），那么经过归一化后，其 softmax 后的值将接近于零。因此，对于这样的块，我们可以跳过三个昂贵的操作：(1) 计算 softmax 的指数，(2) 从 HBM 加载对应的值块，以及 (3) 注意力与值之间的乘法。这个简单的剪枝规则每个块只需要一次比较，并且无缝集成到现有的注意力 API 中，只需要输入一个单一的标量阈值。

为了最大化 BLASST 的实际影响，我们为 Blackwell 和 Hopper 提供了优化过的 CUDA 内核，实现了我们的稀疏注意力算法。我们的内核设计有两个关键目标：(1) 通过重用已计算的统计量，为块跳过决策逻辑引入最小的开销；(2) 战略性地针对每个阶段的瓶颈资源——在计算受限的预填充阶段减少 CUDA 核心和张量核心的使用，在内存受限的解码阶段减少内存带宽消耗。我们的预填充和解码内核针对它们各自不同的计算模式进行了定制。我们的内核在预填充阶段实现了高达 1.52 倍的加速（稀疏度 71.9%），在解码阶段实现了高达 1.48 倍的加速（稀疏度 73.2%），超越了 FlashAttention 基线 (Shah 等人, 2024; Zadouri 等人, 2026)，同时保持数值稳定性并支持常见的注意力变体 (MHA, MQA, GQA, MLA)。

除了核心算法和内核实现，我们还开发了两种关键技术来增强 BLASST 的部署和性能。首先，我们提出了一种自动校准过程，可以确定任意目标稀疏度水平下的最优阈值。我们的校准揭示了阈值与上下文长度 L 之间存在稳健的反比关系 λ = a/L，使得无需手动调整即可在不同场景下可靠部署。其次，我们探索了稀疏感知训练作为自然的扩展，表明模型可以被训练得对稀疏注意力模式天生更具鲁棒性。这种训练方法进一步推动了准确率-稀疏度的前沿，使得在准确率损失最小的情况下达到更高的稀疏度水平。

我们的贡献包括：

1. BLASST 算法，一种即插即用的方法，没有预计算开销和代理分数，实现最小的准确率损失。
2. 自动超参数选择和稀疏感知训练，以实现稳健、灵活和可扩展的部署。
3. 为预填充和解码实现 BLASST 的优化 CUDA 内核，可在 TensorRT-LLM¹ 和 FlashInfer 中获得。 ¹GPU 内核和推理框架支持可在 https://github.com/NVIDIA/TensorRT-LLM 找到。

2 相关工作

有效利用稀疏注意力的特性需要在不产生昂贵的选择开销或重训练的情况下，要么减少对不重要交互的计算，要么减少内存占用（例如 KV 缓存）。与以下相关工作相比，BLASST 以一种无需训练的方式同时解决了这两个方面。表 1 总结了现有工作的格局。

表 1: 稀疏注意力方法的特征对比。BLASST 是唯一一种无需训练或昂贵预计算步骤，即可加速预填充和解码两个阶段的方法。

2.1 计算优化的稀疏性

几种方法通过选择重要的交互来减少注意力的计算。静态模式方法，如 Sparse Transformer (Child 等人, 2019)、LongFormer (Beltagy 等人, 2020) 和 BigBird (Zaheer 等人, 2020)，通过局部或基于块的注意力降低了复杂度。基于检索头的方法 (Wu 等人, 2025; Xiao 等人, 2025) 通过聚焦计算于关键检索头来加速模型解码。动态稀疏方法，如 MInference (Jiang 等人, 2024) 使用预计算的重要性分数，XAttention (Xu 等人, 2025) 对反对角块进行排序，FlexPrefill (Lai 等人, 2025) 提供编译器支持的灵活块模式；虽然这些方法对预填充有效，但其预计算和调度开销可能限制实际加速效果。训练辅助稀疏性，如 SeerAttention (Gao 等人, 2025b)，通过（预）训练门控来诱导高稀疏度，提高了效率但增加了训练成本，并显示出下游模型性能不一。FLASH-D (Alexandridis 等人, 2025) 以与 BLASST 类似的方式利用在线 softmax 的数学特性，但旨在提高自定义硬件加速器上的数值稳定性和并行性。SpargeAttention (Zhang 等人, 2025) 的设计与 BLASST 最为相似。我们在三个关键方面有所不同：(1) BLASST 使用专门的内核优化了预填充和解码，而 SpargeAttention 仅针对预填充；(2) 我们直接使用已计算的统计量做出跳过决策，开销为零，而 SpargeAttention 使用了单独的预测步骤；(3) 我们的解码内核跳过从 HBM 加载 Value，在计算节省的基础上解决了内存瓶颈。此外，我们提供了自动校准和稀疏感知训练。

2.2 内存优化的稀疏性

令牌/KV 稀疏性侧重于减少内存占用和解码时间成本。H2O (Zhang 等人, 2023)、TOVA (Oren 等人, 2024) 和 InfLLM (Xiao 等人, 2024a) 基于查询模式丢弃令牌。StreamingLLM (Xiao 等人, 2024b) 保留初始和最近的令牌以实现一致的延迟和内存使用。Quest (Tang 等人, 2024) 根据当前查询修剪令牌，Rectified Sparse Attention (Sun 等人, 2025) 自适应选择令牌以在高稀疏度下保持准确率，RocketKV (Behnam 等人, 2025) 通过选择性驱逐压缩 KV 缓存，以及最近的用于超大规模扩展的 KV 压缩 (Łańcucki 等人, 2025) 进一步扩展了有效上下文；TidalDecode (Yang 等人, 2026) 通过位置持久模式稳定解码效率。我们进一步区分面向推理的压缩方法，如 RPC (Song 等人, 2025)，它在内存约束下优先保留对推理关键的信息；非驱逐方法如 Loki (Singhania 等人, 2024)，它避免显式 KV 驱逐同时减少有效内存/计算开销。通常，这些方法减少了解码阶段的内存访问，而 BLASST 在无需训练的情况下，在预填充和解码阶段都减少了计算和内存访问。

2.3 新的注意力变体

除了上述方法，替代机制包括滑动窗口注意力 (Beltagy 等人, 2020)、线性或门控注意力 (Qiu 等人, 2025) 以及状态空间模型 (SSM, Gu and Dao, 2024)。原生稀疏注意力 (NSA, Yuan 等人, 2025) 和 DeepSeek 稀疏注意力 (DSA, DeepSeek-AI, 2025) 虽然在某些场景下有效，但通常需要架构改变和训练。相比之下，BLASST 是一种无需训练的方法，无需代理分数或复杂的预计算即可加速预填充和解码，并与 FlashAttention 实现无缝集成。

3 方法论

3.1 利用运行最大值剪枝注意力

BLASST 的核心洞察在于观察到，在 FlashAttention 计算注意力分数的过程中，许多块在 softmax 归一化后对最终输出的贡献可以忽略不计。我们的方法在前向传播过程中动态识别并跳过这些块，无需预计算或代理分数。

3.1.1 关键洞察