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 或訂閱通訊。