FA4 consumer Blackwell (sm_120) integration

flash_attn/cute/README.md

@@ -22,6 +22,45 @@ from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
 out = flash_attn_func(q, k, v, causal=True)
 ```
 
+## Consumer Blackwell (sm_120 / RTX PRO 6000, RTX 50-series)
+
+FA4 runs on consumer Blackwell (compute capability 12.x), which exposes SM80-class
+`mma.sync` tensor cores (no WGMMA/tcgen05/TMEM). Dispatch and tile selection are
+auto-tuned for this arch; no environment variables are required for normal use.
+
+**Supported:** forward and backward for dense, causal, and local/sliding-window
+attention; MHA / GQA / MQA; variable-length (`flash_attn_varlen_func`); paged-KV;
+block sparsity; `score_mod` / `mask_mod`; learnable sink. Head dims 64/96/128/192/256.
+
+**fp8 KV-cache decode (e4m3/e5m2):** for decode (`seqlen_q == 1`) with a quantized
+K/V cache and a bf16/fp16 query, pass fp8 `k`/`v` plus per-`(batch, kv_head)` fp32
+`k_descale`/`v_descale`. This auto-routes to a memory-efficient GEMV decode kernel
+(no env var needed) and is ~1.6–1.9× faster than bf16 at GQA ratios ≤ 4 while halving
+KV-cache bandwidth. Accuracy is within ~2e-3 of an fp8-quantized reference.
+
+**Environment flags:**
+- `FLASH_ATTENTION_SM120_DECODE_KERNEL=1` — opt into the experimental **bf16** GEMV
+  decode kernel for `seqlen_q == 1` (the fp8 decode path above is always on when fp8
+  K/V is supplied). Off by default.
+- `FLASH_ATTENTION_ARCH` — override the detected compute capability (testing/compile).
+
+**Known performance floors vs FA2 on sm_120** (hardware-bound, not bugs):
+- Causal *MHA* (`qhead_per_kvhead == 1`) at `seqlen ≥ 8192` is ~0.95× FA2 — a register/
+  occupancy wall (255 regs/thread → 1 CTA/SM). GQA (the common case) is at parity or faster.
+- fp8 KV-cache decode regresses below bf16 at GQA ratio ≥ 8 (the GEMV loop becomes
+  compute-bound); it still halves KV memory, so it remains the only fp8-cache path.
+- Backward is at ~parity with FA2; fp8 is forward/decode-only (no fp8 backward).
+
+**Feature limitations on sm_120:**
+- `learnable_sink` is incompatible with SplitKV (each split would double-count the sink
+  in the combine step), so SplitKV is disabled when a sink is present — attention runs
+  in a single split (correct, but without the decode SplitKV speedup).
+- Negative-offset sliding windows (`window_size` with a negative bound, e.g. `(None, -X)`
+  or `(-X, None)`) are forward-only: the backward raises `NotImplementedError` (its
+  dK/dV are incorrect for these offset windows). Non-negative windows are fully supported.
+- Deterministic backward (`deterministic=True`) is not supported (the SM80-base backward
+  lacks the dQ-semaphore path).
+
 ## Development
 
 ```sh

flash_attn/cute/block_sparse_utils.py

@@ -1,3 +1,4 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
 """
 Block-sparse runtime utilities for CUTE DSL kernels.
 
@@ -705,6 +706,114 @@ def produce_block_sparse_loads_sm100(
     return kv_producer_state, q_producer_phase
 
 
+@cute.jit
+def run_block_sparse_mainloop_sm80(
+    blocksparse_tensors: BlockSparseTensors,
+    batch_idx,
+    head_idx,
+    m_block,
+    mma_one_n_block,
+    mask_fn,
+    mask_mod,
+    fastdiv_mods,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+    q_subtile_factor: cutlass.Constexpr[int] = 1,
+):
+    """Block-sparse mainloop iteration for SM80/SM120.
+
+    NOTE: This implementation hard-codes the non-varlen 4D indexing pattern
+    (lines below this docstring access blocksparse_tensors with 3D indices into
+    batch/head/m_block_sparse). Varlen + block-sparse on SM80/SM120 would
+    require routing through get_curr_blocksparse_tensors(...) (which handles
+    both 2D varlen and 4D non-varlen layouts) and threading seqlen_info into
+    this function. The SM120 dispatcher in interface.py currently does not
+    support varlen + block-sparse together, so this is intentionally narrow;
+    if you lift that restriction, update this function accordingly.
+
+    Processes mask blocks first (applying mask_mod), then full blocks (seqlen masking only).
+    The first full block always receives seqlen masking regardless of whether mask blocks
+    preceded it, since full blocks may be at higher n positions than mask blocks.
+
+    Mirrors the non-intra-wg-overlap path of consume_block_sparse_loads for SM90/SM100.
+
+    Args:
+        mma_one_n_block: callable with signature
+            (n_block, mask_fn, is_first_n_block) -> None
+        mask_fn: partial of mask.apply_mask with batch/head/m_block/thr_mma already bound.
+            Called as mask_fn(acc_S, n_block=n, mask_mod=..., mask_seqlen=...,
+                              fastdiv_mods=...)
+        mask_mod: the user mask_mod constexpr (None → no mask_mod application)
+        fastdiv_mods: fast-division helpers when mask_mod is not None
+
+    Returns:
+        processed_any: True if at least one block was processed.
+    """
+    # SM80/SM120 only need the first 4 fields; trailing fields are SM100/backward-only.
+    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors
+
+    m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+
+    curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse]
+    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None]
+
+    if const_expr(full_block_cnt is not None):
+        curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse]
+        curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None]
+    else:
+        curr_full_block_cnt = Int32(0)
+        curr_full_block_idx = None
+
+    processed_any = curr_mask_block_cnt + curr_full_block_cnt > 0
+
+    # Process mask blocks: first gets is_first=True and seqlen masking; rest get is_first=False.
+    if curr_mask_block_cnt > 0:
+        n_block = curr_mask_block_idx[curr_mask_block_cnt - 1]
+        mma_one_n_block(
+            n_block=n_block,
+            mask_fn=partial(
+                mask_fn,
+                mask_mod=mask_mod,
+                mask_seqlen=True,
+                fastdiv_mods=fastdiv_mods if const_expr(mask_mod is not None) else None,
+            ),
+            is_first_n_block=True,
+        )
+        for i in cutlass.range(1, curr_mask_block_cnt):
+            n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i]
+            mma_one_n_block(
+                n_block=n_block,
+                mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=False),
+                is_first_n_block=False,
+            )
+
+    # Process full blocks: first full block always gets seqlen masking (it may be at the
+    # highest n position even when mask blocks were present). No mask_mod applied.
+    if const_expr(full_block_cnt is not None):
+        if curr_full_block_cnt > 0:
+            n_block = curr_full_block_idx[curr_full_block_cnt - 1]
+            if curr_mask_block_cnt == 0:
+                mma_one_n_block(
+                    n_block=n_block,
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
+                    is_first_n_block=True,
+                )
+            else:
+                mma_one_n_block(
+                    n_block=n_block,
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
+                    is_first_n_block=False,
+                )
+            for j in cutlass.range(1, curr_full_block_cnt):
+                n_block = curr_full_block_idx[curr_full_block_cnt - 1 - j]
+                mma_one_n_block(
+                    n_block=n_block,
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False),
+                    is_first_n_block=False,
+                )
+
+    return processed_any
+
+
 @cute.jit
 def get_total_block_count(
     blocksparse_tensors: BlockSparseTensors,