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
Reverts the effect of d2c9d61. tcgen05.ld is asynchronous and PTX requires
tcgen05.wait::ld before its destination registers are used. ptxas does
scoreboard the register consumers, so three of the four sites were safe in
practice, but the dQ epilogue arrives on empty_bar_dq right after its two
loads and before any consumer: in SASS the SYNCS.ARRIVE carries an empty
wait mask and precedes the first wait on the loads' barriers (B3/B1). That
hands tmem_dq to the MMA warp while the loads may still be reading it; the
next GEMM overwrites the same columns (tmem_dq aliases tmem_dpt), and the
mixed registers are reduce-added into dqaccum as a valid partial. With the
wait restored ptxas emits `NOP wait={B1}` ahead of the arrive.
Measured cost on GB200: none (0.4-0.9% in favour of the build with waits,
within noise); outputs bitwise identical where the race does not fire.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01G4KbxfFobc4FUXLZV3H4QF
The destination registers of tcgen05.ld are scoreboarded, so the consumers wait on their own; the
four waits (dV/dK merge, S^T, dP^T and dQ loads) are removed, the fence stays. Same change as the
reference kernel in books/. Verified: CPU-reference checks incl. ragged variable_block_sizes,
zero-count kv blocks and STRESS_N=5 bitwise reruns; tests/test_block_sparse_bwd_sm100a.py 15
passed, tests/test_block_sparse_sm100a_dispatch.py 10 passed, tests/test_block_sparse_sm100a.py
37 passed (extension rebuilt from this tree). Interleaved A/B 4k..524k: -1.2..+0.9%, neutral.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Review finding: the binding's per-device static cache of a torch::arange identity array was
neither stream-safe (a larger call replaced the tensor while an earlier backward on another
stream could still read it; the replacement arange was only ordered with its own stream) nor
thread-safe (unguarded static std::vector).
Fix: no identity array at all. decode_workitem reads
`real_item_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id`; the launch leaves
workitem_remap nullptr below ORDER_MIN_KV_BLOCKS (one launch of every item, so the work id is
the item id) and runs the order kernel into order_workspace from there on. Slice launches
(>= 524288 tokens) always sit in the order-kernel regime and keep passing `work_remap + base`.
block_sparse_bwd_supported requires the workspace only when the order kernel will run. The
binding passes nullptr always and allocates the workspace from ORDER_MIN_KV_BLOCKS on, as
before; identity_remap and the cache are gone.
Verified: CPU-reference checks incl. odd top-k, ragged variable_block_sizes and zero-count kv
blocks on the nullptr path; device order == host sort check at 1024 kv blocks;
tests/test_block_sparse_bwd_sm100a.py 15 passed, tests/test_block_sparse_sm100a_dispatch.py
10 passed, tests/test_block_sparse_sm100a.py 37 passed (extension rebuilt from this tree).
Interleaved A/B vs the previous build: neutral (4k..262k within +-1.2%).
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Review finding: the two softmax warps of a TMEM lane group split the 128-column S^T (and
dP^T) tile into fp32 halves, but the col_half == 1 warp packed its bf16x2 P^T / dS^T into
columns 32-63, inside the fp32 range the col_half == 0 warp was still loading. Nothing but
the shared full_bar_st / full_bar_dpt wait ordered the two, so a delayed col_half == 0 warp
could read bf16 pairs where it expected fp32 scores.
Fix, no extra synchronization: a warp's bf16x2 overlay now starts at the same column as the
64 fp32 columns it loads (ST_QBLOCK_COLS = 64 per q64 block), so it only overwrites data it
has itself consumed; the dV / dK TS MMAs read slot p's bf16 atoms from column
p * ST_QBLOCK_COLS instead of 32 * p. The dQ GEMM already overwrote the dP^T tile only after
the dK GEMM consumed dS^T (MMA issue order), unchanged.
Verified: CPU-reference checks incl. odd top-k, ragged variable_block_sizes and zero-count kv
blocks; tests/test_block_sparse_bwd_sm100a.py 15 passed, tests/test_block_sparse_sm100a_dispatch.py
10 passed, tests/test_block_sparse_sm100a.py 37 passed (extension rebuilt from this tree).
Interleaved A/B vs the previous build: neutral (4k..131k within noise).
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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>
