💬 NLP

BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholding

BLASST 是一种无需训练、即插即用的动态稀疏注意力机制，它通过利用在线 Softmax 统计信息设定固定阈值来跳过无关注意力块，从而在保持精度的同时显著加速大语言模型在预填充和解码阶段的推理速度。

原作者： Jiayi Yuan, Cameron Shinn, Kai Xu, Jingze Cui, George Klimiashvili, Guangxuan Xiao, Perkz Zheng, Bo Li, Yuxin Zhou, Zhouhai Ye, Weijie You, Tian Zheng, Dominic Brown, Pengbo Wang, Markus Hoehnerbach

发布于 2026-04-07

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jiayi Yuan, Cameron Shinn, Kai Xu, Jingze Cui, George Klimiashvili, Guangxuan Xiao, Perkz Zheng, Bo Li, Yuxin Zhou, Zhouhai Ye, Weijie You, Tian Zheng, Dominic Brown, Pengbo Wang, Markus Hoehnerbach, Richard Cai, Julien Demouth, John D. Owens, Xia Hu, Song Han, Timmy Liu, Huizi Mao

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一个名为 BLASST 的新方法，它就像是大语言模型（LLM）的“智能过滤器”，能让模型在处理超长文本时变得更快、更省内存，而且几乎不会损失智能。

为了让你轻松理解，我们可以把大语言模型想象成一位正在阅读一本超级厚书（比如整个互联网或几百万字的代码库）的超级天才。

1. 核心痛点：为什么现在的模型“读”得慢？

想象一下，这位天才每读到一个新词（Token），他都要回头去翻整本书，把之前读过的每一个词都重新看一遍，思考它们之间有没有关系。

问题：如果书只有 10 页，这很快。但如果书有 100 万页，他每读一个新词都要翻 100 万次，这简直累死人（计算量呈指数级爆炸）。
现状：现有的技术（如 FlashAttention）虽然优化了“翻书”的动作，让他翻得更快，但他依然试图去翻每一页，并没有真正减少工作量。

2. BLASST 的解决方案：聪明的“略读”策略

BLASST 的核心思想非常直观：并不是每一页都值得你仔细读。

传统做法：不管内容多无关紧要，都强行计算一遍。
BLASST 的做法：它给这位天才装了一个“智能眼镜”。
- 当天才快速浏览时，如果发现某几页的内容（比如“的”、“了”、“啊”这种词，或者完全无关的段落）对当前句子的理解贡献微乎其微，他的“智能眼镜”会立刻告诉他：“跳过！别读了！”
- 跳过什么？ 跳过计算、跳过从内存里读取数据、跳过乘法运算。
- 怎么判断？ 它不需要重新计算，而是利用正在进行的“注意力分数”（可以理解为“重要程度打分”）。如果某一块内容的分数远低于当前已知的最高分，那就直接忽略。

3. 三大亮点：为什么它这么厉害？

A. 无需“预习”或“特训” (Drop-in, Training-free)

比喻：以前的很多加速方法，就像要求这位天才在读书前，先花几个小时写一份“重点摘要”，或者让他重新上几年学来适应新的阅读方式。这太麻烦了，而且容易出错。
BLASST：就像给他递了一副现成的眼镜。他不需要重新学习，也不需要提前做功课。戴上眼镜就能立刻开始用，直接插入现有的系统里就能跑。

B. 自动调节的“灵敏度” (Dynamic Calibration)

比喻：如果书很短（比如 100 页），眼镜的灵敏度可以调低一点，多读几页；如果书很长（比如 100 万页），眼镜就要调得更灵敏，只读最关键的几页。
BLASST：它发现了一个简单的数学规律：书越长，阈值就要越严格。它有一个自动校准程序，能根据书的长度自动设定“跳过多少”的标准。你不需要手动去调参数，它自己会搞定。

C. 两头通吃：既加速“读前”也加速“读中”

比喻：
- Prefill（预填充/读前）：就像在开始写文章前，快速浏览整本书找灵感。BLASST 在这里能跳过大量无关计算，让“找灵感”的速度快 1.5 倍。
- Decode（解码/读中）：就像在写文章时，每写一个字都要回头看。BLASST 在这里能跳过从内存里搬运数据的步骤，让“写字”的速度快 1.5 倍。
大多数旧方法只能加速其中一步，而 BLASST 是两头一起加速。

4. 实际效果：快了多少？

在最新的顶级显卡（如 NVIDIA Blackwell 和 Hopper）上测试：

速度提升：在保持 99% 以上准确度的情况下，处理速度提升了 1.5 倍 左右。
跳过率：它成功跳掉了 70% 以上 的无用计算和内存读取。
结果：原本需要几分钟才能处理完的超长文档，现在可能只需要几十秒。

5. 总结：它意味着什么？

BLASST 就像是给大语言模型装上了**“自动驾驶”和“智能节油”系统**。

以前，让 AI 读一本 100 万字的书，既慢又贵，甚至跑不动。
现在，有了 BLASST，AI 可以忽略那些废话，只关注核心内容，从而跑得飞快，而且不费油（省显存）。

这让那些需要处理超长上下文（比如分析整个公司的代码库、阅读整本小说、进行超长对话）的应用变得真正可行和实用。它不需要你改变模型结构，也不需要你重新训练，就像给旧车换上了新引擎，直接就能跑。

1. 问题背景 (Problem)

大型语言模型（LLM）在处理长上下文（Long Context）推理时面临严峻挑战，主要瓶颈在于自注意力机制（Self-Attention）的二次方计算复杂度（ $O(n^2)$ ）和内存消耗。

现有方案的局限性：
- FlashAttention 及其变体：虽然优化了内存带宽利用（通过分块和内核融合），但仍需计算完整的注意力矩阵，无法从根本上解决二次方复杂度问题。
- 现有稀疏注意力方法：存在五大落地障碍：
  1. 需要昂贵的预计算（Pre-computation）来确定稀疏模式，抵消了加速收益。
  2. 需要微调模型或训练新架构（Training requirements）。
  3. 通常仅针对 Prefill（预填充）或 Decode（解码）阶段优化，缺乏端到端加速。
  4. 缺乏对现代 GPU（如 Blackwell, Hopper）的内核支持。
  5. 难以集成到现有框架中，需要侵入式修改。

2. 核心方法论 (Methodology)

BLASST 提出了一种无需训练、无需预计算的动态稀疏注意力机制。其核心思想是利用 FlashAttention 在线 Softmax 计算过程中已有的统计信息，动态跳过对最终输出贡献可忽略的注意力块。

2.1 核心洞察：基于运行最大值的剪枝

在 FlashAttention 的分块在线 Softmax 计算中，模型维护一个运行最大值（Running Maximum, $m_i$ ）。

判断逻辑：对于当前的注意力块 $S_j$ ，计算其局部最大值 $\tilde{m}_j$ 。如果 $\tilde{m}_j$ 显著小于当前的运行最大值 $m_i$ （即 $\tilde{m}_j - m_i < \ln(\lambda)$ ，其中 $\lambda$ 是阈值），则该块经过 Softmax 归一化后的值将趋近于 0。
跳过操作：一旦满足上述条件，直接跳过该块的以下昂贵操作：
1. 计算节省：跳过 Softmax 中的指数运算（exp）和归一化求和。
2. 计算节省：跳过注意力权重与 Value 矩阵的乘法（MMA 操作）。
3. 内存节省：跳过从 HBM（高带宽内存）加载对应的 Value 块。

2.2 自动阈值校准 (Automated Threshold Calibration)

为了在不同上下文长度下保持稳定的稀疏率和精度，BLASST 提出了一种自动校准机制：

发现：最优阈值 $\lambda$ 与上下文长度 $L$ 呈反比关系： $\lambda = a / L$ 。
原理：由于注意力分数行归一化后总和为 1，序列越长，单个 Token 的平均分数越低，因此需要更小的阈值来维持相同的稀疏率。
实施：通过少量样本（Algorithm 2）拟合参数 $a$ ，使得用户只需设定目标稀疏率，即可自动获得适用于任意长度上下文的阈值，无需手动调整。

2.3 针对 Prefill 和 Decode 的专用内核优化

BLASST 为现代 GPU（Blackwell/Hopper）设计了专用 CUDA 内核，针对两个阶段的不同瓶颈进行优化：

Prefill 阶段（计算受限 Compute-bound）：
- 重点跳过 Softmax 计算和矩阵乘法（MMA）。
- 由于内存带宽不是瓶颈，Value 块仍按顺序加载以维持流水线预测性，但计算单元被释放以加速后续操作。
- 通过重叠计算任务，将执行时间从 18 个时间单位压缩至 14 个。
Decode 阶段（内存受限 Memory-bound）：
- 重点跳过 Value 块从 HBM 的加载（这是 Decode 的主要瓶颈）。
- 挑战：跳过加载需要在计算完 Query-Key 乘积（BMM1）并做出跳过决策后才能进行，这会导致流水线气泡。
- 解决方案：采用批处理加载调度（Batched Load Scheduling）。连续计算多个 Query-Key 乘积，待所有跳过决策完成后，再批量加载通过检查的 Value 块，从而消除流水线气泡并显著减少内存带宽消耗。

2.4 稀疏感知训练 (Sparsity-Aware Training)

作为一种扩展，BLASST 支持在微调阶段引入稀疏性。

方法：在训练的前向传播中应用相同的阈值剪枝逻辑，被跳过的块在反向传播中不接收梯度。
效果：迫使模型将重要信息集中在高分块中，使模型在推理时更能适应稀疏模式，进一步提升精度 - 稀疏率权衡（Pareto Frontier）。

3. 主要贡献 (Key Contributions)

BLASST 算法：一种即插即用（Drop-in）的动态稀疏注意力方法，无需预计算、无需代理分数（Proxy Scores）、无需训练，即可在 Prefill 和 Decode 阶段实现加速。
自动校准与训练扩展：提出了基于上下文长度的自动阈值校准方法，以及稀疏感知训练策略，显著提升了部署的鲁棒性和灵活性。
高性能 CUDA 内核：为 Blackwell 和 Hopper 架构实现了优化的 Prefill 和 Decode 内核，已集成至 TensorRT-LLM 和 FlashInfer 框架。
广泛的兼容性：支持 MHA、GQA、MQA 及 MLA（Multi-head Latent Attention）等多种注意力变体。

4. 实验结果 (Results)

实验在 Llama-3.1-8B 和 Qwen3-8B 等模型上，针对 RULER、LongBench、MATH500 等基准进行了测试。

加速比 (Speedup)：
- Prefill：在约 71.9% 的稀疏率下，Blackwell 架构上实现 1.52× 加速；在 Hopper 上实现 1.52× 加速。
- Decode：在约 73.2% 的稀疏率下，实现 1.48× 加速。
- 端到端：在长上下文场景下，TTFT（首字延迟）和 TPOT（每字延迟）均有约 1.1× 的提升。
精度保持：
- 在 50%-75% 的稀疏率下，精度损失极小（通常 < 0.5%），甚至在某些推理任务（如 MATH500）中，由于去除了冗余干扰，精度略高于稠密基线。
- 在 200K 超长上下文（RepoQA）测试中，BLASST 依然保持有效，Prefill 稀疏率可达 57.5%。
对比优势：
- 相比 MInference、FlexPrefill 等需要预计算的方法，BLASST 在保持精度的同时提供了更稳定的加速。
- 相比 Quest、RocketKV 等仅优化 Decode 的方法，BLASST 实现了端到端加速。
- 校准后的阈值在不同长度下能保持稳定的稀疏率（误差仅 1.2%），而固定阈值会导致稀疏率剧烈波动。

5. 意义与影响 (Significance)

打破长上下文推理瓶颈：BLASST 通过动态跳过无效计算和内存访问，显著降低了长序列推理的延迟和显存占用，使得在消费级或企业级 GPU 上部署 100K+ 甚至 1M Token 上下文成为可能。
降低落地门槛：无需重新训练模型、无需复杂的预计算流程、无需修改模型架构，即可直接集成到现有推理框架中，极大地促进了稀疏注意力技术的实际落地。
硬件感知优化：针对现代 GPU 架构（Hopper/Blackwell）的瓶颈特性（计算 vs 内存）设计了专用内核，充分挖掘了硬件潜力。
未来方向：为构建高效、长上下文的 Agent AI 系统提供了基础，展示了硬件感知稀疏模式与训练结合的巨大潜力。

总结：BLASST 通过巧妙地复用在线 Softmax 统计信息，以极低的开销实现了动态的块级注意力剪枝。它不仅解决了长上下文推理的计算和内存瓶颈，还通过自动校准和专用内核设计，解决了现有稀疏方法难以部署的痛点，是目前长上下文 LLM 推理加速领域的一项突破性工作。