Summary

MiniMax Sparse Attention (MSA) provides highly optimized dense FlashAttention and block-sparse attention kernels specifically for NVIDIA SM100 GPUs. It targets researchers and engineers building large language models, offering significant performance gains through on-demand JIT compilation and support for advanced quantization techniques.

How It Works

MSA employs two distinct Just-In-Time (JIT) compiled stacks within a unified Python package. The csrc JIT stack delivers dense FMHA and sparse top-k selection, compiled from Jinja templates at runtime. Complementing this, the CuTe-DSL stack enables full sparse attention, including paged FP8 decode, compiled via cute.compile. This dual approach allows for runtime optimization tailored to specific hardware and model requirements, leveraging techniques like FlashAttention and advanced quantization.

Quick Start & Requirements

• Install: Clone the repository recursively (git clone --recursive https://github.com/MiniMax-AI/MSA.git), then pip install . or pip install -e ..
• Prerequisites: NVIDIA SM100 GPU (compute capability 10.0), CUDA Toolkit (>= 12.x) with nvcc on PATH, Python (>= 3.10), Linux x86_64 (aarch64 untested).
• Setup: Initial JIT compilation of kernels can take 30 seconds to a few minutes.
• Resources: Algorithmic reference available at docs/MiniMaxSparseAttention.pdf. Kernels library usage example on Hugging Face Hub.

Highlighted Details

• Optimized kernels for dense FlashAttention (FMHA) and sparse top-k attention.
• Runtime JIT compilation for both C++ CUDA and CuTe-DSL based sparse attention.
• Extensive support for quantization: BF16, FP8, NVFP4, and FP4, particularly for sparse attention and decoding.
• Features paged attention and FP8 decode capabilities.

Maintenance & Community

The provided README does not detail specific contributors, sponsorships, or community channels like Discord/Slack. Contribution guidelines are outlined, emphasizing pre-alignment on public API changes.

Licensing & Compatibility

The project is licensed under MIT for its own code. However, it incorporates and derives from third-party components (e.g., NVIDIA CUTLASS, FlashInfer) which retain their original licenses (e.g., BSD-3-Clause, Apache-2.0). Users must comply with all applicable licenses. The MIT license generally permits commercial use.

Limitations & Caveats

Initial JIT compilation can be time-consuming. Support for aarch64 operating systems is untested and may require modifications. Internal implementation details outside the documented stable API entries may change without notice.