From b714911e8f3ff94c8fbc7f218ed4ee696131d431 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Mon, 17 Aug 2026 23:27:38 -0700
Subject: [PATCH 1/7] test(mhas_v2): NaN-poison ragged capacity tails so
 capacity-vs-live-token bugs fail deterministically (GitHub #624)

The f16/bf16 harness filled the entire packed K/V/Q/dO capacity - pad
tokens included - with finite random data, so an engine that reads past
the last ragged offset (e.g. THD zero-host-read binding K/V views to
buffer capacity instead of live token counts) stayed green: the padding
mask turns finite garbage into exact zeros, while real recycled device
memory holds NaN bit patterns and 0 x NaN = NaN poisons whole output
rows. The fp8 harness already NaN-fills its pads via
convert_uniform_to_packed, which is why the only test able to catch
GitHub #624 in the wild was fp8 ragged - and only with a lucky seed.

- packed_token_capacity: capacity is now always strictly greater than
  the packed total (next multiple of 64), so every ragged buffer has a
  poisonable tail; previously an exact-multiple total had no tail and
  the bug class structurally could not fire.
- fp16 harness: NaN-fill the capacity tail of Q/K/V/dO after the random
  fill, matching the fp8 harness contract.
- fp16/fp8 harnesses: use the shared packed_token_capacity instead of
  three duplicated inline formulas.

Validation on SM100 (B200):
- native routing: failure set bit-identical to unmodified develop on
  MHAS_NUM_TESTS=16/32 ragged sweeps (no collateral).
- FROST routing (CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1, checkout with
  the #606 zero-host-read path): old harness 24/24 ragged f16 fwd pass;
  new harness fails 19/24 with the exact #624 signature (o_gpu 24% NaN
  in valid rows).

mxfp8 note: the mxfp8 harness has no ragged/THD path (its random
configs never draw "ragged"), so this class is not yet exercisable
there; fp8 covers the 1-byte-dtype THD engines.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/python/sdpa/fp16.py          | 16 ++++++++++++++--
 test/python/sdpa/fp8.py           |  8 ++++++--
 test/python/sdpa/random_config.py |  9 +++++++--
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py
index 2d2db5e07..e53be63b2 100644
--- a/test/python/sdpa/fp16.py
+++ b/test/python/sdpa/fp16.py
@@ -169,6 +169,18 @@ def allocate_tensors(cfg, rng_data_gen, perf=False):
             allocs[TensorUid.dK] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type, strides=k_strides)
             allocs[TensorUid.dV] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type, strides=v_strides)
             allocs[TensorUid.dO] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type, strides=o_strides, rng=rng_data_gen, mean=0.0, std=0.1, sparse_int=si)
+        # NaN-poison the capacity tail (tokens past the last ragged offset).
+        # No engine may ever read those rows; finite random data there hides
+        # capacity-vs-live-token binding bugs because the padding mask turns
+        # them into exact zeros (0 x finite = 0), while real recycled device
+        # memory holds NaN bit patterns (0 x NaN = NaN) — GitHub issue #624.
+        # This makes f16/bf16 consistent with the fp8 harness, whose
+        # convert_uniform_to_packed always NaN-fills the tail.
+        total_t_q, total_t_kv = sum(cfg.seq_len_q), sum(cfg.seq_len_kv)
+        for uid, total in ((TensorUid.q, total_t_q), (TensorUid.k, total_t_kv), (TensorUid.v, total_t_kv)):
+            allocs[uid][0][total:] = float("nan")
+        if cfg.is_train:
+            allocs[TensorUid.dO][0][total_t_q:] = float("nan")
     else:
         allocs[TensorUid.q] = alloc_tensor(cfg.shape_q, cfg.data_type, strides=cfg.stride_q, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si)
         allocs[TensorUid.k] = alloc_tensor(cfg.shape_k, cfg.data_type, strides=cfg.stride_k, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si)
@@ -733,8 +745,8 @@ def _to_seq_len_ref(seq_len_gpu, cu_seq_len_gpu):
     if cfg.is_ragged and cfg.is_train:
         dO_ref = convert_packed_to_uniform(dO_ref, seq_len_q_ref, cfg.s_q)
 
-    max_t_q = max(64, ((seq_len_q_ref.sum().item() + 63) // 64) * 64) if cfg.is_ragged else None
-    max_t_kv = max(64, ((seq_len_kv_ref.sum().item() + 63) // 64) * 64) if cfg.is_ragged else None
+    max_t_q = packed_token_capacity(seq_len_q_ref.tolist()) if cfg.is_ragged else None
+    max_t_kv = packed_token_capacity(seq_len_kv_ref.tolist()) if cfg.is_ragged else None
 
     attn_scale = 0.125
 
diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py
index bcc16584d..c051dcc09 100644
--- a/test/python/sdpa/fp8.py
+++ b/test/python/sdpa/fp8.py
@@ -27,6 +27,7 @@
     profile_execution,
     note_frost_routing,
 )
+from .random_config import packed_token_capacity
 
 # fmt: off
 
@@ -359,8 +360,11 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle):
     if is_ragged:
         seq_len_q_gpu = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1)
         seq_len_kv_gpu = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1)
-        max_t_q = max(64, ((seq_len_q_gpu.sum().item() + 63) // 64) * 64)
-        max_t_kv = max(64, ((seq_len_kv_gpu.sum().item() + 63) // 64) * 64)
+        # Guaranteed capacity tail (> total tokens); convert_uniform_to_packed
+        # NaN-fills it, so engines that read past the last ragged offset fail
+        # deterministically (GitHub #624).
+        max_t_q = packed_token_capacity(seq_len_q_list)
+        max_t_kv = packed_token_capacity(seq_len_kv_list)
 
         # With the ragged offset multiplier, offsets are stored in coarser units
         # (divided by the per-tensor multiplier; always divides evenly) and the
diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py
index d4f638fb3..65659cff1 100644
--- a/test/python/sdpa/random_config.py
+++ b/test/python/sdpa/random_config.py
@@ -42,8 +42,13 @@ def get_strides_from_indices(shape, indices=[0, 1, 2, 3], gaps=[0, 0, 0, 0], rng
 
 
 def packed_token_capacity(seq_lens):
-    """Token capacity of a packed (ragged) buffer, rounded up to a multiple of 64."""
-    return max(64, ((sum(seq_lens) + 63) // 64) * 64)
+    """Token capacity of a packed (ragged) buffer: total tokens rounded up to the
+    next multiple of 64, always strictly greater than the total. The surplus
+    guarantees every ragged buffer has a capacity tail past the last ragged
+    offset — the harnesses NaN-poison that tail so an engine that reads it
+    (e.g. binding K/V views to capacity instead of live token counts,
+    GitHub #624) fails deterministically instead of only on recycled memory."""
+    return (sum(seq_lens) // 64 + 1) * 64
 
 
 def get_strides_from_layout(shape, layout, gaps=[0, 0, 0, 0], rng_geom=None):

From 22ec16fff09c2dfa9666cf33d1f1d6c2a6b683d3 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 18 Aug 2026 01:04:01 -0700
Subject: [PATCH 2/7] test(mhas_v2): stop drawing silently-ignored padded
 configs for mxfp8
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The sdpa_mxfp8 python API has no seq_len/padding arguments and
exec_sdpa_mxfp8 never reads cfg.seq_len_q/kv, so a "padded" draw ran
dense-full while the repro config claimed padding — inflated coverage.
Draw full-only with a pointer to re-add padded/ragged once the API
grows seq-len support (the shared packed_token_capacity /
convert_uniform_to_packed helpers then provide the NaN-poisoned
capacity tails from the previous commit automatically).

Also documents the analogous fp8 gap in-place: a dense "padded" draw
runs as full because exec_sdpa_fp8 binds seq_len tensors only on the
paged and ragged paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/python/sdpa/mxfp8.py   |  8 ++++++++
 test/python/test_mhas_v2.py | 12 +++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py
index 5c44dd06b..3ad0c4871 100644
--- a/test/python/sdpa/mxfp8.py
+++ b/test/python/sdpa/mxfp8.py
@@ -31,6 +31,14 @@
     quantize_to_mxfp8,
 )  # noqa: F401  (re-exported; mxfp8_ref imports it from here)
 
+# NOTE: this harness is dense-full only. The sdpa_mxfp8 python API exposes no
+# seq_len/padding arguments (and the MXFP8 engines defer THD/varlen), so
+# cfg.is_padding / cfg.seq_len_q / cfg.seq_len_kv are intentionally ignored
+# here. When padding/THD support lands in the API, wire it through the shared
+# packed_token_capacity / convert_uniform_to_packed helpers so the ragged
+# capacity tails come NaN-poisoned (see GitHub issue #624 for why that
+# poisoning is load-bearing).
+
 # fmt: off
 
 def ceil_div(a: int, b: int) -> int:
diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py
index 9e746ae39..2b7c9c3f3 100644
--- a/test/python/test_mhas_v2.py
+++ b/test/python/test_mhas_v2.py
@@ -700,6 +700,10 @@ def test_sdpa_fp8_fwd_L0(env_info, test_no, request, cudnn_handle):
         output_type=RandomChoice({torch.float8_e4m3fn: 1, torch.float8_e5m2: 1, torch.float16: 2}),
         with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10),
         diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
+        # KNOWN GAP: a dense "padded" draw currently runs as full — exec_sdpa_fp8
+        # binds seq_len tensors only for the paged and ragged paths, so the
+        # padding mask is never applied here (only the ragged fp8 suites below
+        # exercise real padding).
         is_ragged_or_padded_or_full=RandomChoice({"ragged": 0, "padded": 1, "full": 1}),
         with_sink_token=RandomChoice({True : 1, False : 2}),
     ) as randomization_ctx:
@@ -949,7 +953,13 @@ def test_sdpa_mxfp8_fwd_L0(env_info, test_no, request, cudnn_handle):
         output_type=RandomChoice({torch.float16: 2, torch.bfloat16: 1}),  # FP16 more often for tighter tolerance testing
         with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10),
         diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
-        is_ragged_or_padded_or_full=RandomChoice({"ragged": 0, "padded": 1, "full": 3}),
+        # full-only: the sdpa_mxfp8 API has no seq_len/padding arguments, so a
+        # "padded" draw would silently run dense-full (exec_sdpa_mxfp8 never
+        # reads seq_len_q/kv) and inflate padded coverage. When the API grows
+        # seq-len support, re-add padded/ragged draws — the shared
+        # packed_token_capacity / convert_uniform_to_packed helpers then give
+        # the NaN-poisoned capacity tails that catch the GitHub #624 class.
+        is_ragged_or_padded_or_full=RandomChoice({"full": 1}),
         with_sink_token=RandomChoice({True : 1, False : 2}),
     ) as randomization_ctx:
         test.cfg = randomization_ctx(rng, data_seed, geom_seed)

From 656d3ab72464396f4903a948fb330c2a98b5ef29 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 1 Sep 2026 07:44:55 -0700
Subject: [PATCH 3/7] frost(sdpa): clamp the f16 THD K/V descriptors to the
 packed total (issue #624)

A THD caller binds K/V at BUFFER CAPACITY, not at the packed total, and the
descriptor extent is the host-derived capacity. Rows in [total, capacity) are
therefore inside the extent and get loaded. They are masked, so P == 0 -- but
0 * NaN == NaN, so an uninitialized capacity tail wipes whole tiles of O.

The FP8/MXFP8 flavors have always handled this: build_thd_meta_o_kv_descs_kernel
copies the K/V base TensorMaps into o_desc_words and patches GLOBAL_DIM ord=2 to
cu_k[B] device-side, so tail rows become TMA-OOB and land as exact zeros -- no
memory touched, nothing written into the caller's buffer. The f16/bf16 flavors
used the O-only setup variant and had no clamp.

Port it: build_thd_meta_o_descs_kernel takes base_k_desc/base_v_desc and writes
the two clamped descriptors after the dead-unit pad slot; all four f16 SM100
kernels (d128, d192_d128, d256, d512) thread o_desc_words into
_tmaldg_warp_group and bind tma_k/tma_v to runtime-descriptor closures under
CFG.THD_VARLEN; _odesc_len and api_dsl's o_desc_slots grow by two slots (which
is now unconditional for THD, not FP8-only).

The clamp block is duplicated rather than factored into a helper on purpose: a
plain Python function called from a @cute.kernel body is not AST-transformed,
and hoisting it turns these tests red.

Adds test_dsl_sm100_thd_nan_capacity_tail: three sequences totalling 397 (not a
multiple of TILE_N, so the last tile steps past it) bound into capacity-sized
Q/K/V buffers whose tails are NaN. Asserts O is finite and matches the
per-sequence reference.

Validation, B200 SM100, cuDNN 9.26.0.33, CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1,
same setup for both columns:

  test_mhas_v2 -k fwd_ragged_L0    without clamp: 19 failed / 27 passed
                                   with clamp:     0 failed / 46 passed
  test_mhas_v2 -k fwd (dense +
    ragged + fp8 + mxfp8)          212 passed, 28 skipped, 0 failed

SM120 (prefill_f16_sm120.py) still lacks the clamp -- it uses the meta-only
build_thd_meta_kernel, so the K/V slots must be introduced rather than switched.
---
 python/cudnn/sdpa/fwd/api_dsl.py              | 20 ++--
 .../fwd/kernels/prefill_d128_f16_sm100.py     | 26 ++++-
 .../kernels/prefill_d192_d128_f16_sm100.py    | 23 ++++-
 .../fwd/kernels/prefill_d256_f16_sm100.py     | 23 ++++-
 .../fwd/kernels/prefill_d512_f16_sm100.py     | 23 ++++-
 python/cudnn/sdpa/fwd/kernels/thd_sm100.py    | 35 ++++++-
 .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py     | 95 +++++++++++++++++++
 7 files changed, 218 insertions(+), 27 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 390ebbd37..11dd93ab9 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/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
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
index 64263f32f..78a70ff7f 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
+++ b/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,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py
index 47125817d..db5f49d06 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py
+++ b/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,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
index b6b491c52..dd02511b8 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
+++ b/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,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
index 15bcf30f1..1df375e06 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
+++ b/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,),
diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
index 556c58369..79651a888 100644
--- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
+++ b/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,
diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
index 8ad146950..e6ae8abe1 100644
--- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
+++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
@@ -958,6 +958,101 @@ def _dense_buf(packed):
     torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2)
 
 
+# Issue #624: a THD caller binds K/V at BUFFER CAPACITY, not at the packed
+# total, and the descriptor extent is the host-derived capacity. Rows in
+# [total, capacity) are therefore inside the extent and get loaded. They are
+# masked, so P == 0 -- but 0 * NaN == NaN, so an UNINITIALIZED capacity tail
+# wipes whole tiles of O. The setup kernel patches the K/V descriptor extent to
+# the packed total cu_k[B] device-side so those rows are TMA-OOB and land as
+# exact zeros instead. Total 397 is deliberately not a multiple of TILE_N, so
+# the last sequence's tail tile steps past it.
+@pytest.mark.L0
+@pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS)
+@pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS)
+@torch_fork_set_rng(seed=0)
+def test_dsl_sm100_thd_nan_capacity_tail(dtype, d):
+    """THD with a NaN-poisoned capacity tail: O must be finite and correct."""
+    _require_dsl()
+    import cudnn
+
+    dev = "cuda"
+    H = 8
+    seq_lens = [200, 150, 47]
+    B = len(seq_lens)
+    S_max = max(seq_lens)
+    T = sum(seq_lens)
+    cu = [0]
+    for s_i in seq_lens:
+        cu.append(cu[-1] + s_i)
+    scale = 1.0 / math.sqrt(d)
+
+    q_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
+    k_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
+    v_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
+
+    stride = (S_max * H * d, d, H * d, 1)
+
+    def _poisoned_buf(packed):
+        """Capacity-sized storage; only [0, T) is live, the tail is NaN."""
+        stor = torch.full((B * S_max * H * d,), float("nan"), device=dev, dtype=dtype)
+        stor[: T * H * d] = packed.reshape(-1)
+        return stor, stor.as_strided((B, H, S_max, d), stride)
+
+    q_stor, q_gpu = _poisoned_buf(q_pk)
+    k_stor, k_gpu = _poisoned_buf(k_pk)
+    v_stor, v_gpu = _poisoned_buf(v_pk)
+    o_stor = torch.zeros(B * S_max * H * d, device=dev, dtype=dtype)
+    o_gpu = o_stor.as_strided((B, H, S_max, d), stride)
+
+    slq = torch.tensor(seq_lens, dtype=torch.int32, device=dev).view(B, 1, 1, 1)
+    slk = slq.clone()
+    cu_t = torch.tensor(cu, dtype=torch.int64, device=dev)
+    ro = (cu_t * H * d).view(B + 1, 1, 1, 1)
+
+    io = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16
+    g = cudnn.pygraph(io_data_type=io, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT)
+    tq = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="q")
+    tk = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="k")
+    tv = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="v")
+    sq = g.tensor_like(slq)
+    skv = g.tensor_like(slk)
+    qro = g.tensor_like(ro)
+    kro = g.tensor_like(ro)
+    vro = g.tensor_like(ro)
+    oro = g.tensor_like(ro)
+    tq.set_ragged_offset(qro)
+    tk.set_ragged_offset(kro)
+    tv.set_ragged_offset(vro)
+    o, _ = g.sdpa(
+        name="sdpa", q=tq, k=tk, v=tv, generate_stats=False, attn_scale=scale, use_causal_mask=True, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv
+    )
+    o.set_output(True).set_dim([B, H, S_max, d]).set_stride(list(stride))
+    o.set_ragged_offset(oro)
+
+    g.validate()
+    g.build_operation_graph()
+    g.create_execution_plans([cudnn.heur_mode.A])
+    _select_engine(g, engine_name())
+    g.check_support()
+    g.build_plans()
+    vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro, kro: ro, vro: ro, oro: ro}
+    g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8))
+    torch.cuda.synchronize()
+
+    o_out = o_stor[: T * H * d].reshape(T, H, d)
+    assert not torch.isnan(o_out).any(), f"{int(torch.isnan(o_out).sum())} NaNs in O -- the capacity tail reached BMM2"
+
+    o_ref = torch.zeros(T, H, d, device=dev, dtype=dtype)
+    for b in range(B):
+        lo, hi = cu[b], cu[b + 1]
+        qb = q_pk[lo:hi].permute(1, 0, 2).unsqueeze(0)
+        kb = k_pk[lo:hi].permute(1, 0, 2).unsqueeze(0)
+        vb = v_pk[lo:hi].permute(1, 0, 2).unsqueeze(0)
+        ob = _ref_sdpa_full(qb, kb, vb, scale=scale, is_causal=True)
+        o_ref[lo:hi] = ob.squeeze(0).permute(1, 0, 2)
+    torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2)
+
+
 @pytest.mark.L0
 @pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS)
 @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS)

From 550e229a18fd8f4ef11a900ecf8b7e2edc45c99f Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 1 Sep 2026 07:45:03 -0700
Subject: [PATCH 4/7] docs(agents): record the CuTeDSL kernel-body and
 editable-install traps

Two traps that each cost real debugging time on issue #624, written down with
the runnable check rather than as prose advice.

python/cudnn/AGENTS.md -- 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 -> ir_loop, if -> scf region); a helper is not transformed, so the
ops it emits can land outside the enclosing region. Hoisting an 11-line block
that ran correctly inline turned 212 passing forward tests into 31 failures.
These break at compile(), not at import, so --collect-only stays green.

test/AGENTS.md -- confirm you are testing the code you edited. `pip install -e .`
installs a sys.meta_path finder whose MAPPING hard-codes one checkout; meta-path
finders beat sys.path, so PYTHONPATH cannot shadow it and edits in a second
clone or a git worktree are silently not under test. The symptom is
indistinguishable from a real result: a deliberately destructive probe leaves
the output bit-identical. Includes the `python -c "import cudnn; print(...)"`
check and a sitecustomize.py recipe for redirecting the finder for one run.
---
 python/cudnn/AGENTS.md | 24 ++++++++++++++++++++++++
 test/AGENTS.md         | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md
index babedb9ff..e9b5cde90 100644
--- a/python/cudnn/AGENTS.md
+++ b/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:
diff --git a/test/AGENTS.md b/test/AGENTS.md
index ff216bfca..68698106d 100644
--- a/test/AGENTS.md
+++ b/test/AGENTS.md
@@ -49,3 +49,36 @@ pytest fe_api/gemm/          # OSS kernel tests
 - **A regression test must be seen RED.** Before trusting one, run it against the unfixed code — restore the old line, confirm it fails, restore the fix. `test_dsl_sm100_thd_interleaved_kv_views` and `test_varlen_backward_does_not_sync` were both checked this way, and both were genuinely red beforehand; a test written for a bug and never seen to fail is asserting an unknown.
 - **Seed before you allocate.** `torch.manual_seed()` after constructing the inputs seeds nothing that matters. Two runs meant to be compared then differ by data, and the assertion fails (or worse, passes) for a reason unrelated to what is under test — if two runs must be comparable, build the inputs once and reuse them.
 - **When you remove a fallback, invert its counter assertion — do not delete it.** Tests that asserted `calls["bwd_cpp"]` incremented had to become "`calls["bwd"]` increments **and** `bwd_cpp` does not", so a silent regression to the old path fails the suite instead of passing it.
+
+### Confirm you are testing the code you edited
+
+`pip install -e .` does **not** put the package on `sys.path`. It installs a
+`sys.meta_path` finder (`__editable___nvidia_cudnn_frontend_*_finder.py`) whose
+`MAPPING` hard-codes an absolute path to the checkout it was installed from.
+Meta-path finders run *before* `sys.path`, so **`PYTHONPATH` cannot shadow it** —
+if you edit a different clone or a git worktree, your changes are silently not
+under test. Symptoms are indistinguishable from a real result: a probe that
+should change the output leaves it bit-identical, and edits appear to do nothing.
+
+Check first, every time you work outside the installed checkout:
+
+```bash
+python -c "import cudnn; print(cudnn.__file__)"   # must be YOUR tree
+```
+
+`conftest.py` prints the same path in its banner (`cuDNN Frontend Path:`) — read
+it rather than assuming. To point the editable install at another tree for one
+run, patch the finder's `MAPPING` from a `sitecustomize.py` on `PYTHONPATH`
+(`site` imports it after processing `.pth` files, so the finder already exists):
+
+```python
+# sitecustomize.py
+import __editable___nvidia_cudnn_frontend_1_28_0_finder as f
+f.MAPPING["cudnn"] = "/path/to/your/worktree/python/cudnn"
+```
+
+The same trap hides *inside* a run: `python/cudnn/frost/template_loader.py`
+loads kernel templates by absolute path via `spec_from_file_location`, so the
+template that serves a config may come from elsewhere too. To find out which
+template a test actually compiles, log `path` at the top of `load_template` —
+do not infer it from `_pick_flavor` by reading the source.

From 8ca5b63d48bd0d426231b628d4759e930bd31935 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 1 Sep 2026 07:51:03 -0700
Subject: [PATCH 5/7] test(sdpa): the undeclared THD path is NaN-safe now too
 (issue #624)

test_dsl_sm100_thd_declared_total_bounds_capacity_tail ended with a control
asserting that WITHOUT max_total_seq_len the capacity tail is still reachable,
so that the clamp assertions above it were testing the clamp rather than a
benign shape. The device-side K/V extent clamp makes that false by design --
the tail is now out of reach on every THD flavor, declared or not -- which is
the case the assertion's own message called out as "needs rethinking".

Invert it: assert the undeclared path is clean, that O is independent of the
tail fill either way, and that declaring the total does not change the result
bit for bit. Declaring is now a host-side tightening (capacity, grid and
workspace sizing) rather than the only thing standing between the kernel and
the tail.

Seen red: with the clamp reverted and these tests kept, 7 failed
(test_dsl_sm100_thd_nan_capacity_tail at 201728/403456/806912 NaNs in O across
d128/d256/d512, plus this test); with the clamp, 7 passed.
---
 .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py     | 27 +++++++++++++------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
index e6ae8abe1..73f8eb2f9 100644
--- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
+++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
@@ -1801,9 +1801,11 @@ def test_dsl_sm100_thd_declared_total_bounds_capacity_tail():
     FINITE. A caller that over-allocates and leaves the tail unwritten
     therefore poisons whole tiles through ``0 * NaN`` (issue #624).
 
-    Declaring the total clamps the TMA extent to the packed total, putting that
-    tail out of reach. Same graph, same data, only the tail fill differs: with
-    the declaration a NaN tail must be inert."""
+    Two mechanisms now put that tail out of reach: declaring the total clamps
+    the host-derived extent, and the THD setup kernel clamps the K/V extent to
+    ``cu_k[B]`` device-side regardless. Same graph, same data, only the tail
+    fill differs -- a NaN tail must be inert either way, and the two paths must
+    agree bit for bit."""
     _require_dsl()
     import cudnn
 
@@ -1877,13 +1879,22 @@ def _run(tail_value, declare_total):
     )
     assert torch.equal(declared_nan, declared_zero), "O must not depend on the capacity tail once the total is declared"
 
-    # Control: the tail is genuinely reachable without the declaration, so the
-    # assertions above are testing the clamp rather than a benign shape.
+    # The undeclared path is NaN-safe too, and deliberately so: the THD setup
+    # kernel now patches the K/V descriptor extent to cu_k[B] device-side on
+    # every flavor, which bounds the tail without any host declaration. Before
+    # that clamp this same call poisoned ~50% of O, and this assertion was
+    # inverted -- it asserted the NaN to prove the tail was reachable.
     undeclared_nan = _run(float("nan"), declare_total=False)
-    assert torch.isnan(undeclared_nan).any(), (
-        "expected the undeclared path to read the capacity tail (issue #624); if this no longer "
-        "holds the extent is exact by other means and this test needs rethinking"
+    undeclared_zero = _run(0.0, declare_total=False)
+    assert not torch.isnan(undeclared_nan).any(), (
+        f"the device-side K/V extent clamp must bound the capacity tail with or without a declared "
+        f"total: {int(torch.isnan(undeclared_nan).sum())} NaNs in O"
     )
+    assert torch.equal(undeclared_nan, undeclared_zero), "O must not depend on the capacity tail"
+    # Declaring the total is now a host-side tightening (capacity, grid and
+    # workspace sizing) rather than the only thing standing between the kernel
+    # and the tail -- so both paths must agree bit for bit.
+    assert torch.equal(declared_nan, undeclared_nan), "declaring the packed total must not change the result"
 
 
 @pytest.mark.L0

From b3606fc975886e5074f8dfcd2d36a4d15bef9d1e Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 1 Sep 2026 09:30:15 -0700
Subject: [PATCH 6/7] test(sdpa): cover the d192/d128 flavor in the
 capacity-tail test; drop the hard-coded finder version

Two CodeRabbit findings on #646, both valid.

test_dsl_sm100_thd_nan_capacity_tail parametrized a single `d` over _FLAVORS,
so it could only select d128/d256/d512 -- it never reached the d192/d128
flavor, which is one of the four kernels this PR changes. Parametrize
(d_qk, d_v) over [(128,128), (192,128), (256,256), (512,512)] instead, with Q/K
at d_qk and V/O at d_v (separate ragged offsets per width). Confirmed the new
case is a real regression guard: with the clamp reverted, dsv3 fails at 201728
NaNs in O for both dtypes; with it, all 8 combos pass.

test/AGENTS.md hard-coded __editable___nvidia_cudnn_frontend_1_28_0_finder in
the sitecustomize recipe, which raises ModuleNotFoundError once __version__
bumps -- and it would fail before reaching the MAPPING patch, i.e. exactly when
the reader is already confused about which checkout is loaded. Discover the
module via pkgutil.iter_modules() instead. The recipe as written is the one
used to validate this PR.
---
 test/AGENTS.md                                | 13 +++-
 .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py     | 60 ++++++++++---------
 2 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/test/AGENTS.md b/test/AGENTS.md
index 68698106d..c3dfd9d16 100644
--- a/test/AGENTS.md
+++ b/test/AGENTS.md
@@ -72,9 +72,16 @@ run, patch the finder's `MAPPING` from a `sitecustomize.py` on `PYTHONPATH`
 (`site` imports it after processing `.pth` files, so the finder already exists):
 
 ```python
-# sitecustomize.py
-import __editable___nvidia_cudnn_frontend_1_28_0_finder as f
-f.MAPPING["cudnn"] = "/path/to/your/worktree/python/cudnn"
+# sitecustomize.py -- the finder module name embeds the installed version, so
+# discover it rather than hard-coding it (it changes when __version__ bumps).
+import importlib, pkgutil
+
+name = next(
+    m.name
+    for m in pkgutil.iter_modules()
+    if m.name.startswith("__editable___nvidia_cudnn_frontend_") and m.name.endswith("_finder")
+)
+importlib.import_module(name).MAPPING["cudnn"] = "/path/to/your/worktree/python/cudnn"
 ```
 
 The same trap hides *inside* a run: `python/cudnn/frost/template_loader.py`
diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
index 73f8eb2f9..5db778feb 100644
--- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
+++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
@@ -967,10 +967,12 @@ def _dense_buf(packed):
 # exact zeros instead. Total 397 is deliberately not a multiple of TILE_N, so
 # the last sequence's tail tile steps past it.
 @pytest.mark.L0
-@pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS)
+# (d_qk, d_v) rather than a single d: d192/d128 is its own kernel flavor, and
+# a single-d parametrization cannot reach it.
+@pytest.mark.parametrize("d_qk,d_v", [(128, 128), (192, 128), (256, 256), (512, 512)], ids=["llama", "dsv3", "qwen", "dsv4"])
 @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS)
 @torch_fork_set_rng(seed=0)
-def test_dsl_sm100_thd_nan_capacity_tail(dtype, d):
+def test_dsl_sm100_thd_nan_capacity_tail(dtype, d_qk, d_v):
     """THD with a NaN-poisoned capacity tail: O must be finite and correct."""
     _require_dsl()
     import cudnn
@@ -984,49 +986,51 @@ def test_dsl_sm100_thd_nan_capacity_tail(dtype, d):
     cu = [0]
     for s_i in seq_lens:
         cu.append(cu[-1] + s_i)
-    scale = 1.0 / math.sqrt(d)
+    scale = 1.0 / math.sqrt(d_qk)
 
-    q_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
-    k_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
-    v_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
+    q_pk = torch.randn(T, H, d_qk, device=dev, dtype=dtype)
+    k_pk = torch.randn(T, H, d_qk, device=dev, dtype=dtype)
+    v_pk = torch.randn(T, H, d_v, device=dev, dtype=dtype)
 
-    stride = (S_max * H * d, d, H * d, 1)
+    def _stride(dd):
+        return (S_max * H * dd, dd, H * dd, 1)
 
-    def _poisoned_buf(packed):
+    def _poisoned_buf(packed, dd):
         """Capacity-sized storage; only [0, T) is live, the tail is NaN."""
-        stor = torch.full((B * S_max * H * d,), float("nan"), device=dev, dtype=dtype)
-        stor[: T * H * d] = packed.reshape(-1)
-        return stor, stor.as_strided((B, H, S_max, d), stride)
+        stor = torch.full((B * S_max * H * dd,), float("nan"), device=dev, dtype=dtype)
+        stor[: T * H * dd] = packed.reshape(-1)
+        return stor, stor.as_strided((B, H, S_max, dd), _stride(dd))
 
-    q_stor, q_gpu = _poisoned_buf(q_pk)
-    k_stor, k_gpu = _poisoned_buf(k_pk)
-    v_stor, v_gpu = _poisoned_buf(v_pk)
-    o_stor = torch.zeros(B * S_max * H * d, device=dev, dtype=dtype)
-    o_gpu = o_stor.as_strided((B, H, S_max, d), stride)
+    q_stor, q_gpu = _poisoned_buf(q_pk, d_qk)
+    k_stor, k_gpu = _poisoned_buf(k_pk, d_qk)
+    v_stor, v_gpu = _poisoned_buf(v_pk, d_v)
+    o_stor = torch.zeros(B * S_max * H * d_v, device=dev, dtype=dtype)
+    o_gpu = o_stor.as_strided((B, H, S_max, d_v), _stride(d_v))
 
     slq = torch.tensor(seq_lens, dtype=torch.int32, device=dev).view(B, 1, 1, 1)
     slk = slq.clone()
     cu_t = torch.tensor(cu, dtype=torch.int64, device=dev)
-    ro = (cu_t * H * d).view(B + 1, 1, 1, 1)
+    ro_qk = (cu_t * H * d_qk).view(B + 1, 1, 1, 1)
+    ro_v = (cu_t * H * d_v).view(B + 1, 1, 1, 1)
 
     io = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16
     g = cudnn.pygraph(io_data_type=io, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT)
-    tq = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="q")
-    tk = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="k")
-    tv = g.tensor(dim=[B, H, S_max, d], stride=list(stride), data_type=io, name="v")
+    tq = g.tensor(dim=[B, H, S_max, d_qk], stride=list(_stride(d_qk)), data_type=io, name="q")
+    tk = g.tensor(dim=[B, H, S_max, d_qk], stride=list(_stride(d_qk)), data_type=io, name="k")
+    tv = g.tensor(dim=[B, H, S_max, d_v], stride=list(_stride(d_v)), data_type=io, name="v")
     sq = g.tensor_like(slq)
     skv = g.tensor_like(slk)
-    qro = g.tensor_like(ro)
-    kro = g.tensor_like(ro)
-    vro = g.tensor_like(ro)
-    oro = g.tensor_like(ro)
+    qro = g.tensor_like(ro_qk)
+    kro = g.tensor_like(ro_qk)
+    vro = g.tensor_like(ro_v)
+    oro = g.tensor_like(ro_v)
     tq.set_ragged_offset(qro)
     tk.set_ragged_offset(kro)
     tv.set_ragged_offset(vro)
     o, _ = g.sdpa(
         name="sdpa", q=tq, k=tk, v=tv, generate_stats=False, attn_scale=scale, use_causal_mask=True, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv
     )
-    o.set_output(True).set_dim([B, H, S_max, d]).set_stride(list(stride))
+    o.set_output(True).set_dim([B, H, S_max, d_v]).set_stride(list(_stride(d_v)))
     o.set_ragged_offset(oro)
 
     g.validate()
@@ -1035,14 +1039,14 @@ def _poisoned_buf(packed):
     _select_engine(g, engine_name())
     g.check_support()
     g.build_plans()
-    vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro, kro: ro, vro: ro, oro: ro}
+    vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro_qk, kro: ro_qk, vro: ro_v, oro: ro_v}
     g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8))
     torch.cuda.synchronize()
 
-    o_out = o_stor[: T * H * d].reshape(T, H, d)
+    o_out = o_stor[: T * H * d_v].reshape(T, H, d_v)
     assert not torch.isnan(o_out).any(), f"{int(torch.isnan(o_out).sum())} NaNs in O -- the capacity tail reached BMM2"
 
-    o_ref = torch.zeros(T, H, d, device=dev, dtype=dtype)
+    o_ref = torch.zeros(T, H, d_v, device=dev, dtype=dtype)
     for b in range(B):
         lo, hi = cu[b], cu[b + 1]
         qb = q_pk[lo:hi].permute(1, 0, 2).unsqueeze(0)

From 7c36f4b523d9e83c2f1565ec337df9dcc44ac7e7 Mon Sep 17 00:00:00 2001
From: Vedaanta Agarwalla <vagarwalla@nvidia.com>
Date: Tue, 1 Sep 2026 09:55:33 -0700
Subject: [PATCH 7/7] test(sdpa): drop the unused packed-storage bindings in
 the capacity-tail test

RUF059: q_stor/k_stor/v_stor were never read after unpacking. The as_strided
view keeps its storage alive on its own, so _poisoned_buf now returns just the
view rather than the pair. o_stor stays bound -- the assertions read the packed
prefix out of it.
---
 test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
index 5db778feb..2db5f5b91 100644
--- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
+++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
@@ -996,14 +996,17 @@ def _stride(dd):
         return (S_max * H * dd, dd, H * dd, 1)
 
     def _poisoned_buf(packed, dd):
-        """Capacity-sized storage; only [0, T) is live, the tail is NaN."""
+        """Capacity-sized storage; only [0, T) is live, the tail is NaN.
+
+        Returns the strided view only -- it keeps the storage alive, so there
+        is no need to bind the flat tensor as well."""
         stor = torch.full((B * S_max * H * dd,), float("nan"), device=dev, dtype=dtype)
         stor[: T * H * dd] = packed.reshape(-1)
-        return stor, stor.as_strided((B, H, S_max, dd), _stride(dd))
+        return stor.as_strided((B, H, S_max, dd), _stride(dd))
 
-    q_stor, q_gpu = _poisoned_buf(q_pk, d_qk)
-    k_stor, k_gpu = _poisoned_buf(k_pk, d_qk)
-    v_stor, v_gpu = _poisoned_buf(v_pk, d_v)
+    q_gpu = _poisoned_buf(q_pk, d_qk)
+    k_gpu = _poisoned_buf(k_pk, d_qk)
+    v_gpu = _poisoned_buf(v_pk, d_v)
     o_stor = torch.zeros(B * S_max * H * d_v, device=dev, dtype=dtype)
     o_gpu = o_stor.as_strided((B, H, S_max, d_v), _stride(d_v))