NVIDIA / NVIDIA/Megatron-LM

Feature Request: MiniMax Sparse Attention (MSA)

Open
#5,322 0 comments 3 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Summary
Explore MiniMax Sparse Attention (MSA) as a long-context attention path. MSA is designed around GQA: a selector scores KV blocks per GQA group, always includes the local block, and uses a fixed sparse budget.

The public `MiniMax-AI/MSA` repo provides an `fmha_sm100` package with dense SM100 FMHA, `sparse_topk_select`, CuTe-DSL sparse attention, CSR metadata builders, BF16/FP8/NVFP4/FP4 support, paged FP8 decode, and an FMHA-style sparse adapter.

## Motivation
Full attention remains a bottleneck for ultra long-context training and inference.

Reported MSA paper results:
- 28.4× per-token attention FLOP reduction at 1M context
- 14.2× prefill and 7.6× decode wall-clock speedups at 1M context on H800
- 109B MoE experiments where MSA stays close to a GQA full-attention baseline

The open-source implementation is not a drop-in Megatron training backend today: its documented public contract is SM100-focused and mostly forward/inference oriented.

## Requested Feature
Investigate adding an experimental MSA-style backend or attention variant in `megatron.core.transformer`, initially behind a feature flag and compatible with existing GQA (`num_query_groups`) models.

## References
- [MiniMax Sparse Attention paper](https://github.com/MiniMax-AI/MSA/blob/main/docs/MiniMaxSparseAttention.pdf)
- [GitHub: MiniMax-AI/MSA](https://github.com/MiniMax-AI/MSA)
- [MSA CuTe-DSL sparse attention README](https://github.com/MiniMax-AI/MSA/blob/main/python/fmha_sm100/cute/README.md)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.