[perf] Add FA4 CuTe backward support for VSA-256 (#1639)

Co-authored-by: Hyunsung Lee <hyunsungl@sizigistudios.com>
Co-authored-by: alexzms <3036648523@qq.com>

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py

@@ -1,4 +1,4 @@
-"""VSA-256 block-sparse attention wrapper.
+"""VSA-128/256 block-sparse attention wrappers.
 
 The default 256-block path is Triton: it expands the logical 256-block map
 to the existing 64-block Triton kernel via a dense 4x4 expansion per logical
@@ -7,8 +7,8 @@ edge ("route A"), and requires no optional dependencies.
 The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is
 *opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to
 :mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates
-on 128-token KV blocks (this wrapper expands the logical 256-block map /
-sizes into that physical 128-block representation). The CuTe kernel
+on 128-token Q/KV blocks (the 256 wrapper expands its logical KV map and
+sizes into that physical representation). The CuTe kernel
 (``flash_attn.cute`` with block-sparsity) is an optional dependency,
 imported lazily only when this fastpath is selected.
 
@@ -35,7 +35,7 @@ _KV_BLOCK_TRITON = 64  # Existing Triton path uses 64-token KV blocks.
 
 
 def _resolve_backend() -> str:
-    """Pick the backend for the 256-block VSA path.
+    """Pick the backend for the 128/256-block VSA paths.
 
     Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in
     via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe
@@ -49,6 +49,26 @@ def _resolve_backend() -> str:
     return "triton"
 
 
+def _expand_mask_and_sizes_128_to_64(
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Expand a [B, H, Qb128, KVb128] map to 64-token Triton tiles."""
+    expanded_mask = logical_mask_128.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
+    sizes_i32 = logical_kv_sizes_128.to(torch.int32)
+    offsets = torch.tensor(
+        [0, _KV_BLOCK_TRITON],
+        dtype=torch.int32,
+        device=sizes_i32.device,
+    )
+    expanded_sizes = torch.clamp(
+        sizes_i32[:, None] - offsets[None, :],
+        min=0,
+        max=_KV_BLOCK_TRITON,
+    ).reshape(-1)
+    return expanded_mask, expanded_sizes
+
+
 def _expand_mask_and_sizes_256_to_128(
     logical_mask_256: torch.Tensor,
     logical_kv_sizes_256: torch.Tensor,
@@ -112,6 +132,63 @@ def _triton_via_route_a(
     return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
 
 
+def _triton_via_route_a_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    from .triton_kernels.index import map_to_index as triton_map_to_index
+
+    mask_64, sizes_64 = _expand_mask_and_sizes_128_to_64(logical_mask_128, logical_kv_sizes_128)
+    q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
+    return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
+
+
+def block_sparse_attn_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, H, S, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        return _triton_via_route_a_128(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
+    return block_sparse_attn_cute_fwd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
+def block_sparse_attn_128_bshd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, S, H, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        out_bhsd, aux = _triton_via_route_a_128(
+            q.transpose(1, 2).contiguous(),
+            k.transpose(1, 2).contiguous(),
+            v.transpose(1, 2).contiguous(),
+            logical_block_map_128,
+            logical_variable_block_sizes_128,
+        )
+        return out_bhsd.transpose(1, 2).contiguous(), aux
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
+    return block_sparse_attn_cute_fwd_bshd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
 def block_sparse_attn_256(
     q: torch.Tensor,
     k: torch.Tensor,

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py

@@ -1,18 +1,18 @@
-"""CuTe-DSL block-sparse attention forward kernel.
+"""FA4 CuTe-DSL block-sparse attention adapter.
 
-Thin wrapper around `flash_attn.cute.interface._flash_attn_fwd` that adapts
-VSA's `(block_map, variable_block_sizes)` inputs into FA4's
-`BlockSparseTensorsTorch` representation and the per-KV-block validity mask.
+This module adapts VSA's ``(block_map, variable_block_sizes)`` inputs into
+FA4's forward and backward ``BlockSparseTensorsTorch`` representations.
+FA4's public ``flash_attn_func`` owns the forward/backward autograd bridge.
 
 Both [B, H, S, D] (BHSD) and [B, S, H, D] (BSHD) entrypoints are provided.
-The BSHD variant is preferred from VSA-256 callers to avoid layout
+The BSHD variant is preferred from VSA-128/256 callers to avoid layout
 round-trips on the hot path.
 
 The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with
 ``block_sparsity``) is an *optional* dependency: it is imported lazily and
-only exercised when the VSA-256 CuTe fastpath is explicitly selected
-(``FASTVIDEO_VSA_CUTEDSL=1``). The default VSA-256 path is Triton and does
-not require it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
+only exercised when the VSA-128/256 CuTe fastpath is explicitly selected
+(``FASTVIDEO_VSA_CUTEDSL=1``). The default path is Triton and does not require
+it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
 """
 
 from __future__ import annotations
@@ -22,13 +22,14 @@ from typing import Tuple
 
 import torch
 
-_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
+_FA4_IMPORT_HINT = ("VSA-128/256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
                     "provides `flash_attn.cute` with block-sparsity support (plus "
                     "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
-                    "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
+                    "dependency; the default path is Triton. Install the FA4 CuTe "
                     "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.")
 
 
+@functools.lru_cache(maxsize=1)
 def _load_fa4_cute():
     """Lazily import the optional FA4 CuTe block-sparse symbols.
 
@@ -38,14 +39,39 @@ def _load_fa4_cute():
     """
     try:
         from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
-        from flash_attn.cute.interface import _flash_attn_fwd
+        from flash_attn.cute.interface import (
+            _flash_attn_bwd,
+            _flash_attn_fwd,
+            flash_attn_func,
+        )
     except ImportError as exc:  # pragma: no cover - optional dependency
         raise ImportError(_FA4_IMPORT_HINT) from exc
-    return BlockSparseTensorsTorch, _flash_attn_fwd
+    return BlockSparseTensorsTorch, flash_attn_func, _flash_attn_fwd, _flash_attn_bwd
 
 
-# Q-side tile size; kv_block_size comes from the caller's VSA logical KV block.
-_M_BLOCK_SIZE_DEFAULT = 128
+# FA4's physical Q tile size; KV block size comes from the VSA caller.
+_FA4_Q_BLOCK_SIZE = 128
+
+
+class _SingleQStageLength(int):
+    """Keep the real length while selecting FA4's one-stage Q128 path.
+
+    On sm_100 FA4 derives ``q_stage`` from ``max_seqlen_q > tile_m``. Its
+    kernel supports one 128-token Q stage, but the fixed-length public wrapper
+    does not expose that choice. VSA-128 must select it explicitly; otherwise
+    adjacent logical Q blocks are merged into a 256-token sparse block.
+    """
+
+    def __mul__(self, other):
+        return type(self)(int(self) * int(other))
+
+    def __rmul__(self, other):
+        return type(self)(int(other) * int(self))
+
+    def __gt__(self, other):
+        if int(other) == _FA4_Q_BLOCK_SIZE:
+            return False
+        return int(self) > int(other)
 
 
 def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -64,12 +90,12 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     return triton_map_to_index(block_map)
 
 
-def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DEFAULT) -> int:
-    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > m_block_size.
+def _choose_q_sparse_block_size(q_len: int, q_tile_size: int = _FA4_Q_BLOCK_SIZE) -> int:
+    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > q_tile_size.
     major, _ = torch.cuda.get_device_capability()
-    if major >= 10 and q_len > m_block_size:
-        return 2 * m_block_size
-    return m_block_size
+    if major >= 10 and q_len > q_tile_size:
+        return 2 * q_tile_size
+    return q_tile_size
 
 
 def _aggregate_q_block_map(
@@ -134,23 +160,35 @@ def _build_vbs_mask_mod(kv_block_size: int):
     return _vbs_mask_mod
 
 
-def _cute_forward(
-    q_bshd: torch.Tensor,
-    k_bshd: torch.Tensor,
-    v_bshd: torch.Tensor,
+def _build_sparse_tensors(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
     *,
+    q_len: int,
     q_block_size: int,
     kv_block_size: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Internal: FA4 CuTe BSA fwd with BSHD inputs."""
-    BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute()
-    q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1])
-    q_sparse_block_size = max(
-        q_block_size,
-        ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
-    )
+    need_backward: bool,
+    force_q_sparse_block_size: int | None = None,
+) -> Tuple[object, object | None]:
+    """Build the Q-owned forward and KV-owned backward sparse metadata.
+
+    ``need_backward`` is False on inference-only calls: the backward metadata
+    is a pair of dense ``[B, H, kv_blocks, q_blocks]`` int32 index tensors that
+    FA4 keeps alive on its autograd ctx until backward runs, so building it
+    when nothing requires grad is pure overhead (~80 MiB per call at Wan-14B
+    720p shape).
+    """
+    BlockSparseTensorsTorch, _, _, _ = _load_fa4_cute()
+    if force_q_sparse_block_size is None:
+        q_sparse_candidate = _choose_q_sparse_block_size(q_len)
+        q_sparse_block_size = max(
+            q_block_size,
+            ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
+        )
+    else:
+        q_sparse_block_size = force_q_sparse_block_size
+        if q_sparse_block_size < q_block_size or q_sparse_block_size % q_block_size != 0:
+            raise ValueError("force_q_sparse_block_size must be a positive multiple of q_block_size")
     sparse_map = _aggregate_q_block_map(
         block_map,
         q_sparse_block_size=q_sparse_block_size,
@@ -158,35 +196,166 @@ def _cute_forward(
     )
     kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1)
     kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1)
-    full_map = sparse_map & kv_full
-    mask_map = sparse_map & kv_partial
 
-    full_block_idx, full_block_cnt = _map_to_index(full_map)
-    mask_block_idx, mask_block_cnt = _map_to_index(mask_map)
+    def from_maps(full_map: torch.Tensor, mask_map: torch.Tensor) -> object:
+        full_block_idx, full_block_cnt = _map_to_index(full_map.contiguous())
+        mask_block_idx, mask_block_cnt = _map_to_index(mask_map.contiguous())
+        return BlockSparseTensorsTorch(
+            full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
+            full_block_idx=full_block_idx.to(torch.int32).contiguous(),
+            mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
+            mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
+            block_size=(q_sparse_block_size, kv_block_size),
+        )
 
-    sparse_tensors = BlockSparseTensorsTorch(
-        full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
-        full_block_idx=full_block_idx.to(torch.int32).contiguous(),
-        mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
-        mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
-        block_size=(q_sparse_block_size, kv_block_size),
+    forward_sparse_tensors = from_maps(
+        sparse_map & kv_full,
+        sparse_map & kv_partial,
     )
 
-    # _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
-    out, lse = _flash_attn_fwd(
+    if not need_backward:
+        return forward_sparse_tensors, None
+
+    # FA4 backward is KV-owned: for each physical KV tile, list the sparse
+    # query tiles that selected it. Full and partial KV tiles stay separate
+    # so the token-level validity mask only runs for padded tiles.
+    backward_sparse_tensors = from_maps(
+        (sparse_map & kv_full).transpose(2, 3),
+        (sparse_map & kv_partial).transpose(2, 3),
+    )
+    return forward_sparse_tensors, backward_sparse_tensors
+
+
+def _cute_attention_q128_forward(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    *,
+    need_backward: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, object | None]:
+    """Run FA4 with one physical Q stage per logical VSA-128 block."""
+    _, _, flash_attn_fwd, _ = _load_fa4_cute()
+    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
+        block_map,
+        variable_block_sizes,
+        q_len=q_bshd.shape[1],
+        q_block_size=_FA4_Q_BLOCK_SIZE,
+        kv_block_size=_FA4_Q_BLOCK_SIZE,
+        need_backward=need_backward,
+        force_q_sparse_block_size=_FA4_Q_BLOCK_SIZE,
+    )
+    out, lse = flash_attn_fwd(
         q_bshd,
         k_bshd,
         v_bshd,
-        tile_mn=(_M_BLOCK_SIZE_DEFAULT, kv_block_size),
-        mask_mod=_build_vbs_mask_mod(kv_block_size),
-        block_sparse_tensors=sparse_tensors,
+        tile_mn=(_FA4_Q_BLOCK_SIZE, _FA4_Q_BLOCK_SIZE),
+        max_seqlen_q=_SingleQStageLength(q_bshd.shape[1]),
+        mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
+        block_sparse_tensors=forward_sparse_tensors,
         aux_tensors=[variable_block_sizes],
         causal=False,
         return_lse=True,
     )[:2]
+    return out, lse, backward_sparse_tensors
+
+
+class _CuteAttentionQ128(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes):
+        out, lse, backward_sparse_tensors = _cute_attention_q128_forward(
+            q_bshd,
+            k_bshd,
+            v_bshd,
+            block_map,
+            variable_block_sizes,
+            need_backward=True,
+        )
+        ctx.save_for_backward(q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes)
+        ctx.backward_sparse_tensors = backward_sparse_tensors
+        ctx.mark_non_differentiable(lse)
+        ctx.set_materialize_grads(False)
+        return out, lse
+
+    @staticmethod
+    def backward(ctx, grad_out, grad_lse):
+        del grad_lse
+        q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes = ctx.saved_tensors
+        if grad_out is None:
+            grad_out = torch.zeros_like(out)
+        _, _, _, flash_attn_bwd = _load_fa4_cute()
+        dq, dk, dv = flash_attn_bwd(
+            q_bshd,
+            k_bshd,
+            v_bshd,
+            out,
+            grad_out.contiguous(),
+            lse,
+            softmax_scale=q_bshd.shape[-1]**-0.5,
+            mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
+            aux_tensors=[variable_block_sizes],
+            block_sparse_tensors=ctx.backward_sparse_tensors,
+        )
+        return dq, dk, dv, None, None
+
+
+def _cute_attention_q128(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
+    if need_backward:
+        return _CuteAttentionQ128.apply(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
+    out, lse, _ = _cute_attention_q128_forward(
+        q_bshd,
+        k_bshd,
+        v_bshd,
+        block_map,
+        variable_block_sizes,
+        need_backward=False,
+    )
     return out, lse
 
 
+def _cute_attention(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Run FA4's autograd-enabled block-sparse attention with BSHD inputs."""
+    _, flash_attn_func, _, _ = _load_fa4_cute()
+    q_block_size = q_bshd.shape[1] // block_map.shape[2]
+    kv_block_size = k_bshd.shape[1] // block_map.shape[3]
+    if q_block_size == kv_block_size == _FA4_Q_BLOCK_SIZE:
+        return _cute_attention_q128(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
+    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
+    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
+        block_map,
+        variable_block_sizes,
+        q_len=q_bshd.shape[1],
+        q_block_size=q_block_size,
+        kv_block_size=kv_block_size,
+        need_backward=need_backward,
+    )
+    return flash_attn_func(
+        q_bshd,
+        k_bshd,
+        v_bshd,
+        mask_mod=_build_vbs_mask_mod(kv_block_size),
+        aux_tensors=[variable_block_sizes],
+        block_sparse_tensors=forward_sparse_tensors,
+        block_sparse_tensors_bwd=backward_sparse_tensors,
+        return_lse=True,
+    )
+
+
 def block_sparse_attn_cute_fwd(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -194,34 +363,25 @@ def block_sparse_attn_cute_fwd(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """CuTe forward-only block-sparse attention with [B, H, S, D] inputs."""
+    """Autograd-enabled CuTe block-sparse attention for [B, H, S, D]."""
     if block_map.dim() == 3:
         block_map = block_map.unsqueeze(0)
-    q_block_size = q.shape[2] // block_map.shape[2]
-    kv_block_size = k.shape[2] // block_map.shape[3]
 
     q_bshd = q.transpose(1, 2).contiguous()
     k_bshd = k.transpose(1, 2).contiguous()
     v_bshd = v.transpose(1, 2).contiguous()
-    out_bshd, lse_bshd = _cute_forward(
+    out_bshd, lse = _cute_attention(
         q_bshd,
         k_bshd,
         v_bshd,
         block_map,
         variable_block_sizes,
-        q_block_size=q_block_size,
-        kv_block_size=kv_block_size,
     )
     out = out_bshd.transpose(1, 2).contiguous()
-    if lse_bshd is None:
-        lse = torch.empty(
-            (q.shape[0], q.shape[1], q.shape[2]),
-            dtype=torch.float32,
-            device=q.device,
-        )
-    else:
-        lse = lse_bshd.transpose(1, 2).contiguous()
-    return out, lse
+    # FA4 already returns lse as [B, H, S], matching the Triton path's aux
+    # contract, so it needs no transpose. Detach before any further op: the
+    # value is informational and callers never backprop through it.
+    return out, lse.detach()
 
 
 def block_sparse_attn_cute_fwd_bshd(
@@ -231,27 +391,16 @@ def block_sparse_attn_cute_fwd_bshd(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """CuTe forward-only block-sparse attention with [B, S, H, D] inputs."""
+    """Autograd-enabled CuTe block-sparse attention for [B, S, H, D]."""
     if block_map.dim() == 3:
         block_map = block_map.unsqueeze(0)
-    q_block_size = q.shape[1] // block_map.shape[2]
-    kv_block_size = k.shape[1] // block_map.shape[3]
 
-    out, lse_bshd = _cute_forward(
+    out, lse = _cute_attention(
         q,
         k,
         v,
         block_map,
         variable_block_sizes,
-        q_block_size=q_block_size,
-        kv_block_size=kv_block_size,
     )
-    if lse_bshd is None:
-        lse = torch.empty(
-            (q.shape[0], q.shape[2], q.shape[1]),
-            dtype=torch.float32,
-            device=q.device,
-        )
-    else:
-        lse = lse_bshd.transpose(1, 2).contiguous()
-    return out, lse
+    # lse is [B, H, S] regardless of the q/k/v layout; see above.
+    return out, lse.detach()

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -2,6 +2,8 @@ import math
 import torch
 from .block_sparse_attn import block_sparse_attn
 from .block_sparse_attn_256 import (
+    block_sparse_attn_128,
+    block_sparse_attn_128_bshd,
     block_sparse_attn_256,
     block_sparse_attn_256_bshd,
 )
@@ -74,12 +76,13 @@ def video_sparse_attn(
 
     Dispatches the sparse branch by ``block_elements = prod(block_size)``:
     - 64  -> existing TK/Triton path (see ``block_sparse_attn_from_indices``).
+    - 128 -> Triton fallback or CuTe FA4 block-sparse attention.
     - 256 -> CuTe FA4 block-sparse attention (see ``block_sparse_attn_256``).
 
     Backend overrides:
     - ``FASTVIDEO_VSA_TRITON=1`` forces Triton in either path.
     - ``FASTVIDEO_VSA_TK=1`` prefers the sm_90 TK kernel in the 64-block path.
-    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 256-block path.
+    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 128/256-block paths.
     """
     if isinstance(block_size, int):
         block_size = (block_size, block_size, block_size)
@@ -119,8 +122,9 @@ def video_sparse_attn(
     # Sparse branch (fused Triton topk mask)
     mask = fused_topk_mask(scores, topk)
 
-    if block_elements == 256:
-        out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
+    if block_elements in (128, 256):
+        attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256
+        out_s = attention(q, k, v, mask, variable_block_sizes)[0]
     else:
         out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
 
@@ -142,14 +146,14 @@ def video_sparse_attn_bshd(
     """VSA entrypoint for [B, S, H, D] tensors.
 
     Avoids the BHSD<->BSHD round-trip that ``video_sparse_attn`` performs on
-    the CuTe 256-block path; the 64-block path still expects BHSD and is not
+    the CuTe 128/256-block paths; the 64-block path still expects BHSD and is not
     supported here.
     """
     if isinstance(block_size, int):
         block_size = (block_size, block_size, block_size)
     block_elements = block_size[0] * block_size[1] * block_size[2]
-    if block_elements != 256:
-        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
+    if block_elements not in (128, 256):
+        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=128 or 256 "
                          f"(got {block_elements}); use video_sparse_attn for the 64-block path.")
 
     batch, q_seq_len, heads, dim = q.shape
@@ -171,19 +175,15 @@ def video_sparse_attn_bshd(
         raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
                          f"got {q_variable_block_sizes.numel()}")
 
-    # Compression branch (BSHD-native: mean over the 256-token axis).
-    token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32)
-    q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1,
-                                                                         1)).view(1, q_num_blocks, block_elements, 1, 1)
-    kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1,
-                                                                        1)).view(1, kv_num_blocks, block_elements, 1, 1)
-
+    # Compression branch (BSHD-native: match fused_block_mean's semantics).
+    # Padding values are expected to be zero; gradients are broadcast across
+    # the full padded block, just like the BHSD fused common path.
     q_c = q.view(batch, q_num_blocks, block_elements, heads, dim)
     k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim)
     v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim)
-    q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
-    k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
-    v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
+    q_c = (q_c.float().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
     q_ch = q_c.permute(0, 2, 1, 3).contiguous()
     k_ch = k_c.permute(0, 2, 1, 3).contiguous()
     v_ch = v_c.permute(0, 2, 1, 3).contiguous()
@@ -195,13 +195,15 @@ def video_sparse_attn_bshd(
 
     # Sparse branch (fused Triton topk mask + CuTe BSHD).
     mask = fused_topk_mask(scores, topk)
-    out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)
+    attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd
+    out_s, _ = attention(q, k, v, mask, variable_block_sizes)
 
-    out = out_s
-    out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
+    # Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its
+    # backward, so mutating it in place invalidates the graph.
+    out_view = out_s.view(batch, q_num_blocks, block_elements, heads, dim)
     if compress_attn_weight is not None:
         gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim)
-        out_view.add_(out_c_blk.unsqueeze(2) * gate_view)
+        out = out_view + out_c_blk.unsqueeze(2) * gate_view
     else:
-        out_view.add_(out_c_blk.unsqueeze(2))
-    return out
+        out = out_view + out_c_blk.unsqueeze(2)
+    return out.view(batch, q_seq_len, heads, dim)