1. 项目概述异步GPU架构下的稀疏矩阵乘法加速稀疏矩阵乘法SpMM是科学计算、图神经网络和大语言模型推理等领域的核心计算瓶颈。简单来说它计算一个稀疏矩阵A和一个稠密矩阵B的乘积C。这里的“稀疏”意味着矩阵A中绝大部分元素都是零直接进行稠密矩阵乘法会浪费大量计算资源在零元素上。因此SpMM算法的核心挑战在于高效地组织非零元素的计算同时最大限度地利用现代硬件的并行计算能力。传统的GPU SpMM优化思路主要围绕如何让线程高效地协作进行数据的加载和计算。然而随着NVIDIA Hopper等新一代GPU架构的引入游戏规则发生了变化。这些架构的核心思想是“异步”和“专用化”将数据搬运由Tensor Memory Accelerator, TMA负责和矩阵计算由Warpgroup Matrix-Multiply-Accumulate, WGMMA负责这两个任务解耦交给不同的硬件单元并行执行。这就像在厨房里以前只有一个厨师CUDA核心既要切菜加载数据又要炒菜计算现在有了专门的配菜师TMA和炒锅Tensor Core他们可以同时工作大大提升了整体效率。AsyncSparse项目正是瞄准了这一架构演进带来的新机遇。它不再仅仅优化线程间的协作而是从系统层面重新设计SpMM内核使其与Hopper的异步执行模型深度协同。项目提出了两种针对不同稀疏模式的内核BCSR块压缩稀疏行格式擅长处理非零元素聚集的“结构化稀疏”而WCSR窗口压缩稀疏行格式则能更灵活地处理“不规则稀疏”。通过将TMA的异步数据搬运、WGMMA的异步张量核心计算以及Warp专用化让不同的线程组专门负责加载或计算等技术有机结合AsyncSparse实现了显著的性能提升。对于从事高性能计算、AI推理系统优化或底层GPU编程的开发者而言理解这套设计思路不仅是掌握一个工具更是洞察未来硬件优化方向的关键。2. 核心设计思路与架构解析2.1 异步执行范式的价值重估在深入代码之前我们必须理解为什么“异步”在Hopper架构上如此重要。在Ampere或更早的架构上执行模型本质上是同步的。一个Warp32个线程发出mma.sync指令进行矩阵乘加时它必须等待数据从寄存器中准备好并且指令执行完成才能继续。数据加载同样需要整个Warp的线程协同进行地址计算和内存访问。这种模式导致计算单元Tensor Core在等待数据时闲置而内存带宽在计算时未被充分利用。Hopper架构通过两项革新打破了这一瓶颈异步张量核心指令WGMMAwgmma.mma_async指令由128个线程组成的Warpgroup集体发出。关键之处在于它可以直接从共享内存中读取操作数并且指令本身是异步的——发出后立即返回计算在Tensor Core单元上后台执行。程序需要通过wgmma.fence、wgmma.commit_group和wgmma.wait_group这一套显式的栅栏协议来确保计算完成的顺序性。张量内存加速器TMA这是一个专用的硬件DMA引擎。以前加载一个数据块到共享内存需要所有线程参与计算全局地址、发出加载指令、将数据写入共享内存并处理存储体冲突。现在只需要一个线程发出cp.async.bulk.tensor指令TMA单元就会自动完成从全局内存到共享内存的整个批量数据传输包括地址计算和避免存储体冲突的数据重排Swizzle。这极大地解放了线程减少了寄存器压力。AsyncSparse的设计核心就是构建一个高效的“生产者-消费者”流水线让TMA生产者和WGMMA消费者这两个强大的异步硬件单元持续、重叠地工作从而压榨出硬件的每一分性能。2.2 双格式策略BCSR与WCSR的权衡没有一种稀疏格式能通吃所有场景。AsyncSparse采用了两种互补的格式来应对不同的稀疏模式这是其高性能的关键。2.2.1 块压缩稀疏行BCSR格式BCSR将稀疏矩阵A在逻辑上划分为固定大小的块例如64x64。它只存储那些包含至少一个非零元素的块并将这些块以稠密矩阵的形式连续存放。通过三个数组来记录结构信息block_row_ptr: 记录每个块行block-row起始的非零块索引。block_col_idx: 记录每个非零块所在的块列block-column索引。blocks: 连续存储所有非零块的稠密值。设计优势TMA友好每个非零块在内存中是连续的二维数据块TMA可以发起一次高效的2D批量加载硬件自动进行Swizzle优化。计算规整每个非零块都可以被视为一个小的稠密矩阵乘法micro-GEMM完美匹配WGMMA指令的计算模式。B矩阵访问连续由于B矩阵是稠密的且BCSR按块列索引访问B的行这些行在内存中也是连续的同样可以用TMA高效加载。设计代价存储与计算浪费块内可能包含零元素。这些零不仅被存储还会参与计算浪费存储空间和算力。我们用一个“填充率”fill_ratio nnz / (nnz_blocks * b_row * b_col)来衡量这种浪费。对于非零元素非常分散的矩阵BCSR的填充率会很低效率不佳。2.2.2 窗口压缩稀疏行WCSR格式WCSR提供了更细粒度的压缩以应对不规则稀疏。它将矩阵的行分组为固定高度的“窗口”例如64行。对于每个窗口它收集该窗口内所有行出现过的列索引的并集然后仅存储这些列对应的值。数据同样由三个数组组织window_row_ptr: 记录每个窗口起始的压缩列索引。window_col_idx: 记录每个压缩列位置对应的原始列索引用-1表示填充的列。values: 按窗口和压缩列存储的稠密值。设计优势高压缩率只存储窗口内实际出现的列避免了块内零值的存储对于高度不规则的稀疏模式更节省内存。计算更精确只计算实际有非零值的列避免了BCSR的块内零值计算。设计挑战B矩阵间接访问这是WCSR最大的性能挑战。由于window_col_idx数组的存在要加载B矩阵中对应的行需要进行一次间接查找。TMA无法处理这种非连续的、依赖索引的加载模式因此必须退回到由线程协作完成的“聚集”Gather操作这会增加开销。负载不均衡不同窗口的非零列数差异可能巨大。一个窗口可能有几百个非零列而另一个可能只有几个。如果简单地将一个窗口分配给一个线程块会造成严重的负载不均。AsyncSparse为这两种格式量身定制了不同的内核实现策略这正是其设计的精妙之处没有追求单一的“银弹”而是根据数据特征选择最优路径。3. BCSR内核Warp专用化生产者-消费者流水线针对BCSR格式规整、TMA友好的特点AsyncSparse设计了一个三阶段Warp专用化流水线这是其性能飞跃的核心。3.1 线程块与Warpgroup的角色分配一个线程块CTA包含384个线程划分为3个Warpgroup每个128线程生产者Warpgroup (WG 0)唯一职责是使用TMA加载数据。由于TMA加载只需一个线程发起该Warpgroup中实际上只有线程0在工作其余127个线程在此阶段空闲。这种“浪费”是刻意为之旨在将数据搬运任务完全卸载给专用硬件。消费者Warpgroup (WG 1 2)唯一职责是执行WGMMA计算。每个消费者Warpgroup负责输出矩阵C的一部分列例如如果总列宽BN256则每个消费者负责128列。它们从共享内存的缓冲区中读取A和B的数据块进行计算。3.2 多阶段环形缓冲区与同步机制为了实现生产与计算的重叠内核在共享内存中维护了一个深度为Q通常为3的环形缓冲区。每个缓冲区槽Stage能存放一个A块64x64和一个B切片64xBN。同步通过两组内存屏障mbarrier数组full[Q]和empty[Q]来管理并配合“相位比特”进行精确跟踪生产者流程生产者等待empty[q]屏障表示槽q为空且可用。然后它通过mbarrier.arrive.expect_tx设置期望的事务字节数即即将加载的数据量接着发起针对A块和B切片的TMA加载。TMA加载完成后硬件会自动触发full[q]屏障。消费者流程消费者等待full[q]屏障表示槽q的数据已就绪。然后它们执行一系列WGMMA指令处理完该槽的数据后触发empty[q]屏障通知生产者该槽可被复用。通过设置Q3可以实现“双缓冲”甚至“三缓冲”当消费者在处理槽i的数据时生产者可以同时为槽i1加载数据。理想情况下当消费者处理完槽i时槽i1的数据已经加载完毕计算单元无需等待。3.3 动态寄存器分配与资源管理三个Warpgroup对寄存器资源的需求很高尤其是消费者Warpgroup需要大量寄存器来存放WGMMA累加器的中间结果Fragment。如果静态分配可能导致每个SM只能驻留一个线程块降低占用率Occupancy。AsyncSparse利用了Hopper的setmaxnregPTX指令进行运行时寄存器池管理生产者Warpgroup执行setmaxnreg.dec.sync.aligned24主动释放一部分寄存器到池中。消费者Warpgroup执行setmaxnreg.inc.sync.aligned240从池中申请更多寄存器。这要求同一个Warpgroup内的所有线程同步执行该指令并且在寄存器调整前后需要显式同步。这种动态分配机制确保了在高寄存器需求下SM上仍能容纳多个线程块更好地隐藏了各种延迟。实操心得屏障与相位管理实现这种异步流水线时最易出错的是屏障和相位管理。务必确保“到达”arrive和“等待”wait操作配对正确并且相位比特在每次循环后正确翻转。一个常见的调试技巧是使用__nanosleep在关键同步点插入微小延迟模拟最坏情况下的执行交错以暴露竞争条件。此外TMA加载的完成信号是通过屏障传递的必须确保在发起加载前正确初始化屏障的期望计数。3.4 性能收益分析从论文的消融实验来看在BCSR内核上逐步启用这些异步特性带来了巨大收益仅使用WGMMAopt1相比纯CUDA核心标量计算基线opt0性能提升约5.3倍。这证明了即使同步使用张量核心的算力也远超通用核心。启用TMAopt2在opt1基础上用TMA替代协作加载性能又提升了约3.7倍相对基线。这省去了地址计算的开销并开启了异步重叠的可能性。启用Warp专用化opt3这是最大的单次提升在opt2基础上又带来了约2.75倍的性能增长相对基线。这完全归功于生产与计算在硬件层面的真正并行。三者叠加贡献了超过98%的总性能提升将性能从远低于cuSPARSE提升到其4.3倍以上。这清晰地表明对于BCSR这类规整稀疏模式构建一个与异步硬件深度匹配的流水线是性能优化的决定性因素。4. WCSR内核间接加载与动态负载均衡WCSR内核面临BCSR所没有的挑战B矩阵的间接访问和窗口间的不均衡负载。因此它的设计哲学与BCSR截然不同。4.1 单Warpgroup混合加载计算模型由于B矩阵的加载需要根据window_col_idx进行聚集操作这需要所有线程参与。因此为加载专门分配一个Warpgroup像BCSR那样收益甚微因为该Warpgroup的所有线程在加载阶段本就处于忙碌状态。因此WCSR内核采用单Warpgroup设计128线程在每个迭代中依次执行混合加载阶段线程0发起TMA加载将当前窗口的A值连续内存加载到共享内存。与此同时所有128个线程协作执行B矩阵的聚集加载每个线程读取一个window_col_idx根据索引从全局内存获取对应的B矩阵行并手动以128字节Swizzle模式写入共享内存。同步等待通过__syncthreads()和TMA完成等待确保A和B的数据在共享内存中就绪。计算阶段同一个Warpgroup集体执行WGMMA指令进行计算。这种设计简化了同步没有跨Warpgroup的屏障但将加载和计算串行化在同一个Warpgroup内无法实现BCSR那样的深度重叠。4.2 基于任务分解的负载均衡策略WCSR格式下不同窗口的非零列数即工作量差异可能极大。直接将一个窗口分配给一个线程块会导致严重的负载不均处理小窗口的线程块早早完工闲置而处理大窗口的线程块成为瓶颈。AsyncSparse的解决方案是任务分解任务生成在启动内核前或在内核初始化阶段遍历所有窗口。如果一个窗口的压缩列数超过一个阈值例如大于MAX_COLS_PER_TASK则将其水平切分成多个固定大小的子任务。任务映射线程块索引blockIdx.x不再直接对应窗口索引而是对应一个全局任务描述符数组的索引。每个任务描述符包含{窗口ID, 列起始偏移, 列数}。原子累加保证正确性当多个线程块处理同一个窗口的不同部分时它们会向输出矩阵C的同一行累加部分结果。这需要通过atomicAdd操作来保证结果的正确性。虽然原子操作有开销但相比因负载不均导致大量SM闲置的代价这通常是更优的选择。这种动态调度策略将不规则的工作负载转化为粒度更均匀的任务更好地利用了GPU上大量的线程块是处理高度不规则稀疏问题的关键。4.3 与BCSR内核的对比与选型下表总结了两种内核的核心区别与适用场景特性BCSR内核WCSR内核核心架构生产者-消费者流水线Warp专用化单Warpgroup混合加载计算数据加载A和B均通过TMA连续加载A通过TMA加载B通过线程协作间接加载负载均衡依赖CUDA运行时调度每个块行工作量固定通过任务分解实现动态负载均衡计算效率高但可能计算块内零值高只计算非零列但间接加载有开销存储效率低存储块内零值填充率是关键高只存储窗口内出现的列最佳场景结构化稀疏非零元素聚集填充率高50%不规则稀疏非零元素分散或窗口内列数差异大性能瓶颈TMA加载带宽WGMMA计算吞吐B矩阵的间接聚集带宽任务调度开销选型建议在实际应用中可以先对稀疏矩阵进行简单分析。计算其在不同块大小下的BCSR填充率。如果填充率较高例如在目标块大小下超过70%则BCSR内核通常是更优选择因为它能充分发挥异步流水线的威力。如果填充率很低或者矩阵本身极度不规则则WCSR内核更能避免存储和计算浪费其动态负载均衡机制也能更好地应对不均匀性。5. 高级优化尝试与其在稀疏场景下的失效分析论文不仅展示了成功的优化还坦诚地分享了在密集矩阵乘法GEMM中常用、但在SpMM中失效的高级优化技术这部分经验极具参考价值。5.1 线程块集群与TMA多播Hopper支持线程块集群Cluster允许最多16个协同调度的线程块通过分布式共享内存DSMEM相互访问。结合TMA多播Multicast一个线程块发起的TMA加载可以广播到集群内所有线程块的共享内存中。理论优势在SpMM中沿N维度划分的相邻线程块需要相同的A矩阵块但B的不同列切片。使用多播只需加载一次A块即可广播给集群内所有块减少L2缓存流量和全局内存带宽压力。实际结果与失效原因实验显示启用2-CTAs集群多播后性能反而出现了倒退。协同调度约束集群要求所有CTAs必须同时在相邻的SM上启动。这限制了CUDA运行时调度器的灵活性。对于稀疏负载不同块行的工作量差异很大灵活的调度本可以更快地完成轻量任务并接手重型任务。集群的刚性约束破坏了这种动态负载均衡能力。跨CTA同步开销在生产者-消费者流水线中消费者需要通知生产者缓冲区已空。在集群模式下这需要跨CTA的屏障操作mbarrier.arrive.shared::cluster其延迟远高于CTA内部的屏障。收益有限对于稀疏矩阵A块的数据复用机会本身就不如密集GEMM频繁。不同块行访问的A块不同集群内共享A块的收益被高昂的同步和调度代价所抵消。避坑指南不要盲目移植密集优化这是一个经典教训在密集计算中行之有效的优化如多播、持久化内核在稀疏计算中可能无效甚至有害。稀疏计算的固有特征是不规则性和负载不均。任何增加调度刚性或引入额外同步的优化都必须谨慎评估其在稀疏负载下的实际收益。在设计稀疏内核时应优先保证调度灵活性和负载均衡。5.2 持久化内核与PID置乱持久化内核是另一种在密集GEMM中提升L2缓存命中率的常用技术。它启动固定数量的线程块通常等于SM数这些线程块持久地驻留在SM上以轮询方式处理所有输出块。理论优势重叠存储与加载在处理当前输出块的同时可以预加载下一个块的数据。PID置乱通过以列优先的顺序重新映射线程块ID确保同时活跃的SM处理空间上相邻的输出块从而提高B矩阵切片在L2缓存中的复用率。实际结果与失效原因在SpMM中持久化内核导致了性能衰退。严重的负载不均这是致命伤。在稀疏矩阵中不同输出块对应不同的A块行需要计算的非零块数量天差地别。采用静态轮询调度的持久化内核中处理稀疏行的线程块很快完工并进入空闲等待而处理密集行的线程块成为瓶颈。这种“木桶效应”抵消了持久化带来的启动开销减少和潜在的缓存优化收益。L2复用收益低PID置乱旨在让相邻块复用B的列。但在SpMM中每个输出块通过block_col_idx访问B中完全不同的行集这些行在内存中可能毫不相邻因此L2缓存复用的机会很少。累加器清零开销持久化内核需要为每个新任务重置累加器。在密集GEMM中每个任务包含数百次迭代清零开销可忽略。但在SpMM中一个任务可能只包含几个非零块清零开销就显得相对较大。最终选择AsyncSparse最终选择了非持久化内核。CUDA运行时的默认调度器本身就具备动态负载均衡能力处理完稀疏行任务的SM会立即被分配新的、更繁重的任务。这种动态性无需额外的原子操作或任务队列开销对于不规则稀疏工作负载而言是最简单高效的策略。6. 关键参数调优与性能分析6.1 分块大小Tile Size的选择在WGMMA指令中M维度固定为64匹配A的块行或窗口高度K维度固定为16用于BF16累加。唯一的自由参数是N维度即WGMMA_N它可以是8到256之间的8的倍数。选择策略更大的WGMMA_N意味着每次加载一个A块后能计算更宽的B切片BN 2 * WGMMA_N因为有两个消费者。这能更好地分摊TMA加载和屏障同步的开销提升计算密度。如图7所示几何平均吞吐量随着WGMMA_N增大而显著提升。填充开销如果BN不能整除稠密矩阵的宽度N内核必须将B填充到BN的整数倍。这些填充的列会参与计算但不贡献有效结果造成浪费。当WGMMA_N较大时填充带来的浪费比例也更大。经验法则选择能整除N的最大WGMMA_N值。例如当N1024时WGMMA_N128BN256是一个好的选择因为它能整除1024且提供了较大的计算粒度。如果N是变量一种稳健的策略是选择一组固定的WGMMA_N值如64, 128, 256并在运行时根据N选择填充开销最小的那个。6.2 实际性能表现与对比在414个SuiteSparse矩阵上的综合评估显示了AsyncSparse的显著优势WCSR全面领先在稠密矩阵宽度N1024、矩阵密度≥1%的子集上WCSR内核达到了23.53 TFLOPS的几何平均吞吐量相比cuSPARSE的BELL格式实现了4.86倍加速相比最新的研究内核AccSpMM和FlashSparse也分别有2.40倍和2.85倍的优势。BCSR在结构化稀疏中表现优异在密度≥0.1%的矩阵上BCSR内核开始超越所有基线。在密度≥1%时其性能达到18.05 TFLOPS虽略低于WCSR但仍远优于其他方案。这证明了在非零元素聚集的场景下其异步流水线的强大威力。扩展性良好随着N从256增加到1024WCSR和BCSR的性能持续提升而AccSpMM等基线性能却有所下降。这表明AsyncSparse的设计能更好地利用宽矩阵带来的计算并行度而传统基于散射-聚集scatter-gather的方法可能在内存带宽上遇到瓶颈。6.3 端到端大语言模型推理加速将BCSR内核集成到Qwen2.5-7B模型的前馈网络FFN中并与稀疏注意力技术MInference结合展示了其在真实应用中的价值FFN层加速在90%块稀疏度下单个FFN投影层gate_proj相比稠密cuBLAS实现获得了1.58倍到1.98倍的加速。在99%稀疏度下加速比可达3.19倍。端到端收益在短序列如1K时FFN是主要计算瓶颈稀疏FFN能带来**~1.4倍的端到端加速。在长序列如64K时注意力计算成为主导稀疏注意力技术MInference能带来1.73倍加速。而将两者结合在64K序列长度、90% FFN稀疏度下实现了2.66倍**的端到端加速。这证明了稀疏化FFN和注意力是互补的优化方向。7. 局限性与未来展望7.1 当前设计的局限粒度固定BCSR和WCSR的块/窗口高度固定为64以匹配WGMMA的M维度。对于某些稀疏模式更细的粒度如32或16可能减少填充开销但这需要重新设计TMA描述符和WGMMA流水线目前不支持。预处理依赖本文未深入探讨矩阵预处理算法如行列重排序。更智能的预处理可以将非零元素更紧密地聚集在一起从而提升BCSR的填充率或改善WCSR的负载均衡这有进一步优化的空间。格式转换开销从通用的CSR或COO格式转换为BCSR/WCSR需要预处理步骤。在动态稀疏性如稀疏注意力的场景下这部分开销需要被考虑在内。7.2 未来优化方向支持FP8精度Hopper及后续架构支持FP8精度的WGMMA其吞吐量是BF16的两倍。支持FP8将直接受益于量化大语言模型推理带来进一步的性能提升。多GPU扩展通过NCCL等通信库将计算扩展到多GPU以处理单卡内存无法容纳的超大规模矩阵。负载均衡算法改进探索更高效的任务窃取Work Stealing或基于代价模型的动态调度策略以应对极端不规则负载。面向新硬件研究如何将TMA和Warp专用化流水线迁移到NVIDIA Blackwell等新一代GPU架构并利用其新的集群启动控制等特性。编译器集成将AsyncSparse的核心思想异步流水线、双格式选择集成到Triton、TACO等高级编译器中降低开发者使用门槛使其能自动应用于更广泛的稀疏张量运算。AsyncSparse的工作揭示了一个重要趋势未来的高性能稀疏计算必须与GPU底层异步、专用化的硬件设计哲学协同。它不仅仅是一组高性能内核更是一份关于如何在新时代的硬件上重新思考稀疏计算范式的设计蓝图。