本文提出了一种面向 N:M 稀疏视觉 Transformer 推理的软硬件协同设计框架，能够在相同目标加速比下实现更准确可靠的加速。

https://t.co/uga3Ba6C1L https://t.co/7rZBHxWQMr

通过搜索与内核协同设计实现可落地的 N:M 稀疏 Transformer 推理

来源：https://arxiv.org/html/2607.12505 11institutetext:中国科学技术大学，合肥，中国22institutetext:中国科学技术大学苏州高等研究院，苏州，中国33institutetext:复杂系统建模与仿真全国重点实验室，中国 33email:[email protected],33email:[email protected],33email:[email protected] Wenqi Lou✉https://orcid.org/0000-0002-2240-6672 Zhiguang Wang https://orcid.org/0009-0002-5829-3339 Zhiwei Ke https://orcid.org/0009-0003-7636-2446 Fengrui Zuo https://orcid.org/0009-0003-3468-9280 Chao Wang https://orcid.org/0000-0002-9403-5575 Xuehai Zhou https://orcid.org/0000-0002-8360-3143

摘要

视觉 Transformer (ViT) 能够实现高准确率，但推理延迟较高。半结构化 N:M 稀疏性可以降低算术成本，但其理论上的节约效果往往无法在现代 GPU 上转化为比例性的端到端加速。这种不匹配的产生，是因为部署延迟不仅取决于算术缩减，还依赖于稀疏条件下的执行规整性和硬件调度。因此，实现实际加速需要在稀疏执行与稀疏配置之间进行协同设计。为此，我们提出了一种面向 N:M 稀疏 ViT 推理的软硬件协同设计框架。在硬件方面，我们设计了 MD-SpMM，一个 N:M 稀疏 CUDA 内核，它将稀疏 GEMM 重新组织为微密集的、与 Tensor Core 对齐的数据流，并使用推理感知的自适应并行性来维持利用率。在软件方面，我们在明确的端到端延迟预算下，执行逐层稀疏性搜索，采用三阶段启发式搜索并结合约束松弛，以避免过早收敛，并实现部署感知的稀疏性分配。在多个 ViT/Swin 模型和 GPU 平台上的实验表明，该框架在保持可比准确率的同时，实现了超过 2.2 倍的延迟加速，并且在相同延迟约束下提供了更优越的准确率。源代码已公开在 https://github.com/liuganhuo/realizable-nm-sparse-transformer。

1 引言

视觉 Transformer (ViT) 利用自注意力机制捕捉长距离依赖关系，在图像分类和语义分割等任务中取得了优异性能[20 (https://arxiv.org/html/2607.12505#bib.bib18)]。然而，这种优势伴随着高昂的推理成本：ViT 严重依赖于投影层和前馈层中的矩阵乘法，带来了巨大的计算和内存需求，使得在实际硬件上的部署日益困难[24 (https://arxiv.org/html/2607.12505#bib.bib20),16 (https://arxiv.org/html/2607.12505#bib.bib15)]。

模型压缩，特别是剪枝，已被广泛研究，旨在通过移除冗余参数同时保持准确率来缓解这一负担[2 (https://arxiv.org/html/2607.12505#bib.bib2)]。在众多剪枝方案中，半结构化 N:M 稀疏性已成为在准确率保持与硬件规整性之间的实用折中方案，并已被应用于 CNN、ViT 及相关加速器中[25 (https://arxiv.org/html/2607.12505#bib.bib21),10 (https://arxiv.org/html/2607.12505#bib.bib10),12 (https://arxiv.org/html/2607.12505#bib.bib12),9 (https://arxiv.org/html/2607.12505#bib.bib9)]。

尽管取得了这些进展，N:M 稀疏性的理论节约效果往往无法在实际推理中转化为比例性的端到端延迟增益，尤其是在 INT8 部署下，第 6 节的实验比较也证实了这一点。在这种情况下，延迟不仅取决于算术缩减，还取决于执行规整性、内存行为以及稀疏条件下的调度效应。这揭示了理论效率与实际性能之间持续存在的差距。这种差距常常被误解为稀疏性“无效”，而更准确地说，是由于在当前的部署栈上难以高效地实现稀疏性。

因此，我们关注的是可落地的加速：N:M 稀疏性能否在实际部署栈上实现目标端到端延迟降低，而不仅仅是改进诸如 FLOPs、全局稀疏率或孤立算子加速比等代理指标[4 (https://arxiv.org/html/2607.12505#bib.bib4),23 (https://arxiv.org/html/2607.12505#bib.bib19),18 (https://arxiv.org/html/2607.12505#bib.bib17)]。从系统角度来看，可落地的加速需要满足两个紧密耦合的条件。首先，稀疏计算必须是执行可落地的：稀疏算子应在真实硬件上提供稳定高效的执行，而不是表现出不规则行为，从而妨碍持续的延迟降低[10 (https://arxiv.org/html/2607.12505#bib.bib10)]。其次，稀疏性分配必须是基于延迟的：当以逐层方式应用 N:M 稀疏性时，应在明确的部署级延迟约束条件下选择每层的稀疏比，而不是基于抽象的稀疏率或 FLOPs 目标，后者难以准确反映端到端成本[8 (https://arxiv.org/html/2607.12505#bib.bib8)]。任何一个条件的不满足都会在实践中导致不可实现的加速，即使理论计算缩减非常显著。然而，现有方法往往只优化其中一方面[7 (https://arxiv.org/html/2607.12505#bib.bib7),9 (https://arxiv.org/html/2607.12505#bib.bib9),10 (https://arxiv.org/html/2607.12505#bib.bib10),13 (https://arxiv.org/html/2607.12505#bib.bib25)]，这促使我们提出一个端到端的协同设计框架，该框架将执行可落地的稀疏内核与基于延迟的逐层配置以及轻量级候选评估结合起来。

为此，我们提出了一个用于可落地 N:M 稀疏推理加速的软硬件协同设计框架。在硬件方面，我们设计了 MD-SpMM，一个 N:M 稀疏 CUDA 内核，它通过微密集重构和自适应并行性，将稀疏计算重组为与 Tensor Core 对齐的执行。在软件方面，我们引入了一个基于延迟的逐层搜索流水线，该流水线结合了基于性能分析的延迟预测器和可行区域感知的进化求解器，以高效评估候选方案。通过将执行可落地的内核与基于延迟的稀疏性搜索集成，该框架能够将稀疏性带来的计算缩减持续转化为实际部署栈上的端到端延迟改进。

总之，本工作做出了以下贡献：

• •我们提出了一个面向可落地 N:M 稀疏 Transformer 加速的搜索–内核协同设计框架，该框架建立在部署延迟约束下的执行可实现性与配置可实现性之上，能够在相同目标加速比下实现更准确可靠的加速。
• •我们设计了MD-SpMM，一个执行可落地的、原生 Tensor Core 的 N:M 稀疏内核，它将稀疏计算转换为以 MMA 为中心的规整数据流，并具有可扩展的推理并行性。在多种矩阵乘法形状和稀疏度水平上，它实现了相较于 cuBLAS 密集 GEMM 平均 2.0 倍的加速，以及相较于 nmSPARSE[10 (https://arxiv.org/html/2607.12505#bib.bib10)] 1.6–4.0 倍的加速。
• •我们提出了一个基于延迟的逐层 N:M 搜索流水线，该流水线集成了一个经过标定的基于查找表 (LUT) 的延迟预测器，以及一个可行区域感知的进化求解器，以高效发现高质量的稀疏配置。通过我们的软硬件协同设计，最终解决方案在保持更高准确率的同时实现了超过 2.2 倍的延迟降低，在 ImageNet-1K 上 Top-1 准确率提升高达 +1.4%，并且在 ViT 和 Swin 推理场景中始终优于先前方法。

2 背景与动机

2.1 视觉 Transformer 与逐层 N:M 配置

视觉 Transformer (ViT) 将图像处理为一系列图像块嵌入，并通过一组 Transformer 块进行前向传播。在实际推理中，时间主要消耗在 Q/K/V 投影、注意力输出投影以及两个 FFN 线性层上。对于常见的视觉任务，这些以矩阵运算为主的算子占据了大部分端到端推理时间，因此成为稀疏化的主要目标[24 (https://arxiv.org/html/2607.12505#bib.bib20)]。

N:M 稀疏性约束每组 M 个权重中保留 N 个非零值。在 ViT 中，它通常被统一应用于所有 Transformer 块，如图 2 所示。然而，尽管具有相似的块结构，不同深度的层由于其深度依赖的角色而对模型准确率和推理延迟的贡献不同[8 (https://arxiv.org/html/2607.12505#bib.bib8),18 (https://arxiv.org/html/2607.12505#bib.bib17),5 (https://arxiv.org/html/2607.12505#bib.bib5)]。这导致各块对稀疏性的敏感度存在差异，使得单一的全局 N:M 比并非最优。因此，ViT 稀疏化更自然地表述为一个逐层 N:M 配置问题。

参见图注 图1: N:M 稀疏性与 ViT 中的统一配置与逐层配置。 参见图注 图2: GPU 上 N:M 稀疏 INT8 推理的 Roofline 分析。

2.2 缺乏高效的稀疏执行机制

尽管 N:M 稀疏性减少了乘加运算次数，但它并不能在现代 Tensor Core GPU 上直接带来高效执行。这种不匹配在 INT8 推理中尤为突出，因为密集 GEMM 已经通过规整的数据流、可预测的控制流和高 MMA 利用率实现了高效率[15 (https://arxiv.org/html/2607.12505#bib.bib14),14 (https://arxiv.org/html/2607.12505#bib.bib13),1 (https://arxiv.org/html/2607.12505#bib.bib1)]。在这种情况下，稀疏加速不仅取决于算术缩减，还取决于执行是否与硬件执行模型兼容。现有的稀疏内核通常依赖于元数据处理、运行时解码和不规则内存访问，这些都会破坏 Tensor Core 的数据流并降低利用率[6 (https://arxiv.org/html/2607.12505#bib.bib6),10 (https://arxiv.org/html/2607.12505#bib.bib10),22 (https://arxiv.org/html/2607.12505#bib.bib24)]。

如图 2 中的 Roofline 分析所示，现有方法表现出明显的执行权衡。nmSPARSE[10 (https://arxiv.org/html/2607.12505#bib.bib10)] 遵循稀疏解码和不规则访问模式，导致算术强度低且利用率有限；而 cuBLAS 通过密集 GEMM 维持了高效率，但无法利用 N:M 结构。这一差距凸显了需要一种稀疏执行机制，既能保持 Tensor Core 友好的规整性，又能发挥 N:M 稀疏性的算术优势。

2.3 基于延迟的逐层配置的必要性

将延迟约束引入稀疏配置搜索是必要的，但也具有挑战性，因为满足目标延迟需要针对不同模型、层和部署环境采用不同的稀疏度水平。传统方法优化代理指标如稀疏率或 FLOPs，往往无法捕捉真实的部署行为。即使总体稀疏度相似，不同的逐层分配也可能导致截然不同的端到端延迟。此外，施加延迟约束会通过丢弃无效配置而碎片化可行空间，削弱搜索空间的连通性，使得探索容易陷入狭窄的可行区域。

为了解决这个问题，稀疏配置优化必须融入部署级的延迟语义。特别是，需要基于目标硬件和部署栈构建的延迟模型来引导搜索，以选择那些既满足延迟约束又能实现可落地端到端加速的配置。这有助于缓解抽象优化目标与实际硬件行为之间的不匹配[3 (https://arxiv.org/html/2607.12505#bib.bib3),4 (https://arxiv.org/html/2607.12505#bib.bib4),19 (https://arxiv.org/html/2607.12505#bib.bib22)]。

3 问题形式化与方法概述

给定一个 Transformer 模型和一个目标端到端延迟预算，问题在于实现既能保持模型质量又能在目标硬件上提供实际加速的 N:M 稀疏推理。

如图 3 所示，我们通过两个协同的组件来解决这个问题。在执行层面，我们设计了 MD-SpMM，一个通过 Tensor Core 对齐的执行来实现 N:M 稀疏性的稀疏内核。在配置层面，我们构建了一个基于延迟的逐层搜索流水线，其中经过标定的基于 LUT 的延迟预测器强制执行部署可行性，而硬件感知的搜索策略则在显式延迟约束下探索碎片化的可行区域。

这些组件共同实现了第 1 节中引入的可落地加速的两个要求。MD-SpMM 解决了执行可实现性问题，而基于延迟的搜索通过识别与端到端部署行为一致的逐层 N:M 分配来解决配置可实现性问题。因此，只有当稀疏执行和稀疏配置被联合优化时，才能实现可落地的加速。

参见图注 图3: 面向 N:M 稀疏 Transformer 加速的搜索-内核协同设计框架总览。

4 MD-SpMM: 面向数据流的 N:M 稀疏内核

4.1 设计方法

从执行可实现性的角度来看，N:M 稀疏性的 INT8 推理不仅取决于算术量的减少，还取决于稀疏执行是否与硬件执行模型匹配。我们识别出两个首要需求：(i) 数据流规整性，以维持 Tensor Core MMA 利用率；(ii) 推理规模并行性，以在不同算子形状下保持高 SM 占用率。任何因引入索引驱动的控制流、不规则内存访问或不足的线程块并行性而导致的需求违反，都可能抹去端到端延迟中的理论稀疏性增益。如图 4 所示，MD-SpMM 使用三种机制：(1) 权重打包，用于解码友好的稀疏存储；(2) 微密集执行，用于以 MMA 为中心的规整块；(3) 自适应并行性，以提高占用率同时控制缩减开销。

4.2 微密集：从 N:M 稀疏性到 Tensor Core 数据流

为了使 N:M 稀疏性与 Tensor Core 执行兼容，我们引入了微密集 (Micro-Dense)，它将 N:M SpMM 重新表述为以 MMA 为中心的规整数据流。在 N:M 稀疏性下，每个大小为 M 的窗口包含 N 个非零值。微密集不将其作为不规则的稀疏操作数暴露给计算阶段，而是仅将这 N 个非零值具体化为固定的、与 MMA 兼容的密集块，这些块与 Tensor Core 片段形状对齐，因此 Tensor Core 仍然对规整的密集片段进行操作，运行时计算完全基于 MMA。

这种设计将稀疏不规则性移出了 MMA 计算路径，移入了加载和解码阶段。为了保持由此产生的开销较低，我们采用了解码友好的打包表示 (B_pack)，它在每个 N:M 窗口内联合编码非零值及其位置。这消除了单独的索引加载，并支持轻量级解码，该解码与共享内存放置和 MMA 发布重叠。

因此，微密集将稀疏性处理限制在打包加载和轻量级解码之内，使得运行时计算