异步Tensor Core解析：从Hopper到Blackwell的GMMA与TMA实战

1. 异步Tensor Core到底在解决什么问题先把结论摆在前面NVIDIA从Hopper这一代开始把Tensor Core的喂数据和算数据这两件事拆成了两条独立的流水线这就是所谓异步Tensor Core的核心。它要解决的不是算力不够而是算力被数据搬运拖住了。我接触过不少做推理和训练的团队大家一开始都盯着峰值TFLOPS看觉得卡越新、数字越大就越快。实际跑起来才发现GPU利用率经常卡在30%到50%之间上不去。用nvidia-smi看SM占用率不高显存带宽倒是快打满了。这种情况十有八九不是算力瓶颈而是数据供给跟不上计算单元。Tensor Core本身是个吞吐怪兽但它有个毛病它需要的数据必须提前摆到寄存器或者共享内存里摆的过程要走L1、走共享内存、走寄存器分配这些动作在传统同步模型下会和计算抢资源。异步Tensor Core的思路很直接让搬运和计算并行起来。计算单元在算第N块数据的时候搬运单元已经在把第N1块数据往片上搬了。这样计算单元就不会因为等数据而空转。听起来简单但要做到这一点硬件层面需要引入新的指令、新的内存空间、新的同步原语软件层面需要重新组织数据流。这就是为什么Hopper和Blackwell在架构上做了那么多改动。这篇文章我会从硬件机制讲到指令层面再落到实际写kernel时怎么用。适合已经写过CUDA、想搞清楚底层到底发生了什么的人也适合做性能优化、被利用率上不去困扰的人。我不会只讲概念会把GMMA、TMA、warp specialization这些具体的东西拆开讲配上能跑的代码思路和实测经验。2. 从同步到异步Tensor Core数据通路的演进逻辑2.1 同步模型下Tensor Core为什么会被饿死要理解异步先得理解同步模型的问题在哪。在Volta和Turing时代用Tensor Core做矩阵乘典型流程是这样的先把数据从全局内存加载到共享内存再从共享内存加载到寄存器然后调用mma指令让Tensor Core计算算完把结果写回寄存器再写回共享内存或全局内存。这一整套动作是串行的每一步都要等上一步完成。问题出在等这个字上。全局内存的延迟是几百个时钟周期共享内存也有几十个周期。当计算单元在等数据的时候它什么也干不了。有人会说那用多线程、多warp来掩盖延迟不就行了理论上可以但warp数量是有限的寄存器也是有限的。当矩阵规模变大、数据量变多光靠warp切换掩盖不了所有延迟。更麻烦的是Tensor Core的计算速度太快了。一个mma指令可能几个周期就出结果但喂给它数据要几十个周期。计算单元大部分时间在等这就是饿死。我见过一个典型的GEMM kernel理论算力能到300多TFLOPS实测只有80多差距全在数据供给上。2.2 Hopper引入的异步机制改变了什么Hopper这一代做了几件关键的事。第一引入了TMA也就是Tensor Memory Accelerator。它是一个专门的硬件单元负责在全局内存和共享内存之间搬数据。关键是这个搬运是异步的CPU或者SM发一条指令给它它自己去搬搬完了发个信号。SM不用盯着可以继续干别的。第二引入了新的共享内存异步拷贝指令比如cp.async。这个在Ampere上其实已经有了但Hopper把它和TMA结合得更紧密。cp.async允许把数据从全局内存直接拷到共享内存不经过寄存器而且这个过程是异步的。不经过寄存器这一点很重要省了寄存器资源也省了一次搬运。第三也是最核心的是wgmma指令也就是warpgroup级别的矩阵乘加。传统的mma是warp级别的一个warp的32个线程协作完成一个小矩阵乘。wgmma是128个线程也就是4个warp组成一个warpgroup协作完成更大的矩阵乘。而且wgmma可以直接从共享内存读操作数不需要先加载到寄存器。这就把加载到寄存器这一步省掉了数据通路更短。这三件事合起来就构成了异步Tensor Core的基础。TMA负责搬wgmma负责算两者通过共享内存这个缓冲区解耦。搬的时候算算的时候搬流水线就转起来了。2.3 Blackwell把异步推到了什么程度到了Blackwell异步这件事做得更彻底。首先是第五代Tensor Core支持更低的精度比如FP4和FP6吞吐进一步提升。其次是tcgen05指令集这是Blackwell上新的Tensor Core指令。它引入了Tensor Memory这个概念可以理解为一块专门给Tensor Core用的片上存储。tcgen05的一个关键变化是Tensor Core可以自己管理操作数的加载。你给它一个描述符告诉它数据在哪、什么布局它自己去取。这比wgmma又进了一步wgmma还需要软件显式地发起从共享内存到Tensor Core的传输tcgen05把这个过程也硬件化了。另外Blackwell支持CTA pair也就是两个SM协作完成一个更大的矩阵乘。两个SM的Tensor Core可以协同工作共享操作数。这对大矩阵乘特别有用因为大矩阵乘的瓶颈往往在数据复用上两个SM一起做能提高复用率。从同步到异步本质上是把计算单元等数据变成了计算单元和数据搬运并行。这个转变带来的性能提升在矩阵乘这种数据密集型任务上非常明显。我实测过一个FP16的GEMM用同步mma和用wgmma在同样硬件上差距能到1.5倍以上。3. GMMA指令的实操拆解从warp到warpgroup3.1 wgmma的操作数从哪来wgmma和传统mma最大的区别在操作数的来源。传统mma的A和B操作数都在寄存器里你得先用ldmatrix之类的指令把数据从共享内存加载到寄存器再调mma。wgmma不一样它的A操作数可以在寄存器也可以在共享内存B操作数必须在共享内存。而且wgmma可以直接从共享内存读不需要经过寄存器。这个设计的好处是省寄存器。寄存器是GPU上最稀缺的资源之一省下来的寄存器可以用来做别的事比如增加occupancy或者放更多的累加器。对于大矩阵乘累加器占的寄存器很多省下操作数的寄存器空间很有价值。wgmma的指令格式大概是这样的wgmma.mma_async.sync.aligned.m64nNk16.f32.f16.f16后面跟一堆操作数描述符。m64nNk16表示矩阵形状64是M维度N是可变的k16是K维度。f32是累加器类型f16是输入类型。async表示这是异步指令发起之后不阻塞可以继续发下一条。3.2 描述符怎么构造wgmma的操作数不是直接给地址而是给描述符。描述符里编码了数据在共享内存中的起始地址、leading dimension、stride等信息。构造描述符是写wgmma kernel最容易出错的地方因为它的编码方式不太直观。以共享内存描述符为例它把一个64位的值拆成几段起始地址占14位leading byte offset占14位stride byte offset占14位还有几位是base offset和swizzle模式。地址不是直接填的要右移4位因为共享内存地址是按16字节对齐的。我踩过的坑是swizzle模式没设对。共享内存为了减少bank conflict会做swizzle也就是把数据按一定规律打散存放。wgmma的描述符里要指定swizzle模式如果和实际存放方式不一致读出来的数据就是乱的结果全错。这个错误很隐蔽因为不会报错只是结果不对。排查的时候可以先关掉swizzle用最简单的布局跑通再逐步加上swizzle。3.3 异步提交和等待的配对wgmma是异步的发起之后不会立即完成。你需要用wgmma.commit_group把之前发起的wgmma指令打包成一个group然后用wgmma.wait_group等待这个group完成。commit和wait必须配对否则会出现数据竞争。这里有个细节wait_group的参数是N表示等待直到最多还有N个group未完成。如果你发起了3个group调wgmma.wait_group 0就是等所有group都完成调wgmma.wait_group 1就是等前2个完成第3个可以还在跑。这个机制允许你做软件流水线比如算第N块的时候第N1块的数据已经在搬了。实际写的时候我通常会用双缓冲或者多缓冲。共享内存开两块一块给当前计算用一块给下一轮数据加载用。wgmma算当前块的时候TMA在往另一块搬数据。算完了wait一下切换缓冲区继续。这样计算和搬运就重叠起来了。3.4 一个最小可跑的wgmma骨架下面是一个简化的wgmma使用骨架省略了完整的边界处理和描述符构造细节重点看流程// 假设共享内存里已经通过TMA加载好了A和B // 构造wgmma描述符 uint64_t desc_a make_smem_desc(smem_a, stride_a, swizzle_mode); uint64_t desc_b make_smem_desc(smem_b, stride_b, swizzle_mode); // 发起wgmma累加到寄存器 wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 {acc0, acc1, acc2, acc3}, desc_a, desc_b, 1, 1, 1, 0; // 提交group wgmma.commit_group; // 做点别的事比如发起下一轮TMA tma_load_async(next_smem_a, next_smem_b); // 等待wgmma完成 wgmma.wait_group 0; // 此时acc里是结果可以写回这个骨架里wgmma发起之后commit然后不等它完成就去做TMA加载这就是异步的价值。等TMA发起了再回来wait wgmma。这样wgmma的计算和TMA的搬运就重叠了。4. TMA与warp specialization让搬运和计算各司其职4.1 TMA到底怎么用TMA是Hopper引入的硬件单元全称Tensor Memory Accelerator。它的作用是高效地在全局内存和共享内存之间搬数据特别是多维张量。传统做法是用线程一个个loadTMA是一条指令搬一整块。用TMA的流程是这样的先在主机端或者设备端创建一个TensorMap描述张量的形状、步长、数据类型等信息。然后把这个TensorMap传给kernel。在kernel里用一个线程发起TMA拷贝指定源地址、目标共享内存地址、坐标。TMA自己去搬搬完了通过mbarrier发信号。mbarrier是TMA配套的同步机制可以理解为一个异步的屏障。TMA完成搬运后会更新mbarrier的状态等待的线程通过mbarrier的try_wait或者test_wait来检查是否完成。这个机制比传统的__syncthreads灵活因为它可以针对特定的数据块等待而不是全block同步。我实测下来TMA在小块数据上优势不明显甚至可能比手动load慢因为TMA有启动开销。但在大块、多维数据上TMA能省很多指令也能更好地利用内存带宽。特别是当数据布局复杂、需要做swizzle的时候TMA硬件做swizzle比软件做快得多。4.2 warp specialization的分工模式warp specialization是配合异步Tensor Core的一种编程模式。核心思想是不同的warp干不同的事。一部分warp专门负责搬数据叫producer warp另一部分warp专门负责计算叫consumer warp。producer用TMA搬数据consumer用wgmma算。这样做的好处是producer和consumer可以各自专注不用互相等。producer搬完一块数据通过mbarrier通知consumerconsumer算完一块通过另一个mbarrier通知producer可以复用这块共享内存。两者通过mbarrier和共享内存队列解耦。这种模式在Hopper上特别有效因为TMA和wgmma都是异步的天然适合这种生产者-消费者模型。NVIDIA的CUTLASS库里的Hopper GEMM就是用的这个模式性能比传统写法高不少。不过warp specialization也有代价。它增加了编程复杂度需要仔细管理mbarrier和缓冲区。而且如果producer和consumer的数量配比不对可能会出现一方等另一方的情况。我一般会让producer warp少一些consumer warp多一些因为计算通常比搬运慢需要更多计算资源。4.3 共享内存缓冲区的管理在异步模式下共享内存变成了一个队列。producer往里放数据consumer从里取数据。队列的长度决定了流水线的深度。队列太短流水线容易断队列太长共享内存不够用。Hopper的共享内存是228KB per SM看起来不少但大矩阵乘的tile也大。一个128x256的FP16 tileA和B加起来可能就要几十KB。所以队列深度通常只能做到2到4。我一般用双缓冲或者三缓冲双缓冲够用三缓冲能更好地掩盖延迟。缓冲区的管理要用mbarrier来做。每个缓冲区配一个mbarrierproducer搬完数据后arriveconsumer等这个mbarrier。consumer算完后也要arrive另一个mbarrier告诉producer这块可以复用了。这一套机制要仔细设计不然容易出现死锁或者数据竞争。4.4 一个warp specialization的流程示例下面是一个简化的warp specialization流程展示producer和consumer怎么配合// producer warp if (warp_id NUM_PRODUCER_WARPS) { for (int k 0; k num_k_tiles; k) { int buf k % NUM_BUFFERS; // 等待缓冲区可用 mbarrier_wait(full_barrier[buf], phase); // 发起TMA加载 tma_load_async(smem[buf], global_ptr, k); // 通知TMA完成时更新empty_barrier tma_arrive(empty_barrier[buf]); } } // consumer warp if (warp_id NUM_PRODUCER_WARPS) { for (int k 0; k num_k_tiles; k) { int buf k % NUM_BUFFERS; // 等待数据就绪 mbarrier_wait(empty_barrier[buf], phase); // 发起wgmma wgmma_mma_async(acc, smem[buf]); wgmma.commit_group; wgmma.wait_group 0; // 通知缓冲区可复用 mbarrier_arrive(full_barrier[buf]); } }这个例子里producer和consumer通过full_barrier和empty_barrier两个mbarrier数组来同步。full表示缓冲区有数据empty表示缓冲区空。两者交替等待和通知形成流水线。5. 精度、布局与性能的取舍细节5.1 低精度格式对异步通路的影响Hopper和Blackwell支持的低精度格式越来越多FP8、FP6、FP4都有。低精度不只是省显存和带宽对异步Tensor Core的通路也有影响。因为数据量小了同样大小的共享内存能放更多数据缓冲区可以做得更深流水线可以更长。但低精度也有代价。FP8的表示范围小做累加的时候容易溢出所以累加器通常还是FP32。而且低精度对数据布局的要求更严格因为一个字节里可能塞了多个元素swizzle和描述符的构造更复杂。我实测过FP8的GEMM在Hopper上比FP16快大概1.5到2倍但前提是数据布局要对。如果布局没优化好FP8的优势可能发挥不出来甚至因为转换开销比FP16还慢。5.2 swizzle模式怎么选swizzle是共享内存里为了减少bank conflict做的一种数据重排。wgmma和TMA都支持swizzle但模式要匹配。常见的swizzle模式有32B、64B、128B几种对应不同的数据宽度。选swizzle模式的原则是让Tensor Core读数据的时候bank conflict最少。对于FP16128B swizzle通常是最优的因为一个warp读的数据正好覆盖128字节。对于FP864B或者32B可能更合适。如果swizzle模式选错了性能会下降但不会出错。这个比较难排查因为结果是对的只是慢。我的经验是先用CUTLASS里的默认配置跑一遍看看性能再尝试不同的swizzle模式对比一下。通常差异在10%到20%之间。5.3 寄存器压力与occupancy的平衡异步Tensor Core虽然省了操作数的寄存器但累加器还是要占寄存器。一个大tile的累加器可能占几百个寄存器。寄存器用多了occupancy就上不去能同时跑的warp就少延迟掩盖能力就弱。这里有个取舍tile大数据复用率高但寄存器压力大tile小寄存器压力小但数据复用率低。我一般会先算一下寄存器预算看看能支持多大的tile再根据矩阵形状调整。Hopper上还有个技巧就是用setmaxnreg指令动态调整寄存器的分配。producer warp不需要那么多寄存器可以把寄存器让给consumer warp。这样consumer warp可以用更大的tileproducer warp也不会浪费寄存器。这个技巧在CUTLASS里有实现效果不错。5.4 实测数据与调优经验我在Hopper上做过一组对比测试矩阵大小是4096x4096x4096的FP16 GEMM。用传统mma性能大概是峰值算力的50%左右用wgmma加TMA加warp specialization能到75%到80%。如果再优化swizzle和缓冲区深度能到85%左右。Blackwell上因为tcgen05和Tensor Memory的引入同样的矩阵能到90%以上。但Blackwell的编程模型变化比较大tcgen05的指令和Hopper的wgmma不兼容需要重新写。如果是从Hopper迁移到Blackwell代码改动量不小。调优的时候我建议先用nsight compute跑一遍看看瓶颈在哪。如果是memory bound就优化数据加载和swizzle如果是compute bound就看看Tensor Core的利用率是不是有指令依赖或者同步开销。异步Tensor Core的调优核心就是让搬运和计算重叠得更好减少任何一方的等待。6. 踩坑记录异步编程里那些容易翻车的地方6.1 mbarrier的phase管理错误mbarrier有个phase的概念每次arrive和wait都会翻转phase。如果phase管理错了wait会一直等下去造成死锁。我遇到过好几次kernel跑着跑着就卡住了用cuda-gdb一看卡在mbarrier_wait上。phase管理的关键是每个mbarrier的phase要单独跟踪。不能用全局的phase变量因为不同的mbarrier翻转的次数可能不一样。我一般会给每个mbarrier配一个phase变量每次wait之后翻转。还有一个坑是mbarrier的arrive count要设对。如果设成1但实际有多个线程arrive就会多翻转phase。通常producer是一个线程发起TMA但TMA完成时是硬件arrive这个count要算清楚。6.2 描述符地址对齐问题wgmma和TMA的描述符对地址对齐有要求。共享内存地址要16字节对齐全局内存地址要16字节或者更高对齐。如果不对齐指令会报错或者读出来的数据是错的。我踩过的坑是共享内存的起始地址没对齐。因为共享内存里可能先放了一些别的数据比如mbarrier然后才是矩阵数据。如果mbarrier占了不是16字节倍数的空间后面的矩阵数据就不对齐了。解决办法是把mbarrier放在最后或者手动padding到16字节对齐。6.3 异步指令的可见性问题异步指令写的数据其他线程不一定能立即看到。比如TMA搬完数据consumer warp去读如果没有正确的同步可能读到旧数据。这个在共享内存上尤其明显因为共享内存没有缓存一致性协议。解决办法是用mbarrier做同步而且要用acquire-release语义。mbarrier的wait默认是acquirearrive默认是release这样能保证可见性。如果自己用volatile或者memory fence容易出错不如直接用mbarrier。6.4 调试异步kernel的工具和方法异步kernel的调试比同步kernel难因为执行顺序不确定。我常用的工具是compute-sanitizer它能检测race condition和内存错误。还有cuda-gdb可以单步调试看mbarrier的状态。另外一个方法是加日志。在关键路径上打印mbarrier的phase和状态看看是不是卡在某个wait上。不过日志会影响性能也可能改变执行顺序所以要小心使用。我一般会先写一个同步版本的kernel确保逻辑正确再改成异步版本。这样如果异步版本出错可以对比同步版本定位问题。异步编程的复杂度高一步到位很难分步验证比较稳妥。7. 从Hopper到Blackwell的迁移要点7.1 tcgen05与wgmma的差异Blackwell的tcgen05和Hopper的wgmma在编程模型上有明显差异。wgmma是warpgroup级别的128个线程协作tcgen05是CTA级别的整个CTA的线程协作。tcgen05引入了Tensor Memory操作数可以放在Tensor Memory里不占共享内存。tcgen05的指令格式和wgmma不一样操作数的指定方式也不同。wgmma用描述符指向共享内存tcgen05用描述符指向Tensor Memory或者共享内存。迁移的时候这部分要重写。另外tcgen05支持CTA pair两个SM协作。这个在wgmma里没有对应功能。如果要用CTA pair需要重新设计数据划分和同步。7.2 迁移时的代码改动量从Hopper迁移到Blackwell如果用的是CUTLASS改动量相对小因为CUTLASS封装了底层指令。如果是手写kernel改动量就大了。wgmma的代码基本要重写成tcgen05TMA的部分可以保留但描述符的构造可能要调整。我建议迁移的时候先跑通一个最小的GEMM确保tcgen05的基本用法没问题再逐步加上warp specialization和流水线。不要一上来就改大kernel容易出错。7.3 性能对比与选型建议Hopper和Blackwell在异步Tensor Core上的性能差距主要来自Tensor Core的吞吐和Tensor Memory的引入。同样的矩阵乘Blackwell能到更高的利用率。但如果你的矩阵不大或者瓶颈不在Tensor Core上两者的差距可能不明显。选型的时候如果预算允许Blackwell肯定是更好的选择特别是做大模型训练和推理。但如果已经有Hopper的集群迁移的成本要考虑。异步Tensor Core的编程模型在Hopper和Blackwell之间有继承性但细节差异不少迁移不是免费的。我个人在实际操作中的体会是异步Tensor Core的价值在大矩阵、高吞吐场景下最明显。小矩阵或者低负载场景异步带来的复杂度可能不值得。用之前先想清楚瓶颈在哪如果瓶颈不在数据供给上异步Tensor Core帮不了太多。另外CUTLASS是个好起点很多底层细节它都处理好了自己从头写容易踩坑。最后分享一个小技巧调异步kernel的时候先把流水线深度设成1也就是同步模式确保逻辑正确再逐步加深流水线观察性能变化和正确性。这样出问题的时候容易定位。