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。
這裡需要注意的另一個重要方麵是，可以通過現有的Triton flash注意實現“免費”提供這種性能增益，無需更改代碼。
圖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
