DSA: fix dK SMEM handoff race and two latent NaN sources in the dense indexer

[ForFishes]

What

Four fixes in the dense DSA warmup path (dense_indexer_score_recompute + dense_attn_score_recompute + dense_indexer_backward). Found while chasing global_grad_norm: nan in a DSA warmup run where the forward loss stayed finite but the indexer gradients blew up.

1. dense_indexer_backward_sm100.py — SMEM handoff race on dK (the main one)

In _reduce_warpgroup_2q, all 128 lanes of the reduce warpgroup scatter their TMEM readback into sdK_reduce, but the cp.reduce.async.bulk that ships the tile is issued by a single thread and reads the whole tile. The only thing between them was fence_proxy("async.shared", space="cta"), which is a proxy fence: it orders the executing thread's own prior generic-proxy SMEM writes against the async proxy, and neither waits for nor publishes any other lane's stores. There is no happens-before edge between lanes 1..127's stores and thread 0's DMA read. sdK_reduce is single-buffered and never zeroed, so a read that overtakes the stores picks up either the previous KV block's dK or uninitialised SMEM.

A second, separate race sits in the same loop: the cp_async_bulk_wait_group that is supposed to stop the next iteration from overwriting the staging tile was executed by all 128 lanes, but bulk groups are per-thread and only thread 0 ever commits one — for the other 127 lanes it is a no-op.

Both are fixed with a warpgroup NamedBarrier (id 5; 0/3/4 are taken by sync_threads, compute_sync_barrier and tmem_alloc_barrier), following the arrive_and_wait / fence_proxy / arrive_and_wait pattern that the dQ epilogue in the same file already uses — including the comment at the dQ site explaining why the second barrier is needed before the staging tile is reused. dQ got it right; dK was missed.

2. _interface_sm100.py — denom_out was never initialised

Allocated with torch.empty and, unlike out, never filled. The tile scheduler only visits q rows below cu_seqlens_q[-1], so in THD every row in [cu_seqlens_q[-1], total_q) comes back as whatever the allocator handed over. Zeroed next to out.fill_, which also covers a caller-supplied buffer. 0 rather than -inf: consumers compute exp(out - denom) and out is already -inf on those rows, so -inf - (-inf) would be a NaN.

Worth noting separately: the varlen validation never asserts cu_seqlens_q[-1] == total_q. It is a caller-side invariant that nothing checks.

3. dense_score_recompute_sm100.py — attn exp2 had no upper clamp

Neither branch bounded the exponent (the qhpkv != 64 branch only clamped the lower end against denormals). A finite but far too small lse — e.g. the additive -1e30 mask constant a caller's own LSE pass can leave on a row whose candidate set degenerated — overflows fp32, giving out = +inf and denom = +inf, so the consumer's score / denom is a silent NaN. Clamped to 120, chosen so the head sum stays finite (2**120 * 64 = 8.5e37 < FLT_MAX).

This bounds each term, not the column sum, so a degenerate lse can still saturate the L1 norm to +inf. That is benign — the consumer's target becomes finite/inf = 0 and the row drops out of the loss — but it is a mitigation, not a guarantee of a correct value.

4. indexer_backward/api.py — reject unsupported head_dim

check_support only asserted heads >= 64. Measured: head_dim 64 and 128 are correct; 96 / 100 / 112 return silently wrong d_index_k (relative error up to 4e2, and head_dim=100 also corrupts d_index_q with NaN and 1e38 values); 192 / 256 abort with cudaErrorInvalidValue. The dK epilogue's bulk-reduce size assumes the global row stride equals head_dim_padded, and the TMEM load atom only tiles the staging buffer completely at the widths it was tuned for. Rejecting is better than returning bad gradients.

How it was verified

For a sum of N terms accumulated in fp32 in any order the error is bounded by (N-1) * 2**-24 * sum|t_i|. The launch grid is (ceil(s_q/2), batch), so thousands of CTAs reduce into the same global dK rows and d_index_k is inherently non-deterministic even with a correct kernel — "only dK is non-deterministic" proves nothing on its own. Run-to-run differences above that bound cannot be explained by re-association, and that is the criterion used here.

Five THD shapes ([256], [512,512], [300,100,612], [1000,24], [1024,1024] with q_causal_offsets), byte-identical inputs, 8 repeats:

shape	spread before	after	elements over bound, before	after
[256]	3.82e-06	9.54e-07	0 / 32768	0
[512,512]	9.16e-05	1.91e-06	711 / 131072	0
[300,100,612]	1.27e-04	3.82e-06	756 / 129536	0
[1000,24]	1.14e-05	1.91e-06	0 / 131072	0
[1024,1024]	3.62e-05	1.91e-06	90 / 262144	0

Before the fix, up to 22.6x over the loosest possible bound. After, zero elements exceed it on any shape and the residual spread sits at 0.08..0.29 of the bound — back inside what the cross-CTA reduction order alone explains. Also holds under a concurrent GPU load used to perturb warp scheduling.

Other checks:

• The patch is surgical. Fully patched vs unpatched: index_score, index_lse, attn_score, attn_l1norm, d_index_q, d_weights are all bit-identical; only d_index_k changes. d_index_q / d_weights do not go through the staging buffer, and stay bit-exact over 20 repeats before and after.
• Both forward changes are bit-identical to the previous behaviour on legitimate inputs — the clamp never fires when lse is a real LSE, and the zeroing does not touch a covered row.
• Accuracy: d_index_k vs an fp32 reference goes from 2.0e-2..5.4e-2 of peak (failing) to 2.75e-3, i.e. the same level as d_index_q / d_weights, which is the bf16 dS staging floor.
• Defects 2 and 3 each have a reproducer that fails before and passes after: the denom_out probe recovers leaked values from the previous call (including an exact 64.0 — a previous attn_l1norm — and fp32 denormals around 1e-43); the degenerate-lse probe goes from 37888 +inf score elements and 37888 NaN targets to zero of both.
• Forward ops are bit-deterministic over 60 repeats on all five shapes.
• Cost: +1.9% on the backward, +1.7% on the forward, measured at s_local=8192 (64 KV blocks per CTA), alternating runs of 30 iterations.
• barrier_id=5 confirmed free: cutlass only constructs NamedBarrier in tmem_allocator (which receives id 4 from this kernel) and in an unrelated gemm template; PipelineTmaStore uses no named barrier; this kernel has no CLC scheduler.

Caveats

• Verified on NVIDIA B30Z, cc 10.3. The manifestation of an SMEM race depends on warp scheduling and SASS codegen, so the quantitative numbers should be discounted across architectures. That a formal data race exists is a static fact and arch-independent, and every before/after comparison above was taken on the same machine.
• This proves the race is real and fixed. It does not prove it was the cause of the training-side NaN that started the investigation. Closing that loop needs a training rerun.

---

Note on this branch (paddle/v1.27.0)

The same change was developed and verified on paddle/v1.26.0; this is the port.

Three of the four files applied with offsets only; _interface_sm100.py was re-derived by hand because it differs by 609/150 lines between the two branches, and the attn hunk carries an extra comment explaining why the existing denom_out.zero_() at allocation time is not sufficient (it is gated behind precision == "mxfp8" and a shape check, so the bf16 path never reaches it).

All three defects were confirmed to still be present on this branch before patching. Stripping comments, the added executable lines are identical between the two ports (23 lines each), and the removed lines match as well — so the verification above transfers, modulo the surrounding-code differences.

Not functionally verified on this branch. The environment used for testing has nvidia-cutlass-dsl 4.4.1, and paddle/v1.27.0 needs a newer one (sparse_attention_backward/dsa_bwd_sm100.py imports cutlass.cute.nvgpu.OperandMajorMode, which 4.4.1 does not provide; the DSA package __init__ loads all symbols eagerly, so this blocks the whole namespace). What was checked here: the hunks land in the correct context, all four files compile, and the executable diff is identical to the verified v1.26.0 port. A run of the test suite on a machine with the required cutlass-dsl is still owed.