Blackwell Programming for the Masses With OpenAI Triton
Blackwell Programming for the Masses With OpenAI Triton
Phil Tillet | OpenAI
dePaul Miller | NVIDIA
目录
- 演讲大纲 (Talk Outline
- Triton 编程语言 (The Triton Programing Language
- Triton 编程模型 (The Triton Programming Model
- 矩阵乘法 (Matrix Multiplication
- 持久化矩阵乘法 (Persistent Matrix Multiplication
- Blackwell
- 再次探讨 Blackwell 上的 Group GEMM
- 寻找唾手可得的优化点 (Finding Low-Hanging Fruits
- 修复网格大小 (Fixing the Grid Size
- 与持久化 GEMM (Persistent GEMM) 的比较
- 调试性能差距：Nsight Systems
- 调试性能差距：Nsight Compute
- NVIDIA Blackwell 架构
- Blackwell 性能 101
- 性能结果
- 进一步优化
- 更进一步 (Going Further
- 致谢
演讲大纲 (Talk Outline)
- 引言 (Introduction)
- Hopper 的问题与 Blackwell 的解决方案 (Hopper Problems; Blackwell Solutions)
- 为 Blackwell 优化 Triton 核函数 (Optimizing Triton Kernels for Blackwell)
Triton 编程语言 (The Triton Programing Language)
动机：CUDA 的权衡 (Motivations: CUDA trade-offs)
- CUDA-C++：以牺牲简单性为代价换取灵活性
-
优点： 极高的优化潜力
- 性能： 开发者可以对计算进行最优调度
- 灵活性： 开发者可以使用最合适的数据结构
-
缺点： 极低的研究迭代速度
- 生产力： 复杂的修改可能非常耗时且不具备可移植性
- 可靠性： 复杂的修改可能容易引入错误（例如，竞争条件）
动机：XLA/Torch 的权衡 (Motivations: XLA/Torch trade-offs)
- 图编译器：以牺牲灵活性为代价换取简单性
-
优点： 极高的研究迭代速度
- 生产力： 某些想法可以非常快速地进行原型设计
- 可靠性： （假设编译器没有错误）错误的表面积很小
-
缺点： 极低的优化潜力
- 性能： 并行化/调度由编译器完成；没有操作空间
- 灵活性： 并非所有的研究想法都能被表达出来
动机：Triton 的权衡 (Motivations: Triton trade-offs)
- Triton：一种折中方案
-
优点： 较高的研究迭代速度；较高的优化潜力
- 生产力： 在 CUDA 中需要数周/数月的工作，在 Triton 中约需数天完成
- 可靠性： 相较于 CUDA，出错空间更小（例如，未暴露簇内并发）
- 性能： 开发者保持对粗粒度并行化和调度的控制
- 灵活性： 开发者保持对粗粒度控制流的控制（例如，树、基数排序...）
-
缺点： 较低的研究迭代速度；较低的优化潜力
- 生产力： 在 Torch 中需要数小时的工作，在 Triton 中约需数天完成
- 可靠性： 相较于 Torch，出错空间更大（例如，暴露了簇间并发）
- 性能： 开发者失去对细粒度并行化和调度的控制
- 灵活性： 开发者失去对细粒度控制流的控制（例如，warp 间的并发）
生产力 / 性能帕累托前沿 (Productivity / Performance Pareto Frontier)
Triton 编程模型 (The Triton Programming Model)
概述 (Overview)
- 可以认为是 Torch，但张量存在于片上内存（on-chip）！
- 存在一些注意事项：
- 形状 (shapes): 张量维度必须是 2 的幂次方
- spmd: 程序在启动时通过不同的 ID 进行批量实例化
- 内存模型 (memory model): HBM 通过指针和张量描述符进行访问
矩阵乘法 (Matrix Multiplication)
概述 (Overview)
以下是使用 Triton 实现矩阵乘法的代码示例。
性能 (Performance) (FP16; M=N=8192; cuBLAS 12.8)
下图比较了 Triton 和 cuBLAS 在矩阵乘法上的性能。
软件流水线 (Software Pipelining)
- 存在大量的“气泡” (bubbles) 👾
- 每次程序实例完成时，软件流水线都会被清空 (flushed)
- 这会损害性能，尤其当归约循环 (reduction loop) 较浅时
- 这个问题只有通过拥有能够处理多个输出瓦片（output tiles）的程序实例（即持久化 persistent）才能缓解
持久化矩阵乘法 (Persistent Matrix Multiplication)
概述 (Overview)
以下是持久化矩阵乘法的 Triton 实现代码。
软件流水线 (Hopper)
- 气泡减少了，但仍不完美！🔮
- Triton 将来自不同外层循环迭代的 I/O 与计算重叠
- 累加器 (Accumulators) 存在于寄存器中 (在 Hopper 架构上)
- Epilogue 被暴露：在寄存器被需要之前无法发出下一个 MMA 指令...
- ...除非使用 warp-specialization + ping-pong schedule (例如 cuBLAS)
H100 性能 (FP16; M=N=8192; cuBLAS 12.8)
下图展示了在 H100 上的性能表现。
软件流水线 (Blackwell)
下图对比了 Hopper 和 Blackwell 架构上的软件流水线。
让我们尝试在 Hopper 上运行 matmul 内核，然后……看看会发生什么？
GB200 性能表现
下图展示了在 GB200 平台上，使用 FP16 精度，当 M=N=8192 时，Triton 和 cuBLAS 12.8 在矩阵乘法上的性能对比。图中显示，Triton 的性能（蓝色曲线）显著低于 cuBLAS（绿色曲线）。cuBLAS 的性能大约在 1400 TFLOPS，而 Triton 的性能大约在 900 TFLOPS。
瓶颈分析
- 注意：流水线图示并非按比例绘制！
- 内核（Kernel）的瓶颈在于加载延迟（load latency）。
从下方的流水线图中可以看出，在加载 Load A₂ 和 Load B₂ 的同时，张量核心（Tensor Cores）正在执行 MMA₀。但在 MMA₀ 完成后，由于 Load A₂ 和 Load B₂ 尚未完成，张量核心进入空闲状态，直到数据准备好才能开始执行 MMA₁。这种空闲状态表明计算单元未被充分利用。
- 解决方案：尝试将流水线深度增加到 num_stages=4，观察其效果。
增加流水线深度
下图对比了 num_stages=3 和 num_stages=4 的情况。
- 当 num_stages=4 时，流水线图看起来更好。通过增加流水线深度，可以更好地隐藏数据加载的延迟，使得 MMA 操作能够连续执行，从而减少或消除了张量核心的空闲时间。
- 让我们运行一下，看看实际结果如何。
Blackwell 平台性能问题... 哦不！
当尝试在 Blackwell 平台上运行时，程序抛出了资源不足的错误。
错误信息如下：
triton.runtime.errors.OutOfResources: out of resource: shared memory. Required: 262176, Hardware limit: 232448. Reducing block sizes or 'num_stages' may help.
这表明增加流水线深度（num_stages）导致共享内存（shared memory）的需求超过了硬件限制。
结尾部分拆分 (Epilogue Splitting)
为了解决共享内存不足的问题，同时保持较高的流水线深度，可以采用“结尾部分拆分”（Epilogue Splitting）的策略。以下是实现该策略的 Triton JIT 代码示例。关键在于 tl.split(acc)，它将累加器 acc 拆分，允许以更小的块进行存储，从而减少对资源的瞬时需求。
应用优化后的 GB200 性能
在应用了结尾部分拆分等优化后，Triton 的性能得到了显著提升。从下图中可以看到，Triton 的性能曲线（蓝色）现在更接近 cuBLAS（绿色），尤其是在 K 值较大时。尽管仍略低于 cuBLAS，但差距已经大幅缩小。
注意事项与未来工作 (Caveats)
- 当前实现仍然较慢... 但并非世界末日！
- 通用人工智能（AGI）的实现不会仅仅来自于 10% 的矩阵乘法速度提升。
- 我们正在努力改进
- 我们仍未使用一些 Blackwell 架构的新特性（例如，2CTA）。也许我们受到了某种限制（throttling）？
- 较长的结尾部分（epilogues）问题仍然部分存在。
下图展示了即使在优化后，结尾部分（Epilogue）仍然可能在流水线中引入气泡（bubble），导致计算单元的短暂空闲。
Blackwell
概述 (Overview)
-
张量内存 (Tensor Memory)！
- Blackwell 张量核心 (tensor cores) 从内存（SHMEM 或 TMEM）中消耗其所有操作数
- 累加器现在可以由我们的编译器透明地进行多重缓冲 (multi-buffered)
-
新的同步原语 (New Synchronization Primitives)！
mbarriers允许不同的 MMA (Matrix Multiply-Accumulate) 进行独立流水线化- 同步由编译器透明地处理
再次探讨 Blackwell 上的 Group GEMM
- 使用 Nsight Compute 运行 Triton Group GEMM 内核，以了解基线性能。
- Nsight Compute 会高亮显示许多 GPU 指标，并给出优化建议。
寻找唾手可得的优化点 (Finding Low-Hanging Fruits)
-
Nsight Compute 的建议包括：
- L2 压缩优化
- 全局访问模式优化
- 修复共享内存库冲突 (Shared Bank Conflicts)
- 改进指令发射槽统计 (Issue Slot Statistics)
- 减少长记分板停顿 (Long Scoreboard Stalls)
- 修复小网格 (Small Grid) 问题
- SM 占用率 (SM Occupancy)
- 工作负载不平衡
-
第一个建议是修复小网格问题。
- 我们总共有 148 个 SM，但只有 128 个被激活！
修复网格大小 (Fixing the Grid Size)
- 修复这个问题就像改变我们的自动调优策略一样简单。
- 修改前:
@triton.autotune(configs=[
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': 128}))
])- 修改后:
def num_sms():
return torch.cuda.get_device_properties('cuda').multi_processor_count
@triton.autotune(configs=[
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
])- 增加了网格大小后，我们获得了 14% 的性能提升。
与持久化 GEMM (Persistent GEMM) 的比较
- 与我们的“光速”基准
matmul_descriptor_persistent内核相比。 - 我们的 Group GEMM 速度仅为该基准的 29%。
调试性能差距：Nsight Systems
- 在 Nsight Systems 中比较两个内核。
- SM 指令部分显示出 Tensor Core 利用率存在巨大差距！
调试性能差距：Nsight Compute
与基准进行比较！
- 通过在 Nsight Compute 中运行两个内核来进一步挖掘。
- 我们可以右键点击 Persistent GEMM，将其设置为比较基准。
- 这证实了 Persistent GEMM 更好地利用了张量核心 (Tensor Cores, TC)。
- 我们需要更好地利用第 5 代张量核心。
NVIDIA Blackwell 架构
-
Blackwell 张量核心 - tcgen05
- 在同等时钟频率下，吞吐量是 Hopper 张量核心的 2 倍。
- 支持新的 8b (MXFP8), 6b (MXFP6) 和 4b (MXFP4) 微缩数据类型。
-
张量内存 (Tensor Memory, TMEM)
- 每个 SM 上用于张量核心输入和输出的新内存。
-
了解更多关于 Blackwell 张量核心的信息：
- "Programming Blackwell Tensor Cores with CuTe and CUTLASS" by Cris Cecka and Mihir Awatramani on Friday
- "How Math Libraries Can Help Accelerate Your Applications on Blackwell GPUs" by Azi Riahi and Babak Hejazi on Wednesday
- "Blackwell TensorCore and Memory Subsystem Optimizations for Applications" by Vishal Meta, Petrick Liu, Allard Hendriksen, Shang Zhang, and Siyuan Fu on Wednesday
- 以及更多！
Blackwell 性能 101
- 第 5 代张量核心支持高达 128x256x16 的 FP16 MMA (矩阵乘加) 运算。
- 通常我们希望 K block size 具有 128B 的连续加载。
- Triton 已经支持所有这些，但需要将其添加到自动调优配置中！
- 因此，让我们增加自动调优的 tile size，以允许这些更大的 tile 配置：
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'NUM_SM': num_sms()}),
性能结果
- 现在我们的速度比基线性能快了 3.2 倍，并且与我们的“光速”内核的性能差距在 19% 以内！
进一步优化
- 我们可以查看 NCU (Nsight Compute) 中的计算工作负载分析部分。
- 看起来 Group GEMM 有更多的 ALU 操作，并且没有使用 TMA！
- TMA (张量内存加速器) 允许从全局内存到共享内存的高效异步张量复制。
- 它在硬件中处理越界索引计算和越界访问预测。
- 它在硬件中处理越界索引计算和越界访问预测。
Triton 中的 TMA
- 在 Triton 中使用 TMA 分为两部分：
- 1. 定义一个描述符：
a_desc = tl.make_tensor_descriptor(
a_ptr,
shape=[group_shape_m, group_shape_k],
strides=[group_stride_m, 1], # row major tensor
block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], # copy block size
)- 2. 执行加载操作：
#
a_desc.load([offs_m, offs_k]) # load with starting coordinate
性能增益
- 对于我们的 8192 x 8192 x 8192 GEMM，使用 TMA 加载可将性能提升 15%。
- Group GEMM现在的性能与我们的“光速”GEMM 内核的性能差距在 7% 以内！
- 当运行更多分组时，我们也能获得更好的性能。
- 例如，对于 8 组 M=256, N=8192, K=8192 的计算，我们的速度是原始内核的 2.4 倍。
- 例如，对于 8 组 M=256, N=8192, K=8192 的计算，我们的速度是原始内核的 2.4 倍。
更进一步 (Going Further)
-
在 Blackwell 上优化现有的 Triton 内核非常简单！
- 自动调优大有帮助！
- 针对 Blackwell 更新了参考教程（包括微缩 MMA）。
- 易于泛化到更多数据类型。
- Group GEMM 现在也增加了 FP8 支持！
-
Nsight 与 Triton 开箱即用，是指导优化的绝佳工具。
- 了解更多关于 Nsight 和优化的信息：
- "It's Easier than You Think - Debugging and Optimizing CUDA with Intelligent Developer Tools" by Mahender Hari
- "What's in Your Developer Toolbox? CUDA and Graphics Profiling, Optimization, and Debugging Tools" by Rafael Campana et al.
- 以及更多！
致谢
- 贡献者致谢
- Open AI: Thomas Raoux, Pawel Szczerbuk, Peter Bell, Jeff Niu
- NVIDIA: Shang Zhang, Samantha Hirsch, Yujia Zhai, Pradeep Ramani, Matthew Brookhart, Masahiro Masuda, Chris Sullivan, Clive Unger, Jason Knight, Wei Liu
