[kernel] sm_100a CUDA backward for VSA block-sparse attention (blk64)

[lpc0220]

Summary

Adds the native GB200 (sm_100a) backward that pairs with the sm_100a block-sparse forward: one
warp-specialized tcgen05 kernel per (batch, head, kv64 block) computing dK, dV and the dQ
partials (dS^T K reduce-added into an fp32 accumulator), plus a preprocess (Delta, Q^T, dO^T,
dqaccum zero, exact-zero dK/dV rows for kv blocks no q block selects) and a postprocess (dQ
unscramble + sm_scale). It consumes the forward's lse (Triton M format) and FastVideo's
invert_indices k2q metadata unchanged and returns (dq, dk, dv) in bf16 with the Triton
backward's scaling.

With FASTVIDEO_VSA_SM100A=1 the autograd backward of the sm_100a forward op now 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. Same
selection rules as the forward, no new environment variable.

Semantics (the points raised on the forward review)

• Per-row k2q counts may differ arbitrarily: one work item is one kv row, nothing is shared
  between rows. Count 0: every warp skips the item identically (no barrier traffic) and the
  preprocess writes exact-zero dK/dV rows.
• variable_block_sizes masks padded kv rows per lane inside the P^T select. Padded q rows rely
  on FastVideo's zero grad_o there, as the Triton backward does (no q-row mask).
• is_supported reads no tensor contents.
• Work-item order: from 1024 kv blocks per sequence on, a device kernel computes a length-binned
  order into a per-call workspace; below that the binding passes a cached torch.arange
  identity (one per device). No host sorting, no sync.

Scope

• sm_100a only: validated on GB200; the sm_103a image is deliberately not built and sm_103a
  devices keep the Triton backward. Enabling it is a one-line gencode + guard change once tested
  on B300/GB300 (separate PR).
• The H3 backend's grad gate ("kernel is forward-only; grad paths keep Triton") is unchanged;
  opening it means switching H3 to the need_lse=True op and is a separate PR.
• 128-token blocks: forward only, backward still raises (unchanged). (separate PR)

Tests (GB200, extension built from this branch with TORCH_CUDA_ARCH_LIST=10.0a)

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   (forward regression)

Performance (GB200, 25% density, D=128; TFLOPS = 2.5 * 4 * D * selected-block pairs / time)

Kernels only, same window (GPU 0 solo, interleaved Triton/blk64/blk64/Triton per size), B=1
H=8, bit-identical inputs, fp32 dQ accumulator (the shipped configuration), warmup 5 / 25 timed
iterations, medians of the two passes. Triton = FastVideo's Triton backward kernels
(tflops_kern); blk64 = this kernel (preprocess + main + postprocess):

S	Triton TFLOPS	sm_100a blk64 TFLOPS	Speedup
4096	198	381	1.93x
8192	313	575	1.84x
16384	389	715	1.84x
32768	433	794	1.84x
65536	448	821	1.83x
131072	452	736	1.63x
262144	452	784	1.73x
524288	437	771	1.76x

[mergify]

Merge Protections

🟠 1 of 1 protections blocking · waiting on 🤖 CI

	Protection	Waiting on
🟠	PR merge requirements	🤖 CI

🟠 PR merge requirements

Waiting for

• check-success=full-suite-passed

Waiting checks: full-suite-passed.

• check-success=full-suite-passed
• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[SolitaryThinker]

Correctness pass over the sm_100a backward. The math and the barrier/phase accounting look right to me (scaling matches the Triton backward, lse consumed correctly in M format, padded kv rows and the clamped tail quad contribute nothing, dqaccum scrambled layout is consistent between epilogue and postprocess, count-0 zeroing covers every row). Two issues below, one in the kernel and one in the binding.

[SolitaryThinker]

Unsynchronized TMEM write/read overlap between softmax warps of the same lane group

Each softmax warp loads its 64 fp32 columns of S^T from tmem_st + tmem_f32_offset, where tmem_f32_offset = tmem_lane_base + col_half * 64, and afterwards stores 32 bf16x2 columns of P^T to tmem_pt_bf16 + tmem_bf16x2_offset, where tmem_bf16x2_offset = tmem_lane_base + col_half * 32 (the store is at L653). Because tmem_pt_bf16 == tmem_st, the warp with col_half == 1 writes columns [32, 64) of its lane group, and those columns lie inside the fp32 range [0, 64) that the col_half == 0 warp of the same lane group is reading. For lane group 0 these are warps 4 and 8; there is no synchronization between warp 4's tcgen05.ld and warp 8's tcgen05.st. Both wait on full_bar_st and then proceed independently. If the col_half == 0 warp is delayed past the col_half == 1 warp's ld + wait::ld + exp2 loop + cvt + st sequence (a few hundred cycles), it reads bf16 P^T bit patterns where it expects fp32 S^T, and the corrupted values propagate into P^T, dS^T, dV, dK and the dQ partials for that work item.

The dS^T store at L681 has the same overlap: tmem_dst_bf16 == tmem_dpt, so the col_half == 1 warp's write to columns [32, 64) overlaps the col_half == 0 warp's fp32 dP^T load from columns [0, 64).

In practice both warps are released by the same mbarrier phase and the load is issued almost immediately, so the window is narrow and the tests pass. It is nevertheless a timing-dependent data race rather than a guaranteed ordering, and it is not covered by the existing barriers (full_bar_pt / full_bar_dst are arrived at only after the store). For comparison, the forward kernel does not have this problem because each warp's P overwrites only the S columns it has itself already loaded (p_tmem_addr = s_tmem_addr); the M=64 fold used here is what introduces the cross-warp overlap.

Suggested fix, which requires no additional synchronization: keep each warp's bf16 output inside the fp32 range it has already consumed.

const uint32_t tmem_bf16x2_offset = tmem_lane_base + (uint32_t)(col_half * HALF_COLS);

and adjust the A-operand addressing in gemm35_dv_dk (L399) so that slot p is read from column 64 * p rather than 32 * p:

const uint32_t tmem_a = tmem_a_base + (uint32_t)(p * HALF_COLS + ki * BF16X2_COLS_PER_K16);

As far as I can tell nothing else depends on the placement of the bf16 P^T / dS^T inside the 128-column st / dpt regions (the dQ GEMM overwrites tmem_dpt only after the dK GEMM has consumed dS^T, which is already ordered by MMA issue order), so these two edits should be sufficient. The alternative is a named barrier over the eight softmax warps between tcgen05_wait_ld() and the tcgen05_st, at the cost of two extra bar.sync per quad.

[SolitaryThinker]

Reproduced on GB200; the race is real and the column-offset fix closes it.

Method: I built the kernel from this PR's head (749a9c5) four ways with torch.utils.cpp_extension (nvcc 13.0, -gencode arch=compute_100a,code=sm_100a -DVSA_BHSD=true) and ran tests/test_block_sparse_bwd_sm100a.py against each build on a GB200 (hpc-rack-3-8, then hpc-rack-3-7), injecting the built module through block_sparse_attn_bwd_sm100a.set_extension.

• base: unmodified.
• delay: the only change is a ~1 ms __nanosleep loop in the col_half == 0 softmax warps between the full_bar_st / full_bar_dpt wait and their fp32 tcgen05.ld. Nothing about what is computed changes, only warp timing, so that the col_half == 1 warps' bf16 tcgen05.st lands before the fp32 load is issued.
• fixdelay: the same delay plus the two-line fix (bf16 offset col_half * HALF_COLS; ts-MMA A operand at p * SUB_COLS_BF16 + ki * BF16X2_COLS_PER_K16).
• fix: the fix alone.

Results (rel_max = max|got - ref| / max|ref| against the test's fp32 reference, case num_blocks=8, topk=4, heads=4, ragged=True):

variant	PR test suite	dq rel_max	dk rel_max	dv rel_max
base	15 passed	4.5e-3	3.1e-3	5.3e-3
delay	13 failed, 2 passed	9.5e-1	9.4e-1	6.3e-1
fixdelay	15 passed	4.5e-3	3.1e-3	5.3e-3
fix	15 passed	4.5e-3	3.1e-3	5.3e-3

The two tests that still pass under delay are the ones that check no numerics (test_invert_indices_torch_matches_q2k, test_unsupported_is_rejected). The corrupted values are finite rather than NaN, which is what one expects from bf16 P^T bit patterns being read as fp32 S^T. The fixed builds produce error statistics identical to base, and their dk/dv are bitwise equal to base on fixed inputs, so relocating the bf16 tiles does not change the arithmetic.

Under natural timing the window did not open: 6000 back-to-back runs of the unmodified kernel on fixed inputs gave zero dk/dv mismatches. So the current layout works today by scheduling margin rather than by ordering; the delay experiment shows the hardware itself provides no ordering between the two warps' TMEM accesses.

Performance of the fix: kernel-only timing (preprocess + main + postprocess), B=1, H=8, D=128, 25% density, medians of two interleaved passes of 25 iterations:

S	base ms	fix ms	fix / base
4096	0.142	0.140	0.991
8192	0.333	0.334	1.006
16384	1.000	0.996	0.996
32768	3.635	3.675	1.011
65536	14.586	14.516	0.995
131072	63.637	63.528	0.998

Neutral within noise, as expected: the fix changes only which TMEM columns hold the bf16 P^T and dS^T tiles, with identical instruction count, TMEM/SMEM traffic, MMA shapes and barrier structure. The alternative fix (a named barrier across the eight softmax warps between tcgen05_wait_ld() and the store) would put two extra bar.sync per quad on the softmax path, so the offset change is the one I would take.

One correction to my comment above: HALF_COLS is a local constexpr of the softmax branch and is not in scope inside gemm35_dv_dk; the A-operand line needs the namespace-scope constant of the same value, i.e.

const uint32_t tmem_a = tmem_a_base + (uint32_t)(p * SUB_COLS_BF16 + ki * BF16X2_COLS_PER_K16);

(the compile fails with identifier "HALF_COLS" is undefined otherwise). The tmem_bf16x2_offset line is inside the softmax branch and can keep HALF_COLS.

[lpc0220]

good catch. I fixed the issue on blk128, but missed blk64. made 3 commits. please test the latest version. 88319b7. The latest test shows perf stays the same, no change. Sorry, 88319b7 is not the latest commit. It is the commit that fixed the issue. Please test the latest, which has 2 more commits to remove the identity cache. Btw, ST_QBLOCK_COLS is the new "HALF_COLS" in the namespace scope.

[SolitaryThinker]

identity_remap cache is neither stream-safe nor thread-safe

The cached identity array is created with torch::arange on whatever stream is current at the first call for a given device and stored in a static per-device vector. Two issues follow:

1. When a call arrives with a larger items than the cached tensor, the old tensor is dropped and its memory is returned to the caching allocator, while a previously launched backward on a different stream may still be reading it through workitem_remap. The replacement arange kernel is likewise only ordered with respect to the allocating stream, so a launch on another stream can read the buffer before it has been filled. Autograd normally replays the backward on the forward's stream, so this is unlikely to bite in FastVideo's current use, but it is a latent cross-stream use-after-free / read-before-write.

2. The static std::vector<torch::Tensor> is read and written without a lock, so concurrent calls from multiple host threads (for example two backward passes on different devices driven from different threads) race on the cache.

Options: allocate the identity array per call (torch::arange(num_items, ...) on the current stream; below ORDER_MIN_KV_BLOCKS the array is at most a few hundred KB and the launch cost is negligible next to the preprocess kernel), or keep the cache but call c10::cuda::CUDACachingAllocator::recordStream on the tensor for the current stream and guard the vector with a mutex.

[lpc0220]

This is a legacy issue of my previous version. I removed this identity cache and remap host call completely. Instead, we will pass nullptr as a signal for the identity assignment on the kernel side. Code is cleaner.

[alexzms]

Request changes. Two things.

1. Cannot reproduce "validated on GB200" on hpc-rack-1-12 (GB200, 152 SMs, driver 580.126.20, VBIOS 97.00.B9.00.99, torch 2.12.1+cu130). tests/test_block_sparse_bwd_sm100a.py fails 15/15 with cudaErrorIllegalAddress, deterministically, on both the PR head and 749a9c5, with nvcc 13.0.88 (same as Will's run) and 13.1.115, CLC on or off, wait::ld dropped or restored. The merged forward passes and is compute-sanitizer-clean in the same build on the same node, so it is not the toolchain or the harness. compute-sanitizer with -lineinfo:

Out-of-range shared or local address
  at tcgen05_mma_ws_f16_ss_1sm_predicated(...)  primitives.cuh:890   (gemm12_st_dpt)
  by thread (384,0,0)                                                (warp 12 = W_MMA)

All-zero k2q counts run fine and return zeros; the smallest non-trivial case (B=H=1, nb=2, one item with count 1) faults on the first .ws GEMM. Repro below. Please run this on a Blackwell node other than 3-7/3-8 before merging — CI has no Blackwell and does not collect top-level tests/, so nothing else will catch it.

2. d2c9d612 (drop tcgen05.wait::ld) should be reverted, at minimum for the dQ epilogue. It landed after the last review pass. In SASS the register consumers do scoreboard-wait on the LDTM barriers, so that part of the commit message holds — but the mbarrier.arrive on empty_bar_dq does not:

LDTM.x64 R68 <- tmem_dq      wr=B4
LDTM.x64 R4  <- tmem_dq      wr=B1
@P0 SYNCS.ARRIVE [empty_bar_dq]   wait={}     <- releases the MMA warp to overwrite tmem_dq
IMAD.SHL ...                      wait={B4}   <- first wait on the loads, after the arrive

tcgen05.fence::before_thread_sync emits no instruction there. With the wait restored the arrive sits behind NOP wait={B1}. That is the same "works by scheduling margin" class as the TMEM race from the earlier review; CUTLASS and FA4 keep wait::ld between a TMEM load and the arrive for exactly this reason ("Need a fence bw TMEM_LOAD and arrive"), and your own A/B puts the cost at noise. tcgen05_wait_ld() in primitives.cuh is dead now, and books/ is not in the repo.

Minor: tests/jit_ext.py referenced in the binding docstring does not exist; work_remap + base on a null remap is pointer UB that only stays safe because ORDER_MIN_KV_BLOCKS*64 < CACHE_WAVE_MIN_SEQ_LEN — worth a static_assert; the nb >= 1024 order-kernel path has no unit test.

The two fixes from the earlier review (TMEM column offsets, nullptr remap) look right.

Repro

Standalone build, same as Will's (TORCH_CUDA_ARCH_LIST=10.0a):

from torch.utils.cpp_extension import load
m = load(name="vsa_bwd", sources=["jit_bind.cpp", "csrc/attention/block_sparse_bwd_sm100a.cu"],
         extra_include_paths=["csrc/attention"], extra_ldflags=["-lcuda"],
         extra_cuda_cflags=["-O3", "-std=c++20", "--use_fast_math", "--expt-extended-lambda",
                            "--expt-relaxed-constexpr", "-DVSA_BHSD=true"])
# jit_bind.cpp: declare block_sparse_sm100a_bwd and m.def() it.
from fastvideo_kernel import block_sparse_attn_bwd_sm100a as bwd
bwd.set_extension(m)

Minimal case (CUDA_LAUNCH_BLOCKING=1 compute-sanitizer --tool memcheck python repro.py):

import torch
B, H, D, nb = 1, 1, 128, 2; S = nb * 64
q, k, v, go, o = (torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16) for _ in range(5))
lse = torch.full((B, H, S), 8.0, device="cuda")
k2q = torch.zeros(B, H, nb, nb, dtype=torch.int32, device="cuda")
k2n = torch.tensor([1, 0], dtype=torch.int32, device="cuda").view(B, H, nb)   # [0, 0] runs and returns zeros
vbs = torch.full((nb,), 64, dtype=torch.int32, device="cuda")
dq, dk, dv = bwd.block_sparse_attn_backward_sm100a_from_k2q(go, q, k, v, o, lse, k2q, k2n, vbs)
torch.cuda.synchronize()

[SolitaryThinker]

Follow-up on @alexzms's review, with the SASS evidence for point 2, a cross-rack attempt at point 1, and a branch carrying the fixes.

Point 2 (d2c9d61, drop tcgen05.wait::ld): confirmed, and it is not just an ISA-compliance nit

I decoded the Volta+ control bits (wait mask / write barrier) of my own builds of d2c9d61 with and without the four waits. ptxas does scoreboard every register consumer of an LDTM, so three of the four sites are safe in practice. The dQ epilogue is the exception, because the mbarrier.arrive on empty_bar_dq is not a register consumer:

d2c9d61 (no wait::ld), dQ epilogue:
  [1844] LDTM.x64 R68, tmem[UR8+0x100]                wr-barrier=B3
  [1846] LDTM.x64 R4,  tmem[UR9+0x140]                wr-barrier=B1
  [1852] @P0 SYNCS.ARRIVE.TRANS64.A1T0 [empty_bar_dq]  wait={}      <-- releases the MMA warp
  [1858] STS.128 [R132+UR24+0x30000], R68             wait={B3}    <-- first wait on the loads

same site with tcgen05.wait::ld restored:
  [1845] LDTM.x64 R4,  tmem[UR9+0x140]                wr-barrier=B1
  [1847] NOP                                          wait={B1}    <-- the wait::ld
  [1853] @P0 SYNCS.ARRIVE.TRANS64.A1T0 [empty_bar_dq]  wait={}

tcgen05.fence::before_thread_sync emits nothing there. So without the wait, the epilogue hands tmem_dq back to the MMA warp while its two loads may still be reading it. The MMA warp is blocked on exactly that barrier, and its next GEMM (the second dQ half, or the next quad's dP^T, since tmem_dq aliases tmem_dpt) writes the same columns. If the tensor-core writes land before the in-flight loads complete, the epilogue's registers hold a mix of old dQ partial and new data, and that goes into dqaccum as a valid partial: silent, finite, nondeterministic dQ corruption, invisible to the test suite unless it happens to fire. dK/dV are unaffected because their path consumes the loaded registers before any arrive. It is the same "works by scheduling margin" class as the TMEM race from the first review, which also looked clean in thousands of small runs and then corrupted every rerun from 32K tokens.

The A/B on GB200 (25 iterations, medians, interleaved) shows no cost for keeping the waits: 4K 1.009x, 16K 1.007x, 64K 1.004x in favour of the build with waits, i.e. noise. Outputs are bitwise identical where the race does not fire. There is no upside to the removal.

Repro of the decode

# decode_ctrl.py <cuobjdump> <object.o>: for each LDTM, walk forward to the first instruction
# that waits on its write barrier and flag any SYNCS.ARRIVE that comes first.
import re, sys, subprocess
out = subprocess.run([sys.argv[1], "-sass", sys.argv[2]], capture_output=True, text=True).stdout.splitlines()
ins, i = [], 0
while i < len(out):
    m = re.match(r"\s+/\*[0-9a-f]+\*/\s+(.*?)\s*;\s*/\*\s*(0x[0-9a-f]+)\s*\*/", out[i])
    m2 = re.search(r"/\*\s*(0x[0-9a-f]+)\s*\*/", out[i + 1]) if m and i + 1 < len(out) else None
    if m and m2: ins.append((m.group(1), int(m2.group(1), 16))); i += 2
    else: i += 1
ctrl = lambda hi: dict(wb=(hi >> 46) & 7, wait=(hi >> 52) & 0x3F)     # bits 105.. of the 128-bit word
for k, (txt, hi) in enumerate(ins):
    if "LDTM" not in txt: continue
    wb = ctrl(hi)["wb"]; print(f"\n[{k}] {txt}  wr-barrier=B{wb}")
    for j in range(k + 1, k + 40):
        t, h = ins[j]; w = ctrl(h)["wait"]
        if "ARRIVE" in t and not (w >> wb & 1): print(f"    [{j}] {t}  wait={{}}  <-- arrive before wait")
        if w >> wb & 1: print(f"    [{j}] {t}  wait={{B{wb}}}"); break

Build the .cu standalone with torch.utils.cpp_extension.load (explicit -gencode arch=compute_100a,code=sm_100a -DVSA_BHSD=true), then python decode_ctrl.py $(python -c "import triton,os;print(os.path.dirname(triton.__file__)+'/backends/nvidia/bin/cuobjdump')") build/block_sparse_bwd_sm100a.cuda.o.

Point 1 (illegal address on hpc-rack-1-12): could not reproduce on four other GB200s

I cannot reach rack 1. On hpc-rack-2-0, 3-6, 3-7 and 3-8 (GB200, 152 SMs, driver 580.82.07, VBIOS 97.00.B9.00.99, torch 2.12.0+cu130, nvcc 13.0.88) the minimal case from the review runs clean and tests/test_block_sparse_bwd_sm100a.py passes 15/15 in fresh processes for every combination I tried: d2c9d61, d2c9d61 with the waits restored, and 749a9c5, each with my flags and with the exact flag set from the review (-O3 -std=c++20 --use_fast_math --expt-extended-lambda --expt-relaxed-constexpr). The one variable I cannot match is the driver (580.126.20). Someone with that driver needs to run it; it stays open from my side.

One harness pitfall worth knowing, because it produces "fails deterministically on every commit" symptoms: torch.utils.cpp_extension._get_cuda_arch_flags silently drops all -gencode flags when any string in extra_cuda_cflags contains the substring arch (my -ccbin .../aarch64-.../g++ triggered it). nvcc then compiles for its default arch, the __CUDA_ARCH_FEAT_SM100_ALL guard empties every kernel, the module still loads, and the suite fails 13/15 with uninitialized-output garbage. cuobjdump -res-usage on the object should show REG:128 ... SHARED:1024 for the main kernel; REG:4 SHARED:0 means the guard ate the body. The pasted recipe does not have that flag, so this probably is not the rack-1 fault, but it is a two-second check.

Branch with the fixes: vsa_bwd_fix (on top of d2c9d61)

• 60165c7 [fix]: restore tcgen05.wait::ld after the TMEM loads (the four sites; comment at the dQ arrive explains why).
• 62f0bc1 [misc]: pin the identity-order/chunked-launch invariant; fix set_extension docstring (a static_assert tying CACHE_WAVE_MIN_SEQ_LEN / sizeof(dq_accum_t) / ORDER_MIN_KV_BLOCKS together plus a runtime guard on the null-remap chunk path; drops the tests/jit_ext.py reference).
• 2d58c9b [test]: sm_100a backward parity and run-to-run stability at 65536 tokens in test_block_sparse_sm100a_dispatch.py: 1024 kv blocks (covers the device-computed order path) at 12.5% density, sm_100a route vs all-Triton within the file's tolerances, and two sm_100a runs within summation-order noise. Bitwise equality is not available through from_indices: invert_indices compacts with atomics, so each kv row's q order differs per call (measured delta 5e-3 rel_max / 7e-6 mean; the TMEM race gave 0.6-1.0 / 4e-2).

Verified on hpc-rack-2-0 with a standalone build of forward (64/128) + backward from the branch: test_block_sparse_bwd_sm100a.py 15 passed, test_block_sparse_sm100a_dispatch.py 11 passed including the new case; the decode above shows the arrive now behind NOP wait={B1} at both dQ sites. Compare view: d2c9d61...vsa_bwd_fix

[lpc0220]

@SolitaryThinker @alexzms Cherry-picked the above 3 commits of vsa_bwd_fix branch. Added a test to cover nb=2 for "Cannot reproduce "validated on GB200" on hpc-rack-1-12". For "Point 1 (illegal address on hpc-rack-1-12): could not reproduce on four other GB200s", i ran your repro code on our racks. It passed. Give you our config:

	rack-1-12 (fails)	ptyche0077 / our GB200 nodes (pass)
driver / GSP	580.126.20	580.178.04
VBIOS	97.00.B9.00.99	97.00.F8.00.1F
torch	2.12.1+cu130	2.13.0+cu130
nvcc	13.0.88 / 13.1.115	13.0 (r13.0)
SASS	identical (sm_100a, offline)	identical
forward (tcgen05.mma)	passes	passes
backward (tcgen05.mma.ws)	cudaErrorIllegalAddress in the first .ws GEMM, all 15 tests	16 tests pass; your nb=2 repro passes; compute-sanitizer memcheck 0 errors (nb=2, nb=8, full suite)

I might need more info re: the error. Can you verify the code again?

Thanks for your time!