MiniMax-AI/MSA

来源：https://github.com/MiniMax-AI/MSA

MiniMax稀疏注意力（MSA）

License: MIT Python GPU 技术栈：CuTe-DSL + Cuda

MSA（fmha_sm100）为 NVIDIA SM100 提供了密集FlashAttention和稀疏top-k注意力核。两个JIT编译栈共享同一个Python包：MSA架构

算法参考：MiniMax稀疏注意力论文

许可证：MIT。 自创文件带有SPDX-License-Identifier: MIT。
参见LICENSE和NOTICE。捆绑/衍生自第三方代码保留其原始许可证 —— 见第三方许可证。

系统要求

• GPU：NVIDIA SM100。
• 工具链：CUDA工具包，需在PATH上存在nvcc（或设置CUDA_HOME / CUDA_PATH）。
• Python：≥ 3.10。
• 操作系统：Linux x86_64（aarch64未经测试；在WSL上JIT构建可能需要小的Makefile修改）。

安装前快速检查：

nvcc --version                    # 期望 ≥ 12.x
nvidia-smi --query-gpu=compute_cap --format=csv | grep "10.0"  # 确认SM100
python -c "import sys; print(sys.version_info[:2])"            # 期望 ≥ (3, 10)

与kernels库一起使用

要快速开始使用MSA核，可以使用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上的核：https://huggingface.co/kernels/kernels-staging/msa

安装

# --recursive 拉取 NVIDIA CUTLASS 子模块（python/fmha_sm100/cutlass/），
# 其头文件是JIT/AOT编译所必需的。
git clone --recursive https://github.com/MiniMax-AI/MSA.git msa
cd msa

# 如果你没有使用 --recursive 克隆：
# git submodule update --init --recursive

pip install .                # 标准安装（也适用于wheel）
# 或者
pip install -e .             # 可编辑模式安装，用于开发

这通过nvidia-cutlass-dsl和quack-kernels引入CuTe-DSL栈；csrc核在首次导入时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

# 稀疏预填充路径的页面大小和top-k
page_size, topk = 128, 16

# 密集代理过程：从廉价的Q切片计算每块最高分。
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 -> 稀疏KV块索引
kv_block_indexes = sparse_topk_select(
    max_score.contiguous(), topk, num_valid_pages=num_pages,
)

# 使用选定的块进行稀疏注意力
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索引器、NVFP4 K/V和分页FP8解码封装，请参阅CuTe-DSL深入指南：

• python/fmha_sm100/cute/README.md

测试

# 快速冒烟测试
python -m pytest tests/smoke -q

# API和端到端集成测试
python -m pytest tests/integration -q
python tests/integration/test_proxy_kv_e2e.py

# 大规模回归测试套件
python tests/regression/test_correctness.py
python tests/regression/test_sparse_attn.py

# CuTe-DSL仅前向稀疏注意力
cd python/fmha_sm100/cute
python -m pytest test_sparse_atten.py -q

基准测试

benchmarks/bench_sparse_attention_ops.py 涵盖密集预填充、分页预填充、稀疏预填充、密集解码、分页解码、稀疏解码，支持fp8和bf16（nvfp4仅用于稀疏预填充）。

python benchmarks/bench_sparse_attention_ops.py --help   # 完整标志列表

常见调用（输出为TSV格式）：

布局

python/fmha_sm100/            Python包
  __init__.py                 公共重导出（CuTe-DSL栈为惰性加载）
  api.py                      fmha_sm100 / fmha_sm100_plan / sparse_topk_select
  jit.py                      运行时JIT（nvcc + ninja），用于csrc栈
  sparse.py                   惰性垫片，加载cute/栈
  sparse_fmha_adapter.py      桥接层：fmha_sm100 API → sparse_atten_func
  csrc/                       CUDA核 + Jinja模板（JIT编译）
  include/                    从FlashInfer / CUTLASS衍生 / TRT-LLM头文件
  cutlass/                    NVIDIA CUTLASS git子模块（include/ + tools/util/include/）
  cute/                       CuTe-DSL稀疏注意力（通过sys.path加载）
tests/                        正确性测试
  smoke/
  integration/
  regression/
scripts/                      预热 + 缓存管理助手
benchmarks/
  bench_sparse_attention_ops.py

技术栈

• csrc JIT — 密集FlashAttention、分页KV和sparse_topk_select索引器。运行时从csrc/*.cu.jinja加上csrc/include/编译。公共入口：fmha_sm100.plan → run。
• CuTe-DSL — 块稀疏预填充、FP8 / NVFP4 / FP4量化、分页FP8解码（SparseDecodePagedAttentionWrapper）、FP4块分数索引器。公共入口：fmha_sm100.sparse_atten_func、fmha_sm100.sparse_decode_atten_func、fmha_sm100.fp4_indexer_block_scores。
• 桥接层 —— sparse_fmha_plan / sparse_fmha 将密集API调用点适配为预填充路径的稀疏后端；当你已经在使用密集核并希望一行代码切换到稀疏时很有用。

第三方许可证

fmha_sm100 捆绑、衍生自或依赖于以下第三方组件。每个组件保留其原始许可证；本节进行总结。官方文本随每个组件一同分发。

随附/衍生源码（存放在本仓库中）

运行时依赖（通过pip安装）

每个已安装包的确切许可证随该包一起分发；有关官方文本，请查阅其元数据（pip show <包名>）。

引用

如果MSA对您的研究有帮助，请引用它。（BibTeX条目将在配套论文/技术报告获得稳定标识符后提供 —— 当前为占位符。）
算法参考位于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。对于内核或运行时合约的更改，请先提出issue以协调公共接口 —— fmha_sm100.api、fmha_sm100.sparse和cute.interface是稳定的入口点；其他所有内容均为内部实现，可能随时更改而不另行通知。