矩阵乘法和注意力机制是现代 AI 工作负载的计算支柱。虽然库如 NVIDIA cuDNN 提供高度优化的实现，而框架如 CUTLASS 提供深度自定义，但许多开发者和研究人员都需要将性能与可编程性结合起来的中间点。
NVIDIA Blackwell 架构 上的 开源 Triton 编译器 通过直观的编程模型展示 Blackwell 的高级功能来满足这一需求。
得益于 NVIDIA 与 OpenAI 的持续合作，Triton 编译器现已支持 NVIDIA Blackwell 架构。这可确保开发者和研究人员可以通过基于 Python 的编译器 (例如 Triton) 轻松使用 Blackwell 架构中最新的出色功能。
NVIDIA Blackwell 的性能提升
NVIDIA Blackwell 架构在原始计算能力和架构创新方面实现了显著提升。NVIDIA 与 OpenAI 的合作重点是通过 Triton 的编译器基础设施透明地利用这些功能，特别是在两个关键领域：
- 矩阵乘法，包括 flash attention
- 全新精度格式
矩阵乘法
NVIDIA Blackwell 架构添加了全新的 Tensor Core，该核心经过全新设计，可提高吞吐量和 能效 。
通过扩展 Triton 的矩阵乘积累加（MMA）流水线机制，我们可以自动利用 NVIDIA Blackwell 的全新 Tensor Cores。这需要仔细分析内存访问模式和复杂的编译器转换，以确保正确高效的计算/数据移动重叠。
因此，开箱即用的 FP8 和 FP16 GEMM 运算性能卓越，这些优化可自动应用到使用 Triton 的 `tl.dot` 基元的任何内核。总体而言，Triton 设法实现了近乎最佳的性能，可与多个关键用例中的库实现相媲美。
图 1 显示，在此针对 GEMM 内核典型生成式 AI 大小的 K 扫描分析中，NVIDIA Blackwell 架构上的 Triton 优化为 FP16 和 FP8 用户带来了硬件性能提升，正如 Triton 教程中所提供的那样。
闪光注意力
Flash Attention 是现代 Transformer 架构中的一个关键基元，通过 Triton 在 NVIDIA Blackwell 上显著加速，与 NVIDIA Hopper GPU 架构相比，FP16 注意力提升高达 1.5 倍。
虽然我们会继续通过 FP8 和其他精度的持续编译器增强来优化绝对性能，但目前的工作可以帮助客户在现有产品发布当天就轻松过渡到 NVIDIA Blackwell。
这里要注意的另一个重要方面是，能够通过现有的 Triton 闪存注意力实现“免费”提供这种性能提升，而无需更改代码。
图 2 显示更复杂的工作负载 (例如 Triton 教程中提供的 flash attention kernel) 再次显示，在使用 Triton 编译器改进进行解锁时，NVIDIA Blackwell 架构可实现巨大的性能提升。这项工作的一些改进也提高了 NVIDIA Hopper 的注意力表现，并且不会显示在此数据中。
全新精度格式
NVIDIA Blackwell 引入了革命性的块级浮点格式，包括 Open Computing Project 的 microscaling 格式 ，Triton 现在为 NVIDIA Blackwell 提供硬件加速功能。
相较于目前 LLM 推理项目中常见的非原生块扩展技术，这些格式可提供更高的平均精度和更高的性能。
对于 OCP 格式支持，Triton 上的 MXFP8 GEMMs 展示了非凡的性能，与本文前面展示的 FP8 GEMMs 加速性能相似，同时原生允许在 Tensor Core 中进行扩展。
同样，MXFP4 在精度性能权衡空间中提供了一个新的操作点，同时提供的硬件加速性能是 FP8 和 MXFP8 GEMMs 的两倍。
如需详细了解新的块级浮点支持，请参阅专门介绍此功能的 新 Triton 教程 。
未来的改进领域
最终用户仍需小心处理 MXFP4 等子字节数据类型格式的布局和打包。我们期待与社区合作，为内核作者改进人体工程学，并实现无缝框架集成。
在 GEMM_K 很小的情况下，所有数据类型中先前引用的矩阵乘法内核的利用率仍然相对较低。这可以通过内核本身中的手动子平铺来缓解，并已在 GEMM 教程中作为示例实施。最终将通过自动 warp-specialization 在编译器中以透明方式解决此问题。
更多信息
在 3 月 17 日举行的 NVIDIA GTC 大会上，Triton 的创建者 Phillippe Tillet 和 NVIDIA 将深入探讨 NVIDIA Blackwell 的这项工作的细节以及由此产生的性能。
此版本为 NVIDIA Blackwell 在 Triton 中的支持奠定了坚实基础，但这还只是开始。以下是您如何帮助塑造未来：
- 在 /triton-lang/triton GitHub 存储库中探索完整的实现。
- 加入充满活力的 GPU MODE Discord 社区。
- 深入了解我们不断增加的 示例和教程 。
立即开始在 NVIDIA Blackwell 上使用 Triton 进行构建，充分发挥 NVIDIA 最新架构的潜力，同时完全掌控您的开发。
是否有想法或遇到问题？请在 Github 上为我们的 NVIDIA 产品经理 Matthew Nicely 添加标签，联系他。
