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

Backward on SM100 (B200/B300) crashed or silently produced wrong dK/dV
when head_dim was not a multiple of 32 (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 to
   match, 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 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
   therefore spilled past the real head_dim and corrupted neighboring dK/dV.
   Route that store through a head-dim-padded scratch buffer, then slice back.

Validated on B300 (SM103) against torch.nn.functional.scaled_dot_product_attention
(fp32 reference): dense + varlen, causal + non-causal, head_dim in
{64, 72, 80, 104, 128}, MHA and GQA(hd=128). 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.

Note: 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, which is
neater for varlen; however in our testing that leaves the dense (TMA-store)
dV path numerically wrong for head_dim % 32 != 0 (large relative error),
whereas padding V to 32 keeps both the dense and varlen paths correct.

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

Author: undefined

flash_attn/cute/flash_bwd_preprocess.py

@@ -13,7 +13,6 @@
 # So the main backward kernel is unchanged; we just replace D with D' = D - dLSE here.
 import math
 import operator
-from functools import partial
 from typing import Callable, Type, Optional
 
 import cuda.bindings.driver as cuda
@@ -365,8 +364,8 @@ def kernel(
             tOpO = None
             if const_expr(self.check_hdim_v_oob):
                 tOpO = copy_utils.predicate_k(tOcO, limit=headdim_v)
-            # Each copy will use the same predicate
-            copy = partial(copy_utils.copy, pred=tOpO)
+            # Each copy uses the predicate sliced to the current m (see loop below).
+            copy = copy_utils.copy
 
             tOrO = cute.make_rmem_tensor_like(tOgO)
             tOrdO = cute.make_rmem_tensor_like(tOgdO)
@@ -378,8 +377,13 @@ def kernel(
                 # Instead of using tOcO, we using t0OcO and subtract the offset from the limit.
                 # This is bc the entries of t0OcO are known at compile time.
                 if t0OcO[0, m, 0][0] < seqlen_limit - tOcO[0][0]:
-                    copy(tOgO[None, m, None], tOrO[None, m, None])
-                    copy(tOgdO[None, m, None], tOrdO[None, m, None])
+                    # The predicate carries a CPY_M mode (broadcast over m); slice it to the
+                    # current m so its shape matches the m-sliced copy source. Otherwise the
+                    # vectorized copy atom's predicate-shape verification fails when head_dim_v
+                    # is not a multiple of the copy-atom width (e.g. 72, 104).
+                    tOpO_cur = tOpO[None, m, None] if const_expr(tOpO is not None) else None
+                    copy(tOgO[None, m, None], tOrO[None, m, None], pred=tOpO_cur)
+                    copy(tOgdO[None, m, None], tOrdO[None, m, None], pred=tOpO_cur)
             # O and dO loads are done; signal that the next kernel can start.
             # Correctness is ensured by griddepcontrol_wait() in bwd_sm90 before it reads our outputs.
             if const_expr(self.use_pdl):

flash_attn/cute/flash_bwd_sm100.py

@@ -68,8 +68,10 @@ def __init__(
         has_aux_tensors: cutlass.Constexpr = False,
         q_subtile_factor: cutlass.Constexpr[int] = 1,
     ):
-        # padding head_dim to a multiple of 16 as k_block_size
-        hdim_multiple_of = 16
+        # Pad head_dim to a multiple of 32 (k_block_size). Must stay a multiple of 32 so
+        # that the dQ accumulator reduce (dQ_reduce_ncol up to 32) tiles evenly and matches
+        # both the preprocess kernel and the interface's head_dim_rounded (also mult of 32).
+        hdim_multiple_of = 32
         self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
         head_dim_v = head_dim_v if head_dim_v is not None else head_dim
         self.same_hdim_kv = head_dim == head_dim_v

flash_attn/cute/interface.py

@@ -1611,6 +1611,28 @@ def _flash_attn_bwd(
     dtype = torch2cute_dtype_map[q.dtype]
     current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
 
+    # Varlen MHA uses the non-TMA dK/dV epilogue store, which writes tile_hdim (a
+    # multiple of 32) columns per row and has no head-dim OOB predication (unlike the
+    # TMA store used elsewhere). When head_dim is not a multiple of 32 the store would
+    # spill past the real head_dim into neighboring memory and corrupt dK/dV. Route the
+    # store through a head-dim-padded scratch buffer, then slice back to head_dim.
+    head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32
+    dkv_hdim_pad = (
+        not dKV_postprocess
+        and cu_seqlens_k is not None
+        and (head_dim != head_dim_rounded or head_dim_v != head_dim_v_rounded)
+    )
+    dk_final, dv_final = dk, dv
+    if dkv_hdim_pad:
+        if head_dim != head_dim_rounded:
+            dk = torch.empty(
+                *dk_final.shape[:-1], head_dim_rounded, dtype=out_torch_dtype, device=device
+            )
+        if head_dim_v != head_dim_v_rounded:
+            dv = torch.empty(
+                *dv_final.shape[:-1], head_dim_v_rounded, dtype=out_torch_dtype, device=device
+            )
+
     if deterministic:
         dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, cluster_size, dtype=torch.int32, device=device)
     else:
@@ -1998,6 +2020,14 @@ def _flash_attn_bwd(
                 cluster_size=cluster_size,
             )
 
+    if dkv_hdim_pad:
+        if not is_fake_mode():
+            if head_dim != head_dim_rounded:
+                dk_final.copy_(dk[..., :head_dim])
+            if head_dim_v != head_dim_v_rounded:
+                dv_final.copy_(dv[..., :head_dim_v])
+        dk, dv = dk_final, dv_final
+
     return dq, dk, dv