在NVIDIA Blackwell上最佳化MiniMax M3稀疏注意力

Fireworks AI團隊在NVIDIA Blackwell（SM100）GPU上為MiniMax M3稀疏注意力開發了一款高效核心。長上下文推理中，注意力機制主導計算和記憶體開銷。MiniMax M3採用的稀疏注意力讓每個查詢只關注最相關的16個128-token KV塊，從而降低開銷。然而，稀疏選擇依賴於資料且隨請求變化，導致不規則訪存模式，常常抵消理論加速。

Fireworks AI的效能團隊構建的Blackwell核心透過KV-固定（KV-stationary）執行路徑彌補了這一差距。該思路源自MiniMax團隊的稀疏注意力論文：在外迴圈中載入每個選中的KV塊一次，然後在內迴圈中處理所有選擇該塊的查詢。基於這一排程，最佳化聚焦於減少聚集載入/分散儲存的記憶體流量，以及改善SM間的負載均衡。

在單塊B200（fp8）上的測試結果令人矚目：注意力核心效率達到約980 TFLOP/s的吞吐量，HBM頻寬約4.1 TB/s。相比查詢固定基線（FlashInfer）實現1.9–2.4倍加速，相比MiniMax開源MSA核心提升約1.6倍。完整模組（含索引對映和組合階段）相比基線提升1.18–1.43倍，相比開源MSA提升1.32–1.41倍。

文章還深入分析了演算法和核心設計空間。M3注意力在分組查詢注意力（GQA）基礎上實現塊狀稀疏：將KV序列分為128-token塊，輕量級索引分支為每個GQA組評分並選擇TopK（=16）塊，主分支僅在選定塊上計算精確注意力。核心結構分為查詢固定（Q-outer）和KV固定（KV-outer）兩種。查詢固定模式下，每個查詢僅讀取其選定的KV塊，但MMA矩陣尺寸小（16×128），張量核心利用率低，且KV塊被多個查詢重複讀取。KV固定則反轉對映：對每個KV頭/塊，收集選擇它的查詢，載入KV塊一次，執行完整128×128矩陣乘法，每個查詢寫入最多16個部分輸出，後續由組合核心進行LSE合併。這樣計算高效，但儲存開銷增加。

I/O分析顯示，兩種結構受不同記憶體瓶頸限制：查詢固定受L2頻寬限制（重複讀取命中L2），KV固定受HBM頻寬限制（每個塊只讀一次，但每個查詢寫入多個部分）。基於B200的峰值頻寬（L2約24 TB/s，HBM約7.4 TB/s）的簡單羅徹斯特線模型表明，當nsb/N < 2.85時KV固定更快（其中N為查詢數，nsb為KV序列長度）。該模型保守假設查詢固定完美命中L2，實際中KV固定更早勝出。

總之，Fireworks AI的KV-固定核心成功最佳化了MiniMax M3稀疏注意力在Blackwell上的效能，為長上下文推理提供了顯著加速。