Command Palette
Search for a command to run...
MiniMax spärliche Aufmerksamkeit
MiniMax spärliche Aufmerksamkeit
Zusammenfassung
Die Fähigkeit zur Verarbeitung extrem langer Kontexte wird für fortschrittliche LLMs unverzichtbar: Agentic Workflows, Code-Reasoning auf Repository-Skala und Persistent Memory erfordern alle, dass das Modell gemeinsam auf Hunderttausende bis Millionen von tokens aufmerksamt, doch die quadratischen Kosten der Softmax-Aufmerksamkeit machen dies im Maßstab des produktiven Einsatzes untragbar. Wir stellen MiniMax Sparse Attention (MSA) vor, eine blockweise sparse Attention, die auf Grouped Query Attention (GQA) aufbaut. Ein leichtgewichtiger Index Branch bewertet Key-Value-Blöcke und wählt unabhängig eine Top-k-Teilmenge für jede GQA-Gruppe aus, wodurch ein gruppenspezifisches sparse Retrieval ermöglicht wird, während eine effiziente blockweise Ausführung aufrechterhalten bleibt; der Main Branch führt anschließend eine exakte block-sparse Attention ausschließlich über die ausgewählten Blöcke durch. MSA ist konsequent auf die Prinzipien der Einfachheit und Skalierbarkeit ausgelegt und bewusst schlank gehalten, was eine effiziente Implementierung auf einer breiten Palette von GPUs erheblich erleichtert. Um die Sparsity in praktische Beschleunigungen zu übersetzen, entwickeln wir MSA gemeinsam mit einem GPU-Ausführungspfad, der eine exp-freie Top-k-Auswahl und KV-outer sparse Attention nutzt, um die Tensor-Core-Auslastung bei blockgranularem Zugriff zu steigern. Auf einem Modell mit 109 Milliarden Parametern und nativem multimodalem Training erreicht MSA eine Leistung, die der von GQA entspricht, und reduziert dabei den pro token erforderlichen Attention-Compute um das 28,4-Fache bei einem Kontext von 1 Million tokens. In Kombination mit unserem gemeinsam entwickelten Kernel erzielt MSA auf H800-GPUs eine 14,2-fache Beschleunigung beim Prefill und eine 7,6-fache Beschleunigung beim Decoding in der Wanduhrzeit. Unser Inference-Kernel steht unter https://github.com/MiniMax-AI/MSA zur Verfügung. Ein produktionsreifes, nativ multimodales Modell, das von MSA angetrieben wird, wurde öffentlich unter https://huggingface.co/MiniMaxAI/MiniMax-M3 veröffentlicht.
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.