flash-msa 0.1.0__py3-none-any.whl

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/_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/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");
+}

flash_msa/flash_msa.py

@@ -0,0 +1,214 @@
+"""Block-sparse top-k GQA attention training entrypoint.
+
+This module owns the Python boundary for the native fused MSA kernels.
+"""
+
+import torch
+
+from flash_msa.msa_select_cutedsl import compute_proxy_lse, select_blocks
+from flash_msa.msa_backward_cutedsl import run_fused_backward
+from flash_msa.msa_forward_cutedsl import run_main_forward
+
+BLOCK_SIZE = 128
+NATIVE_MMA_ROWS_PER_TASK = 128
+
+def _validate_inputs(
+    q_proxy: torch.Tensor,
+    k_proxy: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    top_k: int,
+) -> tuple[int, int]:
+    if q_proxy.ndim != 4 or k_proxy.ndim != 4 or q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError("all attention tensors must have shape (B, H, S, D)")
+
+    b, n_proxy_heads, s, head_dim = q_proxy.shape
+    bk, n_proxy_kv_heads, sk, dk = k_proxy.shape
+    bq, n_heads, sq, dq = q.shape
+    bm, n_kv_heads, sm, dm = k.shape
+    bv, n_v_heads, sv, dv = v.shape
+
+    if not (
+        b == bk == bq == bm == bv
+        and s == sk == sq == sm == sv
+        and head_dim == dk == dq == dm == dv
+        and n_kv_heads == n_v_heads
+    ):
+        raise ValueError("q_proxy, k_proxy, q, k, and v must agree on batch, sequence, head dimension, and KV head count")
+
+    if s % BLOCK_SIZE != 0:
+        raise ValueError(f"sequence length must be divisible by {BLOCK_SIZE}, got {s}")
+    if top_k % BLOCK_SIZE != 0:
+        raise ValueError(f"top_k must be divisible by {BLOCK_SIZE}, got {top_k}")
+    if n_proxy_heads % n_proxy_kv_heads != 0:
+        raise ValueError("n_proxy_heads must be divisible by n_proxy_kv_heads")
+    if n_heads % n_kv_heads != 0:
+        raise ValueError("n_heads must be divisible by n_kv_heads")
+    if n_heads % n_proxy_heads != 0:
+        raise ValueError("n_heads must be divisible by n_proxy_heads")
+    if q_proxy.shape[-1] != 128:
+        raise NotImplementedError("Headdim must be 128")
+
+    num_blocks = s // BLOCK_SIZE
+    top_k_blocks = int(top_k) // BLOCK_SIZE
+    if not 1 <= top_k_blocks <= num_blocks:
+        raise ValueError(f"top_k selects {top_k_blocks} blocks, but sequence has {num_blocks} blocks")
+    if top_k_blocks > 32:
+        raise NotImplementedError("No more than 32 blocks / topk=4096 supported")
+
+    return num_blocks, top_k_blocks
+
+def _run_fused_selected_edge_backward(
+    q_proxy: torch.Tensor,
+    k_proxy: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_indices: torch.Tensor,
+    lse_main: torch.Tensor,
+    o_main: torch.Tensor,
+    grad_o_main: torch.Tensor,
+    grad_kl: torch.Tensor | None,
+    *,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Run the native reverse-index fused selected-edge backward."""
+
+    bsz, n_heads, seq_len, head_dim = q.shape
+    n_proxy_heads = q_proxy.shape[1]
+    n_kv_heads = k.shape[1]
+
+    _main_per_proxy = int(n_heads) // int(n_proxy_heads)
+    query_chunk = NATIVE_MMA_ROWS_PER_TASK // _main_per_proxy
+
+    if NATIVE_MMA_ROWS_PER_TASK % _main_per_proxy  != 0:
+        raise ValueError("Main q heads / proxy q heads ratio must divide NATIVE_MMA_ROWS_PER_TASK evenly, edit NATIVE_MMA_ROWS_PER_TASK to fix")
+
+    lse_proxy = compute_proxy_lse(
+        q_proxy,
+        k_proxy,
+        block_indices,
+        scale=float(scale),
+    )
+    from flash_msa.reverse_index_cuda import build_reverse_index_cuda
+
+    task_meta, task_qids = build_reverse_index_cuda(
+        block_indices,
+        query_chunk=query_chunk,
+    )
+
+    delta_main = (o_main.float() * grad_o_main.float()).sum(dim=-1)
+    grad_kl_scale_value = 0.0
+    if grad_kl is not None:
+        grad_kl_scale_value = float(grad_kl.detach().float().item()) / float(
+            bsz * n_proxy_heads * seq_len
+        )
+
+    return run_fused_backward(
+        q_proxy,
+        k_proxy,
+        q,
+        k,
+        v,
+        grad_o_main,
+        lse_main,
+        lse_proxy,
+        delta_main,
+        task_meta,
+        task_qids,
+        scale=float(scale),
+        grad_kl_scale=grad_kl_scale_value,
+    )
+
+
+class _SparseAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q_proxy: torch.Tensor,
+        k_proxy: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        top_k: int,
+        scale: float,
+    ):
+        b, n_proxy_heads, s, head_dim = q_proxy.shape
+        n_heads = q.shape[1]
+        n_kv_heads = k.shape[1]
+        num_blocks, top_k_blocks = _validate_inputs(q_proxy, k_proxy, q, k, v, int(top_k))
+        block_indices = select_blocks(
+            q_proxy,
+            k_proxy,
+            scale=float(scale),
+            num_blocks=num_blocks,
+            top_k_blocks=top_k_blocks,
+        )
+
+        o_main, lse_main, kl_loss = run_main_forward(
+            q,
+            k,
+            v,
+            block_indices,
+            scale=float(scale),
+        )
+
+        out = o_main.transpose(1, 2).reshape(b, s, -1)
+
+        save_tensors = (
+            q_proxy,
+            k_proxy,
+            q,
+            k,
+            v,
+            block_indices,
+            lse_main,
+            o_main,
+        )
+        ctx.save_for_backward(*save_tensors)
+        ctx.scale = float(scale)
+        return out, kl_loss
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor | None, grad_kl: torch.Tensor | None):
+        q_proxy, k_proxy, q, k, v, block_indices, lse_main, o_main = ctx.saved_tensors
+
+        if grad_out is None:
+            grad_o_main = torch.zeros_like(o_main)
+        else:
+            b, s, _ = grad_out.shape
+            grad_o_main = (
+                grad_out.reshape(b, s, q.shape[1], q.shape[3]).transpose(1, 2).contiguous()
+            )
+
+        dq_proxy, dk_proxy, dq, dk, dv = _run_fused_selected_edge_backward(
+            q_proxy,
+            k_proxy,
+            q,
+            k,
+            v,
+            block_indices,
+            lse_main,
+            o_main,
+            grad_o_main,
+            grad_kl,
+            scale=ctx.scale,
+        )
+        return dq_proxy, dk_proxy, dq, dk, dv, None, None
+
+
+def sparse_attention(
+    q_proxy: torch.Tensor,
+    k_proxy: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    top_k: int,
+    scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Return (attn_out, kl_loss placeholder)
+    Saves Block indices and main attn LSE into HBM for the backward
+    """
+    return _SparseAttentionFunction.apply(q_proxy, k_proxy, q, k, v, int(top_k), float(scale))