矩陣乘法和注意力機制是現代 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 添加標簽，聯系他。
?
