From 2d58c9bfde17bca37e793fe265a159c81f7ba011 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 6 Sep 2026 19:59:19 +0000 Subject: [PATCH] [test]: sm_100a backward parity and run-to-run stability at 65536 tokens The existing backward tests stop at 16 kv blocks. The TMEM race fixed in 88319b7 never fired there but corrupted dk/dv on every rerun from 32K tokens, and the nb >= 1024 device-computed work order had no test at all. This case (1024 kv blocks, 12.5% density) covers both: sm_100a route vs all-Triton within the existing tolerances, and two sm_100a runs within summation-order noise (invert_indices orders each kv row's q list with atomics, so exact bitwise equality is not available through this path; measured delta 5e-3 rel_max / 7e-6 mean, the race gave 0.6-1.0 / 4e-2). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01G4KbxfFobc4FUXLZV3H4QF --- tests/test_block_sparse_sm100a_dispatch.py | 47 ++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py index f0a276553..9bfdc719c 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -181,6 +181,53 @@ def test_backward_uses_sm100a_kernel_when_built(monkeypatch): assert sm100a_backward +def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): + """65536 tokens (1024 kv blocks, 12.5% density): the device-computed work order (nb >= 1024) + and the sequence regime where a TMEM ordering bug corrupted dk/dv while every <= 16-block + test stayed green. sm_100a route vs all-Triton, and two sm_100a runs must agree to within + summation-order noise: invert_indices compacts each kv row's q list with atomics, so the + kernel sums quads in a different order per call (measured run-to-run delta on GB200: + rel_max 5e-3, mean|diff| 7e-6). The TMEM race this guards against gave rel_max 0.6-1.0 + and mean|diff| 4e-2, an order of magnitude past the bounds below on both metrics.""" + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + + torch.manual_seed(0) + block, num_blocks, heads, topk = 64, 1024, 4, 128 + S = num_blocks * block + shape = (1, heads, S, HEAD_DIM) if vsa.BHSD else (1, S, heads, HEAD_DIM) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + scores = torch.rand(1, heads, num_blocks, num_blocks, device="cuda") + idx = scores.topk(topk, dim=-1).indices.sort(dim=-1).values.to(torch.int32) + del scores + num = torch.full((1, heads, num_blocks), topk, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda") + if not vsa_bwd.is_supported(q, vbs): + pytest.skip("sm_100a backward not built for this device") + + def grads(route): + qq, kk, vv = (t.clone().requires_grad_(True) for t in (q, k, v)) + if route == "triton": + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + else: + monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False) + monkeypatch.setenv(ENV, "1") + out, _ = block_sparse_attn_from_indices(qq, kk, vv, idx, num, vbs) + out.float().square().sum().backward() + return [t.grad.float() for t in (qq, kk, vv)] + + ref = grads("triton") + got = grads("sm100a") + again = grads("sm100a") + _assert_grads_close(got, ref) + for g1, g2, name in zip(got, again, "qkv"): + diff = (g1 - g2).abs() + rel_max = diff.max().item() / max(g1.abs().max().item(), 1e-6) + mean_abs = diff.mean().item() + assert rel_max <= 2e-2 and mean_abs <= 1e-4, \ + f"d{name}: two sm_100a runs differ beyond summation-order noise: " \ + f"max|diff|/max|ref|={rel_max:.3e} mean|diff|={mean_abs:.3e}" + + def test_blk128_backward_raises(monkeypatch): """128-token blocks: forward runs, backward refuses (both backwards are 64-block only).""" monkeypatch.setenv(ENV, "1")