Native GB200 (sm_100a) backward pairing with the sm_100a forward: one
warp-specialized tcgen05 kernel per (batch, head, kv64 block) computing dK, dV
and the dQ partials (fp32 accumulator, cp.reduce.async.bulk), plus a preprocess
(Delta, Q^T, dO^T, dqaccum zero, exact-zero dK/dV rows for unselected kv
blocks) and a postprocess (dQ unscramble + sm_scale). Consumes the forward's
lse (Triton M format) and invert_indices' k2q metadata unchanged; returns
(dq, dk, dv) in bf16 with the Triton backward's scaling.
Wiring: with FASTVIDEO_VSA_SM100A=1 the autograd backward of the sm_100a
forward op runs this kernel when block_sparse_attn_bwd_sm100a.is_supported
passes (bf16, head_dim 128, 64-token blocks, even block count, sm_100a device,
op built) and the Triton backward otherwise. No new environment variable.
primitives.cuh gains only the backward's primitives (additions only; the
forward's definitions and SASS are untouched). Built for sm_100a only: the
kernel is validated on GB200, not yet on B300/GB300, so sm_103a devices keep
the Triton backward. The H3 backend's grad gate is unchanged (separate PR).
Semantics: per-row k2q counts may differ arbitrarily (one work item is one kv
row; count 0 is skipped identically by every warp and gets zero dK/dV rows);
variable_block_sizes masks padded kv rows per lane; is_supported reads no
tensor contents. Work-item order: a device kernel sorts by list length from
1024 kv blocks per sequence on; below that a cached identity array is passed.
Tests (GB200, extension built from this branch): tests/test_block_sparse_bwd_sm100a.py
15 passed (fp32 masked-dense autograd reference; ragged vbs, zero-count kv
blocks, top-k 1/2/3/5/7, batch 2), tests/test_block_sparse_sm100a_dispatch.py
10 passed (sm_100a route vs all-Triton grads, ragged vbs, Triton backward
monkeypatched to fail on a supported input), tests/test_block_sparse_sm100a.py
37 passed (forward regression).
Perf (kernels only, same window, B=1 H=8 D=128, 25% density, fp32 accumulator,
TFLOPS on selected blocks) Triton backward vs this kernel:
4k 198/381 (1.93x), 8k 313/575 (1.84x), 16k 389/715 (1.84x), 32k 433/794
(1.84x), 65k 448/821 (1.83x), 131k 452/736 (1.63x), 262k 452/784 (1.73x),
524k 437/771 (1.76x). FastVideo autograd path incl. invert_indices: 1.6-1.7x
at Wan 480P/720P-like shapes.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
tests/local_tests/ holds local-only parity, smoke, and component tests
organized by model family (e.g. tests/local_tests/ltx2/,
tests/local_tests/stable_audio/). They typically require a checked-out
reference repo and/or model weights and are skipped in CI; see each family's
README.md for setup. The top-level
tests/local_tests/README.md is the family index.
The CI-backed test suite still lives in fastvideo/tests/.
