The identity work order (nullptr remap) cannot be offset per chunk, so the
DQ_L2_KEEP chunked launches must only run with a device-computed order.
Today that holds because the L2 transition (256K tokens) lies above the
order-kernel threshold (64K), but nothing tied the two constants together;
a static_assert now does, plus a runtime guard for external callers of the
launch API. Also drops the reference to a tests/jit_ext.py that does not
exist in the repo.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01G4KbxfFobc4FUXLZV3H4QF
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>
