fw-ai/minimax-kernels

fw-ai/minimax-kernels

Description: Kernels for MiniMax models.

Language: Python

License: Apache-2.0

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-07-14T05:55:43Z

Pushed: 2026-07-14T23:09:41Z

Default branch: main

Fork: no

Archived: no

README:

MiniMax Kernels

Open-source GPU kernels for MiniMax models.

This repository contains an NVIDIA Blackwell (SM100) CuTe-DSL implementation of MiniMax M3 block-sparse attention. Its KV-outer (KV-stationary) formulation loads a selected KV block once and processes all queries that selected it, amortizing KV-cache reads when selections are shared.

• [M3 sparse attention technical overview](docs/m3-sparse-attention.md)
• [AOT C++ backend](docs/aot.md)
• Standalone benchmark results: [B200](docs/perf-results-b200.md) ·

[B300](docs/perf-results-b300.md)

• Full Fireworks design and performance blog post

Requirements

• NVIDIA Blackwell GPU with compute capability SM100 or newer
• CUDA 13 toolchain
• A CUDA-enabled PyTorch installation
• nvidia-cutlass-dsl >= 4.5.1 and the CuTe-DSL implementation from

`flash-attn-4`

• A C++17 compiler to build the ahead-of-time extension

The kernels support BF16, FP16, and FP8 inputs, head dimension 128, paged KV caches with 64- or 128-token pages, variable-length batches, MHA/GQA/MQA, and causal or non-causal attention.

Installation

Install the CUDA and PyTorch stack appropriate for the host first, then install this repository from its root:

python -m pip install -v --no-build-isolation .

For an editable developer install:

python -m pip install -v --no-build-isolation -e .

Build isolation is disabled because the extension build needs the CUDA-enabled PyTorch and CuTe-DSL packages from the active environment.

To reproduce the exact public CUDA 13 stack used for correctness and benchmark validation, install PyTorch from its official CUDA index, then apply the repository constraints:

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 .

Quick start

kvouter_attention consumes token-major queries, a paged KV cache, the sparse block selections, and the logical-to-physical page table:

import torch

from minimax_kernels.m3_sparse_attention import kvouter_attention

device = "cuda"
dtype = torch.bfloat16

# One query, 16 query heads, one 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)

# One sequence with four physical pages. The query selects both 128-token
# sparse blocks represented by those pages.
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 has shape [Tq, Hq, 128]. When requested, lse is FP32 with shape [Hq, Tq].

Backends

The installed C++ extension uses ahead-of-time compiled CuTe kernels and is selected by default when available. The Python CuTe-DSL path remains available as a JIT fallback and a numerical reference.

Set MINIMAX_KERNELS_KVOUTER_CPP=0 to force the Python backend, or MINIMAX_KERNELS_KVOUTER_CPP=1 to require the C++ backend. Requiring an unavailable C++ extension raises an error instead of silently falling back.

Tests

GPU tests require a compatible Blackwell GPU. Check GPU availability before running them, then select a free device:

nvidia-smi
CUDA_VISIBLE_DEVICES=0 pytest -q tests/m3_sparse_attention

The suite covers the index builder, load-balancing scheduler, forward and combine kernels against Torch oracles, JIT recompile guards, Python/C++ parity, and end-to-end comparisons with FlashInfer and optionally MiniMax MSA.

Benchmarks

Benchmarks run from a repository checkout. A short smoke run exercises every retained benchmark:

CUDA_VISIBLE_DEVICES=0 \
python benchmarks/m3_sparse_attention/run_benchmark_suite.py --smoke

Run the full sweep with:

CUDA_VISIBLE_DEVICES=0 \
python benchmarks/m3_sparse_attention/run_benchmark_suite.py

Per-case logs are written under the gitignored benchmark_results/ directory. See the [technical overview](docs/m3-sparse-attention.md#tests-and-benchmarks) for focused commands and benchmark coverage.

Acknowledgements

Thanks to the MiniMax team for the MiniMax Sparse Attention algorithm and for helpful discussion. The kernel is built on the FlashAttention4 CuTe-DSL SM100 kernel, and was developed with the help of Claude Code and Cursor.

We also thank Ying Liu and Dmytro (Dima) Dzhulgakov for code reviews and design discussions, and Sofie Yang and Brian Bensch for reviewing and refining the blog post.

Contributors

• Ying Zhang
• Junming Chen

License

The repository is licensed under the [Apache License 2.0](LICENSE). Portions adapted from FlashAttention-4 are distributed under the BSD 3-Clause License; see [NOTICE](NOTICE) and [LICENSES/BSD-3-Clause.txt](LICENSES/BSD-3-Clause.txt).