[CuTe,Bwd,Sm100] Support head_dim not a multiple of 32 in backward

[KaijingOfficial]

Summary

Backward on SM100 (B200/B300) crashes or silently produces wrong dK/dV when head_dim % 32 != 0 (e.g. ViT head_dim=72, 104). Fixes #2492. Three independent root causes:

1. flash_bwd_preprocess.py — the O/dO load predicate (predicate_k) carries a CPY_M mode broadcast over m, but the copy source is sliced per-m (tOgO[None, m, None]) without slicing the predicate. The vectorized 128-bit copy atom's predicate-shape verifier then rejects the extra mode (expects (1,(k)) but got (1,(cpy_m,k))). Slice the predicate per-m, as flash_fwd.py / flash_bwd.py already do.

2. flash_bwd_sm100.py — tile_hdim padded head_dim to a multiple of 16, but the dQ accumulator reduce requires a multiple of 32 (assert tile_hdim % dQ_reduce_ncol == 0), and both the preprocess kernel and the interface already round to 32. head_dim=72 -> 80 tripped the assert. Pad to 32 so all three agree (96); this reuses the already-working head_dim=96 path.

3. interface.py — varlen MHA uses the non-TMA dK/dV epilogue store, which writes tile_hdim (mult of 32) columns per row with no head-dim OOB predication (its 16-wide vectorization cannot mask a 72 boundary). It spilled past the real head_dim and corrupted neighboring dK/dV. Route that store through a head-dim-padded scratch buffer, then slice back.

Test results

Validated on B300 (SM103) against torch.nn.functional.scaled_dot_product_attention (fp32 reference):

• Coverage: dense + varlen, causal + non-causal, head_dim ∈ {64, 72, 80, 104, 128}, MHA and GQA(hd=128).
• Result: dq/dk/dv max relative error 2e-3 – 7e-3 (bf16 precision), matching the multiple-of-32 baselines. No regression on aligned head dims.
• End-to-end: Qwen3.5-VL training (ViT head_dim=72, varlen) trains with a decreasing loss and finite grad norm.

Relation to #2633

This is an alternative to #2633. #2633 keeps V padding at a multiple of 16 and predicates the non-TMA store via a reduced tmem Repetition — neater for varlen. However, in our testing that leaves the dense (TMA-store) dV path numerically wrong for head_dim % 32 != 0 (relative error ~1–13 on dv for hd 72/80/104), while the varlen path is correct. Padding V to 32 keeps both dense and varlen paths correct, at the cost of a little extra dV compute. Happy to combine approaches if maintainers prefer the leaner store.

Test plan

• Numerical parity vs SDPA (fp32 ref), dense + varlen, causal + non-causal, hd {64,72,80,104,128}
• No regression on multiple-of-32 head dims
• End-to-end VLM training smoke (ViT hd=72)
• Maintainer CI matrix

🤖 Generated with Claude Code