fw-ai/minimax-kernels

# MiniMax Kernels 用于 MiniMax 模型的开源 GPU kernel。 本仓库包含一个 NVIDIA Blackwell (SM100) CuTe-DSL 实现， 实现了 MiniMax M3 block-sparse attention。其 KV-outer (KV-stationary) 公式 会加载选定的 KV block 一次，并处理所有选定它的 query， 从而在选择项共享时，分摊 KV-cache 的读取成本。 - [M3 sparse attention 技术概述](docs/m3-sparse-attention.md) - [AOT C++ 后端](docs/aot.md) - 独立基准测试结果：[B200](docs/perf-results-b200.md) · [B300](docs/perf-results-b300.md) - [完整的 Fireworks 设计与性能博客文章](https://fireworks.ai/blog/kernel-optimization-for-minimax-m3-on-nvidia-blackwell) ## 要求 - 具备 SM100 或更新计算能力的 NVIDIA Blackwell GPU - CUDA 13 工具链 - 启用 CUDA 的 PyTorch 安装 - `nvidia-cutlass-dsl >= 4.5.1` 以及来自 [`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) 的 CuTe-DSL 实现 - 用于构建 ahead-of-time 扩展的 C++17 编译器 这些 kernel 支持 BF16、FP16 和 FP8 输入，head dimension 为 128，支持具有 64 或 128 token 页面的 paged KV cache、变长 batch、MHA/GQA/MQA，以及 causal 或 non-causal attention。 ## 安装 请先安装适合宿主机的 CUDA 和 PyTorch 技术栈，然后从其根目录安装本仓库： ``` python -m pip install -v --no-build-isolation . ``` 若要进行可编辑的开发者安装： ``` python -m pip install -v --no-build-isolation -e . ``` 构建隔离被禁用，因为扩展构建需要当前环境中启用了 CUDA 的 PyTorch 和 CuTe-DSL 包。 为了重现用于正确性和基准测试验证的确切公开 CUDA 13 技术栈，请从其官方 CUDA 索引安装 PyTorch，然后应用仓库限制条件： ``` python -m pip install \ --index-url https://download.pytorch.org/whl/cu130 \ "torch==2.9.0+cu130" python -m pip install -c constraints/cu13.txt flashinfer-python fmha-sm100 python -m pip install -v --no-build-isolation -c constraints/cu13.txt . ``` ## 快速开始 `kvouter_attention` 接收 token-major 的 query、paged KV cache、sparse block 选项以及逻辑到物理的 page table： ``` import torch from minimax_kernels.m3_sparse_attention import kvouter_attention device = "cuda" dtype = torch.bfloat16 # 一个 query，16 个 query heads，一个 KV head，head dimension 为 128。 q = torch.randn(1, 16, 128, device=device, dtype=dtype) k_cache = torch.randn(4, 1, 64, 128, device=device, dtype=dtype) v_cache = torch.randn_like(k_cache) # 一个包含四个 physical pages 的 sequence。该 query 选择了这些 pages 所代表的两个 128-token # sparse blocks。 block_tables = torch.arange(4, device=device, dtype=torch.int32).unsqueeze(0) selected = torch.tensor([[[0, 1]]], device=device, dtype=torch.int32) cu_seqlens_q = torch.tensor([0, 1], device=device, dtype=torch.int64) used_kv_lens = torch.tensor([256], device=device, dtype=torch.int32) output, lse = kvouter_attention( q, k_cache, v_cache, selected, block_tables, cu_seqlens_q=cu_seqlens_q, used_kv_lens=used_kv_lens, page_size=64, return_lse=True, ) ``` `output` 的形状为 `[Tq, Hq, 128]`。如果被请求，`lse` 将是形状为 `[Hq, Tq]` 的 FP32。 ## 后端 已安装的 C++ 扩展使用 ahead-of-time 编译的 CuTe kernel，并在可用时作为默认选择。Python CuTe-DSL 路径仍作为 JIT 回退和数值参考提供。 设置 `MINIMAX_KERNELS_KVOUTER_CPP=0` 以强制使用 Python 后端，或设置 `MINIMAX_KERNELS_KVOUTER_CPP=1` 以要求使用 C++ 后端。要求使用不可用的 C++ 扩展会引发错误，而不是静默回退。 ## 测试 GPU 测试需要兼容的 Blackwell GPU。在运行它们之前检查 GPU 可用性，然后选择一个空闲设备： ``` nvidia-smi CUDA_VISIBLE_DEVICES=0 pytest -q tests/m3_sparse_attention ``` 该测试套件涵盖了 index builder、load-balancing scheduler、针对 Torch oracle 的 forward 和 combine kernel、JIT 重新编译防护、Python/C++ 一致性，以及与 FlashInfer 和可选的 MiniMax MSA 的端到端比较。 ## 基准测试 基准测试从仓库检出中运行。简短的冒烟测试将运行每个保留的基准测试： ``` CUDA_VISIBLE_DEVICES=0 \ python benchmarks/m3_sparse_attention/run_benchmark_suite.py --smoke ``` 使用以下命令运行全面扫描： ``` CUDA_VISIBLE_DEVICES=0 \ python benchmarks/m3_sparse_attention/run_benchmark_suite.py ``` 每个用例的日志将写入被 gitignore 的 `benchmark_results/` 目录下。 有关针对性的命令和基准测试覆盖范围，请参阅[技术概述](docs/m3-sparse-attention.md#tests-and-benchmarks)。 ## 贡献者 - Ying Zhang - Junming Chen ## 许可证 本仓库采用 [Apache License 2.0](LICENSE) 授权。改编自 FlashAttention-4 的部分在 BSD 3-Clause License 下分发；请参阅 [NOTICE](NOTICE) 和 [LICENSES/BSD-3-Clause.txt](LICENSES/BSD-3-Clause.txt)。

标签：CUDA, GPU算子, Vectored Exception Handling, 人工智能, 凭据扫描, 大模型, 用户模式Hook绕过, 稀疏注意力, 逆向工具, 高性能计算