[bugfix] make the VSA-256 CuTe backward safe to land

Follow-ups found while reproducing this PR on GB200 (sm_100).

video_sparse_attn_h3.py composed the compression branch onto the attention
output with an in-place addcmul_. On the CuTe backend that output *is* the
tensor FA4's autograd node saved for its backward, so the moment this path
has a backward at all, training dies with "one of the variables needed for
gradient computation has been modified by an inplace operation ... output 0
of FlashAttnFuncBackward". Same defect the BSHD branch of
video_sparse_attn_bshd already fixed here; H3 landed after this branch's base
and so never got it. Now out-of-place.

The KV-owned backward metadata was built unconditionally. It is a pair of
dense [B, H, kv_blocks, q_blocks] int32 index tensors that FA4 keeps alive on
its ctx until backward runs, and inference never reads it. Gated on
requires_grad: at Wan-14B 720p shape (B=1, H=40, D=128, S=92160, topk 45/360)
a no_grad forward drops from 17.23 ms / 1998 MiB transient to 15.92 ms /
1918 MiB. Training keeps the metadata and is unchanged.

The aux (lse) return was transposed to [B, S, H]. FA4 hands back [B, H, S]
already (interface.py builds lse_shape as (batch, num_head, seqlen_q) when qv
is None), which is what the Triton path's aux contract is, so the transpose
made the two backends disagree and cost a 14 MiB fp32 copy per call at the
shape above. Dropped, and detach happens first now.

Tests
- gated compression branch (compress_attn_weight is not None) backward: the
  branch Wan and H3 actually run, and the one the in-place bug lives in.
- partially filled Q tiles, and q_len != kv_len; forward had cross coverage,
  backward had none.
- inference forward is bitwise identical to the training forward, pinning the
  requires_grad gate.
- aux is [B, H, S] on both entrypoints.
- fastvideo/tests/attention/test_vsa_h3_backward.py covers the H3 backend on
  both backends and cross-checks CuTe gradients against Triton.
- Gradient tolerances tightened from avg_abs 2e-2 / max_rel 0.5 to 1e-3 /
  0.25; measured error across every case above is <= 1.2e-4 and <= 0.11.

test_vsa_varlen.py seeds the RNG now. It draws every tensor and every
variable block size from the global RNG and then asserts max_rel < 0.05, so
adding any test that runs before it shifts its inputs and it fails on unlucky
data. That is the "unrelated" CI failure reported on this PR: it reproduces
in a full-suite run and passes in isolation.

Verified on GB200 (sm_100), FA4 82d6441 + nvidia-cutlass-dsl 4.6.0.dev0 +
quack-kernels 0.5.3: fastvideo-kernel/tests/ and fastvideo/tests/attention/
give 318 passed, stable across repeated runs. The one remaining failure,
test_fa4_quantize_op_fake_matches_real, fails identically on an unmodified
main. Reverting either fix above turns the matching new test red.

pre-commit (yapf, ruff, codespell) passes on every changed file; the mypy
hook only errors on the worktree's directory name, and does so identically on
an unmodified main.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XNe3T7jDpv9N5ur7c4S2oX

Author: alexzms

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py

@@ -142,8 +142,16 @@ def _build_sparse_tensors(
     q_len: int,
     q_block_size: int,
     kv_block_size: int,
-) -> Tuple[object, object]:
-    """Build the Q-owned forward and KV-owned backward sparse metadata."""
+    need_backward: bool,
+) -> 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()
     q_sparse_candidate = _choose_q_sparse_block_size(q_len)
     q_sparse_block_size = max(
@@ -174,6 +182,9 @@ def from_maps(full_map: torch.Tensor, mask_map: torch.Tensor) -> object:
         sparse_map & kv_partial,
     )
 
+    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.
@@ -195,12 +206,14 @@ def _cute_attention(
     _, 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]
+    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,
@@ -236,8 +249,10 @@ def block_sparse_attn_cute_fwd(
         variable_block_sizes,
     )
     out = out_bshd.transpose(1, 2).contiguous()
-    lse_bsh = lse.transpose(1, 2).contiguous().detach()
-    return out, lse_bsh
+    # 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(
@@ -258,5 +273,5 @@ def block_sparse_attn_cute_fwd_bshd(
         block_map,
         variable_block_sizes,
     )
-    lse_bsh = lse.transpose(1, 2).contiguous().detach()
-    return out, lse_bsh
+    # lse is [B, H, S] regardless of the q/k/v layout; see above.
+    return out, lse.detach()

fastvideo-kernel/tests/test_vsa256_backward.py

@@ -1,134 +1,224 @@
-"""VSA-256 FA4 CuTe forward/backward parity for BHSD and BSHD APIs."""
+"""VSA-256 FA4 CuTe forward/backward parity for BHSD and BSHD APIs.
+
+Covers the shapes the CuTe backward actually sees in production: the gated
+compression branch (`compress_attn_weight`), partially filled Q tiles,
+and q_len != kv_len. Also pins the inference fast path, which must skip the
+KV-owned backward metadata without changing the forward result.
+"""
 
 from __future__ import annotations
 
+from typing import Tuple
+
 import pytest
 import torch
 
 from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
 
 from .test_vsa256_triton import _metrics, _torch_vsa256_reference
 
+_BLOCK = 256
+_BLOCK_SIZE_3D = (4, 8, 8)  # prod == 256
+
+# Measured on GB200 (sm_100) with bf16 inputs: grads land around 1e-4 avg_abs
+# and <=0.11 max_rel across every case below, so these leave ~10x headroom
+# without being loose enough to hide a real regression.
+_OUT_TOL = (1e-3, 0.2)
+_GRAD_TOL = (1e-3, 0.25)
+
 
 @pytest.fixture(autouse=True)
 def _require_cute_backend(monkeypatch):
     pytest.importorskip(
         "flash_attn.cute.block_sparsity",
         reason="optional FA4 CuTe build (flash_attn.cute) not installed",
     )
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
     monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
     monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
     monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
 
 
+def _zero_pad_tail(x: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
+    """Zero the padded tail of every 256-token tile of a [B, H, S, D] tensor.
+
+    VSA callers scatter into a zeroed tile buffer, so padded slots are zero;
+    both the kernel and the reference rely on that.
+    """
+    bsz, heads, _, dim = x.shape
+    blocks = var.numel()
+    token_idx = torch.arange(_BLOCK, device=x.device, dtype=torch.int32)
+    valid = (token_idx.view(1, -1) < var.view(-1, 1)).view(1, 1, blocks, _BLOCK, 1)
+    valid = valid.expand(bsz, heads, blocks, _BLOCK, dim).reshape_as(x)
+    return x * valid.to(x.dtype)
+
+
+def _make_inputs(
+    q_blocks: int,
+    kv_blocks: int,
+    kv_var: torch.Tensor,
+    q_var: torch.Tensor,
+    heads: int = 2,
+    dim: int = 128,
+    seed: int = 42,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    torch.manual_seed(seed)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
+    q = torch.randn(1, heads, sq, dim, device=device, dtype=dtype)
+    k = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
+    v = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
+    grad_out = torch.randn_like(q)
+    return _zero_pad_tail(q, q_var), _zero_pad_tail(k, kv_var), _zero_pad_tail(v, kv_var), grad_out
+
+
+def _check(tag: str, ref: torch.Tensor, got: torch.Tensor, tol: Tuple[float, float]) -> None:
+    assert torch.isfinite(got).all().item(), f"{tag}: non-finite values"
+    avg_abs, max_rel = _metrics(ref, got)
+    print(f"  {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < tol[0], f"{tag}: avg_abs {avg_abs:.3e} >= {tol[0]:.3e}"
+    assert max_rel < tol[1], f"{tag}: max_rel {max_rel:.3e} >= {tol[1]:.3e}"
+
+
+def _run_bhsd(q, k, v, kv_var, q_var, topk, gate=None):
+    qg, kg, vg = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out = video_sparse_attn(qg, kg, vg, kv_var, q_var, topk, block_size=_BLOCK_SIZE_3D, compress_attn_weight=gate)
+    return out, (qg, kg, vg)
+
+
+def _run_bshd(q, k, v, kv_var, q_var, topk, gate=None):
+    qg, kg, vg = (t.transpose(1, 2).contiguous().requires_grad_(True) for t in (q, k, v))
+    gate_bshd = None if gate is None else gate.transpose(1, 2).contiguous()
+    out = video_sparse_attn_bshd(qg,
+                                 kg,
+                                 vg,
+                                 kv_var,
+                                 q_var,
+                                 topk,
+                                 block_size=_BLOCK_SIZE_3D,
+                                 compress_attn_weight=gate_bshd)
+    return out.transpose(1, 2), (qg, kg, vg)
+
+
+def _reference(q, k, v, q_var, kv_var, topk, gate=None):
+    qr, kr, vr = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out = _torch_vsa256_reference(qr, kr, vr, q_var, kv_var, topk, compress_attn_weight=gate)
+    return out, (qr, kr, vr)
+
+
+def _compare(tag, layout, q, k, v, kv_var, q_var, topk, grad_out, gate=None):
+    runner = _run_bhsd if layout == "bhsd" else _run_bshd
+    out, (qg, kg, vg) = runner(q, k, v, kv_var, q_var, topk, gate=gate)
+    (out * grad_out).sum().backward()
+    grads = [g.grad if g.grad.dim() == 4 and layout == "bhsd" else g.grad for g in (qg, kg, vg)]
+    if layout == "bshd":
+        grads = [g.transpose(1, 2) for g in grads]
+
+    out_ref, refs = _reference(q, k, v, q_var, kv_var, topk, gate=gate)
+    (out_ref * grad_out).sum().backward()
+
+    print(f"[{tag}-{layout}]")
+    _check("out", out_ref, out, _OUT_TOL)
+    for name, ref, got in zip(("dq", "dk", "dv"), refs, grads):
+        _check(name, ref.grad, got, _GRAD_TOL)
+
+
 @pytest.mark.cuda
 @pytest.mark.parametrize("layout", ["bhsd", "bshd"])
 def test_vsa256_cute_forward_backward_vs_torch_ref(layout: str) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA is required")
+    kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var)
+    _compare("vsa256-cute", layout, q, k, v, kv_var, q_var, 2, grad_out)
 
-    torch.manual_seed(42)
-    device = torch.device("cuda")
-    dtype = torch.bfloat16
 
-    bsz, heads, dim = 1, 2, 128
-    q_blocks_256, kv_blocks_256 = 3, 4
-    q_block = 256
-    kv_block = 256
-    topk_logical = 2
-    sq = q_blocks_256 * q_block
-    skv = kv_blocks_256 * kv_block
-
-    q_base = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype)
-    k_base = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
-    v_base = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
-    grad_out = torch.randn_like(q_base)
-
-    q_var = torch.full(
-        (q_blocks_256,), q_block, dtype=torch.int32, device=device
-    )
-    kv_var = torch.tensor(
-        [256, 173, 79, 256], dtype=torch.int32, device=device
-    )
-    token_idx = torch.arange(kv_block, device=device, dtype=torch.int32)
-    kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1)
-    kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1)
-    kv_valid = kv_valid.expand(
-        bsz, heads, kv_blocks_256, kv_block, dim
-    ).reshape(bsz, heads, skv, dim)
-    k_base = k_base * kv_valid.to(k_base.dtype)
-    v_base = v_base * kv_valid.to(v_base.dtype)
-
-    if layout == "bhsd":
-        q = q_base.detach().clone().requires_grad_(True)
-        k = k_base.detach().clone().requires_grad_(True)
-        v = v_base.detach().clone().requires_grad_(True)
-        out = video_sparse_attn(
-            q,
-            k,
-            v,
-            kv_var,
-            q_var,
-            topk_logical,
-            block_size=(4, 8, 8),
-            compress_attn_weight=None,
-        )
-        (out * grad_out).sum().backward()
-        out_bhsd = out
-        dq, dk, dv = q.grad, k.grad, v.grad
-    else:
-        q = q_base.transpose(1, 2).contiguous().requires_grad_(True)
-        k = k_base.transpose(1, 2).contiguous().requires_grad_(True)
-        v = v_base.transpose(1, 2).contiguous().requires_grad_(True)
-        grad_out_bshd = grad_out.transpose(1, 2).contiguous()
-        out = video_sparse_attn_bshd(
-            q,
-            k,
-            v,
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_with_compress_gate(layout: str) -> None:
+    """The gated compression branch is what Wan and MiniMax-H3 actually run.
+
+    It is also the branch that composes the sparse output with the compression
+    output, so it is the one that breaks if that composition mutates FA4's
+    saved output in place.
+    """
+    kv_var = torch.tensor([256, 200, 256, 91], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var, seed=7)
+    gate = torch.randn_like(q) * 0.1
+    _compare("vsa256-cute-gated", layout, q, k, v, kv_var, q_var, 2, grad_out, gate=gate)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_partial_q_blocks(layout: str) -> None:
+    """Q tiles that are not full: only the compression divisor depends on it,
+    but it is the one axis the existing coverage held constant."""
+    kv_var = torch.tensor([256, 128, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.tensor([256, 61, 199], dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 3, kv_var, q_var, seed=11)
+    _compare("vsa256-cute-partial-q", layout, q, k, v, kv_var, q_var, 2, grad_out)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_cross_q_kv(layout: str) -> None:
+    """q_len != kv_len: forward has coverage, backward did not."""
+    kv_var = torch.tensor([256, 143, 256, 256, 88], dtype=torch.int32, device="cuda")
+    q_var = torch.full((2, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(2, 5, kv_var, q_var, seed=13)
+    _compare("vsa256-cute-cross", layout, q, k, v, kv_var, q_var, 3, grad_out)
+
+
+@pytest.mark.cuda
+def test_vsa256_cute_inference_matches_training_forward() -> None:
+    """The KV-owned backward metadata is only built when something requires
+    grad. Skipping it must not perturb the forward result."""
+    kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, _ = _make_inputs(3, 4, kv_var, q_var, seed=5)
+
+    with torch.no_grad():
+        out_infer = video_sparse_attn_bshd(
+            q.transpose(1, 2).contiguous(),
+            k.transpose(1, 2).contiguous(),
+            v.transpose(1, 2).contiguous(),
             kv_var,
             q_var,
-            topk_logical,
-            block_size=(4, 8, 8),
+            2,
+            block_size=_BLOCK_SIZE_3D,
             compress_attn_weight=None,
         )
-        (out * grad_out_bshd).sum().backward()
-        out_bhsd = out.transpose(1, 2)
-        dq = q.grad.transpose(1, 2)
-        dk = k.grad.transpose(1, 2)
-        dv = v.grad.transpose(1, 2)
-
-    q_ref = q_base.detach().clone().requires_grad_(True)
-    k_ref = k_base.detach().clone().requires_grad_(True)
-    v_ref = v_base.detach().clone().requires_grad_(True)
-    out_ref = _torch_vsa256_reference(
-        q_ref, k_ref, v_ref, q_var, kv_var, topk_logical
-    )
-    (out_ref * grad_out).sum().backward()
 
-    tensors = (
-        out_bhsd,
-        dq,
-        dk,
-        dv,
-        q_ref.grad,
-        k_ref.grad,
-        v_ref.grad,
-    )
-    assert all(torch.isfinite(t).all().item() for t in tensors)
-
-    m_out = _metrics(out_ref, out_bhsd)
-    m_dq = _metrics(q_ref.grad, dq)
-    m_dk = _metrics(k_ref.grad, dk)
-    m_dv = _metrics(v_ref.grad, dv)
-    print(
-        f"[vsa256-cute-{layout}] "
-        f"out(avg_abs={m_out[0]:.6e}, max_rel={m_out[1]:.6e}), "
-        f"dq(avg_abs={m_dq[0]:.6e}, max_rel={m_dq[1]:.6e}), "
-        f"dk(avg_abs={m_dk[0]:.6e}, max_rel={m_dk[1]:.6e}), "
-        f"dv(avg_abs={m_dv[0]:.6e}, max_rel={m_dv[1]:.6e})"
-    )
+    out_train, _ = _run_bshd(q, k, v, kv_var, q_var, 2)
+    torch.testing.assert_close(out_infer, out_train.transpose(1, 2).detach(), rtol=0, atol=0)
 
-    assert m_out[0] < 1e-3 and m_out[1] < 0.2
-    assert m_dq[0] < 2e-2 and m_dq[1] < 0.5
-    assert m_dk[0] < 2e-2 and m_dk[1] < 0.5
-    assert m_dv[0] < 2e-2 and m_dv[1] < 0.5
+
+@pytest.mark.cuda
+def test_vsa256_cute_lse_is_bhs() -> None:
+    """The aux return is [B, H, S] on both entrypoints, matching the Triton
+    path's contract."""
+    from fastvideo_kernel.block_sparse_attn_256 import (block_sparse_attn_256, block_sparse_attn_256_bshd)
+
+    device = torch.device("cuda")
+    heads, dim, q_blocks, kv_blocks = 2, 128, 3, 4
+    sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
+    q = torch.randn(1, heads, sq, dim, device=device, dtype=torch.bfloat16)
+    k = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
+    v = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
+    vbs = torch.full((kv_blocks, ), _BLOCK, dtype=torch.int32, device=device)
+    mask = torch.zeros(1, heads, q_blocks, kv_blocks, dtype=torch.bool, device=device)
+    mask[..., :2] = True
+
+    _, lse_bhsd = block_sparse_attn_256(q, k, v, mask, vbs)
+    assert lse_bhsd.shape == (1, heads, sq), lse_bhsd.shape
+
+    _, lse_bshd = block_sparse_attn_256_bshd(
+        q.transpose(1, 2).contiguous(),
+        k.transpose(1, 2).contiguous(),
+        v.transpose(1, 2).contiguous(),
+        mask,
+        vbs,
+    )
+    assert lse_bshd.shape == (1, heads, sq), lse_bshd.shape