路演时代

剖析英伟达 Blackwell 架构——张量核心、PTX指令、SASS、残次品利用、良率

Blackwell 特性

从Hopper到Blackwell，英伟达对架构进行了多项增量改进，并针对MMA相关指令修改了PTX抽象。我们在《英伟达TensorCore演进》一文中涵盖了其中的大部分内容。主要的显著变化包括：

引入了张量内存（TMEM）来存放MMA累加器。线程不再隐式拥有MMA操作的结果，相反，TMEM由软件在MMA作用域内进行显式管理。

t cgen05 操作现在由单个线程代表整个 CTA 发布，而不是像前几代那样在 warp 或 warpgroup 作用域内发布。你可以从 CuTe MMA atom 中看到这一点， Blackwell 现在使用 ThrID = Layout<_1>，而不是 Hopper 中 warpgroup 作用域 MMA 所使用的 ThrID = Layout<_128>。

支持跨协调 CTA 对的 TPC 范围 TMA 和 MMA，在 PTX 中体现为 cta_group::2，在 SASS 中体现为 2CTA。其中组成一个 TPC 的两个 SM 可以对共享操作数执行 t cgen05. mma，通过降低每个 CTA 的 SMEM 带宽需求，从而能够调用具有更高运算强度的 MMA 指令。稍后我们将展示，这种操作数共享对于充分利用现有的 MMA 吞吐量是必不可少的。

原生支持带有微缩放（micro- scaling）的亚字节数据类型

集群启动控制（CLC），作为持久化 CTA 核函数中动态任务调度的硬件支持（将在后续文章中介绍）

Hopper 架构中引入了编程依赖启动（PDL），用于隐藏连续核函数启动和设置的延迟（将在后续文章中介绍）

集群、GPC与残次屏蔽 (Floorsweeping)

自Hopper架构以来，英伟达数据中心GPU支持一项可选特性，该特性有多个名称，如“线程块集群”(thread block clusters)、“CTA集群”和“协作网格数组”(CGAs)，它们指代的都是同一个功能。集群是CTA的逻辑分组，其形状和大小可以在每个内核(kernel）中静态或动态指定。集群在编程模型中以多种实用方式呈现，其中之一是允许向同一集群中的多个CTA进行多播加载（multicast loads）；我们稍后将在TMA多播的章节中讨论这一点。

重要的是，同一集群中的CTA保证会被共同调度到同一个GPC上。在Blackwell架构中，采用“每个SM一个CTA”的“持久化CTA”(persistent CTA）风格内核时，这会产生一个重要后果：如果集群大小不能整除一个GPC中的SM数量，部分SM将处于闲置状态。这种行为可能会让内核开发者感到困惑，如果他们不了解文档中鲜有记载的GPC结构，而在启用集群的情况下盲目启动与SM总数相等的持久化CTA，就会导致部分CTA串行执行。

每个GPC中合格的SM数量并不是固定的，同一芯片上不同GPC之间的SM数量也不尽相同，甚至在同一封装内的不同裸片（die）之间也可能是不对称的。半导体制造过程中会产生缺陷，而这些缺陷可能分布在芯片的任何位置。因此，英伟达必须对其芯片进行设计，以确保这些合格的单元仍能以相对统一的方式呈现给软件。

我们提示Claude编写了一个实用程序，通过启动各种规模的集群并使用PTX %% smid来记录哪些SM出现在同一个GPC中，从而逆向工程SM到GPC的映射关系。结果得到了一个TPC到GPC的逻辑分组列表。该列表比Hopper/Blackwell中存在的8个GPC更长，因为有些TPC似乎占据了它们自己的逻辑GPC，并且从未与其他任何TPC共同调度。

从 SM100 开始，英伟达提供了一个针对这种量化问题的解决方案，使算子（kernel）在利用所有可用 SM 的同时，也能从更大的集群（cluster）中获益。算子启动时可以设置两种集群大小：首选集群大小（preferred cluster size）和备选集群大小（fallback cluster size）。通常情况下，为了利用整个 GPU，备选集群的大小应设为 2 或 1。

参考文献：

Cluster API
Cooperative groups API
CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION
CUTLASS 示例 73

逻辑 GPC 对比物理 GPC

我们上面展示的将 TPC 划分为 GPC 的方式属于逻辑分组。它们代表了软件视角下的 GPC，并未包含每个 GPC 中 20 个实际物理 SM 究竟有哪些被启用，或者每个物理 GPC 位于两个裸片（die）的具体位置等信息。实际上，具有相同逻辑配置的 B200 芯片，其每个 GPC 中产出的物理 SM 并不一定完全相同。对于在软件层面看起来完全相同的 GPU 而言，这是导致性能不确定性的潜在根源。此外，将 SM 逻辑分组为 GPC 的方式，也无法告诉我们 B200 封装中的两个裸片上分别分布着哪些 GPC。

为了进一步了解流式多处理器（SM）的物理布局信息，我们让每个SM遍历一个填满L2缓存的指针追踪（pointer- chase）数组，并测量每次加载的延迟。对于每个地址，我们将每个SM观察到的延迟与其他所有SM观察到的延迟进行对比，从而生成一个SM与SM之间的距离矩阵。其中X轴和Y轴均为SM ID。

我们可以看到两组明显的SM，它们到L2的平均距离相差超过300个周期；这一定是跨芯片（die- to- die）的连接。我们还根据上一节确定的逻辑GPC分组对这些SM进行了标记；有趣的是，单体TPC彼此靠得很近，并且在该基准测试中似乎与GPC0具有很强的相关性，因此可以推测这些TPC在物理上位于GPC0。

基于这些信息，我们可以进一步完善每个GPC中剔除（yielded）TPC的列表，尽管 5+35+3 的配置目前仍仅是推测。

此外，虽然是通过一种间接的方式，但我们可以得出结论：芯片间（die- to- die）的延迟惩罚大约为300个周期。这一点在查看基准测试中单个SM的延迟分布图时也显而易见（该图还包含了大量的L2缓存拥塞情况）。

我们要感谢来自DecartAI的Orian为基准测试提供的灵感。

存储子系统

在本节中，我们将讨论内存子系统：即在计算单元之间移动数据的硬件单元。内存复制指令是使用内存子系统的操作，而较新的世代则具备异步复制指令（异步演进过程请参阅前一篇文章）。在这里，我们将重点关注异步复制指令的两种变体：LDGSTS和TMA（张量内存加速器）。

异步复制

异步复制（PTX：cp. async，SASS：LDGSTS）是在Ampere架构中引入的，该指令异步地将数据从全局内存移动到共享内存。异步复制是非阻塞的，允许内存加载与计算重叠。它还直接写入共享内存而不经过寄存器，从而减轻了寄存器压力。

参考FlashInfer 多头注意力（MHA）算子，我们使用以下配置对异步复制（async copy）进行了基准测试：
1,2,3,4 每个SM的CTA数量：1,2,3,4
阶段数：1,2,4
每个CTA的线程数：64,128,256
加载大小：4B,8B,16B

我们绘制了吞吐量与每个SM在途字节数（即并发内存加载指令正在加载的总字节数）的关系图。

虽然在相同的在途字节数（bytes- in- flight）下，不同的加载大小会趋向于相似的吞吐量，但我们更倾向于使用16字节加载。在相似的在途字节数下，16字节加载能获得略高的吞吐量，同时消耗更少的执行资源。例如，在32 KiB在途数据量时，8字节加载需要使用4个阶段，而16字节加载仅需2个阶段。这节省了2个内存屏障对象（memory barrier objects）的内存空间，并减轻了指令发射压力。

我们的实验表明，增加阶段数可以在更高的在途字节数下实现更高的吞吐量，且在所有配置下，增加每个 CTA 的线程数都能显著提升性能。有趣的是，MLA 使用了 2 个线程束（warp）和 12 个阶段，吞吐量约为 2.2 TB/s。我们认为这是因为执行 Softmax 的线程束需要最多的寄存器，而增加线程束数量会减少每个线程分配到的寄存器数量。

我们对同一组配置的延迟进行了基准测试。我们发现 LDGSTS 的基础延迟约为 600 纳秒，而在在途数据量超过 8 KiB 后，延迟几乎翻倍。这是因为我们需要为 LDGSTS 使用大量线程以实现高在途字节数，这导致大量 Warp 因 MIO（内存输入输出）节流而停顿。

张量内存加速器 (TMA)

TMA (PTX: cp. async. bulk. tensor, SASS: UTMALDg) is Hopper 架构中引入的一种异步数据复制引擎，专门用于将大量数据从全局内存移动到共享内存。单个线程即可启动TMA来执行地址生成、内存交错（swizzling）以及越界处理，从而释放其他线程以执行独立任务。在这里，我们对2D张量版本（cp. async. bulk. tensor. 2d）进行基准测试，以代表典型的TMA用法。

参考 FlashInfer 注意力算子，我们对 TMA 进行了基准测试，每个 SM 仅分配一个 CTA，但每个 CTA 使用 1 到 4 个 Warp，且每个 Warp 仅使用一个线程来发布不同 Box 大小的 TMA 指令。下图显示了每种在途字节数（bytes- in- flight）情况下的最佳吞吐量。

我们使用以下配置对 TMA 进行基准测试：
每个 SM 的 CTA 数量：1
每个 CTA 的线程数：128（4个warp）
TMA 框维度：2D 形状，尺寸从 32x8 增加到 128x128

峰值吞吐量的达到时间远晚于LDGSTS。

异步复制 (Async Copy) 与 TMA 的对比

像 FlashInfer 这样的深度学习算子库会同时使用 TMA 和异步复制来加载数据。TMA 和异步复制具有不同的性能特性：TMA 擅长具有规则访问模式的大批量加载，但延迟较高；而异步复制可以处理不规则的内存访问模式，但存在大小限制。我们将解释在何种条件下应选择其中之一。在此，我们对 FlashInfer 在 MHA 和 MLA 算子中使用的配置进行了基准测试。

我们发现，在吞吐量方面，当在途数据小于 32 字节时，异步复制（async copy）的性能略优于 TMA，但随后 TMA 追赶上来，并能持续扩展至 128 KiB。在延迟方面，我们观察到在在途数据小于 12 KiB 时，异步复制的延迟略低于 TMA，但此后 TMA 的延迟大幅增加。

实际上，Blackwell MLA 算子使用异步复制（async copy）来动态加载页面，而其 MHA 算子则仅使用 TMA。FlashInfer 的大部分 Blackwell MHA 算子是由 TRT- LLM 贡献的，因此我们只能通过研究二进制文件来推测这些算子的行为。我们发现，与 Hopper 类似，所有 Blackwell TRT- LLM 算子都使用 TMA。我们怀疑对于动态页面加载，这些算子遵循了 Hopper 算子的做法，即使用以页面索引作为最后一维的 4D TMA，并在需要时索引到 TensorMap 对象。为了了解这些算子的确切机制，我们敦促 NVIDIA 开源 FlashInfer TRT- LLM 算子，以造福社区。

TMA 多播

TMA 支持多播模式，即通过 CTA 掩码指定，单次加载即可将数据复制到多个 SM 的共享内存中。多播通常用于类 GEMM 模式，其中输入分片在处理不同输出分片的 SM 之间共享。例如，多播对于激活函数 SwiGLU 非常有用，该函数采用双 GEMM 模式，即两个 GEMM 操作共享一个输入矩阵。其主要优势在于减少了 HBM 加载，从而降低了有效带宽占用。它还显著减少了 L2 流量，因为多个 CTA 对共享数据的请求被合并为一个请求。

根据 NCU 的说法，负责处理 TMA 多播请求的单元被称为 L2 请求合并器 (L2 Request Coalescer, LRC):
L2 请求合并器 (LRC) 处理进入 L2 的请求，并尝试在将读取请求转发到 L2 缓存之前对其进行合并。它还负责处理来自 SM 的程序化多播请求，并支持写入压缩。

听起来硬件可能提供了一些多播行为，即使没有明确请求，例如缺失状态保持寄存器（MSHR）。我们通过运行相同的 TMA 多播基准测试来验证这一点，不同之处在于，这次不是由一个 CTA 发布多播加载，而是所有 CTA 都针对相同的数据发布独立的 TMA 加载。

在这里，我们比较三种情况：

每个 SM 加载不同的数据（基准）
TMA 多播（显式）——每个集群（Cluster）中的一个 CTA 向其集群内的所有 CTA 发布多播加载指令
TMA 多播（隐式）——集群（cluster）中的所有 CTA 对相同数据发出普通 TMA 加载指令。

TMA 多播允许以更高的加载带宽填充 SMEM 缓冲区，即使数据尚未存在于 L2 缓存中。对于已知的流量模式，显式 TMA 多播指令可以完美消除 L2 流量，从而实现理想的“1 / 集群大小”的 L2 字节与 SMEM 字节之比。我们还观察到，在这个简单的基准测试中，显式和隐式情况下的 SMEM 填充吞吐量几乎相同。然而，我们可以看到 LRC（本地驻留控制）并不完美；在隐式情况下，L2 接收到的流量略多，尤其是当总量增加时。

在有效内存吞吐量方面，隐式多播与显式多播的表现相当。然而，在减少L2缓存流量方面，当在途数据超过64字节后，隐式多播将失去效力。

DSMEM与SMEM

NVIDIA在Hopper架构中引入了分布式共享内存（DSMEM）。DSMEM允许集群（Cluster）内的线程块（CTA）访问彼此的共享内存。这对于跨CTA归约等模式非常有用。通过DSMEM读取同行CTA内存的吞吐量显著低于SMEM每时钟周期128字节的水平。

我们针对与DSMEM交互的几种不同PTX模式进行了实验。编写DSMEM与SMEM代码的一个重要区别在于，DSMEM加载是像全局加载（GlobalLoad）一样进行分组（Packetized）的，因此其最佳访问模式完全不像局部SMEM那样通过交错访问来避免银行冲突（BankConflict），而更像是对全局内存（GMEM）中连续位置的典型合并访问（CoalescedAccess）。此外，我们观察到，为了获得局部SMEM全额128B/周期的吞吐量，必须使用不带：:cluster修饰符的ld.shared。这是我们在编写基准测试时遇到的一个陷阱，当时我们简单地对局部和远程DSMEM地址都使用了ld.shared::cluster。使用ld.shared时，编译器会发出LDS指令，而不是在使用ld.shared::cluster时发出的通用LD指令，后者似乎无法达到局部SMEM的峰值吞吐量。我们还发现很难通过ld.shared::cluster进一步提升实际吞吐量，直到切换到cp.async.bulk（PTX）/UBLKCP（SASS）以在每条指令中移动更大量的数据后，才在DSMEM上获得了略高一些的吞吐量。

我们使用每种PTX模式所达到的峰值吞吐量如下，以每时钟周期字节数(B/clk)表示，以便与SM本地SMEM中已知的最大可实现值保持一致。

第五代 Tensor Core MMA

MMA 指令是执行矩阵乘法的核心操作。从 Hopper 到 Blackwell, MMA 的性能表现对矩阵形状 (Shape) 的依赖性日益增强。在此, 我们通过对不同形状和数据类型进行遍历测试, 来量化这些性能差异, 并对这一现象展开深入研究。

Blackwell 引入了 2SM MMA, 这是一种新型的 MMA 指令 (. cta_group::2), 其中一个 CTA 对跨越 2 个 SM 协作执行一次 MMA 操作。具体而言, 输入矩阵 A 被复制, 而矩阵 B 和 D 则分片分布在 2 个 SM 上, 且该 CTA 对可以互相访问彼此的共享内存。这使得更大规模的 MMA 形状成为可能。我们研究了 2SM MMA 表现出的是弱扩展性、强扩展性, 还是两者兼有。

我们对 MMA 性能进行了基准测试, 配置范围如下:

格式: MMA 输入和输出数据格式 (BF16, FP8, INT8, FP4, MXFP8, MXFP4)
AB 布局: 矩阵 A 和 B 的数据位置 (SS: 两者均在 SMEM, TS: A 在 TMEM)
CTA 组: 执行一个 MMA 所需的 SM 数量 (1, 2)
MNK: 输入和输出矩阵的形状 (因格式、AB 布局和 CTA 组而异)
流水线深度: 在途 MMA 操作的数量 (因格式而异)

吞吐量

NVIDIA 针对不同的输入数据类型给出了特定的吞吐量性能指标。在此，我们展示了针对每种（格式 + CTA 组）的官方数据，并将其与可达到的最大吞吐量进行了对比。我们证明了 UMMA 在所有格式和 CTA 组中均能实现接近峰值的吞吐量，即使在可能存在协调开销问题的 2SM 版本上也是如此。

对于所有N尺寸下的单SMMMA，我们观察到较小的M=64仅达到理论峰值吞吐量的 50%50% ，而较大的M=128则接近 100%100% 。这证实了M=64仅利用了一半的数据通路。对于双SMMMA，我们发现M=128的吞吐量在N=64时以 90%90% 的峰值开始，并在所有其他N尺寸下接近 100%100% 。M128N64的吞吐量必然受限于不同的硬件单元，如TMEM、L2、SMEM等。同时，M=256在所有配置下均保持接近 100%100% 的峰值吞吐量，这是因为M=256相当于每个SM负责M=128，能够利用完整的数据通路。我们注意到，在相同数据类型位宽的格式下，吞吐量是完全一致的，且微缩放（micro- scaling）数据类型几乎没有额外开销。

MMA支持两种不同的AB布局：两个输入矩阵都存储在SMEM中(SS)，以及矩阵A存储在TMEM且矩阵B存储在SMEM中(TS)。我们观察到，对于M=128，虽然ABLayout=TS达到了接近峰值的吞吐量，但ABLayout=SS在较小的N尺寸下表现不佳，直到N=128时才追赶上米。

我们可以证明，这是因为在SS模式下，当 N<128N<128 时，指令本身受限于共享内存（SMEM）带宽。例如，对于FP16，我们已知硬件每周期每个SM可执行8192次MMA浮点运算（FLOPs），而SMEM带宽为128B/周期（每个SM）。因此，对于 M=128 N=64 K=16M=128 N=64 K=16 ，我们有：
A_bytes = 2MK = 4096; B_bytes = 2NK = 2048; FLOPs = 2MN*K = 262144
SMEM Cycles = (A_bytes + B_bytes) / (128 B/clk) = 48 cycles
Math Cycles = FLOPs / (16384 FLOPs/clk) = 32 cycles

我们针对不断增加的N值进行计算，发现从 N=128N=128 指令开始，终于达到了算力瓶颈（Math limited）。

其他数据类型也是如此——如果两个操作数都在 SMEM 中，当 N<128N<128 时，MMA 指令会受限于 SMEM 带宽。

为了进一步说明这一点，我们绘制了所有形状的 FP8 1SM MMA 的 Roofline 模型图。我们可以清楚地看到，当 N<256N<256 时处于内存受限区域，且斜率约为 128 字节/周期，即 SMEM 带宽。

2SM MMA 在所有格式和形状下都实现了完美的弱扩展性，当使用两倍于 1SM MMA 的计算资源时，速度提升达到了 2 倍。在 A BLayout=SS 的较小形状中，我们观察到了超过 2 倍的加速，这同样是因为在 SS 布局且 N<128 时，该指令受限于共享内存（SMEM）带宽，而 2SM 版本将操作数 B 分散到了两个 SM 之间。

这些实验表明，为了获得最大吞吐量，在给定的SMEM tile大小下，应始终使用可用的最大指令形状。

延迟

我们对单个MMA指令的延迟进行了基准测试，并将对比图绘制如下。在所有配置中，我们观察到延迟随N从64到128呈线性增加，而 N=256N=256 处的峰值可能是由于从128到256的跨度所致。对于单个CTA组的MMA，1SM MMA在 M=64M=64 和 M=128M=128 时在不同N尺寸下具有相似的延迟；而在2SM MMA中， M=256M=256 的延迟增长速度略快于 M=128M=128 ，这符合我们的理论估算。对比数据类型，我们发现1SM的差异很小，但2SM MMA则表现出明显的区分。

我们注意到延迟顺序呈现出一种微小但一致的规律：

S8<BF16=E4M3=F4<MXF8=MXF4

我们认为，整数运算更高的能效比使得S8速度最快，而比例因子的计算则为MXF8和MXF4引入了轻微的额外开销。

不同在途指令数下的吞吐量

在我们的吞吐量基准测试中，我们设置了较高的在途指令数（从256到1024不等），以摊销指令发射和提交等待的开销。然而，内核通常仅使用1到4条在途MMA指令。我们测试了在途指令数为1到10时的吞吐量，并在此讨论吞吐量的变化。

在所有配置中，我们观察到相同的N值和在途（in- flight）MMA指令达到了相近的光速比（SoL）百分比。值得注意的是，只有最大的N值达到了 90%90% 的SoL，而最小的N值仅达到约 70%70% 。对比1SM和2SM的MMA，我们发现1SM的SoL吞吐量比对应的2SM高出约 5%5% 。对于相同的数据格式和CTA组MMA，较大N值的吞吐量始终高于较小的N值。最后，我们观察到4个在途MMA的吞吐量SoL百分比上限在 78%78% 到 80%80% 之间。

实际应用案例：CUTLASS

下面我们将讨论使用算子编写库 CUTLASS 的实际应用案例。我们还将探讨吞吐量、多播（multi- cast）以及芯片布局图（floorplans）。

这里我们来看一个实际的应用案例：分块GEMM的内层循环模式，即沿着内层维度K持续加载分块，以计算输出矩阵的一个分块。

首先，我们使用Nvidia的CUTLASS库实例化一个具有单个持久化CTA（协作线程阵列）的算子，用于计算单个输出分块。我们通过改变DMA- >Math流水线中的加载阶段数，在增加阶段数的同时使用越来越多的SMEM（共享内存），并以此实现更好的延迟掩盖。K维度被设置为极大值，以便我们测量软件流水线在稳态模式下所达到的数学吞吐量，并将其表示为硬件理论极限数学吞吐量的百分比。

对于选定的分块大小（tile size），可以通过使用更多的共享内存（SMEM）来增加流水线级数，从而更好地隐藏延迟。对于给定的流水线深度N，如果其中一级正在执行MMA（矩阵乘累加）计算，则最多可以有N- 1级的A和B缓冲区处于传输过程中。换个角度来看，对于给定的加载阶段，只要其延迟不超过（N- 1)*M（其中M是单级执行MMA所需的时间），就可以被完全隐藏。因此，在分块大小固定的情况下，延迟隐藏能力会随着流水线级数的增加而严格提升。

然而，如果 GEMM 的收尾阶段（Epilogue）消耗了共享内存（SMEM），那么可用于主循环的内存就会减少，因此必须降低所使用的阶段数（Stage Count）。在这里，我们使用了一个不消耗 SMEM 的收尾阶段，因此当阶段缓冲区消耗的总 SMEM 达到每个流式多处理器（SM）的最大 SMEM 容量时，这些线条会在 X 轴方向终止。

因此，共享内存（SMEM）的使用量以及随之而来的流水线级数（stage count）是基于以下因素计算的：

操作数 SMEM 切片（A 和 B）
屏障存储
收尾阶段（Epilogue）的 SMEM 使用

吞吐量与多播

此前，我们研究了单条指令的可实现吞吐量；我们发现某些指令可能本质上受限于共享内存(SMEM)带宽，但在衡量每种指令形状的峰值性能时，我们并未考虑内存系统。在此，我们采用相同的单CTA CUTLASS基准测试，并将其扩展到使用大于1x1的集群(Cluster)规模。此外，对于任何M维度为偶数的集群形状，我们都使用了2SM MMA原子。

图表中显示了一些此类结果，在每个图表中，我们保持每个CTA的切片大小 (tile size) 不变，并针对1SM和2SM分别改变集群大小 (cluster size)。请注意，对于128x128的每CTA切片形状，当集群N维度从1扩展到2(2SM)或从2扩展到4(1SM)时，我们从集群引入的多播 (multicast) 中获得了显著收益。因此我们可以得出结论，小于此尺寸的切片形状受限于存储子系统 / L2带宽。

平面图

Blackwell 与 Blackwell Ultra 的布局相似，如下图所示。正如我们在前一章节中所讨论的，我们可以看到 8 个 GPC、L2 缓存分段以及芯片间互连通道（die-to-die crossings）。

温馨提示：内容源于第三方以及公开平台，仅供用户参考，恕本平台对内容合法性、真实性、准确性不承担责任。如有异议/反馈可与平台客服联系处理（微信：_LYSD_）。