Integrate Fireworks KV-outer sparse prefill (fireworks-msa)

NOTICE

@@ -86,6 +86,23 @@ Where:      Apache-2.0-tagged sources under python/fmha_sm100/csrc/
             sm100_fmha_reduction.hpp, tvm_ffi_utils.h).
 Full text:  https://www.apache.org/licenses/LICENSE-2.0
 
+--------------------------------------------------------------------------------
+Fireworks AI (minimax-kernels KV-outer) — Apache License 2.0
+--------------------------------------------------------------------------------
+
+Copyright (c) 2026 Fireworks AI
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Project:    https://github.com/fw-ai/minimax-kernels
+Where:      python/fmha_sm100/kvouter/ and python/fmha_sm100/csrc/kvouter/
+            (KV-outer sparse attention backend used by sparse_fmha_adapter.py).
+Full text:  https://www.apache.org/licenses/LICENSE-2.0
+
 --------------------------------------------------------------------------------
 NVIDIA TensorRT-LLM — Apache License 2.0
 --------------------------------------------------------------------------------

README.md

@@ -1,13 +1,21 @@
 # MiniMax Sparse Attention (MSA)
 
+> **Fireworks integration branch.** The KV-outer sparse prefill backend on this
+> branch is vendored from
+> [fw-ai/minimax-kernels](https://github.com/fw-ai/minimax-kernels/tree/main),
+> developed by [Fireworks AI](https://fireworks.ai). For upstream docs, the full
+> test matrix, and benchmark results, see that repository. Background on the
+> KV-outer work:
+> [Kernel optimization for MiniMax M3 on NVIDIA Blackwell](https://fireworks.ai/blog/kernel-optimization-for-minimax-m3-on-nvidia-blackwell).
+
 [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
 [![Python](https://img.shields.io/badge/python-≥3.10-blue.svg)](pyproject.toml)
 [![GPU](https://img.shields.io/badge/NVIDIA-SM100-76b900.svg)](#requirements)
-[![Stack: CuTe-DSL + Cuda](https://img.shields.io/badge/stack-CuTe--DSL%20%2B%20Cuda-purple.svg)](#stacks)
+[![Stack: CuTe-DSL + CUDA](https://img.shields.io/badge/stack-CuTe--DSL%20%2B%20CUDA-purple.svg)](#stacks)
 
 **MSA** (`fmha_sm100`) ships dense FlashAttention and sparse top-k attention
-kernels for **NVIDIA SM100**. Two JIT-compiled stacks
-share one Python package:
+kernels for **NVIDIA SM100**. Three kernel stacks plus a routing bridge share one
+Python package; the public `fmha_sm100` / `fmha_sm100_plan` API is unchanged.
 
 ![MSA architecture](docs/architecture.png)
 
@@ -16,26 +24,29 @@ share one Python package:
 | Stack | Path | What it gives you |
 |---|---|---|
 | **csrc JIT** | `python/fmha_sm100/csrc/` | Dense FMHA (`fmha_sm100`, `fmha_sm100_plan`) + `sparse_topk_select` indexer, compiled from Jinja templates by `jit.py` at runtime. |
-| **CuTe-DSL** | `python/fmha_sm100/cute/` | Full sparse attention (forward + paged FP8 decode, BF16 / FP8 / NVFP4 / FP4), compiled at runtime via `cute.compile`. |
-| **Bridge** | `python/fmha_sm100/sparse_fmha_adapter.py` | Adapts the `fmha_sm100` API to call `sparse_atten_func` for sparse prefill paths. |
+| **KV-outer** | `python/fmha_sm100/kvouter/` + `csrc/kvouter/` | Fireworks KV-outer sparse **prefill** (GQA ≥ 8): Python CuTe-DSL + optional C++ AOT extension (`fmha_sm100._C`). Default path when `qhead_per_kv ≥ 8`. |
+| **CuTe-DSL (legacy sparse)** | `python/fmha_sm100/cute/` | CSR sparse prefill fallback (GQA < 8), paged FP8 decode, BF16 / FP8 / NVFP4 / FP4 paths, compiled at runtime via `cute.compile`. |
+| **Bridge** | `python/fmha_sm100/sparse_fmha_adapter.py` | Routes sparse prefill through KV-outer or legacy CuTe based on GQA ratio; adapts the `fmha_sm100` API to both backends. |
 
 > **License: MIT.** Self-authored files carry `SPDX-License-Identifier: MIT`.
 > See [LICENSE](LICENSE) and [NOTICE](NOTICE). Bundled / derived third-party
 > code retains its own license — see [Third-party licenses](#third-party-licenses).
 
 ## Requirements
 
-- **GPU**: NVIDIA SM100.
-- **Toolchain**: CUDA Toolkit with `nvcc` on `PATH` (or `CUDA_HOME` / `CUDA_PATH` set).
+- **GPU**: NVIDIA SM100 (Blackwell).
+- **Toolchain**: CUDA Toolkit **13.x** with `nvcc` on `PATH` (or `CUDA_HOME` / `CUDA_PATH` set). Required for all stacks (csrc JIT, KV-outer, legacy CuTe).
 - **Python**: ≥ 3.10.
+- **PyTorch**: ≥ 2.9 (see `pyproject.toml`).
 - **OS**: Linux x86_64 (aarch64 untested; JIT builds may need small Makefile edits on WSL).
 
 Quick sanity check before installing:
 
 ```bash
-nvcc --version                # expect ≥ 12.x
+nvcc --version                # expect ≥ 13.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)
+python -c "import torch; print(torch.__version__)"                # ≥ 2.9
 ```
 
 ## Using with the `kernels` library
@@ -58,25 +69,42 @@ Check out the kernel on the Hugging Face Hub [here](https://huggingface.co/kerne
 
 ```bash
 # --recursive pulls the NVIDIA CUTLASS submodule (python/fmha_sm100/cutlass/),
-# whose headers are required for JIT/AOT compilation.
+# whose headers are required for dense FMHA JIT / AOT compilation.
 git clone --recursive https://github.com/MiniMax-AI/MSA.git msa
 cd msa
+git checkout fireworks-msa   # this integration branch
 # If you cloned without --recursive:
 #   git submodule update --init --recursive
-pip install .           # standard install (works from a wheel too)
+pip install -e . --no-build-isolation   # editable install; builds fmha_sm100._C (KV-outer)
 # or
-pip install -e .        # editable install for development
+pip install . --no-build-isolation      # standard install
 ```
 
-This pulls in the CuTe-DSL stack via `nvidia-cutlass-dsl` and `quack-kernels`;
-the csrc kernels are JIT-compiled at first import from sources shipped inside
-the package.
+This pulls in `nvidia-cutlass-dsl`, `quack-kernels`, and `flash-attn-4` (see
+`pyproject.toml`). Dense csrc kernels are JIT-compiled on first use; the KV-outer
+C++ extension is built at install time. On first KV-outer call per config, CuTe-DSL
+kernels are AOT-exported (cached under `~/.cache/minfer/`).
+
+**KV-outer backend selection** (optional env vars):
+
+| Variable | Effect |
+|---|---|
+| unset (default) | Use C++ AOT extension when `fmha_sm100._C` is available; else Python CuTe-DSL |
+| `FMHA_SM100_KVOUTER_CPP=1` | Force C++ AOT |
+| `FMHA_SM100_KVOUTER_CPP=0` | Force Python CuTe-DSL |
+| `MINIMAX_KERNELS_KVOUTER_CPP` | Legacy alias for the same flags |
+
+Sparse prefill uses KV-outer when `num_qo_heads // num_kv_heads ≥ 8` (e.g. TP1
+config 64/4); lower GQA ratios still use the legacy CuTe CSR path under `cute/`.
 
 ## Verify
 
 Run a small CUDA smoke test. **The first run JIT-compiles `sparse_topk_select`,
 which takes 30 s – a few minutes on a cold nvcc cache** — this is normal, not
-a hang. Subsequent runs hit the JIT cache and finish in seconds.
+a hang. The KV-outer C++ extension (`fmha_sm100._C`) is built at `pip install`
+time; the first sparse prefill call may additionally AOT-export CuTe-DSL kernels
+(cached under `~/.cache/minfer/`). Subsequent runs hit the caches and finish
+in seconds.
 
 ```bash
 python tests/smoke/test_sparse_topk_forced.py
@@ -124,11 +152,18 @@ out, _ = fmha_sm100(
 )
 ```
 
-For block-sparse prefill with CSR metadata, the FP4 indexer, NVFP4 K/V, and
-the paged FP8 decode wrapper, see the **CuTe-DSL deep dive**:
+For block-sparse prefill with CSR metadata, NVFP4 K/V, and the paged FP8 decode
+wrapper, see the **legacy CuTe-DSL deep dive** (still used for decode, NVFP4, and
+GQA < 8 prefill):
 
 - [`python/fmha_sm100/cute/README.md`](python/fmha_sm100/cute/README.md)
 
+KV-outer internals live under [`python/fmha_sm100/kvouter/`](python/fmha_sm100/kvouter/).
+Upstream docs and benchmark tables:
+[fw-ai/minimax-kernels](https://github.com/fw-ai/minimax-kernels/tree/main).
+Design write-up:
+[Kernel optimization for MiniMax M3 on NVIDIA Blackwell](https://fireworks.ai/blog/kernel-optimization-for-minimax-m3-on-nvidia-blackwell).
+
 ## Test
 
 ```bash
@@ -151,13 +186,18 @@ python -m pytest test_sparse_atten.py -q
 ## Benchmark
 
 `benchmarks/bench_sparse_attention_ops.py` covers dense prefill, paged
-prefill, sparse prefill, dense decode, paged decode, sparse decode, in
-`fp8` and `bf16` (`nvfp4` is sparse-prefill only).
+prefill, sparse prefill, dense decode, paged decode, and sparse decode, in
+`fp8` and `bf16` (`nvfp4` is sparse-prefill only). Sparse prefill on this
+branch uses the KV-outer backend for GQA ≥ 8 (default TP1: `h_q=64`, `h_k=4`).
 
 ```bash
-python benchmarks/bench_sparse_attention_ops.py --help     # full flag list
+FMHA_SM100_KVOUTER_CPP=1 python benchmarks/bench_sparse_attention_ops.py --help
 ```
 
+For three-way KV-outer vs FlashInfer vs MSA comparisons and published latency
+tables, see
+[minimax-kernels benchmarks and perf results](https://github.com/fw-ai/minimax-kernels/tree/main).
+
 Common invocations (output is TSV):
 
 | Goal | Command |
@@ -176,11 +216,14 @@ python/fmha_sm100/                  Python package
   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)
+  sparse_fmha_adapter.py            Bridge: fmha_sm100 API → KV-outer or legacy CuTe
+  kvouter/                          Vendored Fireworks KV-outer (Python + AOT export)
+  csrc/kvouter/                     KV-outer C++ op (fmha_sm100._C)
+  csrc/                             Dense 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)
+  cute/                             Legacy CuTe-DSL sparse attention (loaded via sys.path)
+setup.py                            Builds fmha_sm100._C CUDA extension
 tests/                              Correctness tests
   smoke/  integration/  regression/
 scripts/                            Warmup + cache-management helpers
@@ -191,14 +234,22 @@ benchmarks/                         bench_sparse_attention_ops.py
 
 - **csrc JIT** — dense FlashAttention, page KV, and `sparse_topk_select`
   indexer. Compiled at runtime from `csrc/*.cu.jinja` plus
-  `csrc/include/`. Public entry: `fmha_sm100.plan → run`.
-- **CuTe-DSL** — block-sparse prefill, FP8 / NVFP4 / FP4 quantization, paged
-  FP8 decode (`SparseDecodePagedAttentionWrapper`), FP4 block-score indexer.
+  `csrc/include/`. Public entry: `fmha_sm100_plan` → `fmha_sm100`.
+- **KV-outer** — Fireworks block-sparse prefill for GQA ≥ 8. Vendored from
+  [fw-ai/minimax-kernels](https://github.com/fw-ai/minimax-kernels/tree/main).
+  See the
+  [Blackwell kernel optimization blog post](https://fireworks.ai/blog/kernel-optimization-for-minimax-m3-on-nvidia-blackwell)
+  for background. Index build + forward + combine; C++ AOT path preferred
+  (`fmha_sm100._C`). Public entry: `sparse_fmha` → `kvouter_attention`
+  (via `sparse_fmha_adapter`).
+- **CuTe-DSL (legacy sparse)** — CSR sparse prefill fallback (GQA < 8),
+  FP8 / NVFP4 / FP4 quantization, paged FP8 decode
+  (`SparseDecodePagedAttentionWrapper`), FP4 block-score indexer.
   Public entry: `fmha_sm100.sparse_atten_func`,
   `fmha_sm100.sparse_decode_atten_func`, `fmha_sm100.fp4_indexer_block_scores`.
 - **Bridge** — `sparse_fmha_plan` / `sparse_fmha` adapt the dense-API call
-  site to the sparse backend for prefill paths; useful when you already
-  drive the dense kernel and want a one-line swap to sparse.
+  site to KV-outer or legacy CuTe for prefill; decode and indexer paths are
+  unchanged.
 
 ## Third-party licenses
 
@@ -213,12 +264,14 @@ Authoritative text is shipped with each component.
 | **NVIDIA CUTLASS** | BSD-3-Clause | Git submodule at `python/fmha_sm100/cutlass/` (provides `include/` + `tools/util/include/`), plus BSD-3-tagged headers under `python/fmha_sm100/csrc/include/`. The SM100 MMA descriptor encodings in `python/fmha_sm100/cute/src/common/mma_sm100_desc.py` mirror CUTLASS hardware descriptors. Copyright (c) 2017–2025 NVIDIA CORPORATION & AFFILIATES. |
 | **FlashInfer** | Apache-2.0 | Headers and sources under `python/fmha_sm100/csrc/` and `python/fmha_sm100/csrc/include/` that carry a `Copyright (c) <year> by FlashInfer team` line (e.g. `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`). Project: <https://github.com/flashinfer-ai/flashinfer>. |
 | **NVIDIA TensorRT-LLM + NAVER Corp (CLOVA)** | Apache-2.0 | Portions of `python/fmha_sm100/csrc/include/sparse_topk_select.cuh` — `indexerTopK` histogram-step + insertion-sort derived from `tensorrt_llm/cpp/tensorrt_llm/kernels/indexerTopK.cu`. Copyright (c) 2019–2026 NVIDIA CORPORATION; Copyright (c) 2021 NAVER Corp. The per-file header in `sparse_topk_select.cuh` includes a function-level provenance map. |
+| **Fireworks AI (minimax-kernels KV-outer)** | Apache-2.0 | Vendored under `python/fmha_sm100/kvouter/` and `python/fmha_sm100/csrc/kvouter/`. Upstream: <https://github.com/fw-ai/minimax-kernels>. Copyright (c) 2026 Fireworks AI. See [NOTICE](NOTICE). |
 
 ### Runtime dependencies (installed via pip)
 
 | Package | Upstream | License |
 |---|---|---|
 | `quack-kernels` | <https://github.com/Dao-AILab/quack> | Apache-2.0 |
+| `flash-attn-4` | FlashAttention CuTe SM100 kernels | BSD-3-Clause (see package) |
 | `nvidia-cutlass-dsl` | NVIDIA CUTLASS Python DSL | NVIDIA / BSD-3-Clause (see package) |
 | `apache-tvm-ffi` | Apache TVM FFI | Apache-2.0 |
 | `cuda-python` | NVIDIA | NVIDIA / see package |

pyproject.toml

@@ -1,25 +1,31 @@
 [build-system]
-requires = ["setuptools>=64", "wheel"]
+requires = [
+  "setuptools>=64",
+  "wheel",
+  "torch>=2.9",
+  "nvidia-cutlass-dsl[cu13]>=4.5.1",
+]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "fmha_sm100"
 version = "0.1.1"
-description = "MSA — MiniMax Sparse Attention: dense FlashAttention + sparse top-k indexer (csrc JIT) and CuTe-DSL sparse attention for NVIDIA SM100"
+description = "MSA — MiniMax Sparse Attention: dense FlashAttention + sparse top-k indexer (csrc JIT) and CuTe-DSL sparse attention for NVIDIA SM100 (CUDA 13.x required)"
 readme = "README.md"
 requires-python = ">=3.10"
 # SPDX-License-Identifier: MIT
 license = { text = "MIT" }
 authors = [{ name = "MiniMax" }]
 dependencies = [
-    "torch",
+    "torch>=2.9",
     "apache-tvm-ffi",
     "jinja2",
     "ninja",
     "pybind11",
-    "cuda-python",
-    "nvidia-cutlass-dsl>=4.4.1",
-    "quack-kernels>=0.2.10",
+    "cuda-python>=13,<14",
+    "nvidia-cutlass-dsl[cu13]>=4.5.1",
+    "quack-kernels>=0.4,<0.5",
+    "flash-attn-4[cu13]==4.0.0b15",
 ]
 
 [project.urls]
@@ -37,7 +43,7 @@ package-dir = { "" = "python" }
 # The cute/ sparse sources are loaded via sys.path.insert at runtime by
 # fmha_sm100/sparse.py (not imported as a submodule), so they are shipped as
 # package data.
-packages = ["fmha_sm100"]
+packages = ["fmha_sm100", "fmha_sm100.kvouter"]
 
 [tool.setuptools.package-data]
 fmha_sm100 = [

python/fmha_sm100/csrc/kvouter/bindings.cpp

@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) 2026 Fireworks AI
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <torch/extension.h>
+
+#include "cute_sparse_kvouter.h"
+
+namespace fmha_sm100 {
+
+TORCH_LIBRARY(fmha_sm100, m) {
+  m.def("sparse_kvouter_init", sparse_kvouter_init);
+  m.def("sparse_kvouter_attn", sparse_kvouter_attn);
+}
+
+}  // namespace fmha_sm100
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}