Command Palette
Search for a command to run...
MiniMax スパースアテンション
MiniMax スパースアテンション
概要
最先端の大規模言語モデル（LLM）において、超長文脈処理能力は不可欠となりつつある。agentワークフロー、リポジトリ規模のコード推論、永続的メモリはいずれも、モデルが数十万から数百万のtokensを同時に注意する必要があるが、softmax注意の二次的な計算コストは、本番環境での展開規模においてこれを現実的に不可能にする。本稿では、Grouped Query Attention（GQA）を基盤としたブロック単位のスプース注意機構であるMiniMax Sparse Attention（MSA）を提案する。軽量なIndex Branchはkey-valueブロックのスコアリングを行い、各GQAグループに対して独立してTop-kサブセットを選択する。これにより、効率的なブロックレベルの実行を維持しつつ、グループ固有のスプース検索を可能にする。その後、Main Branchは選択されたブロックのみに対して正確なブロックスプース注意演算を実行する。単純性とスケーラビリティの原則を基盤として設計されたMSAは意図的に簡素化されており、幅広いGPUアーキテクチャに対して効率的な展開を容易にする。スプース性を実際の演算速度向上に結びつけるため、exp-freeなTop-k選択およびKV-outerスプース注意を用い、ブロック粒度のアクセス下におけるtensor-core利用率を向上させるGPU実行パスとMSAを協調設計した。ネイティブなマルチモーダル学習を用いた109Bパラメータモデルにおいて、MSAはGQAと同等の性能を発揮しつつ、1M文脈長においてtokensあたりの注意計算量を28.4倍削減する。本稿で協調設計したカーネルと組み合わせることで、MSAはH800においてprefill処理で14.2倍、デコード処理のウォールクロック時間において7.6倍の速度向上を実現する。本推論カーネルは以下のURLで公開されている：https://github.com/MiniMax-AI/MSA。MSAを搭載した本番環境対応のネイティブマルチモーダルモデルが、以下のURLで一般公開されている：https://huggingface.co/MiniMaxAI/MiniMax-M3。
One-sentence Summary
MiniMax Sparse Attention (MSA) is a Grouped Query Attention variant that employs a lightweight Index Branch to independently select Top-k key-value blocks per group, enabling exact block-sparse computation that matches standard performance while reducing per-token compute by 28.4x at one million tokens and delivering 14.2x prefill and 7.6x decoding speedups on H800 GPUs through a co-designed kernel leveraging exp-free Top-k selection and KV-outer sparse attention.
Key Contributions
- MiniMax Sparse Attention (MSA) is a blockwise sparse attention mechanism built on Grouped Query Attention that employs a lightweight index branch to independently score and select a top-k subset of key-value blocks for each GQA group prior to exact block-sparse attention.
- A co-designed GPU execution path implements exp-free top-k selection and KV-outer sparse attention to maximize tensor-core utilization under block-granular memory access patterns.
- Evaluations on a 109B-parameter multimodal model demonstrate that the approach maintains performance parity with standard Grouped Query Attention while reducing per-token attention compute by 28.4x at a one-million-token context, yielding 14.2x prefill and 7.6x decoding wall-clock speedups on H800 hardware.
Introduction
Long-context modeling in transformer-based language models requires efficient attention mechanisms to mitigate the quadratic computational and memory overhead of dense softmax attention. Prior approaches typically substitute attention with linear or recurrent alternatives, apply fixed content-agnostic sparse patterns, or implement adaptive sparsification that either inherits full-attention training costs or suffers from fragmented memory access and unoptimized inference kernels. The authors leverage a per-GQA-group Top-k sharing strategy combined with block-level selection to preserve contiguous KV cache reads while maintaining adaptive context awareness. They further accelerate the framework by adapting the FlashAttention algorithmic skeleton with loop ordering specifically tuned to this access pattern, effectively translating theoretical FLOP reductions into measurable wall-clock speedups.
Experiment
Two 109B-scale experiments validate replacing dense attention with a sparse mechanism by either training from scratch or continuing pretraining from a full-attention checkpoint. The native sparse approach demonstrates that the model can stably adapt its representations to learn essential attention structures without hard-coded constraints, while the continued pretraining route validates a practical and stable conversion pathway from dense checkpoints. Both methods maintain competitive performance across language, multimodal, and long-context benchmarks despite a strict key-value token budget. Ultimately, the sparse architecture delivers substantial computational efficiency and sustained long-context capabilities, establishing it as a scalable alternative to dense attention.