@KyrieBlunders：用 CuTeDSL 手写了一个 Blackwell GEMM，在 B200 上从 164 提升到 1401 TFLOP/s，约为 cuBLAS 的 97%

手写一个 Blackwell GEMM：从 164 提升到 1400 TFLOP/s

我想弄清楚 Blackwell 上一次快速矩阵乘法到底是由什么构成的。不是 API 调用，而是这些调用背后的理由。于是我在单张 B200 上从零手写了一个，用的是 CUTLASS 的 Python DSL（CuTeDSL），只用底层原语（tcgen05、TMA、mbarrier、TMEM），完全不用 CUTLASS 仓库里附带的那些辅助抽象层。那些辅助层很好，但它们恰好把我最想搞懂的东西藏了起来。

第一个算出正确 GEMM 结果的版本跑到了 164 TFLOP/s，只有 cuBLAS 的 11%。最后一个版本在 4096 立方规模的问题上跑到了 1401 TFLOP/s，达到 cuBLAS 的 97%。总共八次改动把它带到了这个位置。改动清单本身很无聊。有意思的地方在于：其中三次改动毫无效果，一次几乎完成了全部工作，而哪次改动重要取决于你跑的是 prefill 还是 decode。

4096 立方规模下每个内核版本的 GEMM 耗时

4096 立方规模下每个内核版本的 GEMM 耗时

ps：本文中的实验全部由 Claude 完成，不是我。

是两种工作负载，不是一种

大多数 GEMM 文章都用方阵做基准测试，这几乎说明不了推理服务器实际在跑什么。当 N 和 K 很大时，GEMM 每字节内存访问大约对应 M 次浮点运算，而这一个数字就决定了内核的瓶颈在哪里。在 prefill 中，M 是几千个 prompt token，所以瓶颈是 tensor core。在 decode 中，M 是 batch size，几十到两百；在 M=256 时，内核大约是 250 FLOP/B，正好卡在约 200 的临界点附近（此时 cuBLAS 峰值接近 1500 TFLOP/s，HBM 接近 7.7 TB/s）。

我从 Llama 3 70B 的线性层中取了 N 和 K。prefill 用 M=8192，对应 attention 输出投影（N=8192, K=8192）和 MLP up 投影（N=28672, K=8192）。decode 用 M=256，对应 attention 输出、融合 QKV（N=10240, K=8192）、MLP up 和 MLP down（N=8192, K=28672）。我还保留了 4096 立方作为熟悉的参照，并在每一步都用 Nsight Compute 做了 profiling；其他形状则在部分版本上做了剖析。

一个局限：最终的内核要求 M 是 256 的倍数，因为一对 CTA 覆盖 256 行。真实的 decode batch 往往是 1 到 64，在 8192×8192 权重上 cuBLAS 只需要 23 到 24 微秒，而我的内核会 padding 到 256，大约需要 48 微秒。所以这里的 decode 数据都是我的内核所能支持的最小 M 值。

我是怎么测量的

计时方式是在 20 到 50 次迭代前后用 CUDA event 包裹，torch.matmul 在同一进程、同一 GPU、同样的张量上运行。在 fp16 CUDA 张量上，它会走 cuBLAS。同一个内核重复运行的耗时波动大约在 1.7 微秒，所以 2% 以内的差异都算噪声，我直接把这类改动记为“无变化”。

光看耗时隐藏了太多信息，所以我还用 Nsight Compute 做了剖析：耗时、tensor pipe 利用率、DRAM 吞吐、L1/TEX 吞吐、L2 命中率、每个 SM 的 wave 数，以及 warp 停顿的细分。这篇博文的早期版本说 ncu 在 Modal 的 B200 上用不了。其实是可以的，只要加 --clock-control none，因为容器无法锁定 GPU 时钟。这意味着报告中会带上时钟未固定的警告。我会跨内核版本比较百分比，而不把它们当作绝对值。

每次运行都会拿输出和 torch 做校验。除了 MLP down 之外所有形状的最大误差都是 0，而 MLP down 在全部八个版本上误差都是 0.5——所以这来自在 fp16 下累加 K=28672，而不是某次特定改动造成的。

在开始这串阶梯之前，你需要了解四件事

TMA

TMA，即 Tensor Memory Accelerator，随 Hopper 一起到来。一个线程把一个描述符和一组坐标交给硬件，硬件就会自己在全局内存和共享内存之间搬运整个 tile。这里有两个重要后果。搬运是异步的，发起它的 warp 可以立刻继续往下走。而且地址计算（stride、边界、边界谓词）都在描述符里，线程不会在上面浪费寄存器。

有 load，从 GMEM 到 SMEM，也有 store，方向相反。load 用一个 mbarrier 做同步。store 之前需要一个 proxy fence，因为数据是普通线程写的，而 TMA 是通过一条叫 async proxy 的路径来读的。

mbarrier

__syncthreads() 的意思是块里的每个线程都在这里停下，直到所有人都到齐。这对 TMA 毫无用处。一个线程发起搬运，而数据是稍后由 TMA 单元写入的，所以没有线程知道数据什么时候到位。所有线程都可能在数据还在传输途中时就通过了 __syncthreads()。

mbarrier 是共享内存中的一个 64 位对象，同时统计线程和字节。它保存着：一个预期到达数（单个生产者线程时为 1，一个 warpgroup 时为 128）、一个每次 arrive 都会递减的待到达数、一个仍然等待的字节交易数，以及一个 phase 位。当到达数归零且交易数归零时，phase 位翻转，计数器重置，所有等待者被放行。等待者不是在等某个值，而是在等 phase 变化——这就是为什么每次等待调用都要带一个 phase 参数。

协议很短。生产者 arrive 并声明预期的字节数，把这个 barrier 附在 TMA 搬运上发起，然后硬件会在数据落地时递减字节数。当最后一个字节到位，phase 翻转，消费者被唤醒，并保证数据在共享内存中可见。

TMEM

Tensor Memory 是 Blackwell 上的片上内存，独立于共享内存和寄存器。每个 SM 有 256 KB，布局为 128 lane × 512 列 × 32 位单元。一个 128×N 的 fp32 累加器 tile 占用 N 列。

在 Hopper 上，wgmma 把累加器放在寄存器里。这吃掉了寄存器堆的大部分，限制了 tile 大小，强迫整个 warpgroup 拥有这个 MMA，还让 epilogue 和主循环争抢同一批寄存器。Blackwell 把累加器移进了 TMEM。tensor core 拥有结果，一个线程就可以发起 MMA，寄存器得以释放。

使用方式类似 malloc。一个 warp 以 2 的幂的列数（至少 32）调用 tcgen05.alloc，并在结束时释放。发起 MMA 的 warp 调用 tcgen05.mma 并指向 TMEM 地址，然后 tcgen05.commit 加上一个 mbarrier 表示累加器已完成。epilogue warp 用 tcgen05.ld 把结果读回来，每 warp 32 个 lane，做类型转换，再写入全局内存——通常经过共享内存和一次 TMA store。

持久化内核（persistent kernel）

这个我在最后一步它真正出场时再解释，因为只有在看到促成它的那份 profile 之后，它的理由才说得通。

起点 [PR #3]

第一步之前的内核很小。每个 CTA 通过 cute.copy() 在 mbarrier 背后把 A 和 B 的一个 MMA tile 从全局内存加载到共享内存，通过 cute.gemm() 发起一次 UMMA，然后把累加器写出。K 固定为 16，只有一条 MMA 指令的深度，所以它还不是一个通用的 GEMM。两个 mbarrier 各只使用一次，所以每次等待都处于 phase 0。共享内存布局是朴素的，epilogue 是串行的。

它的 profile 里已经藏着一个教训。在每 40.1 个周期的已发指令间隔中，warp 有 19.5 个周期停滞在 L1TEX scoreboarding 依赖上，占总量的 48.8%。我一开始把它读成了内存延迟。其实不是。mbarrier 等待是对一个共享内存地址的轮询，而共享内存访问走的是 L1TEX 管线，所以这里的“停滞在 L1TEX”只意味着线程正卡在那两个 mbarrier_wait 调用上，而 TMA 单元和 tensor core 正在干活。这个内核没有别的事可以让这些线程去做。所有人都同时停滞。这是一个重叠（overlap）的问题，而前两步要做的，就是造出一些可以重叠的东西。

第 1 步：K 循环（基线）[PR #4]

之前每个 CTA 只做一次 MMA，这只有在 K 能装进单个大小为 16 的 tile 里时才成立。这一步让 K 以 16 为步长推进，每个步做一次 MMA，并累加到同一个 TMEM tile 上。第一次 MMA 必须覆盖累加器，之后每次必须累加，而 cute.gemm 不会替你切换这个模式。

C_{tile} = A_{0}B_{0}^{T} + A_{1}B_{1}^{T} + \dots + A_{K/16}B_{K/16}^{T}

这不算优化。它是第一个真正意义上的 GEMM 版本，后面所有东西都以它为参照来测量。使用单缓冲时，这个循环严格是 load、wait、MMA、wait，每步 K 约 950 ns，搬运 12 KiB 并发起一次 MMA。现在 mbarrier 每次迭代都会被复用，所以 phase 会在 0、1、0、1 之间交替，而这个弄错的话会在第 1 次迭代就挂死。

profile 显示出一种新的停顿。已发指令之间的周期数从 40 跳到 167.6，其中 73.8%（123.7 个周期）是 warp 在 CTA barrier 处等待兄弟 warp。整个循环都放在 if warp_idx == 0 之下，所以 warp 1 到 3 无所事事，而 warp 0 要独立跑完全部 256 次迭代。tensor pipe 利用率只有 7.5%，而在 warp 0 的循环内部增加重叠是更便宜的修复手段，所以先做了这个。

每一步之后我都会重新计时四种形状：4096 立方、prefill 的 MLP up（M=8192, N=28672, K=8192）、decode 的 attention 输出（M=256, N=8192, K=8192）以及 decode 的 MLP up（M=256, N=28672, K=8192）。表中给出耗时以及相对 cuBLAS 的速度倍数。

第 2 步：4 级流水线 [PR #5]

num_stages 是共享内存缓冲区的数量，它并不对 K 做切分。当 K=4096、tile 为 16 时，循环仍然跑 256 次，只是四个缓冲区在底下轮转。内核预先加载 tile 0 到 3，然后第 k 次迭代等待自己的槽位、发起 MMA，并用 tile k+4 重新填充这个槽位。这些 MMA 仍然是一个接一个执行的，因为它们要累加到同一个 TMEM tile 上，所以只有 load 发生了重叠。

现在每个 barrier 每绕环一圈只被使用一次，所以当 stage 索引回绕时 load 的 phase 翻转，而 MMA 完成的 phase 仍然每次迭代都翻转。在 4096 立方上这快了 1.87 倍（840 到 448 微秒），在 prefill 形状上快 1.95 倍，在 decode 的 MLP up 上快 1.93 倍。在所有形状上，把 load 和计算重叠起来都有帮助。

每条已发指令的周期数从 167.6 降到 59.6，但 L1/TEX 吞吐率跳到了 98.9%，而 DRAM 只有 2.3%，L2 命中率 95.6%。一条几乎忙满却几乎没有字节来自 HBM 的管线，是在处理大量独立的小请求，而不是大量数据。尺寸很小、分布很差的共享内存访问看起来正是这个样子。我把它记为 bank conflict 嫌疑，然后继续去处理 barrier 停顿。

第 3 步：warp 特化（无变化）[PR #7]

CTA barrier 停顿仍然有 43 个周期，warp 1 到 3 空闲，而 warp 0 在按顺序做全部事情。于是 warp 0 变成了生产者，等待空槽位并发起 TMA；warp 1 变成了消费者，等待满槽位、发起 MMA，并向空 barrier 提交。这些 barrier 在初始化时就预先就绪，所以生产者的第一圈不会阻塞，前置代码也就消失了。warp 2 和 3 一直空闲到 epilogue。

结果：449 微秒对 448 微秒，什么都没发生，prefill 和 decode 的形状也纹丝没动。profile 显示这个改动确实做了我要求的事——CTA barrier 停顿从 43.3 降到 37.1 个周期。但它无济于事，因为 L1/TEX 已经以 98% 成了决定性的瓶颈，削掉 warp 2 和 3 的空闲时间并不能让 warp 0 和 1 跑得更快。

我也误判了特化能消除什么。它消除的是发射的串行化，但 cute.copy、cute.gemm 和 tcgen05.commit 都是异步的，所以第 2 步里的那个单 warp 从来没有卡在发射上。它卡在等待数据上，而两个 warp 只是等待同一份数据。我保留了这个改动，因为后面需要每个 stage 一对“满/空” barrier，而这个零结果迫使我不再从代码出发去推理，转而开始拆解内核的各个部分。

第 4 步：128B swizzle（几乎全部的故事）[PR #8]

profile 中耗时从 477 微秒降到 116 微秒，计时运行中从 448 降到 110 微秒。一次改动带来 4.1 倍。tensor pipe 利用率从 13.7% 升到 63.7%。L1/TEX 吞吐率从 98.8% 降到 68.8%。DRAM 吞吐率从 2.3% 升到 9.4%，因为数据终于快到足以让 HBM 发挥作用了。

这些消融实验是在第 3 步的零结果之后做的。每一行去掉或改变一件事，在 4096 立方上计时：

所以 tensor core 一直没问题，load 才是全部的问题，而流水线的任何改动都没能改变这一点。那就只剩下了共享内存中数据的布局。

共享内存是 32 个 4 字节的 bank，每 128 字节回绕一次。当 tile_k = 64 时，A 的一行是 64 个 fp16 值，正好 128 字节，而行 stride 也是 128 字节。所以每一行都从 bank 0 开始。TMA 一次写入多行，它们以相同的顺序全都堆到相同的 bank 上，访问就开始排队。

swizzle 把每个 128 字节的行切成八个 16 字节的块，并用块索引对行索引做 XOR。第 0 行保持原来的块顺序，第 1 行交换相邻块，第 2 行交换成对的块，以此类推。行与行之间没有任何东西移动，这只是每行内部的一个置换，地址计算就是一个 XOR。现在跨八行读取同一个逻辑列，会触及八组不同的 bank。

跨行读取同一列时，使用与不使用 128B swizzle 的对比

跨行读取同一列时，使用与不使用 128B swizzle 的对比

后续一个基准测试单独隔离了 load 路径，完全不含 MMA。字节数相同、搬运动作相同，只有布局不同。不用 swizzle 时是 3.85 TB/s，用上之后是 18.5 TB/s，相差 4.8 倍。4 路冲突理论上预测约 4 倍的差距，机制和效应大小是一致的。MMA 的读取从来不是问题：只有 MMA 的消融实验在朴素布局下就已经跑出了 cuBLAS 的速度。冲突发生在 TMA 写入上。

有两个细节消耗了我不少时间。128B swizzle 需要 tile_k >= 64，因为它需要 64 个连续的 fp16 值，而在 tile_k = 16 时布局会被直接拒绝。单独把 tile_k 设为 64 是不起作用的（上表中是 448.5 微秒）。只有两处改动同时生效才行，这也是为什么直到第 4 步才找到。单独测试其中任何一半看起来都是一条死胡同。其次，swizzle 加在指针上，而不是布局上：make_fragment_A 会拒绝组合布局，所以 swizzle 那部分交给 recast_ptr，而仿射部分携带 stage 模式。基址指针需要 1024 字节对齐。

它在所有形状上都有帮助，而且幅度差不多。4096 立方上是 4.1 倍，8192 立方的 prefill 形状上是 3.55 倍，decode 的 MLP up 形状上是 3.56 倍。load 路径上的 bank conflict 不在乎 M 是多少。

第 5 步：流水线化 epilogue（无变化）[PR #9]

epilogue 以四个分块把累加器搬出去，每个分块都要经过 TMEM 到寄存器、类型转换、再到全局内存。前四个版本的所有分块复用同一组寄存器，看起来像是存在写后读（WAR）冒险，于是我给寄存器做了双缓冲，并在转换第 i 个分块之前先发起第 i+1 个分块的读取。结果是没有变化：三次运行分别是 108.8、110.4 和 110.5 微秒，而之前是 109.3，profile 中则从 116.2 降到 112.9 微秒，都在噪声范围之内。

这个冒险根本不存在。DSL 会降到 SSA 形式，所以每次 cute.copy 都产生一个新的值，循环也被完全展开，而 ptxas 早已把这四条独立的链给流水线化了。我手动写出了编译器已经做好的东西。

我本该先看看奖品有多大。整个 epilogue 在 109 微秒的内核里最多只占 13 微秒，所以一个完美版本最多也只能省下 12%。在优化一个阶段之前先测一下它占多久，就能省掉这一步。

第 6 步：TMA multicast（无变化，而且在 8192 上更差）[PR #10]

在 4096 立方时，grid 是 32×16 个 CTA，所以共享同一行的那 16 个 CTA 各自独立地取相同的 A tile。multicast 让一个 TMA 请求投递到一个 cluster 中多个 CTA 的共享内存里，所以在 2×2 的 cluster 下，A 和 B 从 L2 请求的字节数都减半。Thien Tran 告诉过我这不会有帮助，并让我自己得出结论。

Thien Tran：“根据我的经验，TMA multicast 用处不大，但你绝对应该试试，得出你自己的观察！”

Thien Tran：“根据我的经验，TMA multicast 用处不大，但你绝对应该试试……”