[bugfix] fastvideo-kernel: fix Triton block-sparse backward logit scaling (bf16 K pre-scaling) (#1730)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py

@@ -237,7 +237,12 @@ def _attn_bwd_dkdv(
         # Load m before computing qk to reduce pipeline stall.
         offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
         m = tl.load(M + offs_m)
-        qkT = tl.dot(k, qT)
+        # Recompute logits exactly as the forward does: raw bf16 operands into
+        # the dot, fp32 scale after accumulation. A bf16 pre-scaled K perturbs
+        # the recomputed logits relative to the saved M by an error
+        # proportional to |logit|, which exp2 amplifies into arbitrarily wrong
+        # probabilities at large activations.
+        qkT = tl.dot(k, qT) * (sm_scale * 1.4426950408889634)
         pT = tl.math.exp2(qkT - m[None, :])
         mask = tl.arange(0, BLOCK_N1) < block_size
         pT = tl.where(mask[:, None], pT, 0.0)
@@ -268,6 +273,7 @@ def _attn_bwd_dq(
         do,
         m,
         D,
+        sm_scale,
         # shared by Q/K/V/DO.
         q2k_index,
         q2k_num,
@@ -315,7 +321,7 @@ def _attn_bwd_dq(
         block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
-        qk = tl.dot(q, kT)
+        qk = tl.dot(q, kT) * (sm_scale * 1.4426950408889634)
         p = tl.math.exp2(qk - m)
         offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
         mask = offs_in_block < block_size
@@ -324,8 +330,7 @@ def _attn_bwd_dq(
         dp = tl.dot(do, vT).to(tl.float32)
         ds = p * (dp - Di[:, None])
         ds = ds.to(tl.bfloat16)
-        # Compute dQ.
-        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        # Compute dQ (kT is raw; the caller applies sm_scale once at the end).
         dq += tl.dot(ds, tl.trans(kT))
         # Increment pointers.
     return dq
@@ -453,6 +458,7 @@ def _attn_bwd(
         do,
         m,
         D,  #
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -470,7 +476,7 @@ def _attn_bwd(
     )
     # Write back dQ.
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq *= LN2
+    dq *= sm_scale
     tl.store(dq_ptrs, dq)
 
 
@@ -591,6 +597,7 @@ def _attn_bwd_dq_kernel(
         Q,
         K,
         V,
+        sm_scale,
         DO,  #
         DQ,
         M,
@@ -663,6 +670,7 @@ def _attn_bwd_dq_kernel(
         do,
         m,
         D,
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -680,7 +688,7 @@ def _attn_bwd_dq_kernel(
     )
 
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
+    dq_acc *= sm_scale
     tl.store(dq_ptrs, dq_acc)
 
 
@@ -748,9 +756,11 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
     dv = torch.empty_like(v)
     BATCH, N_HEAD = q.shape[:2]
     BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
-    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+    # K stays raw: the backward kernels apply sm_scale in fp32 after the dot,
+    # matching the forward's rounding exactly. (A bf16 pre-scaled K perturbs
+    # the recomputed logits vs the saved M; exp2 turns that into unboundedly
+    # wrong probabilities at large activations.)
     arg_k = k
-    arg_k = arg_k * (sm_scale * RCP_LN2)
     PRE_BLOCK = 64
     assert Tq % PRE_BLOCK == 0
     pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
@@ -813,6 +823,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
         q,
         arg_k,
         v,
+        sm_scale,
         do,
         dq,
         M,

fastvideo-kernel/tests/test_vsa_triton_backward_scale.py

@@ -0,0 +1,104 @@
+"""Regression: Triton block-sparse backward gradient parity at realistic activation scale.
+
+The backward used to fold ``sm_scale / ln(2)`` into K in bf16 before the
+exp2-based logit recompute. The bf16 rounding error on the pre-scaled K grows
+proportionally to |logit| and exp2 amplifies it into exponentially wrong
+probabilities, so dQ/dK/dV were correct at unit scale (every pre-existing test)
+but off by orders of magnitude at real activation magnitudes.
+
+This test sweeps the input scale and checks the Triton kernel's gradients
+against an fp32 masked-dense SDPA reference. The unit-scale case is the
+control (it passed even with the broken kernel); the large-scale cases are
+the regression.
+"""
+
+import pytest
+import torch
+
+from fastvideo_kernel.block_sparse_attn import _map_to_index, block_sparse_attn_triton
+
+from .utils import generate_block_sparse_mask_for_function
+
+BLOCK = 64
+
+
+@pytest.fixture(autouse=True)
+def _seed_rng():
+    """Pin the RNG so these cases do not depend on what ran before them.
+
+    Same convention as test_vsa_varlen.py: every tensor here comes from the
+    global torch RNG and the checks use tight thresholds, so an unseeded run
+    would shift inputs whenever an earlier test file draws a different number
+    of randoms.
+    """
+    torch.manual_seed(0)
+    torch.cuda.manual_seed_all(0)
+
+
+def _dense_reference(q, k, v, block_mask):
+    """fp32 masked-dense SDPA over the token-expanded block mask.
+
+    q/k/v: [B, H, S, D]; block_mask: [B, H, S // BLOCK, S // BLOCK] bool.
+    """
+    qf, kf, vf = q.float(), k.float(), v.float()
+    token_mask = block_mask.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)
+    logits = torch.matmul(qf, kf.transpose(-2, -1)) * (q.shape[-1]**-0.5)
+    logits = logits.masked_fill(~token_mask, float("-inf"))
+    return torch.matmul(logits.softmax(dim=-1), vf)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("scale", [1.0, 4.0, 16.0])
+def test_triton_backward_grad_parity_across_input_scales(scale: float) -> None:
+    """Kernel dQ/dK/dV must stay within a few percent of the fp32 reference
+    regardless of input magnitude.
+
+    With the bf16 K pre-scaling bug, scale<=4.0 passes at this geometry while
+    scale=16.0 fails (measured on GB200: dq relative L2 error 5.9e-1 vs 6.9e-3
+    fixed); at larger geometries and real activation magnitudes the broken
+    kernel is off by orders of magnitude. The passing unit-scale case is
+    exactly how the bug survived the original test suite.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    batch, heads, dim = 1, 4, 128
+    num_blocks = 8
+    seq = num_blocks * BLOCK
+
+    q = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype) * scale
+    k = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype) * scale
+    v = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype)
+    grad_out = torch.randn_like(q)
+
+    block_mask = generate_block_sparse_mask_for_function(heads, num_blocks, num_blocks, k=3,
+                                                         device=device).unsqueeze(0)
+    q2k_idx, q2k_num = _map_to_index(block_mask)
+    variable_block_sizes = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device=device)
+
+    q_ker, k_ker, v_ker = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out_ker, _ = block_sparse_attn_triton(q_ker, k_ker, v_ker, q2k_idx, q2k_num, variable_block_sizes)
+    (out_ker.float() * grad_out.float()).sum().backward()
+
+    q_ref, k_ref, v_ref = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out_ref = _dense_reference(q_ref, k_ref, v_ref, block_mask)
+    (out_ref * grad_out.float()).sum().backward()
+
+    # Forward is exact at any scale; this pins the harness itself.
+    fwd_rel = ((out_ker.float() - out_ref).norm() / out_ref.norm()).item()
+    assert fwd_rel < 2e-2, f"scale={scale}: forward rel err {fwd_rel:.3e}"
+
+    for name, g_ker, g_ref in (
+        ("dq", q_ker.grad, q_ref.grad),
+        ("dk", k_ker.grad, k_ref.grad),
+        ("dv", v_ker.grad, v_ref.grad),
+    ):
+        assert torch.isfinite(g_ker).all().item(), f"scale={scale}: non-finite {name}"
+        ref_norm = g_ref.float().norm()
+        rel = ((g_ker.float() - g_ref.float()).norm() / ref_norm.clamp_min(1e-12)).item()
+        ratio = (g_ker.float().norm() / ref_norm.clamp_min(1e-12)).item()
+        print(f"scale={scale} {name}: rel_l2={rel:.4e} norm_ratio={ratio:.4f}")
+        assert rel < 5e-2, f"scale={scale}: {name} rel l2 err {rel:.3e} >= 5e-2"
+        assert 0.98 < ratio < 1.02, f"scale={scale}: {name} grad-norm ratio {ratio:.4f}"