在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上的性能，為長上下文推理提供了顯著加速。