Matrix Multiplication on NVIDIA's Blackwell: A Four-Part Deep Dive

发表时间: 2026-02 · Blog post by Modular (modular.com)

原文（四篇合并）:

作者/机构：Modular

速读

一句话结论 本文通过在 Mojo 语言中逐步应用 NVIDIA Blackwell 架构的新硬件特性，构建了一个全异步的矩阵乘法内核，在 B200 GPU 上实现了超越官方 cuBLAS 库的算力性能。

要解决什么问题 大语言模型超过 83% 的运行时间被矩阵乘法占据，但要发挥最新 GPU 的理论峰值算力面临多重微观机制卡点。首先是全局内存访问极慢，且多线程并发读取时极易引发共享内存的 Bank 冲突。其次，在 Hopper 等前代架构中，张量核心的计算结果必须存入寄存器，导致严重的寄存器压力，并引发张量核心与通用算术逻辑单元之间的资源争夺。此外，传统的线程块启动会产生巨大的内核启动开销；即便使用软件层面的持久化内核，在多内核并发或流式多处理器饥饿时也会导致调度次优。最后，如果内存加载、矩阵乘加计算和结果写回是串行执行的，会导致大量硬件单元处于闲置等待状态，算力利用率无法提升。

怎么做的 方法的核心思路是构建一个深度融合 Blackwell 专属硬件指令的“全异步、Warp 专用化”流水线，将数据搬运与计算完全重叠。首先，利用张量内存加速器（TMA）实现全局内存到共享内存的异步拷贝。为了消除多线程并发读取共享内存时产生的 Bank 冲突，引入了 128 字节的 Swizzling 技术，其核心是通过异或运算打乱地址映射：

$$ \text{Address}_{\text{swizzled}} = \text{Offset} \oplus ((\text{Offset} \ \& \ \text{Mask}_{Y}) \gg \text{Shift}) $$

为了释放寄存器压力，内核调用了 Blackwell 第五代张量核心的 tcgen05.mma 指令，将计算结果直接累加到专用的 256KB 张量内存（TMEM）中。其次，为了减少共享内存的占用，设计了基于集群的分布式共享内存机制。通过 TMA 多播技术，集群内的线程块可以共享加载的输入切片；同时启用 2xSM MMA 指令，让两个流式多处理器协同处理一个大型矩阵乘加任务，使操作数矩阵的共享内存流量减半。最关键的设计是 Warp 专用化与多级流水线。内核将不同的 Warp 分别专门用于 TMA 加载、MMA 计算和结果写回。通过在共享内存和 TMEM 中构建环形缓冲区，各阶段通过内存屏障进行同步。例如，TMEM 被划分为两半，MMA Warp 在其中一半累加新一轮结果时，写回 Warp 正从另一半将上一轮结果异步存入全局内存。屏障相位的翻转控制着流水线的推进：

$$ \text{Phase}_{\text{next}} = \text{Phase}_{\text{current}} \oplus 1 $$
最后，为了消除线程块启动开销，引入了集群启动控制（CLC）。硬件级的调度器 Warp 会自动向空闲流式多处理器分配工作块坐标，并通过共享内存发出信号，这使得下一波的 TMA 加载能与当前波的写回完美重叠。同时，通过线程块 Swizzling 技术，让工作块在 N 维度上呈“之”字形移动，最大化了 L2 缓存的命中率。

效果如何 实验在 NVIDIA B200 GPU 上进行，数据类型采用 BFloat16（累加时使用 FP32 以保证精度）。对比的基线方法是 NVIDIA 官方的 cuBLAS 库，代表了当前闭源、高度优化的最高工艺水平路线。在 4096x4096x4096 的标准方阵测试中，最初的朴素实现仅有 5.6 TFLOPs，而在叠加所有硬件特性与流水线优化后，该内核的性能飙升至 1772.9 TFLOPs，达到了 cuBLAS 性能的 100.6%。此外，作者还在真实的生产环境设置下进行了测试（例如 Gemma 3 27B 模型中 512x8192x5376 的非规则张量形状）。由于非规则形状难以通过经验手动计算最优的流水线深度和 MMA 形状，作者结合了 Mojo 内置的 kbench 自动调优框架来动态选择最优参数。在这种生产负载下，该内核的性能超越了 cuBLAS 基线最高达 6%。代价方面，这种极致的优化高度依赖于 Blackwell 独有的底层 PTX 指令（如 tcgen05 和 CLC 机制），且需要极其繁琐的内存屏障和异步状态管理。

主要贡献

大语言模型（LLM）（如Llama、Qwen、ChatGPT等）的核心计算高度依赖矩阵乘法（matmul）。以运行在2xB200上的Llama 8B模型（FP8）为例，超过83%的模型运行时间被矩阵乘法的各种变体（如线性层、注意力机制、MLP等）所占据。因此，即使是10%的矩阵乘法性能提升，也能带来约8%的端到端加速，这对于大规模部署而言意味着巨大的成本节约。
Llama 8B执行时间中矩阵乘法占比超过80%

本文旨在展示如何利用Mojo语言，在NVIDIA最新的Blackwell（SM100）GPU上编写出性能可与NVIDIA官方cuBLAS库相媲美的高性能GPU内核。文章提供了一套针对Blackwell GPU架构进行矩阵乘法优化的完整蓝图，通过逐步引入并利用Blackwell的新硬件特性（如第五代Tensor Cores、张量内存TMEM、TMA、2SM MMA、集群启动控制CLC等），最终实现了超越当前SOTA（State-of-the-Art）的性能表现。
性能一览

背景知识/关键Observation/设计原则

矩阵乘法的定义与CPU局限性
给定两个稠密矩阵$A$（维度为$M \times K$）和$B$（维度为$K \times N$），矩阵乘法$C = A \cdot B$本质上是跨$K$维度的内积计算。
矩阵乘法的简化图示

for row in range(M):
  for col in range(N):
    C[row][col] = 0
    for inner in range(K):
      C[row][col] += A[row][inner]*B[inner][col]

CPU受限于核心数量（通常在32到128核之间），而现代GPU（如B200）可同时处理超过100,000个线程，是执行矩阵乘法这种高度重复的数据并行操作的理想硬件。
作为内积集合的矩阵乘法

GPU硬件架构视角
GPU包含多个流式多处理器（SM）、被SM共享的L2缓存和全局内存。线程被组织成线程块（Thread Blocks或CTAs），每个线程块被调度到单个SM上执行。SM内的线程可以访问私有寄存器和共享内存（Shared Memory）。
GPU架构
协作线程阵列(CTA)布局

张量核心（Tensor Cores）的工作原理
与一次只能处理一个标量乘加操作的CUDA核心不同，张量核心可以在一条指令中执行整个矩阵乘加（MMA）操作（例如$64 \times 128$的分块）。Blackwell的第五代张量核心（tcgen05）甚至可以在2个SM上扩展至$256 \times 256 \times 16$的MMA操作，并引入了张量内存（Tensor Memory）以减少寄存器压力。
将MxK矩阵与KxN矩阵相乘
单个CUDA核心的矩阵乘法
单个张量核心的矩阵乘法

GPU编程范式与线程调度
GPU采用单指令多线程（SIMT）模型。程序员启动一个内核时，指定网格（Grid）大小和块（Block）大小。每个块被映射到一个SM上，块内的线程进一步被划分为包含32个线程的Warp。SM内部有4个Warp调度器，当某个Warp等待数据时，调度器会切换到其他就绪的Warp，从而隐藏内存延迟。
CTA被分配到SM
CTA划分为块和Warp
优化的Warp调度示例

线程块集群（Thread Block Clustering）
在Hopper和Blackwell架构中，多个CTA可以组成一个集群（Cluster）。集群保证其块被调度到同一GPU处理集群（GPC）内物理连接的SM上，允许不同SM上的块直接访问彼此的共享内存。
线程块集群
线程块集群调度

NVIDIA GPU架构演进及优化策略
* Ampere架构（2020）：引入异步复制指令（cp.async），允许在单个CTA内重叠数据加载和MMA操作。问题在于CTA启动开销大。
Ampere架构
Ampere优化
* Hopper架构（2022）：引入张量内存加速器（TMA）和异步Warp组MMA（WGMMA），催生了持久化内核（Persistent Kernels）技术，消除了内核启动开销。问题是WGMMA消耗大量寄存器。
Hopper架构
Hopper优化
* Blackwell架构：引入tcgen05指令和张量内存（TMEM）。MMA结果直接存入专用硬件TMEM，打破了对寄存器的依赖，形成了加载（TMA）、计算（Tensor Cores写入TMEM）、存储（从TMEM写入全局内存）的三级流水线。
Blackwell架构
Blackwell优化

前Ampere时代的优化
在Ampere之前，内存操作会阻塞计算，必须使用双缓冲和每个SM多个CTA来重叠数据传输和计算。
前Ampere时代的优化

朴素内核实现（Kernel 1）与数据类型
使用Mojo编写的4行代码内核，每个线程计算输出矩阵$C$的一个元素。

acc = Float32(0)
for k in range(K):
  acc += a[global_idx.y, k] * b[global_idx.x, k]
c[row, col] = acc

所有输出的并发点积
为了平衡内存和精度，使用BFloat16（BF16）作为输入和输出，但在计算时转换为Float32（FP32）进行累加，以防止舍入误差。该朴素内核实现了5 TFLOPs的性能。
BF16数据类型
朴素GPU内核的性能

方法细节

优化路线图与基础设定
为了简化说明，后续优化均假定矩阵$A$为$M \times K$，矩阵$B$为$K \times N$（已转置），结果矩阵$C$为$M \times N$，且$M=N=K=4096$。核心挑战在于避免或隐藏缓慢的全局内存加载和存储。
第2部分的性能改进路线图
不同操作的延迟图示
4行代码矩阵乘法的线程内存访问

共享内存与循环分块（Loop Tiling）
为了减少全局内存的冗余加载，采用循环分块技术。将矩阵划分为$BM \times BK$和$BN \times BK$的块（例如$64 \times 64$），并加载到Blackwell SM提供的228KB高速共享内存中。在$K/BK$次循环中，逐步加载块并执行MMA，将中间结果累加，最后只需将最终结果写入全局内存一次。
矩阵乘法的循环分块
使用循环分块的矩阵乘加
将最终结果写入全局内存

Kernel 2：使用TMA和张量核心（Tensor Cores）
此内核结合了分块和张量核心。首先使用Mojo API在主机端创建TMA（张量内存加速器）块（TMA Tile），该块包含张量的形状、步长和全局内存地址。

a_tma_op = create_tma_tile[
  a_type, 2, Index(BM, BK)
](ctx, a_global_mem_address)

在内核中，由单个线程（elect_one_thread）发起异步TMA复制（async_copy），将全局内存数据加载到共享内存中。由于TMA是异步的，必须使用内存屏障（tma_mbar）来阻塞线程，直到复制完成。屏障的相位（phase）会随着预期字节数的到达而翻转。
线程阻塞与TMA屏障示意图
TMA屏障相位翻转
手动切换TMA屏障相位
关于TMA提取的底层细节：张量核心只能理解"核心矩阵"（Core matrices，即$8 \times 8$元素的块）。为了满足K-Major布局，TMA硬件会将逻辑上的$64 \times 64$块划分为8个$64 \times 8$的子块，并逐列进行8次提取。
张量核心矩阵
复制8列核心矩阵8次

发出MMA指令与张量内存（TMEM）
Blackwell引入了256KB的张量内存（TMEM），专门用于存储tcgen05 MMA指令的输入或输出，从而释放了通用寄存器。
张量内存(TMEM)
MMA指令（tcgen05.mma）同样是异步执行的，由单线程发起，并使用mma_arrive与屏障（mma_mbar）链接。由于硬件指令限制K维度必须为32B（BF16下为16个元素），因此对于$BK=64$，需要循环4次发出MMA指令（num_k_mmas）。
MMA需要4次迭代(BK=64)

从TMEM转移到寄存器
数据必须先从TMEM移动到寄存器，才能写入全局内存。使用tcgen05_ld指令执行此操作。TMEM中的数据按特定的$16 \times 256$位模式加载，整个Warp组（4个Warp）读取64个通道。每个线程最终在寄存器数组中持有32个元素。
张量内存的布局组织和访问模式

从寄存器转移到全局内存（GMEM）
利用Mojo的LayoutTensor.tile()方法确定输出块的位置，并进一步为每个Warp分配子块。
输出第3行第3个块
为每个Warp划分输出矩阵
元素到线程的映射(tcgen05.ld.16x256 PTX指令)
为了处理复杂的线程到数据的映射，使用.distribute[Layout.row_major(8, 4)]方法，给每个线程提供其需要写入数据的精确视图，完美映射NVIDIA的底层布局要求。最后通过循环将数据写回全局内存。
16x64块可被视为由2值向量组成的16x32块
子矩阵映射
将distribute操作与NVIDIA的映射进行比较
迭代 n_vec=0, m_vec=0, 单个warp
迭代 n_vec=0, m_vec=1, 单个warp
迭代 n_vec=1, m_vec=0, 单个warp
最终迭代 n_vec=7, m_vec=1
将CTA的块映射到全局内存中的C矩阵

设置共享内存与Kernel 2性能
通过动态共享内存分配（external_memory）设置基础地址，并为A块、B块及屏障分配偏移量。Kernel 2达到了155.0 TFLOPS，是朴素内核的28倍，但仅为cuBLAS性能的8.7%。
async_copy()操作将块复制到共享内存
TMA和张量核心优化的性能

Kernel 3：Swizzling（内存交错）解决Bank冲突
共享内存由32个连续的4B宽的Bank组成。当同一个Warp中的线程试图访问映射到同一个Bank的不同地址时，会发生Bank冲突，导致执行停滞。在128B标准布局下，核心矩阵的8行会映射到相同的Bank 0-3，引发8路冲突。
32个连续4B宽Bank的共享内存布局
多个线程在同一周期服务不同Bank
两个线程访问同一Bank导致冲突
8路Bank冲突
Swizzling通过按位异或（XOR）交换索引，确保数据分布在不同的Bank中。
第0行
第1行
第2行
第3行
采用<3, 4, 3>的128字节Swizzling模式，消除了核心矩阵的Bank冲突。在Mojo中，只需在LayoutTensor和描述符中指定swizzle_mode即可。该优化使性能提升至288.3 TFLOPs。
128字节swizzling模式
128字节swizzled共享内存块
无swizzling的内存访问
有swizzling的内存访问
加入swizzling后的性能提升

Kernel 4：在共享内存中打包输出并使用TMA存储
为了减少全局内存存储指令，Kernel 4利用TMA进行存储。首先必须将寄存器中的数据（FP32）转换为BF16，并使用stmatrix指令将其打包到共享内存中。stmatrix需要4次迭代来完成打包，并同样使用128B Swizzling来避免Bank冲突。
4次迭代的stmatrix
打包完成后，使用fence_async_view_proxy确保内存可见性，随后通过async_store发出TMA存储指令，并使用commit_group()和wait_group[N]()进行同步管理。由于受限于全局内存带宽，此内核性能维持在293.6 TFLOPs，但其异步特性为后续流水线优化奠定了基础。
使用两个线程映射两个TMA存储以最大化并行度
TMA和ST_Matrix优化后的性能
来自NCU的计算和内存吞吐量配置文件
向85% SOTA性能进发的路线图

Kernel 5：多播（Multicast）与2xSM MMA
为了减少多CTA加载相同数据带来的冗余，利用Hopper起支持的分布式共享内存特性。将4个CTA组成一个$2 \times 2$集群，每个CTA只需从全局内存加载一半的块，然后通过TMA多播（async_multicast_load）广播给相邻的CTA。这需要配置一个16位的a_multicast_mask。
使用4个CTA计算两个256x256矩阵乘法
CTA块加载
CTA多播A矩阵的行块
CTA多播B矩阵的列块
进一步，为了解决分布式共享内存中块重复存储的问题，Blackwell引入了2xSM MMA指令（tcgen05.mma.cta_group::2）。成对的两个CTA各加载一半数据，2xSM指令协调两个SM上的张量核心共同完成一个大型MMA操作，从而将共享内存和TMEM的流量减半。通过elect_one_cta确保只有偶数ID的领导CTA发出指令。此优化将性能提升至360.2 TFLOPs。
分布式共享内存中的块重复
张量内存布局
CTA内存多播和2xSM MMA达到20% SOTA性能

Kernel 6：2SM流水线（Pipelining）与Warp特化
此前的内核由于数据依赖，导致TMA和MMA无法并行，硬件处于闲置状态。
CTA的数据依赖争用
通过在共享内存中引入5个阶段的环形缓冲区（Circular Buffer），实现了TMA和MMA的流水线重叠。
5阶段环形缓冲区
使用环形缓冲区重叠通信与计算
为了让TMA和MMA并行执行，采用了Warp特化（Warp specialization）技术。专门分配一个Warp处理TMA加载，另一个处理MMA，它们通过内存屏障（如tma_mbar和mma_mbar）相互通信缓冲区的状态。输出操作也分配了独立的Warp。此流水线策略将性能大幅提升至1429 TFLOPs。
2SM流水线和Warp特化达到81% SOTA性能

Kernel 7：写出阶段的双缓冲（Double-buffering）
为了打破TMEM到寄存器、寄存器到共享内存、共享内存到全局内存的顺序执行瓶颈，在共享内存中声明了两个输出块（双缓冲）。当TMA异步存储第一个块时，系统立即开始将TMEM的数据加载到第二个缓冲区中。通过精确控制wait_group，实现了TMA存储与TMEM加载的重叠。这一优化不仅提高了并行度，还节省了共享内存空间以加深流水线，最终达到1493.0 TFLOPs。
双缓冲流水线的第1次迭代
双缓冲流水线的第2次迭代
写出阶段双缓冲达到85% SOTA性能
CTA启动带来的开销
突破SOTA的路线图

Kernel 8：CLC持久化内核（Persistent Kernel）
为了消除Wave（指一次分配给所有可用SM的一批线程块）之间共享内存和屏障重新初始化的开销，引入了持久化内核。传统的软件静态调度存在无法感知SM空闲状态的问题。
持久化内核与Wave调度
静态持久化内核的调度
Blackwell硬件层面引入了集群启动控制（CLC）。调度器Warp尝试将工作块分配给空闲SM，并将坐标写入集群内所有CTA的共享内存中，通过16字节数据到达信号通知。为消除CLC调度本身的开销，对CLC提取过程进行了流水线化处理，使其与下一Wave的TMA加载重叠。
CLC调度器的流水线化
重叠CLC调度开销
此外，为了解决TMEM单一地址导致的Warp闲置和顺序输出问题，将TMEM也设计为环形缓冲区。通过accum_full和accum_empty屏障，MMA Warp和输出Warp可以并发操作TMEM的不同区域。这些综合优化使性能达到1772.9 TFLOPs。
Blackwell内核各阶段的异步执行
CLC持久化内核达到100.6% cuBLAS性能

Kernel 9：线程块Swizzle（Thread Block Swizzle）
在处理特定形状（如$8192 \times 8192 \times 8192$）时，为了提高L2缓存命中率，采用了线程块Swizzle技术。通过在N维度上采用Z字形（zig-zag）移动模式（block_swizzle_size），使得被调度到同一Wave中的工作块能够最大程度地共享矩阵A和B的全局内存加载，从而减少内存驱逐，进一步提升性能。
6x5集群的28个工作块
分配给每个Wave的工作块
N维度上大小为2的Swizzle移动
采用Swizzle后的Wave 0和1工作块分配

实验环境

实验结果

以下是各阶段内核的实验内容、结果及结论（基于$4096 \times 4096 \times 4096$方形矩阵）：

  1. Ultra Naive (no Tensor Cores)：4行代码的朴素实现。结果：5.6 TFLOPs（cuBLAS的0.3%）。结论：未使用硬件加速器，性能极低。
  2. Tensor Cores：引入共享内存分块、TMA和张量核心。结果：155.0 TFLOPs（cuBLAS的8.8%）。结论：大幅提升（28倍），但受限于共享内存Bank冲突。
  3. Swizzling：应用128字节Swizzling解决Bank冲突。结果：295.6 TFLOPs（cuBLAS的16.8%）。结论：Bank冲突使性能减半，解决后性能显著回升。
  4. TMA + ST_Matrix：在共享内存中打包输出并使用TMA异步存储。结果：293.6 TFLOPs（cuBLAS的16.7%）。结论：受限于全局内存带宽，性能持平，但为异步重叠奠定基础。
  5. 2SM MMA：引入CTA多播和2xSM MMA指令。结果：360.2 TFLOPs（cuBLAS的20.4%）。结论：减少了共享内存重复和流量，但计算仍需等待数据传输。
  6. Pipelining：引入5阶段环形缓冲区和Warp特化。结果：1429.6 TFLOPs（cuBLAS的81.1%）。结论：TMA与MMA的重叠极大隐藏了内存延迟。
  7. Double Buffering Output：输出阶段的双缓冲流水线。结果：1493.0 TFLOPs（cuBLAS的84.7%）。结论：重叠了输出阶段的各级内存传输。
  8. CLC Persistent：引入硬件级CLC调度和TMEM环形缓冲区。结果：1772.9 TFLOPs（cuBLAS的100.6%）。结论：消除了内核启动开销和Warp闲置，达到并微超cuBLAS性能。
  9. 实际生产形状（如Gemma 3 27B）：利用Mojo内置的kbench自动调优框架选择最佳参数（MMA形状、流水线阶段、Swizzle模式等）。结果：在Gemma 3工作负载上，最高超过SOTA性能6%。结论：该实现不仅在方形矩阵上有效，在实际多变维度的生产模型中同样具备卓越的性能泛化能力。
    Gemma 3 27B在Blackwell系统上的性能对比

结论

本文成功展示了如何使用Mojo语言在NVIDIA Blackwell GPU上编写出超越SOTA（cuBLAS）的矩阵乘法内核。通过深入剖析并逐步应用Blackwell架构的先进硬件特性（如TMA、TMEM、2xSM MMA、CLC调度器等），结合高级软件调度策略（流水线、Warp特化、持久化内核、Swizzling），最终实现了高达1772 TFLOPs的性能，并在实际LLM（如Gemma 3）的部署形状中超越SOTA达6%。未来的工作将展示帮助编写高性能代码的工具，探讨Mojo如何在保持强大编程人体工程学的同时实现极致性能，并讨论如何使这些代码模式在不同硬件之间保持可移植性。

附录

描述符（Descriptors）
tcgen05.mma使用描述符来指定共享内存中输入数据的布局、指令形状和数据类型等。在Mojo中通过MMASmemDescriptor创建。
其中最重要的两个参数是：
* LBO（领先维度字节偏移量）：K维度上两个相邻核心矩阵之间的字节数。
* SBO（步长维度字节偏移量）：M/N维度上两个相邻核心矩阵之间的字节数。
在未使用Swizzling的Kernel 2中，A矩阵的aSBO=128（核心矩阵大小$8 \times 16B = 128B$），aLBO=1024（两列核心矩阵之间的距离$BM \times 16B = 1024B$）。
UMMA描述符idesc为32位，编码了稀疏性、数据类型、是否转置等信息。
描述符图示

Swizzling的数学原理
给定的Swizzle定义为Swizzle(bits, base, shift)，其数学逻辑如下：

## A generic Swizzle functor
# 0bxxxYYYxxxxZZZxxxx
#                ^--^  Base is the number of least-sig bits to keep constant
#      ^-^    ^-^      Bits is the number of bits in the mask
#        ^------^      Shift is the distance to shift the YYY mask
1) ZZZ is the first mask, extracted right after the base
2) YYY is the second mask, extracted shift after the base
3) We XOR these two, to get AAA=YYY XOR ZZZ
4) We place this new substring in place of the first mask, ZZZ
5) Final answer becomes:
# 0bxxxYYYxxxxAAAxxxx

底层的Mojo实现通过位掩码和移位操作提取地址的第7-9位与第4-6位进行异或（XOR），从而生成避免Bank冲突的交错模式。
Kernel 3的Swizzling模式