diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py index 71159b43b..0f64fd114 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py +++ b/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, diff --git a/fastvideo-kernel/tests/test_vsa_triton_backward_scale.py b/fastvideo-kernel/tests/test_vsa_triton_backward_scale.py new file mode 100644 index 000000000..51493efde --- /dev/null +++ b/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}"