MiniMax-AI/MSA
GitHub: MiniMax-AI/MSA
MSA 是 MiniMax 开源的面向 NVIDIA SM100 GPU 的 FlashAttention 与 block-sparse attention kernel 库，旨在加速大模型长序列推理中的注意力计算。
Stars: 385 | Forks: 46
# MiniMax Sparse Attention (MSA)
[](LICENSE)
[](pyproject.toml)
[](#requirements)
[](#stacks)
**MSA** (`fmha_sm100`) 为 **NVIDIA SM100** 提供了 dense FlashAttention 和 sparse top-k attention
kernel。两个 JIT 编译的技术栈
共享同一个 Python 包：
| 技术栈 | 路径 | 提供的功能 |
|---|---|
| **csrc JIT** | `python/fmha_sm100/csrc/` | Dense FMHA (`fmha_sm100`, `fmha_sm100_plan`) + `sparse_topk_select` indexer，在运行时由 `jit.py` 从 Jinja 模板编译而来。 |
| **CuTe-DSL** | `python/fmha_sm100/cute/` | 完整的 sparse attention（前向传播 + paged FP8 decode，BF16 / FP8 / NVFP4 / FP4），在运行时通过 `cute.compile` 编译。 |
| **Bridge** | `python/fmha_sm100/sparse_fmha_adapter.py` | 适配 `fmha_sm100` API，以便在 sparse prefill 路径中调用 `sparse_atten_func`。 |
## 环境要求
- **GPU**: NVIDIA SM100。
- **工具链**：PATH 中包含带有 `nvcc` 的 CUDA Toolkit（或已设置 `CUDA_HOME` / `CUDA_PATH`）。
- **Python**: ≥ 3.10。
- **操作系统**：Linux x86_64（aarch64 未测试；在 WSL 上进行 JIT 构建可能需要微调 Makefile）。
安装前的快速检查：
```
nvcc --version # expect ≥ 12.x
nvidia-smi --query-gpu=compute_cap --format=csv | grep "10.0" # confirm SM100
python -c "import sys; print(sys.version_info[:2])" # ≥ (3, 10)
```
## 配合 `kernels` 库使用
要快速开始使用 MSA kernel，你可以使用 [`kernels` 库](https://github.com/huggingface/kernels)：
```
# 请确保已安装 `kernels`：`pip install -U kernels`
from kernels import get_kernel
kernel_module = get_kernel("MiniMaxAI/msa", version=0)
sparse_atten_func = kernel_module.sparse_atten_func
sparse_atten_func(...)
```
在 Hugging Face Hub 上[查看该 kernel](https://huggingface.co/kernels/kernels-staging/msa)。
## 安装
```
# --recursive 会拉取 NVIDIA CUTLASS submodule (python/fmha_sm100/cutlass/)，
# 其 headers 是进行 JIT/AOT compilation 所必需的。
git clone --recursive https://github.com/MiniMax-AI/MSA.git msa
cd msa
# 如果您在 clone 时未使用 --recursive：
# git submodule update --init --recursive
pip install . # standard install (works from a wheel too)
# 或者
pip install -e . # editable install for development
```
这会通过 `nvidia-cutlass-dsl` 和 `quack-kernels` 引入 CuTe-DSL 技术栈；
csrc kernel 会在首次导入时，使用包内附带的源码进行 JIT 编译。
## 验证
运行一个小型的 CUDA 冒烟测试。**首次运行会 JIT 编译 `sparse_topk_select`，
在清空的 nvcc 缓存下需要 30 秒到几分钟的时间** —— 这是正常现象，并非
卡死。随后的运行会命中 JIT 缓存并在几秒钟内完成。
```
python tests/smoke/test_sparse_topk_forced.py
```
## 用法
```
import torch
from fmha_sm100 import fmha_sm100, fmha_sm100_plan, sparse_topk_select
# sparse prefill path 的 page size 和 top-k。
page_size, topk = 128, 16
# Dense proxy pass：从低成本的 Q slice 计算 per-block max score。
proxy_plan = fmha_sm100_plan(
qo_lens, kv_lens, proxy_q.shape[1],
num_kv_heads=1,
page_size=page_size,
output_maxscore=True,
)
_, max_score = fmha_sm100(
proxy_q, proxy_k_pages, proxy_v_pages, proxy_plan,
kv_indices=kv_indices,
output_o=False,
output_maxscore=True,
)
# max_score -> sparse KV block indexes。
kv_block_indexes = sparse_topk_select(
max_score.contiguous(), topk, num_valid_pages=num_pages,
)
# 对所选 blocks 执行 Sparse attention。
sparse_plan = fmha_sm100_plan(
qo_lens, kv_lens, q.shape[1],
num_kv_heads=k_pages.shape[1],
page_size=page_size,
kv_block_num=topk,
)
out, _ = fmha_sm100(
q, k_pages, v_pages, sparse_plan,
kv_indices=kv_indices,
kv_block_indexes=kv_block_indexes,
)
```
有关带有 CSR 元数据、FP4 indexer、NVFP4 K/V 和 paged FP8 decode wrapper 的 block-sparse prefill，请参阅 **CuTe-DSL 深入解析**：
- [`python/fmha_sm100/cute/README.md`](python/fmha_sm100/cute/README.md)
## 测试
```
# 快速 smoke tests。
python -m pytest tests/smoke -q
# API 和 end-to-end integration tests。
python -m pytest tests/integration -q
python tests/integration/test_proxy_kv_e2e.py
# 大型 regression suites。
python tests/regression/test_correctness.py
python tests/regression/test_sparse_attn.py
# CuTe-DSL forward-only sparse attention。
cd python/fmha_sm100/cute
python -m pytest test_sparse_atten.py -q
```
## 基准测试
`benchmarks/bench_sparse_attention_ops.py` 涵盖了 dense prefill、paged
prefill、sparse prefill、dense decode、paged decode、sparse decode，支持
`fp8` 和 `bf16` 格式（`nvfp4` 仅用于 sparse-prefill）。
```
python benchmarks/bench_sparse_attention_ops.py --help # full flag list
```
常见调用方式（输出为 TSV 格式）：
| 目标 | 命令 |
|---|---|
| FP8 全面扫描 | `python benchmarks/bench_sparse_attention_ops.py --dtype fp8 --sections all --output_mode o -o /tmp/msa_fp8.tsv` |
| BF16 全面扫描 | `python benchmarks/bench_sparse_attention_ops.py --dtype bf16 --sections all --output_mode o -o /tmp/msa_bf16.tsv` |
| NVFP4 sparse prefill | `python benchmarks/bench_sparse_attention_ops.py --dtype nvfp4 --sections sparse_prefill --output_mode o -o /tmp/msa_nvfp4.tsv` |
| 快速 CI 冒烟测试 | `python benchmarks/bench_sparse_attention_ops.py --dtype fp8 --sections prefill,decode,sparse_decode --seqs 8192,16384 --tp 1,4 --decode-k 8192,131072 --decode-b 32 --dry-run-ms 50 --repeat-ms 200 -o /tmp/msa_smoke.tsv` |
| 输出模式检查 (dense/paged) | `--output_mode maxscore` 或 `--output_mode full` |
## 布局
```
python/fmha_sm100/ Python package
__init__.py Public re-exports (lazy for the CuTe-DSL stack)
api.py fmha_sm100 / fmha_sm100_plan / sparse_topk_select
jit.py Runtime JIT (nvcc + ninja) for the csrc stack
sparse.py Lazy shim that loads the cute/ stack
sparse_fmha_adapter.py Bridge: fmha_sm100 API → sparse_atten_func
csrc/ CUDA kernels + Jinja templates (JIT-compiled)
include/ Vendored FlashInfer / CUTLASS-derived / TRT-LLM headers
cutlass/ NVIDIA CUTLASS git submodule (include/ + tools/util/include/)
cute/ CuTe-DSL sparse attention (loaded via sys.path)
tests/ Correctness tests
smoke/ integration/ regression/
scripts/ Warmup + cache-management helpers
benchmarks/ bench_sparse_attention_ops.py
```
## 技术栈
- **csrc JIT** — dense FlashAttention、page KV 和 `sparse_topk_select`
indexer。在运行时从 `csrc/*.cu.jinja` 以及
`csrc/include/` 编译而来。公开入口：`fmha_sm100.plan → run`。
- **CuTe-DSL** — block-sparse prefill、FP8 / NVFP4 / FP4 量化、paged
FP8 decode (`SparseDecodePagedAttentionWrapper`)、FP4 block-score indexer。
公开入口：`fmha_sm100.sparse_atten_func`、
`fmha_sm100.sparse_decode_atten_func`、`fmha_sm100.fp4_indexer_block_scores`。
- **Bridge** — `sparse_fmha_plan` / `sparse_fmha` 将 dense-API 调用
点适配到 prefill 路径的 sparse backend；当你已经
在使用 dense kernel 并希望用一行代码切换到 sparse 时非常有用。
## 第三方许可证
`fmha_sm100` 捆绑、派生自或依赖于以下第三方组件。
每个组件均保留其原始许可证；本节仅对其进行概述。
具有权威性的文本随各组件一起分发。
### 捆绑 / 派生源码（内置于本仓库）
| 组件 | 许可证 | 位置 |
|---|---|---|
| **NVIDIA CUTLASS** | BSD-3-Clause | 位于 `python/fmha_sm100/cutlass/` 的 Git 子模块（提供 `include/` + `tools/util/include/`），以及 `python/fmha_sm100/csrc/include/` 下标记为 BSD-3 的头文件。`python/fmha_sm100/cute/src/common/mma_sm100_desc.py` 中的 SM100 MMA 描述符编码映射了 CUTLASS 硬件描述符。Copyright (c) 2017–2025 NVIDIA CORPORATION & AFFILIATES。 |
| **FlashInfer** | Apache-2.0 | 位于 `python/fmha_sm100/csrc/` 和 `python/fmha_sm100/csrc/include/` 下带有 `Copyright (c) by FlashInfer team` 声明的头文件和源码（例如 `allocator.h`、`exception.h`、`utils.cuh`、`cutlass_utils.cuh`、`fmha_cutlass_sm100.cuh`、`sparse_topk_select.cuh`、`plan.cuh`、`sm100_fmha_reduction.hpp`、`tvm_ffi_utils.h`）。项目地址：。 |
| **NVIDIA TensorRT-LLM + NAVER Corp (CLOVA)** | Apache-2.0 | `python/fmha_sm100/csrc/include/sparse_topk_select.cuh` 的部分内容 — `indexerTopK` 直方图步骤 + 插入排序派生自 `tensorrt_llm/cpp/tensorrt_llm/kernels/indexerTopK.cu`。Copyright (c) 2019–2026 NVIDIA CORPORATION；Copyright (c) 2021 NAVER Corp。`sparse_topk_select.cuh` 中的每个文件头都包含一个函数级别的来源映射表。 |
### 运行时依赖（通过 pip 安装）
| 包 | 上游来源 | 许可证 |
|---|---|---|
| `quack-kernels` | | Apache-2.0 |
| `nvidia-cutlass-dsl` | NVIDIA CUTLASS Python DSL | NVIDIA / BSD-3-Clause（详见包说明） |
| `apache-tvm-ffi` | Apache TVM FFI | Apache-2.0 |
| `cuda-python` | NVIDIA | NVIDIA / 详见包说明 |
| `torch` | | BSD-3-Clause |
| `jinja2` | | BSD-3-Clause |
| `ninja` | | Apache-2.0 |
| `pybind11` | | BSD-3-Clause |
每个已安装包的确切许可证均随该包一起分发；
请查阅其元数据（`pip show `）以获取具有权威性的文本。
## 引用
如果 MSA 对你的研究有所帮助，请引用它。（待配套论文 / 技术报告获得稳定标识符后，将提供 BibTeX 条目 — 目前为占位符。）
算法参考文档随附于
[`docs/MiniMaxSparseAttention.pdf`](docs/MiniMaxSparseAttention.pdf)。
```
@software{msa2026,
title = {MiniMax Sparse Attention (MSA): FlashAttention and block-sparse
attention kernels for NVIDIA SM100},
author = {{MiniMax}},
year = {2026},
url = {https://github.com/MiniMax-AI/MSA}
}
```
## 贡献
欢迎在[问题追踪器](https://github.com/MiniMax-AI/MSA/issues)上提交 Issue 和 PR。对于 kernel 或
运行时契约的更改，请先提交一个 Issue 以便就公共
接口达成一致 — `fmha_sm100.api`、`fmha_sm100.sparse` 和
`cute.interface` 是稳定的入口点；其他所有内容
均为内部实现，可能会在不另行通知的情况下发生更改。
标签：CUDA, Vectored Exception Handling, 人工智能, 凭据扫描, 大模型, 注意力机制, 用户模式Hook绕过, 稀疏注意力, 算子库, 逆向工具