[kernel] sm_100a CUDA backward for VSA block-sparse attention (64-token blocks)

Native GB200 (sm_100a) backward pairing with the sm_100a forward: one warp-specialized tcgen05 kernel per (batch, head, kv64 block) computing dK, dV and the dQ partials (fp32 accumulator, cp.reduce.async.bulk), plus a preprocess (Delta, Q^T, dO^T, dqaccum zero, exact-zero dK/dV rows for unselected kv blocks) and a postprocess (dQ unscramble + sm_scale). Consumes the forward's lse (Triton M format) and invert_indices' k2q metadata unchanged; returns (dq, dk, dv) in bf16 with the Triton backward's scaling. Wiring: with FASTVIDEO_VSA_SM100A=1 the autograd backward of the sm_100a forward op runs this kernel when block_sparse_attn_bwd_sm100a.is_supported passes (bf16, head_dim 128, 64-token blocks, even block count, sm_100a device, op built) and the Triton backward otherwise. No new environment variable. primitives.cuh gains only the backward's primitives (additions only; the forward's definitions and SASS are untouched). Built for sm_100a only: the kernel is validated on GB200, not yet on B300/GB300, so sm_103a devices keep the Triton backward. The H3 backend's grad gate is unchanged (separate PR). Semantics: per-row k2q counts may differ arbitrarily (one work item is one kv row; count 0 is skipped identically by every warp and gets zero dK/dV rows); variable_block_sizes masks padded kv rows per lane; is_supported reads no tensor contents. Work-item order: a device kernel sorts by list length from 1024 kv blocks per sequence on; below that a cached identity array is passed. Tests (GB200, extension built from this branch): tests/test_block_sparse_bwd_sm100a.py 15 passed (fp32 masked-dense autograd reference; ragged vbs, zero-count kv blocks, top-k 1/2/3/5/7, batch 2), tests/test_block_sparse_sm100a_dispatch.py 10 passed (sm_100a route vs all-Triton grads, ragged vbs, Triton backward monkeypatched to fail on a supported input), tests/test_block_sparse_sm100a.py 37 passed (forward regression). Perf (kernels only, same window, B=1 H=8 D=128, 25% density, fp32 accumulator, TFLOPS on selected blocks) Triton backward vs this kernel: 4k 198/381 (1.93x), 8k 313/575 (1.84x), 16k 389/715 (1.84x), 32k 433/794 (1.84x), 65k 448/821 (1.83x), 131k 452/736 (1.63x), 262k 452/784 (1.73x), 524k 437/771 (1.76x). FastVideo autograd path incl. invert_indices: 1.6-1.7x at Wan 480P/720P-like shapes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-04 19:39:36 -07:00
parent 7bb76b5ec9
commit 749a9c5d97
11 changed files with 2239 additions and 24 deletions
@@ -372,13 +372,65 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co
 # ---------------------------------------------------------------------------
 # Data-center Blackwell backend custom op (index-native; legacy sm100a API name)
 #
-# Forward runs the sm_100a/sm_103a CUDA extension; backward reuses the Triton kernels.
+# Forward runs the sm_100a/sm_103a CUDA extension. Backward runs the sm_100a CUDA
+# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64-token blocks,
+# sm_100a device, extension built with the op) and the Triton kernels otherwise.
 # The native forward emits lse in exactly Triton's M format (max*log2e +
-# log2(l)), so the pairing needs no conversion. The Triton backward is
+# log2(l)), so either pairing needs no conversion. Both backwards are
 # hardcoded to 64-token blocks, hence the block-size assert below.
 # ---------------------------------------------------------------------------


+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_backward_sm100a",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_backward_sm100a(
+    grad_o: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    lse: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    from fastvideo_kernel.block_sparse_attn_bwd_sm100a import (
+        block_sparse_attn_backward_sm100a_from_k2q, )
+
+    num_kv_blocks = variable_block_sizes.numel()
+    k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks)
+    dq, dk, dv = block_sparse_attn_backward_sm100a_from_k2q(
+        grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), o.contiguous(),
+        lse.contiguous(), k2q_idx, k2q_num, variable_block_sizes)
+    return dq, dk, dv
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm100a")
+def _block_sparse_attn_backward_sm100a_fake(
+    grad_o: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    lse: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def _sm100a_backward_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
+    try:
+        from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd_sm100a
+    except ImportError:  # pragma: no cover - extension not built
+        return False
+    return vsa_bwd_sm100a.is_supported(q, variable_block_sizes)
+
+
@torch.library.custom_op(
    "fastvideo_kernel::block_sparse_attn_sm100a",
    mutates_args=(),
@@ -428,11 +480,15 @@ def _backward_sm100a(ctx, grad_o, grad_M):
    block = q.shape[2] // variable_block_sizes.numel()
    if block != 64:
        raise RuntimeError(
-            "block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with the "
-            f"Triton backward, which is hardcoded to 64-token blocks; got {block}. "
+            "block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with a "
+            f"backward that is hardcoded to 64-token blocks; got {block}. "
            "Run 128-token-block metadata without grad, or use the Triton forward.")
-    dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx,
-                                                   q2k_num, variable_block_sizes)
+    if _sm100a_backward_is_supported(q, variable_block_sizes):
+        dq, dk, dv = block_sparse_attn_backward_sm100a(grad_o, q, k, v, o, M, q2k_idx,
+                                                       q2k_num, variable_block_sizes)
+    else:
+        dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx,
+                                                       q2k_num, variable_block_sizes)
    return dq, dk, dv, None, None, None


@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+"""sm_100a (Blackwell) CUDA block-sparse VSA backward.
+
+Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the
+Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and
+FastVideo's k2q index metadata, returns ``(dq, dk, dv)`` in bf16 with the inputs' layout and the
+Triton backward's scaling (dq and dk carry sm_scale, dv does not). 64-token blocks only; any
+other configuration falls back to Triton via ``is_supported``.
+"""
+
+from typing import Tuple
+
+import torch
+
+try:
+    # The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it
+    # (its __init__ is empty, so hasattr on the package fails with the kernel built and present).
+    from fastvideo_kernel._C import fastvideo_kernel_ops as _C
+    _BWD = getattr(_C, "block_sparse_sm100a_bwd", None)
+    _HAS_VSA_BWD_SM100A = _BWD is not None
+except ImportError:  # pragma: no cover - extension not built
+    _C = None
+    _BWD = None
+    _HAS_VSA_BWD_SM100A = False
+
+_SM100 = (10, 0)
+HEAD_DIM = 128
+BLOCK = 64
+# Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true).
+BHSD = True
+
+
+def set_extension(module) -> None:
+    """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd``.
+
+    The in-repo tests build the binding standalone (tests/jit_ext.py) and inject it here, so
+    the backend can be exercised without a fastvideo_kernel install.
+    """
+    global _C, _BWD, _HAS_VSA_BWD_SM100A
+    _C = module
+    _BWD = getattr(module, "block_sparse_sm100a_bwd", None)
+    _HAS_VSA_BWD_SM100A = _BWD is not None
+
+
+def _seqlen(q: torch.Tensor) -> int:
+    return q.shape[2] if BHSD else q.shape[1]
+
+
+def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
+    """True iff this build can run these tensors; otherwise the caller uses Triton.
+
+    Static facts only (shapes, dtypes, arch, layout), never tensor contents, so it is cheap
+    enough for a per-layer dispatch path. The kernel is fixed at 64-token blocks with
+    head_dim 128 and needs seqlen == 64 * num_blocks with an even num_blocks (its preprocess
+    works in 128-token blocks). Per-row k2q counts may be anything in [0, num_q_blocks],
+    including 0: unselected kv blocks get exactly-zero dk/dv rows.
+    """
+    if not _HAS_VSA_BWD_SM100A or not q.is_cuda:
+        return False
+    if torch.cuda.get_device_capability(q.device) != _SM100:
+        return False
+    if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM:
+        return False
+    if not q.is_contiguous():
+        return False
+    # Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving.
+    if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32,
+                                                                              torch.int64):
+        return False
+    num_blocks = variable_block_sizes.numel()
+    if num_blocks == 0 or num_blocks % 2 != 0:
+        return False
+    if _seqlen(q) != BLOCK * num_blocks:
+        return False
+    return True
+
+
+def block_sparse_attn_backward_sm100a_from_k2q(
+    grad_o: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    lse: torch.Tensor,
+    k2q_idx: torch.Tensor,
+    k2q_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Backward from k2q metadata already in hand (``invert_indices`` layout).
+
+    ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q64
+    block ids, ``k2q_num`` ``[B, H, num_kv_blocks]``; entries past a row's count are never read.
+    """
+    sm_scale = 1.0 / (q.shape[-1]**0.5)
+    idx = k2q_idx.to(torch.int32).contiguous()
+    num = k2q_num.to(torch.int32).contiguous()
+    vbs = variable_block_sizes.to(torch.int32).contiguous()
+    res = _BWD(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(),
+               o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale)
+    return res[0], res[1], res[2]
+
+
+def block_sparse_attn_backward_sm100a(
+    grad_o: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    lse: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Backward pass from the forward's q2k metadata. Returns ``(dq, dk, dv)``.
+
+    Mirrors ``block_sparse_attn_backward_triton``: the k2q inversion is recomputed here with
+    FastVideo's Triton ``invert_indices`` rather than saved by the forward.
+    """
+    from fastvideo_kernel.triton_kernels.index import invert_indices
+
+    num_kv_blocks = variable_block_sizes.numel()
+    batch = q.shape[0]
+    heads = q.shape[1] if BHSD else q.shape[2]
+    idx = q2k_idx.to(torch.int32).contiguous()
+    num = q2k_num.to(torch.int32).contiguous()
+    if idx.dim() != 4:
+        idx = idx.view(batch, heads, -1, idx.shape[-1])
+    if num.dim() != 3:
+        num = num.view(batch, heads, -1)
+    k2q_idx, k2q_num = invert_indices(idx, num, num_kv_blocks)
+    return block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num,
+                                                      variable_block_sizes)
@@ -4,10 +4,12 @@
 The historical ``sm100a`` module and symbol names are retained for compatibility, but the
 extension carries native sm_100a and sm_103a images and supports both device generations.

-A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it
-returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward``
-writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so
-``block_sparse_attn_backward_triton`` runs against it unchanged.
+A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This module is the
+forward: it returns ``(out, lse)`` with ``lse`` in exactly the form
+``triton_block_sparse_attn_forward`` writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]``
+fp32 -- so both ``block_sparse_attn_backward_triton`` and the sm_100a CUDA backward
+(``block_sparse_attn_bwd_sm100a``, 64-token blocks, sm_100a devices only) run against it
+unchanged.

 The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
 (tile volumes 64 and 128 in ``build_vsa_metadata``); the block size is inferred from the