越南程序员Thien Tran手写Blackwell tcgen05 kernel，实现98% CuBLAS性能，详解TMA/mbarrier/Tensor Memory等核心技术。
黑客级手把手教学：从零实现98% CuBLAS性能的Blackwell矩阵乘法核
今天最顶级的大模型训练速度越来越快？背后不仅靠的是数据和算法，还有GPU底层硬件的疯狂进化！NVIDIA最新发布的Blackwell架构，已经悄然把计算性能推到了一个恐怖的新高度。
但问题是——官方连一个像样的PTX+CUDA C++教程都没给！别慌，越南天才程序员Thien Tran出手了，他硬是用纯手写kernel，在Blackwell B200上复现了接近CuBLAS 98%性能的矩阵乘法（GEMM）核心！
谁是Thien Tran？一个敢在Modal云上硬刚CuBLAS的男人
先来认识下作者Thien Tran——一位活跃在GitHub和GPU底层开发社区的硬核工程师，他不仅深度研究过Ampere、Hopper，这次更是在Blackwell刚发布不久就上手B200，用Modal云平台跑通了完整的tcgen05 matmul kernel。
他的风格极其务实：拒绝“只会调库”，坚持“从PTX指令看硬件真相”；他不屑Copy-Paste式教程，而是手把手重建所有细节，包括Tensor Memory、TMA、mbarrier、swizzling、pipelining、warp specialization、2-SM MMA、persistent kernel……每一个环节他都亲自验证、调优，最终实现1475 TFLOPS，逼近CuBLAS的1506 TFLOPS！这不仅是一个技术突破，更是对CuBLAS黑盒神话的一次优雅挑战。
为什么Blackwell上的tcgen05如此重要？因为它彻底改变了GEMM编程范式
Blackwell（SM100）引入了一套全新的PTX指令集——tcgen05，专门用于驱动新一代Tensor Core。和Ampere的mma、Hopper的wgmma不同，tcgen05直接在共享内存（shared memory）中读取A/B矩阵，并将结果写入一种全新的“张量内存”（Tensor Memory），完全跳过了传统寄存器加载的步骤（比如不再需要ldmatrix）。
这意味着什么？开发者不再需要像过去那样精细控制线程与寄存器映射，而是可以在“Tile层面”思考问题，大大简化了高性能kernel的开发逻辑。
但代价是：你必须完全理解TMA（张量内存加速器）、mbarrier（内存屏障）、Tensor Memory分配/释放机制、以及tcgen05的内存描述符（descriptor）格式。
张量内存加速器TMA：Blackwell的“快递小哥”，让数据搬运快到飞起
TMA（Tensor Memory Accelerator）最早在Hopper引入，但在Blackwell上才真正大放异彩。
传统cp.async一次最多搬16字节，而TMA单线程就能发起任意尺寸的异步内存拷贝！
这意味着，一个线程块（threadblock）里，你只需一个线程调用TMA，就能把整块A/B矩阵从全局内存搬到共享内存，极大节省了寄存器资源和地址计算开销。但TMA不能随便用！它运行在“异步代理”（async proxy）上，所以必须配合mbarrier做同步。
简单来说：TMA像快递员，把包裹（数据）送到共享内存，mbarrier就像签收单，告诉计算单元“货到了，可以开工了”。
mbarrier：Blackwell上的同步神器，搞定TMA与tcgen05的“握手协议”
mbarrier是NVIDIA在Blackwell上引入的一种64位共享内存同步原语。
它有两个关键计数器：arrival count（有多少线程/操作到达）和tx-count（预期传输多少字节）。
在TMA场景中，你通常只用一个线程发起拷贝，所以arrival count初始化为1；同时你告诉mbarrier“我预期传输X字节”（用mbarrier.arrive.expect_tx），TMA每完成一部分传输，就会自动减少tx-count。
只有当arrival count和tx-count都归零，mbarrier才算完成当前阶段。而tcgen05.mma执行完后，也需要通过tcgen05.commit向mbarrier“打卡签到”。
整个过程就像一场精密的交接仪式，缺一不可。
初代kernel：手写tcgen05.mma，却只跑出254 TFLOPS？
作者的第一版kernel（matmul_v1.cu）虽然结构完整，包括TMA加载、tcgen05.mma计算、Tensor Memory分配/释放、以及从Tensor Memory读回结果的epilogue，但性能惨不忍睹——只有254 TFLOPS，不到CuBLAS的20%！
问题出在哪？原来，Blackwell的tcgen05对共享内存布局有极其严苛的要求：它内部以8x16字节的“核心矩阵”（Core Matrix）为单位读取数据，要求每个8行x16字节的块必须是物理连续的。
而作者最初的TMA布局是按普通二维Tile加载的，导致tcgen05无法高效读取。更致命的是，他没用任何swizzling（交织），也没做pipelining（流水线），计算单元大部分时间都在“等饭吃”。
128字节swizzling：不是为了防冲突，而是为了喂饱TMA！
很多人以为swizzling是为了避免共享内存bank conflict（银行冲突），但在Blackwell上，128B swizzling（CU_TENSOR_MAP_SWIZZLE_128B）其实是为了告诉TMA：“我要用128字节宽的Tile加载数据”。
因为tcgen05内部MMA_K=32字节（BF16下为k16），但TMA一次至少要搬128字节才能达到最大吞吐。所以作者把TMA的boxDim从{8, BLOCK_M}改成{64, BLOCK_M}（64个BF16 = 128字节），并启用128B swizzling。
神奇的是，他完全不需要手动计算swizzle地址——只要在tensor map编码时指定swizzle类型，TMA和tcgen05就会自动对齐！这一改，性能直接飙到695 TFLOPS，暴涨2.7倍！
流水线pipelining：让TMA和tcgen05永远不闲着
有了正确的内存布局，下一步就是让计算和内存传输重叠。
作者采用经典的N-stage pipelining：预加载N-1个Tile到不同的共享内存缓冲区，然后在计算第k个Tile的同时，后台加载第k+N个Tile。这样，Tensor Core几乎不会空闲。实现的关键是为每个stage分配独立的mbarrier，确保TMA加载和tcgen05计算互不干扰。
这一招让性能从695 TFLOPS提升到939 TFLOPS，再涨35%！
Warp特化：1个warp专管TMA，1个warp专管MMA，其他warp干别的
Blackwell的异步执行特性允许不同warp独立运行。
于是作者把128线程的block拆成：warp0专职TMA加载，warp1专职tcgen05.mma计算，剩下4个warp负责epilogue（从Tensor Memory写回结果）。这样就不用在每个循环里检查warp_id，避免了分支发散。
更妙的是，TMA warp和MMA warp可以使用各自的phase变量跟踪mbarrier状态，避免死锁。结果？性能冲到1208 TFLOPS，逼近CuBLAS的80%！
2-SM MMA：跨两个流处理器协同计算，MMA_M翻倍！
Blackwell还支持“2-SM MMA”——两个相邻的线程块（CTA）组成一个cluster，协同完成一个更大的MMA（M=256, N=256, K=16）。每个CTA提供一半的A和B数据，B矩阵甚至需要从对方的共享内存读取！这需要启用cluster_dims(2,1,1)，并通过%cluster_ctarank判断自己的角色。
同步也更复杂：TMA要向主CTA的mbarrier汇报，tcgen05.commit要通过multicast同时通知两个CTA。虽然只带来8%的提升（1302 TFLOPS），但这打开了多SM协同的大门。
持久化Kernel：让每个SM持续工作，彻底榨干硬件！
最后的杀手锏是“持久化Kernel”（Persistent Kernel）：不再为每个输出Tile启动新线程块，而是只启动148个block（等于B200的SM数量），每个block循环处理多个Tile。这样，epilogue（写回）可以和下一个Tile的TMA/MMA完全重叠！
作者甚至为Tensor Memory也做了双缓冲，确保MMA warp永远有空闲的输出缓冲区。结果？1475 TFLOPS，98% of CuBLAS！这几乎榨干了Blackwell的所有潜力。
代码保留：关键PTX指令与C++封装
以下是作者封装的TMA加载函数：
cpp
device inline
void tma_2d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int mbar_addr) {
asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];"
:: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(mbar_addr) : "memory");
}
cpp
auto make_desc = [](int addr) -> uint64_t {
const int SBO = 8 * 128; // 8行x128字节
return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
};
cpp
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
: "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
"=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
: "r"(addr));
总结：Blackwell GEMM开发已进入“Tile时代”
Thien Tran的教程证明：在Blackwell上写高性能GEMM，不再是寄存器和线程的“微观舞蹈”，而是Tile布局和异步流水的“宏观调度”。只要你理解TMA、mbarrier、Tensor Memory、tcgen05 descriptor这四大支柱，就能写出接近CuBLAS的kernel。
随着FlashAttention 4等混合kernel的出现，这种“人机协同”的底层开发将变得越来越重要。