MoonMath AI 開源 AMD MI300X 的 HIP 注意力內核，在各項指標和舍入模式下均超越 AITER v3

MoonMath AI 團隊近日發佈了針對 AMD MI300X GPU 的 bf16 前向注意力內核。該內核使用 HIP 編寫（而非手寫彙編），以 MIT 許可證開源發佈。MoonMath.ai 團隊聲稱，該內核在測試的所有形狀上均超越了 AMD 自己的優化內核 AITER v3。裸機訪問由 AMD 雲服務提供商 HotAisle 提供。

注意力機制是每個 Transformer 中的融合 softmax(QKᵀ/√d)·V 操作。MI300X 是 AMD 的 CDNA3 數據中心 GPU，目標 ISA 為 gfx942，該內核僅在此硬件上運行。

內核是一款直接在 GPU 眾多核心上運行的小程序，專門用於儘可能快地執行注意力運算。該內核僅在 MI300X 上計算 bf16 前向注意力，支持 BSHD 或 BHSD 佈局，無需轉置，頭維度固定為 128，支持任意序列長度包括交叉注意力。

目前存在一些限制：不支持因果掩碼、GQA 或變長批處理。輸出為 bf16，僅在 gfx942 硬件上運行。數值控制嚴格：所有三種舍入模式均與 AITER 的逐模式舍入規則匹配，每個有限輸出與 AITER 的偏差在 1 bf16 ULP 以內，NaN 和 Inf 處理完全一致，結果具有確定性。

核心技巧：單指令彙編包裝器 內核的核心理念規避了常見的困境：編譯器內聯函數保持代碼整潔但可能允許編譯器重排序或重命名操作數；而原始內聯彙編提供控制但需手動管理寄存器和地址。MoonMath 團隊將單條指令包裝在 device forceinline 函數中，使用擴展 asm 約束描述操作數。團隊選擇操作碼，編譯器仍分配寄存器並跟蹤數據流。

例如，asm mfma 函數中，"+v"(c) 約束將累加器輸入和輸出綁定到同一 VGPR，無需發出複製指令。這使得內核接近普通的 HIP 代碼，但仍能逐條指令控制機器。

架構：八波流水線，兩組，兩個屏障 一個 CDNA3 計算單元有四個 SIMD 單元。傳統方法是四個波前，但 MoonMath 在每個塊中使用八個波前，分為兩組，每組四個。兩組執行相同的 Q*K、softmax、O += P*V 序列，但相位偏移。當一組飽和矩陣核心時，另一組執行 softmax 併發出加載指令，然後交換，確保矩陣核心永不空閒。每次迭代有兩個 s_barrier：一個在相位交接處，一個在迭代邊界。其餘同步由逐計數器等待處理。這類似於 FlashAttention-3 的矩陣乘法和 softmax 交替模式，但未採用其生產者和消費者波前拆分。在 CDNA3 上，所有內存移動已是異步的，因此無需專用的生產者波前。

數據存放位置及 16×16×16 的選擇 大多數加速來自內存佈局：K 從 HBM 流入 LDS，雙緩衝，由所有八個波前共享；V 保持在 L1 中，每次 PV 矩陣乘法時讀取；Q 和累加器位於寄存器中。團隊選擇了 16×16×16 MFMA 而非 32×32×8，兩者吞吐量相同，但更小的瓦片每個通道累加 4 個 fp32 元素（而非 16 個），降低了累加器壓力，為更深的預取和第三個 Q 瓦片留出空間。

後續兩個優化進一步縮小差距：第三個 Q 瓦片（3Q）提高了每個加載的 K 和 V 瓦片的數據複用率；類似 Flash-Decoding 的尾部 KV 拆分解決了 MI300X 的 304 個 CU 上剩餘分數的循環問題。這些優化層層遞進：將 V 移到 L1 釋放了 LDS，然後第三個 Q 瓦片填充了它。

基準測試 測試在 MI300X 上以 bf16 格式執行，頭維度 128。每個形狀在三種舍入模式下測量：RTNE（舍入到最近偶數）、RTNA（舍入到最近，遠離零）、RTZ（向零截斷）。結果顯示 MoonMath 的內核全面領先：幾何平均相對 AITER 為 1.18×（RTNE）、1.15×（RTNA）、1.08×（RTZ）；相對 Modular MAX 的幾何平均加速比為 1.44× 到 1.49×，單形狀最高達 1.59×。

實際用例 該內核通過 pip 安裝，提供簡潔的 API，在調用者流上啓動，從而與大流水線重疊。具體用例包括視頻擴散：團隊為 SGLang 擴散添加了 LiteAttention 支持並提交了 PR，在 Wan2.1-T2V-1.3B-Diffusers 上，將注意力從 AITER 切換到 liteattention_rocm，端到端生成性能提升了 1.23×，且無可見質量退化。BSHD 佈局直接適用於擴散張量，交叉注意力支持任意 KV 長度且無需填充。

關鍵要點

• 內核為 MI300X 的 bf16 前向注意力，HIP 編寫，MIT 許可。
• 在每種形狀和舍入模式下優於 AITER v3，幾何平均 1.18×/1.15×/1.08×。
• 單指令彙編包裝器在保留操作碼控制的同時讓編譯器分配寄存器。
• 內存佈局是主要加速來源：K 在 LDS，V 在 L1，Q 在寄存器。
• 實際 SGLang PR 加速 Wan2.1 視頻擴散 1.23× 且無質量損失。

更多技術細節請查看原文。歡迎關注我們的 Twitter、加入 150k+ ML SubReddit 或訂閲通訊。