[kernel] sm_100a CUDA backward for VSA block-sparse attention, 128-token blocks (#1866)

Co-authored-by: Pengcheng Li <pengchengl@ptyche0203.ptyche.clusters.nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -52,13 +52,13 @@ def _force_tk() -> bool:
 def _force_sm100a() -> bool:
     """True iff the data-center Blackwell forward is explicitly opted into.
 
-    Opt-in only (same legacy-named env the H3 backend honors): the extension is
-    forward-only, so this routing pairs it with the Triton backward -- its lse
-    is already in Triton's M format. Honored only when
-    ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token
-    metadata falls through to the default selection; unsupported 128-token
-    metadata raises because Triton has no compatible fallback.
-    ``FASTVIDEO_VSA_TRITON`` still wins.
+    Opt-in only (same legacy-named env the H3 backend honors). Honored only when
+    ``block_sparse_attn_sm100a.is_supported`` passes; the backward then runs the
+    sm_100a/sm_103a CUDA backward when ``block_sparse_attn_bwd_sm100a.is_supported``
+    passes and Triton otherwise (the forward's lse is already in Triton's M
+    format). Unsupported 64-token metadata falls through to the default
+    selection; unsupported 128-token metadata raises because Triton has no
+    compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins.
     """
     return _env_bool("FASTVIDEO_VSA_SM100A")
 
@@ -389,11 +389,12 @@ 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 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 either pairing needs no conversion. Both backwards are
-# hardcoded to 64-token blocks, hence the block-size assert below.
+# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64- or 128-token
+# blocks, data-center Blackwell device, extension built with that block's op) and
+# the Triton kernels otherwise. The native forward emits lse in exactly Triton's M
+# format (max*log2e + log2(l)), so either pairing needs no conversion. The Triton
+# backward is hardcoded to 64-token blocks, so unsupported 128-token metadata
+# raises instead of falling back.
 # ---------------------------------------------------------------------------
 
 
@@ -493,16 +494,17 @@ def _setup_context_sm100a(ctx, inputs, output):
 
 def _backward_sm100a(ctx, grad_o, grad_M):
     q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
-    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 a "
-            f"backward that is hardcoded to 64-token blocks; got {block}. "
-            "Run 128-token-block metadata without grad, or use the Triton forward.")
     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:
+        block = q.shape[2] // variable_block_sizes.numel()
+        if block != 64:
+            raise RuntimeError(
+                "block_sparse_attn_sm100a backward: no sm_100a/sm_103a backward for "
+                f"{block}-token blocks on this build/device, and the Triton backward is "
+                "hardcoded to 64-token blocks. Run this metadata without grad, or build "
+                "the extension with block_sparse_sm100a_blk128_bwd.")
         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
@@ -536,8 +538,9 @@ def block_sparse_attn_from_indices(
 
     # Backend resolution:
     # - FASTVIDEO_VSA_TRITON forces Triton everywhere.
-    # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward
-    #   (Triton backward). The environment name is retained for compatibility.
+    # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward (CUDA
+    #   backward when supported, else Triton). The environment name is retained
+    #   for compatibility.
     #   Unsupported 64-token metadata falls through; unsupported 128-token
     #   metadata raises because Triton cannot consume it.
     # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py

@@ -4,8 +4,10 @@
 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``.
+Triton backward's scaling (dq and dk carry sm_scale, dv does not). Two kernels, picked by the
+metadata's block size: 64-token blocks (``block_sparse_sm100a_bwd``) and 128-token blocks
+(``block_sparse_sm100a_blk128_bwd``); any other configuration falls back to Triton via
+``is_supported``.
 """
 
 from typing import Tuple
@@ -16,45 +18,64 @@ 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
+    _BWD_BY_BLOCK = {
+        64: getattr(_C, "block_sparse_sm100a_bwd", None),
+        128: getattr(_C, "block_sparse_sm100a_blk128_bwd", None),
+    }
+    _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values())
 except ImportError:  # pragma: no cover - extension not built
     _C = None
-    _BWD = None
+    _BWD_BY_BLOCK = {}
     _HAS_VSA_BWD_SM100A = False
 
 _SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)}
 HEAD_DIM = 128
-BLOCK = 64
+BLOCKS = (64, 128)
 # 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``.
+    """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd`` and / or
+    ``block_sparse_sm100a_blk128_bwd``.
 
-    A standalone build of ``block_sparse_bwd_sm100a.cu`` (for example through
+    A standalone build of the binding .cu files (for example through
     ``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so
     the backend can be exercised without rebuilding the fastvideo_kernel wheel.
     """
-    global _C, _BWD, _HAS_VSA_BWD_SM100A
+    global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_SM100A
     _C = module
-    _BWD = getattr(module, "block_sparse_sm100a_bwd", None)
-    _HAS_VSA_BWD_SM100A = _BWD is not None
+    _BWD_BY_BLOCK = {
+        64: getattr(module, "block_sparse_sm100a_bwd", None),
+        128: getattr(module, "block_sparse_sm100a_blk128_bwd", None),
+    }
+    _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values())
 
 
 def _seqlen(q: torch.Tensor) -> int:
     return q.shape[2] if BHSD else q.shape[1]
 
 
+def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
+    """The block size the metadata implies (seqlen / num_blocks), 0 when it is not integral.
+
+    A plain int: under dynamic-shape torch.compile the shapes are SymInts, which cannot key
+    ``_BWD_BY_BLOCK``; ``int()`` specializes the value with a shape guard.
+    """
+    num_blocks = variable_block_sizes.numel()
+    seqlen = _seqlen(q)
+    return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks)
+
+
 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.
+    enough for a per-layer dispatch path. Both kernels take head_dim 128 and seqlen == block *
+    num_blocks; the 64-token kernel also needs an even num_blocks (its preprocess works in
+    128-token blocks); batch * heads is the launch grid's y extent, so it must be <= 65535.
+    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
@@ -64,14 +85,16 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
         return False
     if not q.is_contiguous():
         return False
+    if q.shape[0] * (q.shape[1] if BHSD else q.shape[2]) > 65535:
+        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:
+    block = _block_size(q, variable_block_sizes)
+    if _BWD_BY_BLOCK.get(block) is None:
         return False
-    if _seqlen(q) != BLOCK * num_blocks:
+    if block == 64 and variable_block_sizes.numel() % 2 != 0:
         return False
     return True
 
@@ -89,15 +112,21 @@ def block_sparse_attn_backward_sm100a_from_k2q(
 ) -> 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.
+    ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q
+    block ids (64- or 128-token blocks, as the metadata's block size), ``k2q_num``
+    ``[B, H, num_kv_blocks]``; entries past a row's count are never read.
     """
+    block = _block_size(q, variable_block_sizes)
+    backward = _BWD_BY_BLOCK.get(block)
+    if backward is None:
+        raise RuntimeError(f"block_sparse_attn_backward_sm100a: no kernel for {block}-token "
+                           f"blocks (built: {sorted(b for b, f in _BWD_BY_BLOCK.items() if f)})")
     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)
+    res = backward(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]
 
 

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py

@@ -8,7 +8,7 @@ A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This mo
 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
+(``block_sparse_attn_bwd_sm100a``, 64- and 128-token blocks, sm_100a / sm_103a) run against it
 unchanged.
 
 The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
@@ -43,9 +43,10 @@ BHSD = True
 
 
 def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
+    # A plain int: dynamic-shape torch.compile hands SymInt shapes, which cannot key _FWD_BY_BLOCK.
     num_blocks = variable_block_sizes.numel()
     seqlen = q.shape[2] if BHSD else q.shape[1]
-    return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks
+    return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks)
 
 
 def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: