🤖 machine learning

Accelerating Sparse Transformer Inference on GPU

本文提出了 STOF，这是一个 GPU 框架，通过利用分析建模实现高效的多头注意力映射，并采用两阶段搜索策略动态优化算子融合，从而加速稀疏 Transformer 推理，在多头注意力计算和端到端推理中分别实现了高达 1.6 倍和 1.4 倍的加速。

原作者： Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

发布于 2026-05-20

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

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

想象一下，你正在试图阅读一个庞大的图书馆（即大型语言模型）以回答一个问题。这个图书馆被组织成一个个名为Transformer的房间，每个房间里都有一位图书管理员（即多头注意力机制），他必须扫描成千上万页内容，以找到与你问题相关的具体句子。

问题在于，对于许多问题而言，大多数页面都是无关的。图书管理员浪费时间在翻阅空白页面或无关页面上。这正是稀疏性发挥作用的地方：它就像在无关页面上贴上“请勿阅读”的标签。

然而，当前的图书管理员（现有软件）并不擅长使用这些标签。他们仍然会走过标有“请勿阅读”的页面，或者当标签以奇怪、随机的模式放置时感到困惑。此外，图书馆还有其他任务（如摘要或格式化），这些任务通常单独执行，从而在不同任务之间增加了额外的往返时间。

现在，STOF登场了，这是研究人员提出的一种新系统。可以将 STOF 想象为一个专为这些“稀疏”图书馆设计的超高效、智能的图书馆管理系统。以下是其工作原理，分解为简单部分：

1. 智能图书管理员（统一的多头注意力内核）

研究人员意识到，不同的“请勿阅读”模式需要不同的策略。

问题：有些模式是整齐的标签行（如滑动窗口），而另一些则是随机散布的（如彩票）。旧系统试图使用一种“一刀切”的方法，导致速度缓慢。
STOF 解决方案：STOF 扮演一位智能图书管理员的角色，为每项工作选择最佳工具。
- 如果标签位于整齐的小簇中，图书管理员会使用**“逐行”**方法：一次性抓取一整行书籍并快速扫描。
- 如果标签分散或图书馆规模巨大，他们则采用**“分块”**方法：将书籍划分为小块、易于管理的部分，仅打开那些带有有效标签的特定块。
结果：通过完全跳过“请勿阅读”的页面，而不仅仅是忽略它们，图书管理员的工作速度大大提升。

2. 流水线（算子融合）

在普通图书馆中，图书管理员可能读完书后，走到另一张桌子进行文本摘要，然后再走到另一张桌子对答案进行格式化。这种走动（在内存和处理器之间移动数据）非常缓慢。

问题：现有系统通常只组合简单的任务。它们将繁重的任务（如复杂数学运算）留给单独的步骤，导致交通拥堵。
STOF 解决方案：STOF 构建了一条定制的流水线。它审视整个流程并问道：“我们能否将这些步骤合并？”
- 它不仅仅是将两个简单任务粘合在一起，而是找出将复杂数学任务与格式化任务结合的完美方式。
- 它使用一个“搜索引擎”来尝试不同的任务组合方式（如同尝试不同的流水线布局），以找到最适合你所阅读图书馆规模的最快方案。

3. 自动驾驶（分层搜索）

你无法为每一种书籍大小和问题类型手动设计完美的流水线；组合方式太多了。

STOF 解决方案：STOF 拥有一个自动驾驶系统，能够即时学习。
- 第一阶段（地图）：它查看图书馆结构，并绘制出“请勿阅读”标签大致位置的粗略地图。
- 第二阶段（优化）：它执行两步搜索。首先，它扩展流水线边界，以查看其能延伸多远。其次，它根据先前尝试的效果，微调工人（参数）的速度。
- 它会记住哪些方法有效（缓存），从而避免浪费时间重新测试相同的缓慢想法。

结果：快了多少？

研究人员在强大的图形处理器（GPU）上使用流行的 AI 模型（如 BERT、GPT 和 LLaMA）对 STOF 进行了测试。

速度：与现有最佳方法相比，STOF 使核心阅读任务（多头注意力）的速度提高了高达1.6 倍。
整体速度：当着眼于回答问题的整个流程（端到端）时，速度提高了高达1.4 倍。
大型图书馆：图书馆越大（文本序列越长），STOF 的表现就越出色，因为它跳过了大量无用工作。

总结

可以将STOF视为一个系统，它阻止 AI 浪费时间阅读不需要的页面，并阻止其在桌子之间来回走动。它采用一种智能、自适应的策略来跳过垃圾内容，并将有用的步骤合并为一个流畅、快速的动作。这使得 AI 模型运行速度显著加快，尤其是在处理长文本或复杂文本时。

技术摘要：加速 GPU 上的稀疏 Transformer 推理

问题陈述

大型语言模型（LLM）严重依赖 Transformer 架构，其中多头注意力（MHA）是核心计算组件。虽然诱导稀疏性的掩码层（如因果掩码、滑动窗口、随机掩码）被广泛用于减少计算量，但它们给 GPU 带来了显著的性能挑战：

碎片化计算：稀疏掩码导致碎片化的数据访问模式，加剧了内存带宽瓶颈，并阻碍了 GPU 资源的有效利用。
优化灵活性不足：现有高性能 MHA 实现（如 FlashAttention）通常仅支持特定的连续掩码模式（如因果掩码或滑动窗口）。它们难以处理离散或非结构化稀疏性（如 Bigbird、随机注意力），或者退化为低效方法，如在计算后重置分数矩阵。
静态算子融合：当前深度学习框架中的算子融合方案大多是静态或基于规则的。它们无法适应多样化的应用场景、变化的序列长度以及不同的模型超参数。此外，融合算子的最优配置往往与单独调优的算子之和存在显著差异，使得顺序调优并非最优。

方法论：STOF 框架

作者提出了 STOF，这是一个旨在通过将灵活的掩码支持与自适应算子融合相结合，从而加速 GPU 上稀疏 Transformer 推理的框架。STOF 包含两个主要模块：

1. 统一 MHA 模块

该模块解决了在不同掩码模式下高效计算稀疏注意力的挑战。

两级存储格式：为了处理任意掩码模式（包括结构化和非结构化），STOF 引入了一种结合块压缩稀疏行（BSR）和位图的新型存储格式。
- 外层瓦片（OT）：表示全局跳过的块。通过指针数组追踪“完整”OT（所有元素有效）和“部分”OT（混合有效/无效元素）。
- 内层瓦片（IT）：表示块内元素分布。 “部分”OT 内的有效元素使用 64 位位图编码（每个 8x8 块对应一个 uint64），从而在不实例化完整稀疏矩阵的情况下实现精确掩码。
双核策略：STOF 基于分析模型实现了两种不同的内核类型：
- 行级内核：针对高稀疏度、短序列场景进行优化。它将查询（ $Q$ ）切片为行，利用 warp 级洗牌，并消除 warp 间同步，以最大化数据局部性。
- 块级内核：适用于更广泛场景的通用内核。它将 $Q$ 、 $K$ 和 $V$ 划分为子块，利用共享内存（SMEM）和异步数据复制，使内存传输与计算重叠。
内核选择：一个分析模型根据掩码稀疏度和序列长度决定使用行级内核还是块级内核，以确保最高效的执行路径。

2. 算子融合模块

该模块解决了超越 MHA 层的自适应融合需求。

融合方案转换：STOF 使用哈希编码将融合方案转换为数值表示。计算图中的相邻算子被映射为二进制数组（表示融合边界），随后被解码为特定的编译模板（通过 Triton 和 TileLang 实现）。
分层搜索引擎：STOF 采用两阶段搜索机制来确定最优融合边界和内核参数，而非依赖静态规则：
1. 初始化：利用神经哈希和预定义规则识别频繁的子图模式，并建立初始融合方案。
2. 两阶段调优：
  - 阶段 1（融合扩展）：使用基于奖励的算法（扩展、夺取、竞争）迭代扩展融合边界，将性能提升与先前结果的缓存进行对比测试。
  - 阶段 2（参数采样）：一旦边界固定，便采样内核参数（如块大小、流水线阶段），以找到特定融合方案的最优配置。

主要贡献

全面分析：作者分析了多样化掩码模式和推理配置的影响，揭示了 MHA 计算和下游算子融合中的优化机会。
统一 MHA 模块：一个新颖的模块，实现了行级和块级内核，采用独特的存储格式（BSR + 位图）和优化手段，并由分析选择模型进行管控。
自适应算子融合：一个模块，通过数值解码将融合方案转化为编译模板，并利用两阶段搜索引擎动态确定最优设置。
STOF 框架：一个完整的推理系统，在 GPU 上实现了灵活的掩码模式和最优算子融合。

实验结果

作者在 NVIDIA RTX 4090 (Ada) 和 A100 (Ampere) GPU 上评估了 STOF，对比的基线包括 PyTorch Native、FlashAttention-2 (FA2)、FlexAttention、ByteTransformer、MCFuser 和 Bolt 等最先进方案。

MHA 性能：在 MHA 计算中，STOF 相比现有最佳 MHA 实现（如 FlexAttention）实现了最高 1.6× 的加速。它在各种掩码模式（因果、滑动窗口、Longformer、Bigbird）和序列长度（高达 4,096）下均表现出一致的优越性。
端到端推理：在 BERT、GPT-2、LLaMA、T5 和 ViT 等模型的端到端推理中，STOF 实现了最高 1.4× 的加速。
调优效率：STOF 的调优过程显著快于竞争对手（MCFuser 和 Bolt），在大规模场景下快 6.7× 至 6.9×，这归功于其基于奖励的采样和缓存机制。
可扩展性：在超长序列（高达 16k）实验中，STOF 显示出显著的性能提升（相比 PyTorch Compile 最高达 16.8×）和更好的内存效率，在基线方案出现内存溢出（OOM）错误的序列长度下仍能正常运行。

意义与主张

本文声称 STOF 是首个系统，能够同时为稀疏 Transformer 场景启用灵活的掩码模式和多样化的算子融合方案。其意义在于弥合了手工调优的稀疏专用内核与自动稠密算子融合之间的差距。通过提供一个将自定义 MHA 内核与生成式编译模板集成在内的完整技术栈，STOF 建立了此前静态或基于规则的方法无法实现的更广泛的优化机会。作者强调，他们的方法可扩展至新兴的 LLM 架构（如混合专家模型 MoE）和更新的 GPU 代际，为大规模模型推理的持续性能提升提供了一条途径。