OpenAI Triton on NVIDIA Blackwell提高AI性能和可编程性
矩阵乘法和注意力机制是现代人工智能工作负载的计算支柱。虽然像NVIDIA cuDNN这样的库提供了高度优化的实现，而像CUTLASS这样的框架提供了深度定制，但许多开发人员和研究人员需要一个将性能与可编程性相结合的中间地带。
NVIDIA Blackwell架构上的开源Triton编译器通过直观的编程模型暴露Blackwell的高级功能，解决了这一需求。
由于NVIDIA与OpenAI的持续合作，Triton编译器现在支持NVIDIA Blackwell架构。这确保了开发人员和研究人员可以在Triton等基于python的编译器中轻松地使用Blackwell架构的最新和最强大的功能。
NVIDIA Blackwell的性能进步
NVIDIA Blackwell架构在原始计算能力和架构创新方面都有了实质性的改进。NVIDIA与OpenAI的合作重点是通过Triton的编译器基础设施透明地利用这些功能，特别是在两个关键领域：
- 矩阵乘法包括Flash注意
- 新的精度格式
矩阵乘法
NVIDIA Blackwell架构增加了全新的张量核心，旨在提高吞吐量和能源效率。
通过扩展Triton的Matrix Multiply-Accumulate （MMA）流水线机制，我们已经能够自动开发NVIDIA Blackwell的新张量核心。这需要仔细分析内存访问模式和复杂的编译器转换，以确保正确和有效的计算/数据移动重叠。
结果是FP8和FP16 GEMM操作的卓越性能开箱即用，这些优化自动应用于使用Triton ‘ tl的任何内核。点的原始。总的来说，Triton设法实现了接近最优的性能，可以与跨几个关键用例的库实现相媲美。
图1所示。性能改进与Triton在NVIDIA Blackwell
图1显示了Triton在NVIDIA Blackwell架构上的优化为FP16和FP8的用户带来了硬件性能的提高，在这个针对典型的GEMM内核生成AI大小的K扫描分析中，如Triton教程所提供的。
Flash的关注
Flash注意力是现代变压器架构中至关重要的基本要素，通过Triton， NVIDIA Blackwell的速度显著提高，FP16的注意力比NVIDIA Hopper GPU架构提高了1.5倍。
当我们继续通过FP8和其他精度上的编译器增强来优化绝对性能时，目前的工作可以帮助客户在现有产品的第0天轻松过渡到NVIDIA Blackwell。
图2。为更复杂的工作负载提供更大的性能提升
图2显示了更复杂的工作负载，例如Triton教程中提供的flash注意力内核，在使用Triton编译器改进解锁后，再次显示了NVIDIA Blackwell架构的巨大性能提升。这项工作的一些改进也提高了NVIDIA Hopper的注意力表现，但没有在本数据中显示出来。
新的精度格式
NVIDIA Blackwell引入了革命性的块缩放浮点格式，包括开放计算项目的微缩放格式，Triton现在为NVIDIA Blackwell驱动的硬件加速解锁。
这些格式提供了更高的平均精度和更高的性能，而不是今天在LLM推理项目中经常模拟的非本地块缩放技术。
对于OCP格式的支持，Triton上的MXFP8 GEMMs表现出了类似于FP8 GEMMs性能加速的优异性能，同时允许在张量核心中进行原生扩展。
同样，MXFP4在精度性能权衡领域提供了一个新的操作点，同时提供了两倍于FP8和MXFP8 gem的硬件加速性能。
要了解有关新的块缩放浮点支持的更多信息，请查看专门用于此功能的新的Triton教程。
有待改进的领域
像MXFP4这样的子字节数据类型格式的布局和打包仍然需要最终用户的注意。我们期待与社区合作，为内核作者和无缝框架集成改进人机工程学。
当GEMM_K很小时，前面提到的跨所有数据类型的矩阵乘法内核的利用率仍然相对较低。这可以通过在内核本身中手动进行子平铺来缓解，GEMM教程中已经实现了这个示例。它最终将在编译器中通过自动翘曲专门化透明地解决。
更多的信息
Triton的创造者philippe Tillet和NVIDIA将在3月17日的NVIDIA GTC会议上深入研究NVIDIA Blackwell工作的细节和由此产生的性能。
您可以在这里注册：
https://www.nvidia.com/gtc/pricing/?nvid=nv-int-bnr-346123
