flash-msa 0.1.0__tar.gz

flash_msa-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2026 flash-msa authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

flash_msa-0.1.0/MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include LICENSE
+recursive-include flash_msa/csrc *.cu

flash_msa-0.1.0/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.4
+Name: flash-msa
+Version: 0.1.0
+Summary: CuTeDSL Flash-MSA training kernels for sparse attention.
+Author: Flash-MSA contributors
+License-Expression: MIT
+Keywords: attention,cuda,cutedsl,pytorch,sparse-attention
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: cuda-python
+Requires-Dist: nvidia-cutlass-dsl[cu13]
+Requires-Dist: packaging
+Requires-Dist: torch==2.12.0
+Requires-Dist: ninja
+Provides-Extra: fa4
+Requires-Dist: flash-attn-4[cu13]; extra == "fa4"
+Provides-Extra: dev
+Requires-Dist: numpy; extra == "dev"
+Requires-Dist: psutil; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Dynamic: license-file
+
+# Flash-MSA
+Flash-MSA is written in CuTeDSL for Hopper and Blackwell GPUs (eg H100, B200) on CUDA 13.
+
+These kernels implement training for the MiniMax Sparse Attention paper:
+https://arxiv.org/abs/2606.13392
+
+Briefly, MSA is a style of sparse attention fitted to GQA that uses a small proxy attention layer to select blocks of keys to provide to the main attention layer. This offers a massive speedup to inference by slashing the memory-bandwidth bottleneck of loading the full KV cache from HBM.
+
+The proxy heads are trained via a KL-divergence loss between the main attention layer's attention scores over the sparsely selected blocks. The proxy heads are assigned groups of main attention heads to select keys for & average scores over for KL-teaching.
+
+This library also includes MSA warmup kernels, which run the main attention densely and train the proxy attention on the full sequence.
+
+More information is included in the [blog post](https://nanduruganesh.github.io/flash-msa).
+
+# Installation
+
+flash-msa depends on [flash-attn](https://github.com/dao-ailab/flash-attention). Try to configure your CUDA/Python/Torch versions to match one of the flash-attn wheels for a fast installation, but if you must build from source, set `MAX_JOBS=<max jobs>` to avoid `pip install flash-msa[attn]` bricking your CPU.
+
+You will also need Python headers, e.g. `apt-get install python3.12-dev`, for whichever python version you are using.
+
+```
+pip install flash-msa
+```
+
+From source:
+
+```
+python setup.py install
+```
+or
+```
+uv pip install -e . --no-build-isolation
+```
+# Usage
+```
+from flash_msa import flash_msa_func
+attn_out, kl_loss = flash_msa_func(Q_proxy, K_proxy, Q, K, V, top_k, head_dim ** -0.5)
+```
+or
+```
+from flash_msa import flash_msa_warmup
+attn_out, kl_loss = flash_msa_warmup_func(Q_proxy, K_proxy, Q, K, V, top_k, head_dim ** -0.5)
+```
+
+Note that kl_loss in the forward is just a torch.zeros placeholder, but after adding it to the main model loss, calling backward() will activate the on-the-fly gradient calcs equivalent to the actual proxy KL loss signal.
+
+# Caveats
+
+1. Flash-MSA only supports headdims 128, block size 128.
+2. Flash-MSA does not currently return fully materialized KL div. loss term in the fwd/bwd (see [blog](https://nanduruganesh.github.io/flash-msa) for explanation).
+3. No support for quantized training (fp8, nvfp4, mxfp4).
+4. No support for attn temps / oai-style softmax bias.
+5. Proxy Q is grouped by Main KV so Q_p <= KV heads for now.
+
+These are not ridiculous to implement though so if there is demand or if someone makes a PR, I will update the repo to include these features.
+
+# Testing
+
+Test sparse MSA correctness against an eager implementation of MSA: `python tests/test_eager_match.py [args]`
+
+Test warmup MSA correctness against an eager implementation of MSA: `python tests/test_warmup_eager_match.py [args]`
+
+# Training
+
+An MSA training example is implemented in this [Megatron-LM fork](https://github.com/nanduruganesh/Megatron-LM). 
+
+Notably, you must add the kl_loss returned by MSA kernels to the model's main CE loss before backward to train the proxy attention. The kl_loss is currently treated as a torch.zeros` placeholder and calculated on-the-fly in the backward, so logging the kl_loss will not reflect how proxy training is actually going. Some solutions to get some signal on proxy training are logging grad/update norms of proxy weights, or patching the forward kernel to calculate and accumulate KL div, but only doing this once every n steps to amortize how slow this would make the forward.
+
+In general if you are going to train with this it is highly recommended to follow tips from [the paper](https://arxiv.org/abs/2606.13392), use MSA warmup before turning on MSA sparse training, and replicate any transformations to the main attention queries and keys (RoPE, QK norm, QK clip, etc) to the proxy queries and keys to improve proxy convergence.
+
+# Inference
+See MiniMax's [official repo](https://github.com/MiniMax-AI/MSA) for MSA inference kernels.

flash_msa-0.1.0/README.md

@@ -0,0 +1,72 @@
+# Flash-MSA
+Flash-MSA is written in CuTeDSL for Hopper and Blackwell GPUs (eg H100, B200) on CUDA 13.
+
+These kernels implement training for the MiniMax Sparse Attention paper:
+https://arxiv.org/abs/2606.13392
+
+Briefly, MSA is a style of sparse attention fitted to GQA that uses a small proxy attention layer to select blocks of keys to provide to the main attention layer. This offers a massive speedup to inference by slashing the memory-bandwidth bottleneck of loading the full KV cache from HBM.
+
+The proxy heads are trained via a KL-divergence loss between the main attention layer's attention scores over the sparsely selected blocks. The proxy heads are assigned groups of main attention heads to select keys for & average scores over for KL-teaching.
+
+This library also includes MSA warmup kernels, which run the main attention densely and train the proxy attention on the full sequence.
+
+More information is included in the [blog post](https://nanduruganesh.github.io/flash-msa).
+
+# Installation
+
+flash-msa depends on [flash-attn](https://github.com/dao-ailab/flash-attention). Try to configure your CUDA/Python/Torch versions to match one of the flash-attn wheels for a fast installation, but if you must build from source, set `MAX_JOBS=<max jobs>` to avoid `pip install flash-msa[attn]` bricking your CPU.
+
+You will also need Python headers, e.g. `apt-get install python3.12-dev`, for whichever python version you are using.
+
+```
+pip install flash-msa
+```
+
+From source:
+
+```
+python setup.py install
+```
+or
+```
+uv pip install -e . --no-build-isolation
+```
+# Usage
+```
+from flash_msa import flash_msa_func
+attn_out, kl_loss = flash_msa_func(Q_proxy, K_proxy, Q, K, V, top_k, head_dim ** -0.5)
+```
+or
+```
+from flash_msa import flash_msa_warmup
+attn_out, kl_loss = flash_msa_warmup_func(Q_proxy, K_proxy, Q, K, V, top_k, head_dim ** -0.5)
+```
+
+Note that kl_loss in the forward is just a torch.zeros placeholder, but after adding it to the main model loss, calling backward() will activate the on-the-fly gradient calcs equivalent to the actual proxy KL loss signal.
+
+# Caveats
+
+1. Flash-MSA only supports headdims 128, block size 128.
+2. Flash-MSA does not currently return fully materialized KL div. loss term in the fwd/bwd (see [blog](https://nanduruganesh.github.io/flash-msa) for explanation).
+3. No support for quantized training (fp8, nvfp4, mxfp4).
+4. No support for attn temps / oai-style softmax bias.
+5. Proxy Q is grouped by Main KV so Q_p <= KV heads for now.
+
+These are not ridiculous to implement though so if there is demand or if someone makes a PR, I will update the repo to include these features.
+
+# Testing
+
+Test sparse MSA correctness against an eager implementation of MSA: `python tests/test_eager_match.py [args]`
+
+Test warmup MSA correctness against an eager implementation of MSA: `python tests/test_warmup_eager_match.py [args]`
+
+# Training
+
+An MSA training example is implemented in this [Megatron-LM fork](https://github.com/nanduruganesh/Megatron-LM). 
+
+Notably, you must add the kl_loss returned by MSA kernels to the model's main CE loss before backward to train the proxy attention. The kl_loss is currently treated as a torch.zeros` placeholder and calculated on-the-fly in the backward, so logging the kl_loss will not reflect how proxy training is actually going. Some solutions to get some signal on proxy training are logging grad/update norms of proxy weights, or patching the forward kernel to calculate and accumulate KL div, but only doing this once every n steps to amortize how slow this would make the forward.
+
+In general if you are going to train with this it is highly recommended to follow tips from [the paper](https://arxiv.org/abs/2606.13392), use MSA warmup before turning on MSA sparse training, and replicate any transformations to the main attention queries and keys (RoPE, QK norm, QK clip, etc) to the proxy queries and keys to improve proxy convergence.
+
+# Inference
+See MiniMax's [official repo](https://github.com/MiniMax-AI/MSA) for MSA inference kernels.

flash_msa-0.1.0/flash_msa/__init__.py

@@ -0,0 +1,14 @@
+"""Public API for Flash-MSA."""
+
+from flash_msa.flash_msa import sparse_attention as flash_msa_func
+from flash_msa.warmup.flash_msa_warmup import sparse_attention_warmup
+
+flash_msa_func_warmup = sparse_attention_warmup
+flash_msa_warmup_func = sparse_attention_warmup
+
+__all__ = [
+    "flash_msa_func",
+    "flash_msa_func_warmup",
+    "flash_msa_warmup_func",
+    "sparse_attention_warmup",
+]

flash_msa-0.1.0/flash_msa/_flash_attn_compat.py

@@ -0,0 +1,76 @@
+""" Patch to call the right varlen flash func depending on if FA2,3, or 4 was installed"""
+
+from __future__ import annotations
+
+import torch
+
+
+def flash_attn_varlen_forward(
+    *,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    softmax_scale: float,
+    causal: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return ``(out, lse)`` for either FlashAttention 2 or 3."""
+
+    try:
+        from flash_attn.flash_attn_interface import _flash_attn_varlen_forward
+
+        out, lse, *_ = _flash_attn_varlen_forward(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            dropout_p=0.0,
+        )
+        return out, lse
+    except ModuleNotFoundError as exc:
+        if exc.name not in {"flash_attn", "flash_attn.flash_attn_interface"}:
+            raise
+
+    try:
+        from flash_attn_interface import flash_attn_varlen_func
+
+        out, lse = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            return_attn_probs=True,
+        )
+        return out, lse
+    except ModuleNotFoundError as exc:
+        if exc.name != "flash_attn_interface":
+            raise
+
+    from flash_attn.cute.interface import flash_attn_varlen_func
+
+    out, lse = flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        return_lse=True,
+    )
+    return out, lse

flash_msa-0.1.0/flash_msa/csrc/reverse_index_cuda.cu

@@ -0,0 +1,211 @@
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cstdint>
+
+#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
+#define CHECK_INT32(x) TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x); CHECK_INT32(x)
+
+namespace {
+
+constexpr int kThreads = 256;
+
+__global__ void count_edges_kernel(
+    int const* __restrict__ block_indices,
+    int* __restrict__ counts,
+    int B,
+    int Hp,
+    int S,
+    int Kb,
+    int NB)
+{
+    int64_t E = (int64_t)B * Hp * S * Kb;
+    int64_t stride = (int64_t)blockDim.x * gridDim.x;
+    for (int64_t e = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; e < E; e += stride) {
+        int slot = (int)(e % Kb);
+        int64_t tmp = e / Kb;
+        int q = (int)(tmp % S);
+        tmp /= S;
+        int p = (int)(tmp % Hp);
+        int b = (int)(tmp / Hp);
+
+        int key_block = block_indices[((int64_t)(b * Hp + p) * S + q) * Kb + slot];
+        if ((unsigned)key_block >= (unsigned)NB) {
+            continue;
+        }
+        int bucket = (b * Hp + p) * NB + key_block;
+        atomicAdd(counts + bucket, 1);
+    }
+}
+
+__global__ void scan_fill_meta_kernel(
+    int const* __restrict__ counts,
+    int* __restrict__ bucket_offsets,
+    int* __restrict__ task_meta,
+    int B,
+    int Hp,
+    int NB,
+    int query_chunk,
+    int padded_tasks)
+{
+    if (blockIdx.x != 0 || threadIdx.x != 0) {
+        return;
+    }
+
+    int task = 0;
+    int buckets = B * Hp * NB;
+    for (int bucket = 0; bucket < buckets; ++bucket) {
+        bucket_offsets[bucket] = task;
+        int count = counts[bucket];
+        int chunks = (count + query_chunk - 1) / query_chunk;
+        int key_block = bucket % NB;
+        int tmp = bucket / NB;
+        int proxy_head = tmp % Hp;
+        int batch = tmp / Hp;
+
+        for (int c = 0; c < chunks && task < padded_tasks; ++c) {
+            int valid = count - c * query_chunk;
+            valid = valid > query_chunk ? query_chunk : valid;
+            int64_t row = (int64_t)task * 4;
+            task_meta[row + 0] = batch;
+            task_meta[row + 1] = proxy_head;
+            task_meta[row + 2] = key_block;
+            task_meta[row + 3] = valid;
+            ++task;
+        }
+    }
+}
+
+__global__ void scatter_edges_kernel(
+    int const* __restrict__ block_indices,
+    int const* __restrict__ bucket_offsets,
+    int* __restrict__ write_counts,
+    int* __restrict__ task_qids,
+    int B,
+    int Hp,
+    int S,
+    int Kb,
+    int NB,
+    int query_chunk,
+    int padded_tasks)
+{
+    int64_t E = (int64_t)B * Hp * S * Kb;
+    int64_t stride = (int64_t)blockDim.x * gridDim.x;
+    for (int64_t e = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; e < E; e += stride) {
+        int slot = (int)(e % Kb);
+        int64_t tmp = e / Kb;
+        int q = (int)(tmp % S);
+        tmp /= S;
+        int p = (int)(tmp % Hp);
+        int b = (int)(tmp / Hp);
+
+        int key_block = block_indices[((int64_t)(b * Hp + p) * S + q) * Kb + slot];
+        if ((unsigned)key_block >= (unsigned)NB) {
+            continue;
+        }
+
+        int bucket = (b * Hp + p) * NB + key_block;
+        int local = atomicAdd(write_counts + bucket, 1);
+        int task = bucket_offsets[bucket] + local / query_chunk;
+        if (task >= padded_tasks) {
+            continue;
+        }
+        int lane = local - (local / query_chunk) * query_chunk;
+        task_qids[(int64_t)task * query_chunk + lane] = q;
+    }
+}
+
+} // namespace
+
+void run_build_reverse_index(
+    torch::Tensor block_indices,
+    torch::Tensor counts,
+    torch::Tensor write_counts,
+    torch::Tensor bucket_offsets,
+    torch::Tensor task_meta,
+    torch::Tensor task_qids,
+    int64_t block_size,
+    int64_t query_chunk)
+{
+    CHECK_INPUT(block_indices);
+    CHECK_INPUT(counts);
+    CHECK_INPUT(write_counts);
+    CHECK_INPUT(bucket_offsets);
+    CHECK_INPUT(task_meta);
+    CHECK_INPUT(task_qids);
+
+    TORCH_CHECK(block_indices.dim() == 4, "block_indices must have shape [B, Hp, S, Kb]");
+    TORCH_CHECK(task_meta.dim() == 2 && task_meta.size(1) == 4, "task_meta must have shape [T, 4]");
+    TORCH_CHECK(task_qids.dim() == 2 && task_qids.size(1) == query_chunk, "task_qids must have shape [T, query_chunk]");
+    TORCH_CHECK(block_size > 0 && query_chunk > 0, "block_size and query_chunk must be positive");
+
+    int B = (int)block_indices.size(0);
+    int Hp = (int)block_indices.size(1);
+    int S = (int)block_indices.size(2);
+    int Kb = (int)block_indices.size(3);
+    int NB = S / (int)block_size;
+    int buckets = B * Hp * NB;
+    int padded_tasks = (int)task_meta.size(0);
+
+    TORCH_CHECK(S % block_size == 0, "S must be divisible by block_size");
+    TORCH_CHECK(counts.numel() == buckets, "counts has wrong size");
+    TORCH_CHECK(write_counts.numel() == buckets, "write_counts has wrong size");
+    TORCH_CHECK(bucket_offsets.numel() == buckets, "bucket_offsets has wrong size");
+    TORCH_CHECK(task_qids.size(0) == padded_tasks, "task_meta/task_qids row mismatch");
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    C10_CUDA_CHECK(cudaMemsetAsync(counts.data_ptr<int>(), 0, counts.numel() * sizeof(int), stream));
+    C10_CUDA_CHECK(cudaMemsetAsync(write_counts.data_ptr<int>(), 0, write_counts.numel() * sizeof(int), stream));
+    C10_CUDA_CHECK(cudaMemsetAsync(task_meta.data_ptr<int>(), 0, task_meta.numel() * sizeof(int), stream));
+    C10_CUDA_CHECK(cudaMemsetAsync(task_qids.data_ptr<int>(), 0xff, task_qids.numel() * sizeof(int), stream));
+
+    int64_t edges = (int64_t)B * Hp * S * Kb;
+    int blocks = (int)std::min<int64_t>((edges + kThreads - 1) / kThreads, 65535);
+    blocks = std::max(blocks, 1);
+
+    count_edges_kernel<<<blocks, kThreads, 0, stream>>>(
+        block_indices.data_ptr<int>(),
+        counts.data_ptr<int>(),
+        B,
+        Hp,
+        S,
+        Kb,
+        NB);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    scan_fill_meta_kernel<<<1, 1, 0, stream>>>(
+        counts.data_ptr<int>(),
+        bucket_offsets.data_ptr<int>(),
+        task_meta.data_ptr<int>(),
+        B,
+        Hp,
+        NB,
+        (int)query_chunk,
+        padded_tasks);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    scatter_edges_kernel<<<blocks, kThreads, 0, stream>>>(
+        block_indices.data_ptr<int>(),
+        bucket_offsets.data_ptr<int>(),
+        write_counts.data_ptr<int>(),
+        task_qids.data_ptr<int>(),
+        B,
+        Hp,
+        S,
+        Kb,
+        NB,
+        (int)query_chunk,
+        padded_tasks);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("run_build_reverse_index", &run_build_reverse_index, "Build MSA reverse index on CUDA");
+}