MiniMax 發表 MSA 稀疏注意力機制：突破長文本運算瓶頸，讓 AI 處理超長序列更高效

解決長文本運算的二次方成本瓶頸

大型語言模型在處理長文本時，傳統的 softmax 注意力機制會帶來巨大的運算負擔。MiniMax 研究團隊近期發表了 MiniMax Sparse Attention (MSA)，這是一種基於分組查詢注意力（Grouped Query Attention, GQA）架構的稀疏注意力方法，旨在解決長上下文運算中的效率瓶頸。

MSA 的核心邏輯在於將注意力運算拆分為兩個階段：索引分支（Index Branch）與主分支（Main Branch）。索引分支負責決定每個查詢（Query）應讀取哪些鍵值（Key-Value）區塊，而主分支則僅針對這些被選中的區塊執行精確的 softmax 注意力運算。

區塊化篩選與運算機制

與傳統逐 token 篩選不同，MSA 採用「區塊粒度」進行篩選，預設區塊大小為 128 個 token。每個查詢與 GQA 分組被限制保留 16 個區塊，這意味著每個查詢的運算預算被固定在 2,048 個 token（16 x 128）。隨著上下文長度（N）增加，傳統 GQA 的運算成本會隨之增長，但 MSA 的成本卻能維持固定，從而大幅提升長文本處理的擴展性。

為了訓練這個非微分（non-differentiable）的 Top-k 選擇過程，MiniMax 採用了 KL 散度對齊損失（KL alignment loss），將索引分支的分佈與主分支的注意力模式進行匹配，並透過梯度分離（Gradient Detach）與索引器預熱（Indexer Warmup）等機制確保訓練穩定。

專為硬體優化的推理核心

理論上的稀疏性若缺乏硬體配合，難以轉化為實際效能。MiniMax 為此設計了兩項關鍵的 Kernel 優化：

• 無指數（exp-free）Top-k 選擇：跳過了 softmax 中的指數與求和步驟，直接對原始分數進行排序，在 128K 上下文長度下，比 torch.topk 快 5.1 倍。
• KV 外積稀疏注意力：透過將查詢位置打包至分數矩陣中，提升算術強度。

目前 MiniMax 已開源相關的 fmha_sm100 推理核心，支援 NVIDIA SM100 GPU，並相容 BF16、FP8、NVFP4 及 FP4 等多種精度。

應用場景與效能表現

MSA 特別適用於上下文長度受限的部署環境，例如：

• 長程代理（Long-horizon agents）：處理數百個推理與行動步驟的長紀錄。
• 儲存庫規模的程式碼推理：僅路由至相關的程式碼區塊，忽略無關檔案。
• 長影片理解：利用其多模態訓練背景，處理長序列視覺 token。

在 3T token 的訓練預算下，MSA 在 MMLU、GSM8K 及 VideoMME 等基準測試中，表現與全注意力（Full-Attention）模型相當，證明了在大幅降低運算成本的同時，模型仍能保持高效能。