sdpa(#624): clamp the f16 THD K/V descriptor extents to the packed total, and NaN-poison ragged capacity tails in test_mhas_v2

python/cudnn/AGENTS.md

@@ -218,6 +218,30 @@ python/cudnn/gemm/
 
 Shared helpers (schedulers, metadata utils, e.g. `gemm/cutedsl/grouped/moe_*.py`) stay internal to the family package — never exported through `cudnn`.
 
+## CuTeDSL kernel bodies
+
+**Do not factor code out of a `@cute.kernel` body into a plain Python helper.**
+The DSL AST-transforms only the decorated function's own source: `for` becomes
+an `ir_loop`, `if` becomes an `scf` region. A helper called from the kernel is
+not transformed, so the ops it emits can land outside the enclosing region.
+
+Hoisting an 11-line block that ran correctly inline into a
+`write_clamped_kv_descs(...)` helper — called from inside
+`if nvvm.elect_sync() and tidx < 32:` — turned 212 passing forward tests into
+31 failures (`Error building ...`, traceback through `ir_loop` →
+`scf_execute_dynamic`). Unrolling the helper's own loop did not help; the
+helper *call* was the problem. Duplicating the block across flavors is the
+correct trade here. Factor only host-side code, or code you can mark
+`@cute.jit`.
+
+Related: inside a kernel body, `for x in (a, b)` over a Python tuple is
+rewritten into a dynamic `ir_loop` and cannot iterate heterogeneous objects
+(e.g. `GridConstant[TensorMap]`). Unroll it, or use `cutlass.range_constexpr`.
+
+**Detector.** These break at `compile()`, not at import — `python -c "import ..."`
+and `pytest --collect-only` both stay green. After any refactor of a kernel
+body, run that flavor's own tests.
+
 ## The APIBase contract (`api_base.py`)
 
 Every OSS kernel API extends `APIBase` and implements:

python/cudnn/sdpa/fwd/api_dsl.py

@@ -1333,13 +1333,12 @@ def scratch_workspace_bytes(self) -> int:
             # head-major (H, head_stride)); without one it compiles with
             # has_lse=False and no LSE buffer exists at all. No slq/slk
             # copies either: the metadata is built DEVICE-side by the setup
-            # kernel (issue #552). o_desc: 16 int64 per
-            # sequence + 16 spare, the per-sequence O TMA descriptors the
-            # builder kernel fills.
-            # o_desc: 16 int64 per sequence + the dead-unit pad slot; the
-            # FP8/MXFP8 flavors carry two more slots for the packed-total-
-            # clamped K/V runtime descriptors (see the kernels' THD closures).
-            o_desc_slots = b + (3 if self._fp8 else 1)
+            # kernel (issue #552).
+            # o_desc: 16 int64 per sequence + the dead-unit pad slot + two
+            # slots for the packed-total-clamped K/V runtime descriptors the
+            # setup kernel writes (see the kernels' THD closures). Every THD
+            # flavor carries those two now, not just FP8/MXFP8 (issue #624).
+            o_desc_slots = b + 3
             return ws_align((4 * b + 4) * 4) + ws_align(o_desc_slots * 16 * 8) + (0 if self.has_sink else ws_align(qh * 4))
         if self._fp8 and self.split_kv == 1:
             return 0  # dense FP8/MXFP8: no per-execute scratch (dummies are cached one-time)
@@ -1675,9 +1674,10 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
         # before any consumer read, so stale bytes never survive — a fill
         # here is a wasted kernel launch on the execute hot path (Rule 1).
         with _torch_stream_context(current_stream, dev):
-            # +2 slots on the FP8/MXFP8 flavors: the packed-total-clamped K/V
-            # runtime descriptors the setup kernel writes after the pad slot.
-            o_desc_slots = b + (3 if self._fp8 else 1)
+            # +2 past the pad slot: the packed-total-clamped K/V runtime
+            # descriptors the setup kernel writes. Every THD flavor carries
+            # them now, not just FP8/MXFP8 (issue #624).
+            o_desc_slots = b + 3
             o_desc = carver.take(o_desc_slots * 16, torch.int64) if carver is not None else torch.empty(o_desc_slots * 16, dtype=torch.int64, device=dev)
         # The PLAN-TIME envelope grid — dead units exit by kernel contract.
         # PERSISTENT THD grid: cap the launch at what the device can hold

python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py

@@ -601,6 +601,7 @@ def _kernel(
             n_q_supers=n_q_supers,
             n_qh=n_qh,
             n_batch=n_batch,
+            o_desc_words=o_desc_words,
             qh_per_kh=qh_per_kh,
             is_leader=is_leader,
             cta_in_pair=cta_in_pair,
@@ -672,6 +673,7 @@ def _tmaldg_warp_group(
     n_q_supers,
     n_qh,
     n_batch,
+    o_desc_words,
     qh_per_kh,
     is_leader,
     cta_in_pair,
@@ -693,8 +695,19 @@ def _tmaldg_warp_group(
     mb_q_reload = bars.mb_q_o_alias if cutlass.const_expr(IS_QO_ALIAS) else bars.mb_q_empty
 
     tma_q = GmemTileTma(tma_q_desc)
-    tma_k = GmemTileTma(tma_k_desc)
-    tma_v = GmemTileTma(tma_v_desc)
+    if cutlass.const_expr(CFG.THD_VARLEN):
+        # THD: K/V ride the setup kernel's packed-total-clamped runtime
+        # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last
+        # sequence's tile-tail lands as exact zeros instead of reading the
+        # buffer's capacity tail (issue #624). Same closure shape as the dense
+        # GmemTileTma, so every load site below stays branch-free.
+        _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords)  # noqa: E731
+        tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords)  # noqa: E731
+    else:
+        tma_k = GmemTileTma(tma_k_desc)
+        tma_v = GmemTileTma(tma_v_desc)
 
     q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split(
         sched.bidx_init,
@@ -2199,6 +2212,8 @@ def _tma_swz(byte_w: int):
         _build_thd_meta_o_descs_kernel(
             o_tensor,
             tma_o_desc,
+            tma_k_desc,
+            tma_v_desc,
             o_desc_words,
             seq_kv_lens_tensor,
             thd_q_lens_tensor,
@@ -2412,9 +2427,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
         if CFG.SEQ_Q_LENS_PRESENT
         else None
     )
-    # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot;
-    # dummy 1-elem when THD off (kernel never reads it).
-    _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
+    # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot
+    # + 2 slots for the packed-total-clamped K/V runtime descriptors the setup
+    # kernel writes (issue #624); dummy 1-elem when THD off (never read).
+    _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
     fake_o_desc = cute.runtime.make_fake_compact_tensor(
         cutlass.Int64,
         (_odesc_len,),

python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py

@@ -634,6 +634,7 @@ def _kernel(
             n_q_supers=n_q_supers,
             n_qh=n_qh,
             n_batch=n_batch,
+            o_desc_words=o_desc_words,
             qh_per_kh=qh_per_kh,
             is_leader=is_leader,
             cta_in_pair=cta_in_pair,
@@ -685,6 +686,7 @@ def _tmaldg_warp_group(
     n_q_supers,
     n_qh,
     n_batch,
+    o_desc_words,
     qh_per_kh,
     is_leader,
     cta_in_pair,
@@ -706,8 +708,19 @@ def _tmaldg_warp_group(
     mb_q_reload = bars.mb_q_o_alias if cutlass.const_expr(IS_QO_ALIAS) else bars.mb_q_empty
 
     tma_q = GmemTileTma(tma_q_desc)
-    tma_k = GmemTileTma(tma_k_desc)
-    tma_v = GmemTileTma(tma_v_desc)
+    if cutlass.const_expr(CFG.THD_VARLEN):
+        # THD: K/V ride the setup kernel's packed-total-clamped runtime
+        # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last
+        # sequence's tile-tail lands as exact zeros instead of reading the
+        # buffer's capacity tail (issue #624). Same closure shape as the dense
+        # GmemTileTma, so every load site below stays branch-free.
+        _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords)  # noqa: E731
+        tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords)  # noqa: E731
+    else:
+        tma_k = GmemTileTma(tma_k_desc)
+        tma_v = GmemTileTma(tma_v_desc)
 
     q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split(
         sched.bidx_init,
@@ -2266,6 +2279,8 @@ def _tma_swz(byte_w: int):
         _build_thd_meta_o_descs_kernel(
             o_tensor,
             tma_o_desc,
+            tma_k_desc,
+            tma_v_desc,
             o_desc_words,
             seq_kv_lens_tensor,
             thd_q_lens_tensor,
@@ -2471,7 +2486,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
     )
     # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot;
     # dummy 1-elem when THD off (kernel never reads it).
-    _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
+    # +2 slots beyond the pad: the packed-total-clamped K/V runtime
+    # descriptors the setup kernel writes (issue #624).
+    _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
     fake_o_desc = cute.runtime.make_fake_compact_tensor(
         cutlass.Int64,
         (_odesc_len,),

python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py

@@ -447,6 +447,7 @@ def _kernel(
             n_q_supers=n_q_supers,
             n_qh=n_qh,
             n_batch=n_batch,
+            o_desc_words=o_desc_words,
             qh_per_kh=qh_per_kh,
             is_leader=is_leader,
             cta_in_pair=cta_in_pair,
@@ -493,6 +494,7 @@ def _tmaldg_warp_group(
     n_q_supers,
     n_qh,
     n_batch,
+    o_desc_words,
     qh_per_kh,
     is_leader,
     cta_in_pair,
@@ -502,8 +504,19 @@ def _tmaldg_warp_group(
     kv_state = PipelineState.start(phase=1)
 
     tma_q = GmemTileTma(tma_q_desc)
-    tma_k = GmemTileTma(tma_k_desc)
-    tma_v = GmemTileTma(tma_v_desc)
+    if cutlass.const_expr(CFG.THD_VARLEN):
+        # THD: K/V ride the setup kernel's packed-total-clamped runtime
+        # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last
+        # sequence's tile-tail lands as exact zeros instead of reading the
+        # buffer's capacity tail (issue #624). Same closure shape as the dense
+        # GmemTileTma, so every load site below stays branch-free.
+        _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords)  # noqa: E731
+        tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords)  # noqa: E731
+    else:
+        tma_k = GmemTileTma(tma_k_desc)
+        tma_v = GmemTileTma(tma_v_desc)
 
     q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split(
         sched.bidx_init,
@@ -1770,6 +1783,8 @@ def _tma_swz(byte_w: int):
         _build_thd_meta_o_descs_kernel(
             o_tensor,
             tma_o_desc,
+            tma_k_desc,
+            tma_v_desc,
             o_desc_words,
             seq_kv_lens_tensor,
             thd_q_lens_tensor,
@@ -1960,7 +1975,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
         if CFG.SEQ_Q_LENS_PRESENT
         else None
     )
-    _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
+    # +2 slots beyond the pad: the packed-total-clamped K/V runtime
+    # descriptors the setup kernel writes (issue #624).
+    _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
     fake_o_desc = cute.runtime.make_fake_compact_tensor(
         cutlass.Int64,
         (_odesc_len,),

python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py

@@ -625,6 +625,7 @@ def _kernel(
             n_q_supers=n_q_supers,
             n_qh=n_qh,
             n_batch=n_batch,
+            o_desc_words=o_desc_words,
             qh_per_kh=qh_per_kh,
             is_leader=is_leader,
             cta_in_pair=cta_in_pair,
@@ -1644,6 +1645,7 @@ def _tmaldg_warp_group(
     n_q_supers,
     n_qh,
     n_batch,
+    o_desc_words,
     qh_per_kh,
     is_leader,
     cta_in_pair,
@@ -1655,8 +1657,19 @@ def _tmaldg_warp_group(
     o_empty_for_v_state = PipelineState.start(phase=1)
 
     tma_q = GmemTileTma(tma_q_desc)
-    tma_k = GmemTileTma(tma_k_desc)
-    tma_v = GmemTileTma(tma_v_desc)
+    if cutlass.const_expr(CFG.THD_VARLEN):
+        # THD: K/V ride the setup kernel's packed-total-clamped runtime
+        # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last
+        # sequence's tile-tail lands as exact zeros instead of reading the
+        # buffer's capacity tail (issue #624). Same closure shape as the dense
+        # GmemTileTma, so every load site below stays branch-free.
+        _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
+        tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords)  # noqa: E731
+        tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords)  # noqa: E731
+    else:
+        tma_k = GmemTileTma(tma_k_desc)
+        tma_v = GmemTileTma(tma_v_desc)
 
     q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split(
         sched.bidx_init,
@@ -1975,6 +1988,8 @@ def _tma_swz(byte_w: int):
         _build_thd_meta_o_descs_kernel(
             o_tensor,
             tma_o_desc,
+            tma_k_desc,
+            tma_v_desc,
             o_desc_words,
             seq_kv_lens_tensor,
             thd_q_lens_tensor,
@@ -2161,7 +2176,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
         if CFG.SEQ_Q_LENS_PRESENT
         else None
     )
-    _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
+    # +2 slots beyond the pad: the packed-total-clamped K/V runtime
+    # descriptors the setup kernel writes (issue #624).
+    _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
     fake_o_desc = cute.runtime.make_fake_compact_tensor(
         cutlass.Int64,
         (_odesc_len,),

python/cudnn/sdpa/fwd/kernels/thd_sm100.py

@@ -231,8 +231,12 @@ def build_thd_meta_o_kv_descs_kernel(
     n_batch: cutlass.Int32,
     o_row_stride: cutlass.Int32,
 ) -> None:
-    """``build_thd_meta_o_descs_kernel`` + packed-total-clamped K/V descriptors
-    (the FP8/MXFP8 THD flavors).
+    """THD setup for the FP8/MXFP8 flavors: metadata, per-batch O descriptors
+    and the packed-total-clamped K/V descriptors.
+
+    Same body as ``build_thd_meta_o_descs_kernel`` minus the persistent
+    scheduler's live-unit total and claim counter, which these flavors do not
+    launch with. Both kernels clamp K/V (issue #624).
 
     The K/V loads tile in TILE_N rows, so the LAST sequence's tile steps past
     the packed KV total into the buffer's capacity tail — caller-owned bytes
@@ -305,6 +309,8 @@ def build_thd_meta_o_kv_descs_kernel(
 def build_thd_meta_o_descs_kernel(
     o_tensor: cute.Tensor,
     base_o_desc: cutlass.GridConstant[tmap.TensorMap],
+    base_k_desc: cutlass.GridConstant[tmap.TensorMap],
+    base_v_desc: cutlass.GridConstant[tmap.TensorMap],
     o_desc_words: cute.Tensor,
     meta_t: cute.Tensor,
     q_lens_t: cute.Tensor,
@@ -316,7 +322,9 @@ def build_thd_meta_o_descs_kernel(
     cga_tile_m: cutlass.Int32,
     n_clusters: cutlass.Int32,
 ) -> None:
-    """Per-execute THD setup, one elected thread (issue #552, D2H removal):
+    """Per-execute THD setup for the f16/bf16 flavors, one elected thread —
+    ``build_thd_meta_o_kv_descs_kernel`` plus the persistent scheduler's
+    live-unit total and claim counter (issue #552, D2H removal):
     build the [seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1)] metadata
     buffer DEVICE-side from the caller's length tensors — ``(B,)`` per-batch
     lengths (serial cumsum; B is small) or the ``(B+1,)`` cu prefix-sum form
@@ -364,6 +372,27 @@ def build_thd_meta_o_descs_kernel(
                 new_value=s_i,
                 ord=2,
             )
+        # Packed-total-clamped K/V runtime descriptors (issue #624). K/V load
+        # in TILE_N rows, so the LAST sequence's tile steps past the packed KV
+        # total into the buffer's capacity tail — caller-owned bytes that may
+        # never have been written. Masked S columns are NaN-safe (the mask is
+        # a select), but BMM2's P·V is not: 0 · NaN == NaN wipes every valid
+        # row of the tile. Patching the seq extent (GLOBAL_DIM ord=2) to
+        # cu_k[B] makes those rows TMA-OOB, so they land as EXACT ZEROS
+        # without touching memory — no fill kernel, and nothing written into
+        # the caller's buffer. Mirrors build_thd_meta_o_kv_descs_kernel, which
+        # the FP8/MXFP8 flavors have used for this since they were written.
+        t_kv = cutlass.Int32(meta[cutlass.Int32(3) * n_batch + cutlass.Int32(1)])  # cu_k[B]
+        k_dptr = desc_base + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)
+        k_src = Pointer(base_k_desc.get_ptr(), dtype=cutlass.Int64)
+        for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS):
+            (k_dptr + i).store((k_src + i).load())
+        nvvm.tensormap_replace(nvvm.TensormapField.GLOBAL_DIM, k_dptr, new_value=t_kv, ord=2)
+        v_dptr = desc_base + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)
+        v_src = Pointer(base_v_desc.get_ptr(), dtype=cutlass.Int64)
+        for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS):
+            (v_dptr + i).store((v_src + i).load())
+        nvvm.tensormap_replace(nvvm.TensormapField.GLOBAL_DIM, v_dptr, new_value=t_kv, ord=2)
         nvvm.fence_proxy_release(
             nvvm.MemScope.GPU,
             from_proxy=nvvm.Proxy.GENERIC,