随着 NVIDIA Blackwell 架构的推出，GEMM 性能优化迎来了新的机遇与挑战。本文深入探讨了如何利用线程块cluster和 2-SM UMMA 等关键技术，通过 CUTLASS 库实现极致性能。它不仅解释了底层原理，还结合具体示例，为开发者提供了在 Blackwell GPU 上编写高效 GEMM 核心的实用指南。
智能速览
线程块cluster能将相邻SM分组，促进高级协作。
TMA多播通过协作加载显著降低全局内存访问压力。
2-SM UMMA允许一对CTA协同处理更大的MMA分片。
利用位掩码和mbarrier可实现精准的CTA间同步。
CUTLASS/CuTe库抽象了复杂的底层索引逻辑。
精华内容
要充分释放 Blackwell 架构在 GEMM 上的潜能，必须掌握线程块cluster与UMMA的协同工作方式。下面将深入剖析其核心实现机制与优化细节。
线程块cluster基础
线程块cluster是一种将物理上相近的SM进行分组的结构，确保其中的线程块被共同调度到同一个GPC上。这一特性最早在Hopper架构中引入，为开发人员提供了新的硬件协作层次。cluster中的线程块可以访问彼此的共享内存，即分布式共享内存（DSMEM），为协作加载和同步奠定了基础。Cluster大小以dim3元组定义，最大可移植大小为8，部分Blackwell GPU可扩展至16。
TMA多播加速
TMA多播是一种加速数据传输的关键特性，它能通过一次操作将相同的张量分片加载到同一cluster内的多个CTA的共享内存中。在GEMM场景下，当多个CTA需要相同的操作数分片时，此特性可将全局内存访问量成倍减少。例如，一个4个CTA的cluster，每个CTA仅需加载数据的四分之一，相比传统方式总数据量降低4倍。这种协作式加载极大缓解了内存带宽瓶颈，但其实现依赖于精准的CTA间同步机制。
精准同步机制
实现TMA多播和UMMA协作的关键在于精准的同步。同步主要通过mbarrier和位掩码完成。位掩码是一个16位整数，用于指定cluster中参与特定操作的CTA。例如，在<4,4,1>的cluster中，CTA 0的行操作数多播掩码为0x1111，表示与同行CTA共享数据。对于TMA，mbarrier通过事务计数（tx-count）跟踪总加载字节数，确保所有数据加载完毕后才继续执行。这种细粒度的同步机制避免了不必要的等待，提升了流水线效率。
2-SM UMMA协同
2-SM UMMA（或称Pair-UMMA）是Blackwell架构引入的另一项重要功能，它允许一对物理上相邻的CTA协同处理同一个MMA操作。在这一模式下，每个CTA负责加载操作数分片的一半，并在其张量内存中保存一半的累加器。与让两个CTA分别执行两个独立的MMA相比，这种协同方式在完成相同计算量（FLOP）的同时，将操作数数据传输量减半，显著提升了计算效率。CTA的配对基于其在cluster内的索引，通常按最左维度的奇偶性进行。
Pair-UMMA的实现
在实现Pair-UMMA时，同步和位掩码逻辑变得更加复杂。由于每个CTA只处理MMA分片的一半，TMA多播只需将数据发送给具有相同奇偶性的CTA，减少了通信开销。然而，MMA的完成信号需要覆盖整个CTA对。为此，Blackwell引入了新的同步指令和地址计算方法，通过修改mbarrier地址的第24位，一个CTA可以直接访问其对等CTA的同步屏障。CUTLASS通过`umma_arrive_multicast_2x1SM`等封装函数简化了这一过程，并利用`make_tma_atom_sm100`自动适配多播维度，极大降低了开发者的编程复杂度。
通过线程块cluster、TMA多播和2-SM UMMA的组合，NVIDIA Blackwell架构为GEMM性能带来了质的飞跃。掌握这些技术，意味着能够充分压榨硬件潜能。未来，随着对低精度和块缩放支持的深入，Blackwell的计算能力还将进一步释放，开发者们准备好了吗？
