Mojo 在 NVIDIA Blackwell 上实现矩阵乘法 85% SOTA 性能：multicast、2xSM MMA 与流水线优化全解析

Mojo 在 NVIDIA Blackwell 上实现矩阵乘法 85% SOTA 性能：multicast、2xSM MMA 与流水线优化全解析

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo

本文基于 Mojo 官方设计文档《Matrix Multiplication on Blackwell: Part 3——The Optimizations Behind 85% of SOTA Performance》整理，聚焦于 Modular 平台 Mojo 语言在 NVIDIA Blackwell（SM100）GPU 上把矩阵乘法从约 300 GFLOPS 提升至 SOTA 85% 的三代内核优化路径：CTA 级多播（multicast）与 2xSM MMA、2SM 流水线（pipelining）+ warp specialization，以及写回（write-out）双缓冲。读者将掌握@__llvm_metadata设定 CTA cluster、TMA 多播掩码计算、tcgen05.mma.cta_group::2指令、循环缓冲（circular buffer）与内存屏障（mbar）协作等实战技术，并能在仓库对应的迭代式内核源码中逐一对照验证。

系列背景与本文定位

这是 Mojo 官方 Blackwell 矩阵乘法系列博客的第三篇。前两篇分别对应 Part 1（Blackwell 架构与朴素内核）和 Part 2（共享内存分块、TMA、tcgen05.mma、swizzling、TMA store）。在前两篇结束时，4 行朴素内核的性能仅为 cuBLAS 的 0.3%，经过分块 + 张量核 + swizzling + TMA store 后达到约 293.6 TFLOPS（8.7% cuBLAS），但性能仍被全局内存吞吐所限制。

本篇的目标非常明确：利用 Blackwell 的两项高级特性（CTA 多播与 2xSM MMA），并结合流水线与 warp specialization，将性能提升约 5 倍，达到 SOTA 的 85%。整篇文章围绕三个依次递进的内核展开：

• Kernel 5：引入 CTA cluster、TMA multicast 与tcgen05.mma.cta_group::2（2xSM MMA），性能达到 360.2 TFLOPS（约 20% SOTA）；
• Kernel 6：基于 2SM MMA 实现共享内存循环缓冲流水线，并引入 warp specialization，性能跃升至 1429 TFLOPS（81% SOTA）；
• Kernel 7：对写回路径做双缓冲（double-buffering），额外获得 64 TFLOPS，最终达到 85% SOTA。

该系列的完整推导过程在仓库中有对应的可运行/可基准测试的迭代式内核源码：matmul_blackwell_iterative 目录，其中1_naive_sm100.mojo到8_clc_tmem_ping.mojo与博客系列的 Kernel 1~8 一一对应，BUILD.bazel声明了这些测试仅在有 B200 GPU 的机器上运行。

Kernel 5：multicast 与 2xSM MMA

从 Hopper 代开始，NVIDIA GPU 的流处理器（SM）可以被分组，同一 SM 组（CTA cluster）内的协作线程数组（CTA）能够互相访问彼此的共享内存（即分布式共享内存访问）。基于这一能力，Blackwell 上出现了两个高级优化手段：

• TMA multicasting（Hopper 起支持）：多个 SM 协作，把一个 tile 加载进共享内存；
• 2xSM MMA：2 个 SM 的张量核协作完成一次大型 MMA 运算，输入直接来自各 SM 的共享内存。

声明 CTA cluster：@__llvm_metadata

使用这些特性的第一步，是在编译期把 cluster 的尺寸告诉编译器。在 Mojo 中通过@__llvm_metadata装饰器完成：

@__llvm_metadata(`nvvm.cluster_dim`=cluster_shape) def blackwell_tma_pair_umma_kernel[ a_type, ...other parameters..., cluster_shape ]

上述代码在编译期设置了内核的 "CTA cluster" 发射形状。同一 cluster 内的 CTA 可以互相访问彼此的共享内存。在仓库源码 5_2sm.mojo 中可以看到完整的装饰器写法，cluster 形状通过cluster_shape: StaticTuple[Int32, 3] = StaticTupleInt32, 3作为内核的编译期参数传入，默认(1, 1, 1)表示不使用 cluster，Kernel 5 的测试与基准分别使用cluster_shape=(2, 1, 1)（见 5_2sm.mojo 的 benchmark_blackwell_matmul）。

CTA 内存多播（CTA memory multicasting）

为了直观理解多播，文档使用了一个简单例子：A、B、C 均为 256x256 矩阵，发射 4 个 CTA，每个 CTA 计算一个 128x128 的 tile。若 4 个 CTA 各自独立地从全局内存加载自己所需的 A、B tile，会存在大量冗余加载——这与系列博客开头提到的元素级冗余类似，但这次冗余发生在tile 粒度。

解决办法是把这 4 个 CTA 组成一个 2x2 cluster：由于 SM 可以访问彼此的共享内存，可以让每个 SM 只从全局内存加载 tile 的一半，然后广播（broadcast）给邻居。例如，同一行的两个 CTA 各自只加载 A 矩阵 tile 的一半行，另一半由对端提供；B 矩阵同样按列切分广播。这样每个 CTA 的全局内存加载量减半，却仍能获得完整的 tile。该技术称为 multicasting，且可扩展——例如 4x4 cluster 时，每个 CTA 只需加载自己 tile 的四分之一。

主机侧：按 cluster 维度切分 TMA tile

在 Mojo 代码中，首先在主机侧声明 TMA 算子。由于每个 CTA 会切分 A、B 的共享内存 tile 并广播给对应 peer CTA，因此 TMA tile 的尺寸需要按 cluster 维度缩小：

a_tma_op = create_tma_tile[ (BM // cluster_shape[1], BK), swizzle_mode=a_swizzle ](ctx, a) b_tma_op = create_tma_tile[ (BN // cluster_shape[0], BK), swizzle_mode=b_swizzle ](ctx, b)

在 5_2sm.mojo 的主机封装函数blackwell_kernel_5中可以看到等价实现：A 的 TMA tile 行数为BM // cluster_shape[1]（列维切分），B 的 TMA tile 行数为BN // (cluster_shape[0] // cta_group)（行维按 cluster 与 cta_group 双重切分）。

设备侧：async_multicast_load与多播掩码

在内核中，多播加载使用async_multicast_load：

if elect_one_thread: .... a_tma_op.async_multicast_load( a_smem_slice, tma_mbar[0], (UInt(i * BK), UInt(a_gmem_slice_coord)), a_multicast_mask, )

该 API 与上一篇文章中的async_load类似，但有三个关键差异：

1. 多播掩码a_multicast_mask：这是一个 16 位值，描述参与本次加载的 CTA 索引。每个 bit 代表一个 CTA，一个 cluster 最多 16 个 CTA。例如在上述 2x2 cluster 的例子中，CTA 0 和 CTA 2 负责多播 A tile，因此掩码都是0b101。Mojo 中的计算方式为：

var rank_m = block_id_in_cluster.x # CLUSTER_M 和 CLUSTER_N 是 cluster 的维度 comptime for i in range(CLUSTER_N): a_multicast_mask |= 1 << (i * CLUSTER_M) a_multicast_mask <<= rank_m

仓库 5_2sm.mojo 中同时计算了 A、B 两个掩码：A 掩码按rank_m位移，B 掩码还需结合peer_cta_coord[0]（peer CTA 在 cluster 中的行内序号）与rank_n * CLUSTER_M位移。

2. 共享内存切片a_smem_slice：从原始 tile 的布局张量中按基础指针偏移切出：

alias a_tma_load_size = a_desc_layout.size() var rank_n = block_id_in_cluster.y var a_smem_slice = type_of(a_smem_tile)( a_smem_tile.ptr + rank_n * a_tma_load_size )

3. 全局内存坐标相应更新：

# 每个切片的行数 alias a_tma_rows = a_desc_layout.shape[0].value() a_gmem_slice_coord = block_idx.x * BM + Int(rank_n) * a_tma_rows

2xSM MMA：tcgen05.mma.cta_group::2

多播与分布式共享内存虽然减少了从全局内存到共享内存的数据量，但 tile 在分布式共享内存中仍是重复存储的。例如在前面的例子里，CTA 0 和 CTA 1 在共享内存中各保留一份BN x BK的 B tile 副本——既然每个 CTA 都能访问对方的共享内存，保存两份显然是浪费。

Blackwell 的 2xSM MMA 指令tcgen05.mma.cta_group::2正是为解决此问题而设计：CTA 0 与 CTA 1 组成一对（pair），各自只加载 B tile 的一半；2xSM MMA 指令能同时看到共享内存中的两半数据，协调两个 SM 上的张量核，完成一次相当于两个单 SM MMA 之和的大型 MMA。对比左右两图：两个 SM 计算的是相同的2*BM x MMA_N x BK工作量（其中MMA_N是 MMA 的 N 维尺寸）、产生相同结果，但 2xSM 指令把 B tile 的共享内存占用减半。

与"多播加载 + 单 SM MMA"相比，2xSM MMA 从两个层面减少共享内存流量：其一，CTA 0 和 1 仍各自加载 B tile 的一半（与 Figure 5 的多播加载一致），但这里只是普通 TMA 传输，省略了多播步骤；其二，从共享内存到张量内存（TMEM）的 B tile 流量也减半，因为分布式共享内存中只保存一份副本。

发起 2xSM MMA 的代码

启动 2xSM MMA 的代码与单 SM MMA 非常相似，区别在于使用elect_one_cta：由于两个 SM 协作完成一次 MMA，只需其中一个 CTA（每对 CTA 中 ID 为偶数的那个，即 leader CTA）发起指令：

if elect_one_cta: # 等待数据到达共享内存 ... if elect_one_thread: comptime for j in range(num_k_mmas): var c_scale = 0 if i == 0 and j == 0 else 1 alias idx = IntTuple(0, MMA_K * j) alias a_offset = a_smem_layout(idx) * sizeof[a_type]() alias b_offset = b_smem_layout(idx) * sizeof[b_type]() mmacta_group mma_arrive_multicastcta_group

其中mma在cta_group=2时调用tcgen05.mma.cta_group::2指令；到达（arrive）函数也从上一篇的mma_arrive换为mma_arrive_multicast，它接收cta_group参数，并在cta_group=2时向 leader CTA 的内存屏障发信号。在仓库 5_2sm.mojo 中，2xSM MMA 被封装进MmaOpSM100_SS，并以cta_group=2参数化；elect_one_cta = block_rank_in_cluster() % 2 == 0（L204）用于选出 leader CTA。

Tensor memory（TMEM）布局

使用 2xSM MMA 时，两个 CTA 平分 M 维（BM = MMA_M / 2），每个 CTA 在张量内存中持有形状为BM x MMA_N的一半结果。指令支持MMA_M取 128 与 256 两种值：

• 当MMA_M=256时，每一半的 TMEM 布局与单 SM MMA 相同：上半部分存储于 leader CTA，下半部分存储于其配对 CTA，最终写回全局内存的方式与上一篇的 Kernel 4 一致；
• 当MMA_M=128时布局不同，文档省略了细节（生产环境较少使用），可直接查阅源码。

性能小结

采用多播与 2xSM MMA 后，内核达到360.2 TFLOPS，约 SOTA 目标的 20%。从 profile 看，即使使用了高级 MMA 指令，性能仍受限于全局内存吞吐——计算仍在等待数据传输完成。下一个优化将移除这个障碍，增大计算与内存传输的重叠。

Kernel 6：2SM 流水线（pipelining）

为了保持代码整洁，Kernel 6 把"将 tile 加载进共享内存"的逻辑封装为load_AB()，把"发起 MMA"封装为consume_AB()，写回结果封装为store_C()。以此粒度观察，前一个内核的循环是：

for i range(K/BK): # 发起异步 TMA 加载 load_AB() # 等待屏障直到数据到达 tma_mbar[0].wait(tma_phase) # 发起异步 MMA consume_AB() # 等待屏障直到计算完成 mma_mbar[0].wait(mma_phase) store_C() # 写出结果

在任意时刻，有一半的硬件处于空闲：要么张量核在等数据到达，要么 TMA 在等 MMA 结束以便共享内存缓冲区可复用。数据依赖使得一个 CTA 无法同时利用两类硬件单元。

流水化 MMA 与 TMA：循环缓冲

克服空闲的经典算法是：在共享内存中引入多个缓冲区，通过流水线让 MMA 与 TMA 重叠。首先要核算共享内存预算。

在 Blackwell GPU 上，内核最多可访问227 KB 共享内存（共享内存与 L1 缓存合计提供 228 KB，其中 1 KB 必须留给 L1）。当前配置下（BF16、最大 2xSM MMA 指令形状 256x256x16）的占用为：

• 一个 A tile：BM x BK x 2B = (MMA_M / 2) * 64 * 2B = 16 KB
• 一个 B tile：BN x BK x 2B = (MMA_N / 2) * 64 * 2B = 16 KB
• 一个 C tile：BM x MMA_N x 2B = 64 KB

合计还不到容量的一半。因此引入5 个流水级（stages），组织成循环缓冲：

随后让 TMA 与 MMA并行地遍历这些缓冲区：当 MMA 消费一个缓冲区时，另一个 TMA 可以预取（prefetch）后续计算所需的数据到另一个缓冲区，从而显著提高计算与通信的重叠度。

仓库 6_2sm_pipelined.mojo 的blackwell_kernel_6给出了共享内存的自动预算逻辑：以 B200 上total_smem_size_available = 233472字节为上限，扣除 C tile 后，按每级 A + B + 32 字节屏障开销计算max_pipeline_stages，作为num_pipeline_stages传入内核；smem_size再按 stage 数整体放大。这是文档中"5 级流水"的工程化版本——流水级数由共享内存预算动态推导。

Warp specialization（warp 专业化）

要实现上述流水模式，还需要另一个重要概念——warp specialization。此前所有内核只使用 4 个 warp（源于从 TMEM 加载数据的需要），且 TMA 与 MMA 都由线程 0 发起。为了让 TMA 与 MMA并行发起并作用于不同的流水级，需要为不同 warp 分配不同任务：

if WarpRole.is_main_load(): for i in range(num_iters): ... mma_mbar[stage].wait(phase) # 发起 TMA 加载并信号 tma_bar ... if WarpRole.is_mma(): for i in range(num_iters): ... tma_mbar[stage].wait(phase) # 发起 2xSM MMA 并信号 mma_mbar ... mma_arrive_multicastcta_group

即：一个 warp 专门负责发起 TMA，另一个 warp 专门负责发起 MMA，二者并发地遍历各 tile。warp 之间通过内存屏障（memory barrier）互相告知"tile 已到达"或"tile 已被消费、底层缓冲区可写入新数据"：MMA warp 等待 TMA 屏障（由 TMA warp 发信号）以确认输入就绪；TMA warp 等待 MMA 屏障以确认缓冲区可写。

仓库 6_2sm_pipelined.mojo 定义了WarpRole结构体：MainLoad = Self(4)、Mma = Self(5)、Epilogue = Self(3)，并提供is_main_load()、is_mma()、is_epilogue()静态方法；内核启动 6 个 warp（block_dim=(32 * 6)，见 L1007），其中 1 个 TMA warp、1 个 MMA warp、4 个输出（epilogue）warp。生产者与消费者分别使用独立的PipelineState（producer_phase初始为(0, 1, 0)，consumer_phase初始为(0, 0, 0)）追踪循环缓冲的 stage 与 phase。

输出阶段：compute_barrier

写回 TMEM 仍需要 4 个 warp。warp specialization 后，为输出专门分配 4 个 warp；同时需要一个新的内存屏障，在 MMA warp 与输出 warp 之间传递"MMA 结果已就绪"的信号。整体结构演变为：

if WarpRole.is_main_load(): for i in range(num_iters): ... mma_mbar[stage].wait(phase) # 发起 TMA 加载并信号 tma_bar ... if WarpRole.is_mma(): for i in range(num_iters): ... tma_mbar[stage].wait(phase) # 发起 2xSM MMA 并信号 mma_mbar ... mma_arrive_multicastcta_group # 信号输出 warp if elect_one_sync(): mma_arrive_multicastcta_group if WarpRole.is_epilogue(): compute_barrier[].wait() # 将结果存储到全局内存 ...

MMA warp 额外向新屏障compute_barrier发信号，输出 warp 等待它以确保结果就绪，随后执行与之前相同的存储逻辑。仓库 6_2sm_pipelined.mojo 的kernel_6主循环完整实现了这一结构，mma_complete_mask由self_mask | peer_mask计算（1 << block_rank_in_cluster()与1 << (block_rank_in_cluster() + 1)）。

性能小结

基准测试显示该流水线策略将性能提升至1429 TFLOPS，即 SOTA 的 81%。

Kernel 7：写回双缓冲（double-buffering the write-out）

此前优化一直聚焦于"把数据加载进 matmul"，其实"把结果存储到输出"同样可以优化。先回顾当前把 TMEM 结果写回全局内存的方式：

def store_c(): # 把整个张量内存搬进寄存器 registers = tcgen05_ld parameters # 把寄存器全部搬进共享内存 for block_offset in range(BN/TMA_BN): for st_matrix_offset in range(TMA_BN//16): st_matrix[c_smem_tile, registers] # 把整个共享内存 tile 搬进全局内存 c_tma_op.async_store( c_tma_tile, ((block_idx.x * **MMA_N** + thread_idx.x * TMA_BN), (block_idx.y * BM)), )

存储 C 数据包含三个主要步骤：TMEM → 寄存器 → 共享内存 → 全局内存。由于数据依赖，这些传输原本串行执行；但 TMA store 是异步的，这种串行并非必要。复用此前的思路，可以对 TMA store 与 MMA 结果搬运做流水化。

具体做法：在共享内存中声明两个输出 tile（称为 double-buffer），形状为BM x StageN，在二者之间 ping-pong 切换。为简单起见，先从stageN = 32开始（该值可调）。双缓冲流水策略如下：stmatrix把第一个输出 tile 写入共享内存后，立即发起 TMA store 但不阻塞，转而从 TMEM 加载数据、用另一个缓冲区写共享内存。这样 TMA store 就与下一个 tile 的 "TMEM → 寄存器 → 共享内存" 传输重叠起来。

在 Mojo 中实现时，把输出按MMA_N / StageN = 8次迭代拆分，每次处理 TMEM 中的stageN = 32列（TMEM 加载与stmatrix代码与之前内核基本一致，只是操作更小的 tile）：

# 每次处理 32 列 alias stageN = 32 alias num_stages = MMA_N // stageN comptime for stage in range(num_stages): # 加载 TMEM ... # 用 stmatrix 存入共享内存 ... if elect_one_thread: # 发起 TMA store ... c_tma_op.commit_group() @__parameter # 保持一个 TMA store 在途 if stage < num_stages - 1: c_tma_op.wait_group[1]() # 最后一级等待所有 TMA store 完成 else: c_tma_op.wait_group[0]()

代码的关键在于如何同步 TMA store：commit_group把此前发出的 TMA store 提交为一组，wait_group[N]等待直到只剩最后 N 组仍在途（第 N+1、N+2…组已完成）。上述代码中，第一次迭代不等待（可以直接使用第二个缓冲区）；从第二次到倒数第二次迭代，等待只剩 1 组 TMA store 在途，从而让 TMA store 与下一轮的 TMEM 加载、stmatrix重叠；最后一轮则等待所有 TMA store 完成。

仓库 7_double_buf_writeout.mojo 的store_C完整实现了这套逻辑：stageN = c_smem_layout.shape[1].value()（默认 32）、num_stages = MMA_N // stageN，C 的共享内存迭代器c_iter.next(stage % 2)实现双缓冲 ping-pong；named_barrier用于保证共享内存写入完成、TMA 读共享内存完毕（在 stage 1 至 num_stages-2 之间守卫前一缓冲区），wait_group[1]()/wait_group[0]()的同步逻辑与文档一致。同时，内核参数num_output_stages: Int = 2、output_tile_shape: IndexList[2] = Index(128, 32)（L494-L495）直接暴露了双缓冲配置。

一个"免费"的额外收益

流水化带来一个附带好处。回顾 Kernel 5 的共享内存占用：

• A、B tile 的 5 份流水副本：5 * (BM + BN) * BK * 2B = 160 KB
• 输出 C tile：BM * BN * 2B = 64 KB
• 少量用于内存屏障、TMEM 地址的空间

输出占用约 40% 的共享内存。而 Kernel 7 的输出只占2 * BM * StageN * 2B = 16 KB，省下的 48 KB 可以用来加深流水线，为 TMA 与 MMA 提供更多重叠空间。

性能小结

该优化额外带来64 TFLOPS的提升，内核达到SOTA 的 85%。

下一步：通向最后 15%

至此还剩 15% 的差距。尽管 Kernel 6 已通过重叠 TMA 加载与 MMA 隐藏了全局内存加载开销，但写回全局内存的开销仍然突出（尽管 Kernel 7 做了一定优化）；同时，后续 CTA 调度之间还存在启动开销（launch overhead）：

系列下一篇将用**持久内核（persistent kernel）**配合 Blackwell 的另一项高级特性——cluster launch control（CLC）——同时解决这两个问题，最终拉满 SOTA 差距。仓库中的 8_clc_tmem_ping.mojo 正是这一演进方向的工程实现（结合 TMEM ping-pong 与 CLC），可作为继续研读的入口。

仓库源码地图

本篇文章涉及的核心可验证资源（均位于当前仓库内）：

这些测试文件均可通过mojo_test规则在有 B200 GPU 的环境（仓库 Bazel 配置中//:has_gpu与//:b200_gpu约束）下运行，并支持--benchmark参数复现文档中的 TFLOPS 数据。例如 Kernel 5 的基准入口benchmark_blackwell_matmul会输出 "Average time" 与 "Performance (TFLOPS)"（5_2sm.mojo），Kernel 6、7 的基准入口结构与之类似，可直接对照文档中的 360.2 / 1429 / 1493（85% SOTA）三组数据。

总结

本篇围绕 Blackwell 矩阵乘法性能优化的三条主线展开：数据加载侧用 CTA 多播消除 tile 级冗余加载、用 2xSM MMA 消除分布式共享内存中的重复副本；计算调度侧用循环缓冲 + warp specialization 让 TMA 与 MMA 真正并行、并引入 compute_barrier 衔接 MMA 与输出 warp；写回侧用双缓冲让 TMA store 与下一轮结果搬运重叠。三者叠加，把内核性能从约 300 GFLOPS 一路推到 SOTA 的 85%。这些技巧在 Mojo 中均有对应的高层 API（async_multicast_load、mma_arrive_multicast、PipelineState、commit_group/wait_group等），且在仓库迭代式内核源码中一一可查，是编写生产级 Blackwell 张量内核的绝佳范本。

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo