From cb29d59d9b39e89dba7290093f50e0bbd01c4762 Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Thu, 20 Aug 2026 12:36:04 -0700
Subject: [PATCH 1/9] frost(sdpa): persistent THD scheduler + batch remap on
 SM100/SM120

- SM120 THD: longest-first batch remap and a persistent claim scheduler, so
  the grid is sized to the machine rather than to the plan-time S_q envelope
  (832-2048 CTAs -> 84). 0 to -6.8% on ragged causal shapes.
- Fix a latent K/V pipeline bug: compute_one_kv_tile arrives at
  bar_k/v_consumed once per tile while the load warp only syncs between
  tiles, leaving an unmatched arrival per pass. Harmless for a single-shot
  CTA, fatal once a CTA runs a second tile range.
- Push the SM100 scheduler payload as scalar DSMEM stores; the tuple form
  of store_async_dsmem does not lower.
- Widen the THD metadata to [seq_kv | cu_q | cu_k | remap | live | ctr] on
  SM120 and add a test covering CTAs that claim more than one unit.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/cudnn/frost/tile_dsl/scheduler.py      |  78 ++++
 python/cudnn/sdpa/fwd/api_dsl.py              |  43 +-
 .../cudnn/sdpa/fwd/kernels/_common_sm100.py   |   7 +-
 .../fwd/kernels/prefill_d128_f16_sm100.py     |  39 +-
 .../fwd/kernels/prefill_d128_fp8_sm100.py     |  10 +-
 .../fwd/kernels/prefill_d128_fp8_sm107.py     |  10 +-
 .../fwd/kernels/prefill_d128_mxfp8_sm100.py   |  12 +-
 .../kernels/prefill_d192_d128_f16_sm100.py    |  13 +-
 .../fwd/kernels/prefill_d256_f16_sm100.py     |   8 +-
 .../fwd/kernels/prefill_d512_f16_sm100.py     |   8 +-
 .../sdpa/fwd/kernels/prefill_f16_sm120.py     | 403 +++++++++++++-----
 .../sdpa/fwd/kernels/prefill_fp8_sm120.py     |  31 +-
 python/cudnn/sdpa/fwd/kernels/thd_sm100.py    | 168 +++++++-
 .../sdpa/frost/test_sdpa_fwd_dsl_sm120.py     |  24 ++
 14 files changed, 691 insertions(+), 163 deletions(-)

diff --git a/python/cudnn/frost/tile_dsl/scheduler.py b/python/cudnn/frost/tile_dsl/scheduler.py
index a5c7a2ad7..26fbf3fb7 100644
--- a/python/cudnn/frost/tile_dsl/scheduler.py
+++ b/python/cudnn/frost/tile_dsl/scheduler.py
@@ -11,6 +11,9 @@
 import cutlass.cute as cute
 from cutlass._mlir.dialects import arith
 
+from cutlass._mlir.dialects import arith
+from cutlass.base_dsl.typing import Pointer
+
 from .barrier import PipelineState, advance, wait, arrive_expect_tx
 
 
@@ -107,6 +110,81 @@ class Sched(NamedTuple):
     bidz_init: object
 
 
+@cute.jit
+def scheduler_warp_loop_persistent(
+    sched,
+    sched_stages: int,
+    is_cga_first_cta,
+    meta_t,
+    ctr_off,
+    live_off,
+    cga_size: int,
+):
+    """Persistent tile scheduler over a LIVE-ONLY unit range (THD).
+
+    CLC sizes the grid to the work list, which for THD means the plan-time
+    envelope and therefore dead clusters. Here the grid is occupancy-sized and
+    the bound is a DEVICE value (``meta[live_off]``, written by the setup
+    launch), so no unit past the live total is ever handed out.
+
+    The cluster lead claims one unit with a global atomic and pushes the
+    payload into every CTA's ``tile_id_smem`` over DSMEM, then arrives each
+    peer's scheduler mbarrier -- the same shape as ``read_tile_id_arrive``.
+    Multicast is not available here: it is a clusterlaunchcontrol facility, so
+    dynamic claiming needs an explicit peer write.
+    """
+    meta = cutlass.make_array_view(meta_t)
+    ctr_ptr = Pointer(meta_t.iterator.raw_ptr(), dtype=cutlass.Int32) + ctr_off
+    state = PipelineState.start()
+    is_valid = cutlass.Int32(1)
+
+    while is_valid > cutlass.Int32(0):
+        wait(sched.mb_read_tile_id.subview(state.idx), state.phase)
+
+        # Every CTA expects the 16-byte payload on its own mbarrier, exactly as
+        # the CLC path did; the lead's remote store delivers it and completes
+        # the barrier through the transaction count.
+        if nvvm.elect_sync():
+            arrive_expect_tx(sched.mb_scheduler.subview(state.idx), 16)
+
+        if nvvm.elect_sync() and is_cga_first_cta:
+            uid = cutlass.Int32(nvvm.atomicrmw(nvvm.AtomicOp.ADD, ctr_ptr, cutlass.Int32(1)))
+            live = cutlass.Int32(meta[live_off])
+            valid = cutlass.Int32(arith.select((uid < live).ir_value(), cutlass.Int32(1).ir_value(), cutlass.Int32(0).ir_value()))
+            linear = uid * cutlass.Int32(cga_size)
+            # store_async_dsmem wants CuTe pointers; the smem arrays hand out
+            # base-DSL ones, so convert through the raw addresses.
+            _tile_ptr = cute.make_ptr(
+                cutlass.Int32,
+                sched.tile_id_smem.subview(state.idx * cutlass.Int32(8)).data_ptr().toint(cutlass.Int32),
+                cutlass.AddressSpace.smem,
+                assumed_align=16,
+            )
+            _mbar_ptr = cute.make_ptr(
+                cutlass.Int64,
+                sched.mb_scheduler.subview(state.idx).data_ptr().toint(cutlass.Int32),
+                cutlass.AddressSpace.smem,
+                assumed_align=8,
+            )
+            # One word at a time rather than a v4 payload: store_async_dsmem
+            # accepts a 2/4-tuple per its contract, but lowers whatever it was
+            # handed through Int32(value), so a tuple raises at trace time. Four
+            # scalar stores carry the same 16 bytes and so satisfy the same
+            # transaction count the arrive above expects.
+            _payload = (linear, cutlass.Int32(0), valid, cutlass.Int32(0))
+            for i in cutlass.range_constexpr(cga_size):
+                for w in cutlass.range_constexpr(4):
+                    cute.arch.store_async_dsmem(_tile_ptr + w, _payload[w], _mbar_ptr, i)
+
+        nvvm.bar_warp_sync(cute.arch.FULL_MASK)
+
+        wait(sched.mb_scheduler.subview(state.idx), state.phase)
+        validity = (sched.tile_id_smem.subview(state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()
+        is_valid = validity & cutlass.Int32(1)
+
+        state = advance(state, sched_stages)
+
+
 @cute.jit
 def scheduler_warp_loop(sched, sched_stages: int, is_cga_first_cta):
     state = PipelineState.start()
diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 85ba57e6c..7cfc39921 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -1203,7 +1203,7 @@ def scratch_workspace_bytes(self) -> int:
             # 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)
-            return ws_align((3 * b + 2) * 4) + ws_align(o_desc_slots * 16 * 8) + (0 if self.has_sink else ws_align(qh * 4))
+            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)
         if self.split_kv > 1:
@@ -1520,7 +1520,7 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
         d_qk, d_v = self.head_dim_qk, self.head_dim_v
         carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), label) if workspace is not None else None
         with _torch_stream_context(current_stream, dev):
-            meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev)
+            meta = carver.take(4 * b + 4, torch.int32) if carver is not None else torch.empty(4 * b + 4, dtype=torch.int32, device=dev)
         q_lens_dev = self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens")
         kv_lens_dev = self._checked_cu_seq_lens(seq_kv_lens, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")
         lens_form = (1 if self.cu_seq_q_lens else 0) | (2 if self.cu_seq_kv_lens else 0)
@@ -1544,7 +1544,21 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
             o_desc_slots = b + (3 if self._fp8 else 1)
             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.
-        units = self._thd_unit_envelope()
+        # PERSISTENT THD grid: cap the launch at what the device can hold
+        # resident (one cluster per CTA_MMA SMs) instead of the plan-time
+        # envelope. The kernel pulls units from a device-bounded counter, so
+        # the grid no longer has to cover the work list -- which is what made
+        # 38-84% of clusters dead.
+        _env = self._thd_unit_envelope()
+        _cta_mma = int(getattr(self._k_mod, "CTA_MMA", 1))
+        _sms = torch.cuda.get_device_properties(q_buf.device).multi_processor_count
+        if getattr(self._k_mod, "THD_PERSISTENT", False):
+            units = min(_env, max(1, _sms // max(1, _cta_mma)))
+            _dbg = int(os.environ.get("FROST_THD_CLUSTERS", "0"))  # debug override
+            if _dbg > 0:
+                units = min(_env, _dbg)
+        else:
+            units = _env  # CLC path: the grid must BE the work list
 
         Q = self._thd_view(q_buf, self.q_desc, t_q)
         O = self._thd_view(o_buf, self.o_desc, t_q)
@@ -2778,6 +2792,7 @@ def execute(
             None,  # thd_q_lens / thd_kv_lens / thd_lens_form: THD-only, folded out
             None,
             None,
+            cutlass.Int32(0),  # thd_n_ctas: THD-only persistent grid extent
             current_stream,
         )
         if self.split_kv > 1:
@@ -2944,6 +2959,7 @@ def _execute_fp8(
             pack.q_lens_dev if pack is not None else None,
             pack.kv_lens_dev if pack is not None else None,
             cutlass.Int32(pack.lens_form) if pack is not None else None,
+            cutlass.Int32(0),  # thd_n_ctas: unused on this path (no persistent grid)
             current_stream,
         )
         # Both of these consume what the kernel just wrote, so they belong on
@@ -3025,7 +3041,7 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa
         # validation — AGENTS.md Rule 3; cu prefixes are normalized by the
         # setup kernel).
         with _torch_stream_context(current_stream, dev):
-            meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev)
+            meta = carver.take(4 * b + 4, torch.int32) if carver is not None else torch.empty(4 * b + 4, dtype=torch.int32, device=dev)
         q_lens_dev = self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens")
         kv_lens_dev = self._checked_cu_seq_lens(seq_kv_lens, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")
         lens_form = (1 if self.cu_seq_q_lens else 0) | (2 if self.cu_seq_kv_lens else 0)
@@ -3177,9 +3193,26 @@ def _execute_thd(
             pack.q_lens_dev,
             pack.kv_lens_dev,
             cutlass.Int32(pack.lens_form),
+            cutlass.Int32(self._thd_persistent_ctas(pack.Q.device)),
             current_stream,
         )
 
+    def _thd_persistent_ctas(self, device) -> int:
+        """CTA count for the persistent THD grid.
+
+        Sized to the MACHINE, not to the work: the live unit total is a
+        device-side quantity (issue #552), a CTA with nothing left to claim just
+        retires, and one with more work loops. So over-launching is harmless and
+        under-launching only costs parallelism. One CTA per SM matches the
+        kernel's ``min_blocks_per_mp``.
+        """
+        forced = int(os.environ.get("FROST_THD_CTAS", "0"))
+        if forced > 0:
+            return forced
+        sms = torch.cuda.get_device_properties(device).multi_processor_count
+        per_sm = int(os.environ.get("FROST_THD_CTAS_PER_SM", "1"))
+        return max(1, sms * max(1, per_sm))
+
     def scratch_workspace_bytes(self) -> int:
         if self.thd:
             # [meta(seq_kv, cu_q, cu_k)].
@@ -3192,7 +3225,7 @@ def scratch_workspace_bytes(self) -> int:
             # on sinks. No O-descriptor chunk: SM120 stores O with plain
             # guarded GMEM stores, so THD needs no per-sequence tensor maps.
             b = self.batch_size
-            return ws_align((3 * b + 2) * 4)
+            return ws_align((4 * b + 4) * 4)
         if self.split_kv > 1:
             # Split-major partial slabs (see the SM100 sibling): O_s in the O
             # dtype (half) + lse_s fp32, carved from the caller's workspace.
diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
index 0ca6320fa..52e4a5d66 100644
--- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
@@ -757,12 +757,17 @@ def _thd_decode(linear_cta, seq_kv_lens_t, n_batch, n_qh, cta_in_pair):
         f_head = cutlass.Int32(0)
         f_qc = cutlass.Int32(0)
         done = cutlass.Int32(0)
-        for b in cutlass.range(0, n_batch, 1, unroll=1):
+        # Sequences are visited LONGEST FIRST through batch_remap (built by the
+        # THD setup launch), so the ragged tail of the grid is short sequences.
+        remap0 = cutlass.Int32(3) * n_batch + cutlass.Int32(2)
+        for i in cutlass.range(0, n_batch, 1, unroll=1):
+            b = cutlass.Int32(cu[remap0 + i])
             s_i = cutlass.Int32(cu[cuq0 + b + cutlass.Int32(1)]) - cutlass.Int32(cu[cuq0 + b])
             cb = (s_i + cutlass.Int32(cga_tile_m - 1)) // cutlass.Int32(cga_tile_m)
             units_b = cb * n_qh
             in_rng = (done == cutlass.Int32(0)) & (u < acc + units_b)
             local = u - acc
+            # Natural order within a sequence (head-major, ascending rows).
             f_batch = cutlass.Int32(arith.select(in_rng.ir_value(), b.ir_value(), f_batch.ir_value()))
             f_head = cutlass.Int32(arith.select(in_rng.ir_value(), (local // cb).ir_value(), f_head.ir_value()))
             f_qc = cutlass.Int32(arith.select(in_rng.ir_value(), (local % cb).ir_value(), f_qc.ir_value()))
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 e8c29b3c2..27a988f04 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
@@ -118,6 +118,7 @@
 from cudnn.frost.tile_dsl.scheduler import (
     Sched,
     scheduler_warp_loop,
+    scheduler_warp_loop_persistent,
     read_tile_id_arrive,
     SCHED_NATURAL,
     SCHED_LPT,
@@ -197,6 +198,10 @@
 
 
 CGA_TILE_M = CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA
+CTA_MMA = CFG.CTA_MMA
+# THD uses a persistent grid + device-bounded claim counter (not the CLC
+# envelope). Only this flavor is converted so far.
+THD_PERSISTENT = True
 
 
 # SM100 llama is always cga2 (CTA_MMA=2) → the SCHED_LPT/LPT_L2 reverse-row
@@ -216,9 +221,10 @@
 # factory; O-descriptor builder + TENSOR_MAP_QWORDS from the shared
 # kernels/dsl/common/sdpa/thd.py.  Gated by CFG.THD_VARLEN (folds out otherwise).
 # Supported at cga1 and cga2 (TILES_Q=2 → two Q slabs / O stores per tile).
-# seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2):
+# seq_kv_lens overloaded as the THD metadata buffer (int32 len 4B+4):
 #   [0..B-1]=seq_kv_lens  [B..2B]=cu_q(B+1)  [2B+1..3B+1]=cu_k(B+1)
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS
+#   [3B+2..4B+1]=batch_remap(B)  [4B+2]=live units  [4B+3]=claim counter
+from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS, THD_SETUP_THREADS
 
 _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS
 # The setup kernel builds the THD metadata buffer DEVICE-side from the
@@ -623,7 +629,26 @@ def _kernel(
         # try_cancel.multicast::cluster::all — only (0,0,0) CTA issues; at cga1
         # cta_id_x == 0 always, so flag is 1 unconditionally.
         is_cga_first_cta = cta_id_x == cutlass.Int32(0)
-        scheduler_warp_loop(sched, CFG.SCHEDULER_STAGES, is_cga_first_cta)
+        if cutlass.const_expr(CFG.THD_VARLEN):
+            # THD: persistent grid + device-bounded claim counter, so no unit
+            # past the live total is ever handed out (the CLC path would need
+            # the grid to BE the work list, i.e. the plan-time envelope).
+            # n_batch is a kernel argument -- do NOT re-derive it from the
+            # metadata tensor's layout. seq_kv_lens_tensor is not always a flat
+            # 1-D view (the llama THD configs bind a nested one, where shape[0]
+            # is itself a tuple), so the arithmetic below has to use the value
+            # the caller already threads through every other THD helper.
+            scheduler_warp_loop_persistent(
+                sched,
+                CFG.SCHEDULER_STAGES,
+                is_cga_first_cta,
+                seq_kv_lens_tensor,
+                cutlass.Int32(4) * n_batch + cutlass.Int32(3),
+                cutlass.Int32(4) * n_batch + cutlass.Int32(2),
+                CGA_SIZE,
+            )
+        else:
+            scheduler_warp_loop(sched, CFG.SCHEDULER_STAGES, is_cga_first_cta)
 
 
 # === TMA-LDG warp ===
@@ -2179,7 +2204,9 @@ def _tma_swz(byte_w: int):
             cutlass.Int32(QH),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
-        ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+            cutlass.Int32(CGA_TILE_M),
+            n_thd_units,
+        ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
         # Grid Python-folds on Cfg constant (avoids DSL if staging).
@@ -2353,8 +2380,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
     )
     # seq_kv_lens always part of the ABI; read only when CFG.SEQ_KV_LENS_PRESENT == 1
     # (compile-time fold).  THD overloads it as the [seq_kv_lens(B)|cu_q(B+1)|
-    # cu_k(B+1)] metadata buffer (length 3B+2).
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    # cu_k(B+1)|batch_remap(B)|live|ctr] metadata buffer (length 4B+4).
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_len,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
index 17aa904b1..a190613ed 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
@@ -199,8 +199,11 @@
 # dequant scales, no block-scale SF).  Gated by CFG.THD_VARLEN (folds out:
 # _thd_tma_offsets is (0, 0, batch_idx) dense — TMA coords byte-identical).
 # TILES_Q=2: q_seq_off applies to BOTH Q slabs + both O-store slabs.
-# seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2):
+# seq_kv_lens overloaded as the THD metadata buffer (int32 len 4B+4):
 #   [0..B-1]=seq_kv_lens  [B..2B]=cu_q(B+1)  [2B+1..3B+1]=cu_k(B+1)
+#   [3B+2..4B+1]=batch_remap(B)  [4B+2]=live units  [4B+3]=claim counter
+# The decode walks batch_remap on EVERY THD flavor, so the setup kernel must
+# fill it; the trailing two words are read only by the persistent schedulers.
 # The setup kernel builds it DEVICE-side from the caller's length tensors and
 # the adapter launches the plan-time envelope grid (issue #552) — no length
 # ever reaches the host.
@@ -2349,8 +2352,9 @@ def compile(  # noqa: A001
         assumed_align=16,
     )
     # Always part of the ABI; unread when CFG.SEQ_KV_LENS_PRESENT == 0.  THD
-    # overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)] (len 3B+2).
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    # overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)|batch_remap(B)|
+    # live|ctr] (len 4B+4).
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_len,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
index d46696616..bfaa6bbbd 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
@@ -237,8 +237,11 @@
 # dequant scales, no block-scale SF).  Gated by CFG.THD_VARLEN (folds out:
 # _thd_tma_offsets is (0, 0, batch_idx) dense — TMA coords byte-identical).
 # TILES_Q=2: q_seq_off applies to BOTH Q slabs + both O-store slabs.
-# seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2):
+# seq_kv_lens overloaded as the THD metadata buffer (int32 len 4B+4):
 #   [0..B-1]=seq_kv_lens  [B..2B]=cu_q(B+1)  [2B+1..3B+1]=cu_k(B+1)
+#   [3B+2..4B+1]=batch_remap(B)  [4B+2]=live units  [4B+3]=claim counter
+# The decode walks batch_remap on EVERY THD flavor, so the setup kernel must
+# fill it; the trailing two words are read only by the persistent schedulers.
 # The setup kernel builds it DEVICE-side from the caller's length tensors and
 # the adapter launches the plan-time envelope grid (issue #552) — no length
 # ever reaches the host.
@@ -2428,8 +2431,9 @@ def compile(  # noqa: A001
         assumed_align=16,
     )
     # Always part of the ABI; unread when CFG.SEQ_KV_LENS_PRESENT == 0.  THD
-    # overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)] (len 3B+2).
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    # overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)|batch_remap(B)|
+    # live|ctr] (len 4B+4).
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_len,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py
index f45059f44..10b700bbb 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py
@@ -246,8 +246,11 @@ def _round_up(a: int, b: int) -> int:
 # is (0, 0, batch_idx) and _thd_sf_tile_bases (0, 0) dense — TMA coords
 # byte-identical).  Supported at cga1 and cga2 (TILES_Q=2 → two Q slabs /
 # O stores per tile).  seq_kv_lens overloaded as the THD metadata buffer
-# (int32 len 3B+2): [0..B-1]=seq_kv_lens [B..2B]=cu_q(B+1) [2B+1..3B+1]=cu_k(B+1)
-# — built DEVICE-side by the setup kernel (issue #552).
+# (int32 len 4B+4): [0..B-1]=seq_kv_lens [B..2B]=cu_q(B+1) [2B+1..3B+1]=cu_k(B+1)
+# [3B+2..4B+1]=batch_remap(B) [4B+2]=live units [4B+3]=claim counter
+# — built DEVICE-side by the setup kernel (issue #552).  The decode walks
+# batch_remap on EVERY THD flavor, so the setup kernel must fill it; the
+# trailing two words are read only by the persistent schedulers.
 # MXFP8-only: _thd_sf_tile_bases returns the per-sequence SF-tile prefix bases
 # (cu_sf_q_base / cu_sf_k_base) for the packed scale-factor layout.
 from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_kv_descs_kernel as _build_thd_meta_o_kv_descs_kernel, TENSOR_MAP_QWORDS
@@ -2755,8 +2758,9 @@ def compile(  # noqa: A001
         assumed_align=16,
     )
     # seq_kv_lens always part of the ABI; unread when CFG.SEQ_KV_LENS_PRESENT
-    # == 0.  THD overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)] (len 3B+2).
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    # == 0.  THD overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)|
+    # batch_remap(B)|live|ctr] (len 4B+4).
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_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 87da77e16..dae127c62 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
@@ -200,9 +200,10 @@ def _bounds_for_tile(
 # factory; O-descriptor builder + TENSOR_MAP_QWORDS from the shared
 # kernels/dsl/common/sdpa/thd.py.  Gated by CFG.THD_VARLEN (folds out otherwise).
 # Supported at cga1 and cga2 (TILES_Q=2 → two Q slabs / O stores per tile).
-# seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2):
+# seq_kv_lens overloaded as the THD metadata buffer (int32 len 4B+4):
 #   [0..B-1]=seq_kv_lens  [B..2B]=cu_q(B+1)  [2B+1..3B+1]=cu_k(B+1)
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS
+#   [3B+2..4B+1]=batch_remap(B)  [4B+2]=live units  [4B+3]=claim counter
+from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS, THD_SETUP_THREADS
 
 _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS
 # The setup kernel builds the THD metadata buffer DEVICE-side from the
@@ -2274,7 +2275,9 @@ def _tma_swz(byte_w: int):
             cutlass.Int32(QH),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
-        ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+            cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),
+            n_thd_units,  # CLC path: envelope units; the counter goes unused
+        ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
         # Grid Python-folds on Cfg constant (avoids DSL if staging).
@@ -2438,8 +2441,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
     )
     # seq_kv_lens always part of the ABI; read only when CFG.SEQ_KV_LENS_PRESENT == 1
     # (compile-time fold).  THD overloads it as the [seq_kv_lens(B)|cu_q(B+1)|
-    # cu_k(B+1)] metadata buffer (length 3B+2).
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    # cu_k(B+1)|batch_remap(B)|live|ctr] metadata buffer (length 4B+4).
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_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 b5f17f6f1..e0f6ecc84 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
@@ -121,7 +121,7 @@
 _resolve_seqlen_q = _sdpa_h.resolve_seqlen_q
 
 
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS
+from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS, THD_SETUP_THREADS
 
 _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS
 # The setup kernel builds the THD metadata buffer DEVICE-side from the
@@ -1779,7 +1779,9 @@ def _tma_swz(byte_w: int):
             cutlass.Int32(QH),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
-        ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+            cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),
+            n_thd_units,  # CLC path: envelope units; the counter goes unused
+        ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
         # KV split rides the BATCH axis: z = batch + split*B.  The decode
@@ -1931,7 +1933,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
         stride_order=(0,),
         assumed_align=16,
     )
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_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 73b2e718a..4483bfa42 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
@@ -312,7 +312,7 @@ class KernelTmemLayout:
 _decode_payload = _sdpa_h.decode_payload
 
 
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS
+from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS, THD_SETUP_THREADS
 
 _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS
 # The setup kernel builds the THD metadata buffer DEVICE-side from the
@@ -1984,7 +1984,9 @@ def _tma_swz(byte_w: int):
             cutlass.Int32(QH),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
-        ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+            cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),
+            n_thd_units,  # CLC path: envelope units; the counter goes unused
+        ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
         # KV split rides the BATCH axis: z = batch + split*B.
@@ -2132,7 +2134,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
         stride_order=(0,),
         assumed_align=16,
     )
-    _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
+    _skv_len = (4 * b + 4) if CFG.THD_VARLEN else b
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (_skv_len,),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
index 52897950b..5060fd9d1 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
@@ -59,7 +59,12 @@
 )
 from cudnn.frost.tile_dsl.mma import mma_m16n8k16_f32
 from cudnn.frost.tile_dsl.swizzle import swizzle_xor
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_kernel as _build_thd_meta_kernel
+from cudnn.sdpa.fwd.kernels.thd_sm100 import (
+    build_thd_meta_kernel as _build_thd_meta_kernel,
+    thd_claim_next,
+    thd_decode_unit,
+    THD_SETUP_THREADS,
+)
 from cudnn.sdpa.fwd.config_sm120 import (
     HEAD_TILE_GRANULE,
     SEQ_KV_TILES as _SEQ_KV_TILES,
@@ -140,6 +145,15 @@ def pack_to_i32(
 # L2 working-set budget used by the LPT_L2 group sizing.
 _SCHED_L2_BUDGET_BYTES = 50 * 1024 * 1024
 
+# THD only: pull units from a device-side counter over a machine-sized grid
+# instead of launching the plan-time envelope as a padded rectangle. The
+# envelope scales with the DECLARED S_q, so on ragged batches most of that
+# rectangle is tiles no sequence owns.
+#
+# This is what made the K/V consumer barriers have to balance per unit (see the
+# drain in _run_unit): a CTA here runs the tile range of several units in turn.
+THD_PERSISTENT = True
+
 
 def ceil_div(a: int, b: int) -> int:
     """Return the ceiling division of a by b."""
@@ -794,8 +808,11 @@ def compute_one_kv_tile(
         """
 
         # The K/V loop walks tile indices in reverse order. The mbarrier parity
-        # still follows the load iteration count: 0, 1, 0, 1, ...
-        tma_phase = (num_kv_tiles - 1 - kv_tile_idx) & cutlass.Int32(1)
+        # still follows the load iteration count: 0, 1, 0, 1, ... Under a
+        # persistent grid the mbarriers are armed once for the CTA and reused by
+        # every unit it claims, so the count continues across units — phase_base
+        # is how many tiles this CTA already loaded (0 on the one-unit paths).
+        tma_phase = (basic_params.phase_base + num_kv_tiles - 1 - kv_tile_idx) & cutlass.Int32(1)
         while not prims.mbarrier_try_wait_parity(basic_params.k_tma_mbar, tma_phase):
             pass
 
@@ -818,95 +835,57 @@ def compute_one_kv_tile(
         self.mma_pv(basic_params, mma_params, p_regs)
         prims.barrier_cta_arrive(self.bar_v_consumed, self.threads_kv_pipeline)
 
-    @cute.kernel
-    def kernel(
+    @cute.jit
+    def _run_unit(
         self,
-        q: cute.Tensor,
-        k: cute.Tensor,
-        v: cute.Tensor,
-        o: cute.Tensor,
-        lse: Optional[cute.Tensor],
-        sinks: Optional[cute.Tensor],
-        seq_q_lens: cute.Tensor,
-        seq_kv_lens: cute.Tensor,
-        tma_k_desc: cutlass.GridConstant[cuda.TensorMap],
-        tma_v_desc: cutlass.GridConstant[cuda.TensorMap],
-        softmax_scale_log2: cutlass.Float32,
-        n_q_tiles: cutlass.Int32,
-    ) -> None:
-        """SM120 FMHA prefill kernel.
-
-        :param q: Query tensor.
-        :param k: Key tensor.
-        :param v: Value tensor.
-        :param o: Output tensor.
-        :param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed
-            token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per
-            ``thd_lse_head_major``) under ``thd_varlen``; or ``None`` to
-            compile the LSE store out (the DSL specializes on ``None``).
-        :param sinks: ``(H,)`` fp32 per-Q-head sink logits; ``None`` iff the
-            kernel is configured without ``has_sink``.
-        :param seq_q_lens: Per-batch query lengths, or an unused dummy tensor.
-        :param seq_kv_lens: Per-batch key/value lengths, or an unused dummy tensor.
-        :param tma_k_desc: Tensor map descriptor for K.
-        :param tma_v_desc: Tensor map descriptor for V.
-        :param softmax_scale_log2: ``softmax_scale * log2(e)``, pre-folded host-side.
+        q,
+        k,
+        v,
+        o,
+        lse,
+        sinks,
+        seq_q_lens,
+        seq_kv_lens,
+        tma_k_desc,
+        tma_v_desc,
+        softmax_scale_log2,
+        sKV,
+        sK,
+        sV,
+        k_tma_mbar,
+        v_tma_mbar,
+        lane,
+        warp,
+        phase_base,
+        split_idx,
+        o_batch_idx,
+        q_tile_idx,
+        batch_idx,
+        head_idx,
+    ) -> cutlass.Int32:
+        """One unit of work: the (Q tile, sequence, head) triple named by the
+        last three arguments.
+
+        Split out of ``kernel`` so the persistent THD path can call it once per
+        claimed unit. SMEM and the per-warp register budget belong to the CTA,
+        not to the unit, so they stay in the caller and arrive here as
+        arguments — the same convention ``load_one_kv_tile`` already uses.
         """
-        tidx, _, _ = cute.arch.thread_idx()
-        q_tile_idx, batch_idx, head_idx = cute.arch.block_idx()
-        # KV split rides the BATCH axis: grid.y = batch + split*B.  Q/K/V and the
-        # per-batch seqlens must use the REAL batch; only the O/LSE partial slot
-        # uses the composite.  Folds away entirely at split_kv == 1.  NATURAL-only
-        # (config_sm120 enforces it): LPT / LPT_L2 flatten the grid to 1-D and
-        # derive the batch from the linear tile id, leaving no axis to ride.
-        split_idx = cutlass.Int32(0)
-        o_batch_idx = batch_idx
-        if cutlass.const_expr(self.split_kv > 1):
-            n_batch_real = cutlass.Int32(q.shape[0])
-            split_idx = batch_idx // n_batch_real
-            batch_idx = batch_idx % n_batch_real
-            o_batch_idx = split_idx * n_batch_real + batch_idx
-        if cutlass.const_expr(self.sched_policy != SCHED_NATURAL):
-            _n_qh = cutlass.Int32((q.shape[2] // self.qh_per_kh if self.pack_gqa else q.shape[2]))
-            _n_batch = cutlass.Int32(self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0])
-            # Host-computed (see __call__): the grid is sized from the same
-            # value, so the decode cannot disagree with the launch geometry.
-            _q_tiles = n_q_tiles
-            if cutlass.const_expr(self.sched_policy == SCHED_LPT_L2):
-                q_tile_idx, head_idx, batch_idx = lpt_l2_tile_coords(
-                    q_tile_idx,
-                    _n_qh,
-                    _n_batch,
-                    _q_tiles,
-                    _n_qh // cutlass.Int32(k.shape[2]),
-                    cutlass.Int32(k.shape[1]),
-                    (self.head_tile_qk + self.head_tile_v) * self.in_dtype.width // 8,
-                    _SCHED_L2_BUDGET_BYTES,
-                )
-            else:
-                q_tile_idx, head_idx, batch_idx = lpt_tile_coords(q_tile_idx, _n_qh, _n_batch, _q_tiles)
-            o_batch_idx = batch_idx
-        elif cutlass.const_expr(self.is_causal):
-            # Diagonal-bounded work grows with the Q tile. Launch long tiles
-            # first to avoid leaving a few expensive CTAs in the final
-            # scheduler waves (right-band graphs share the causal shape).
-            grid_q, _, _ = cute.arch.grid_dim()
-            q_tile_idx = grid_q - q_tile_idx - 1
         q_seq_idx = q_tile_idx * (self.q_tile // self.qh_per_kh if self.pack_gqa else self.q_tile)
 
-        lane = tidx % cute.arch.WARP_SIZE
-        warp = cute.arch.warp_idx()
-
         seqlen_q = cutlass.Int32(q.shape[1])
         seqlen_k = cutlass.Int32(k.shape[1])
         q_row_base = cutlass.Int32(0)
         kv_row_base = cutlass.Int32(0)
         if cutlass.const_expr(self.thd_varlen):
-            # THD: seq_kv_lens is the metadata tensor [seq_kv(B) | cu_q(B+1) | cu_k(B+1)].
+            # THD: seq_kv_lens is the metadata tensor
+            # [seq_kv(B) | cu_q(B+1) | cu_k(B+1) | remap(B) | live | ctr].
             # Per-sequence lengths come from the prefix sums; the bases offset
             # every packed (1, T, H, D) access below.
-            n_batch = (seq_kv_lens.shape[0] - 2) // 3
+            n_batch = (seq_kv_lens.shape[0] - 4) // 4
             meta = cutlass.make_array_view(seq_kv_lens)
+            # batch_idx is a real batch index here: whichever dispatch produced
+            # it has already resolved the longest-first batch_remap.
             q_row_base = cutlass.Int32(meta[n_batch + batch_idx])
             seqlen_q = cutlass.Int32(meta[n_batch + batch_idx + 1]) - q_row_base
             kv_row_base = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx])
@@ -986,35 +965,10 @@ def kernel(
             min_kv_tile = _lo
             num_kv_tiles = _lo + _per + _extra
         has_kv_work = num_kv_tiles > 0 and (num_kv_tiles - 1) >= min_kv_tile
-
-        # Shared-memory layout:
-        #   sK: one kv_tile x head_tile_qk K tile
-        #   sV: one kv_tile x head_tile_v V tile
-        # The epilogue later aliases this storage as the q_tile x head_tile_v
-        # sO staging tile after compute warps finish consuming the final K/V
-        # tile.
-        sKV = cutlass.Array(
-            k.dtype,
-            max(self.k_tile_elems + self.v_tile_elems, self.o_tile_elems),
-            space=cutlass.AddressSpace.smem,
-            alignment=128,
-        )
-        sK = sKV
-        sV = sKV.subview(self.k_tile_elems)
-        tma_mbar = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8)
-        k_tma_mbar = tma_mbar
-        v_tma_mbar = tma_mbar.subview(1)
-
-        # Initialize the TMA completion barriers before any load or compute warp
-        # can touch the K/V pipeline.
-        if warp == self.load_warp_id:
-            if prims.elect_sync():
-                prims.prefetch_tensormap(tma_k_desc.get_ptr())
-                prims.prefetch_tensormap(tma_v_desc.get_ptr())
-                prims.mbarrier_init(k_tma_mbar, 1)
-                prims.mbarrier_init(v_tma_mbar, 1)
-        prims.fence_mbarrier_init()
-        prims.barrier_cta_sync(0)
+        # Tiles the load warp will push through the K/V mbarriers for this unit.
+        # Branchless on purpose: when has_kv_work is false the difference is
+        # already <= 0, so this is exactly zero.
+        tiles_loaded = cute.math.max(cutlass.Int32(0), num_kv_tiles - min_kv_tile)
 
         # /////////////////////////////////////////////////////////////////////////////
         #  LOAD K/V
@@ -1087,6 +1041,18 @@ def kernel(
                         envelope=v_envelope,
                     )
                     kv_seq_idx -= self.kv_tile
+                # Consume the LAST tile's consumer arrivals. compute_one_kv_tile
+                # arrives at both barriers for every tile it processes (N), but
+                # the loop above only syncs between tiles (N-1), so without this
+                # each pass leaves one unmatched arrival pending on each barrier.
+                # That is invisible while a CTA runs a single tile range -- the
+                # leftover dies with the CTA -- but a CTA that starts a SECOND
+                # range finds the barrier already part-signalled, satisfies its
+                # first sync immediately, and issues the next
+                # mbarrier.arrive.expect_tx while the previous transaction is
+                # still outstanding.
+                prims.barrier_cta_sync(self.bar_k_consumed, thread_count=self.threads_kv_pipeline)
+                prims.barrier_cta_sync(self.bar_v_consumed, thread_count=self.threads_kv_pipeline)
         # /////////////////////////////////////////////////////////////////////////////
         #  COMPUTE
         # /////////////////////////////////////////////////////////////////////////////
@@ -1118,6 +1084,7 @@ def kernel(
                 o_regs[i] = 0.0
 
             basic_params = SimpleNamespace(
+                phase_base=phase_base,
                 seqlen_q=seqlen_q,
                 seqlen_k=seqlen_k,
                 head_dim_qk=head_dim_qk,
@@ -1376,6 +1343,200 @@ def kernel(
         else:
             prims.setmaxregister(40, prims.SetMaxRegisterAction.DECREASE)
 
+        return tiles_loaded
+
+    @cute.kernel
+    def kernel(
+        self,
+        q: cute.Tensor,
+        k: cute.Tensor,
+        v: cute.Tensor,
+        o: cute.Tensor,
+        lse: Optional[cute.Tensor],
+        sinks: Optional[cute.Tensor],
+        seq_q_lens: cute.Tensor,
+        seq_kv_lens: cute.Tensor,
+        tma_k_desc: cutlass.GridConstant[cuda.TensorMap],
+        tma_v_desc: cutlass.GridConstant[cuda.TensorMap],
+        softmax_scale_log2: cutlass.Float32,
+        n_q_tiles: cutlass.Int32,
+    ) -> None:
+        """SM120 FMHA prefill kernel.
+
+        :param q: Query tensor.
+        :param k: Key tensor.
+        :param v: Value tensor.
+        :param o: Output tensor.
+        :param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed
+            token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per
+            ``thd_lse_head_major``) under ``thd_varlen``; or ``None`` to
+            compile the LSE store out (the DSL specializes on ``None``).
+        :param sinks: ``(H,)`` fp32 per-Q-head sink logits; ``None`` iff the
+            kernel is configured without ``has_sink``.
+        :param seq_q_lens: Per-batch query lengths, or an unused dummy tensor.
+        :param seq_kv_lens: Per-batch key/value lengths, or an unused dummy tensor.
+        :param tma_k_desc: Tensor map descriptor for K.
+        :param tma_v_desc: Tensor map descriptor for V.
+        :param softmax_scale_log2: ``softmax_scale * log2(e)``, pre-folded host-side.
+        """
+        tidx, _, _ = cute.arch.thread_idx()
+        lane = tidx % cute.arch.WARP_SIZE
+        warp = cute.arch.warp_idx()
+
+        # Shared-memory layout:
+        #   sK: one kv_tile x head_tile_qk K tile
+        #   sV: one kv_tile x head_tile_v V tile
+        # The epilogue later aliases this storage as the q_tile x head_tile_v
+        # sO staging tile after compute warps finish consuming the final K/V
+        # tile.
+        sKV = cutlass.Array(
+            k.dtype,
+            max(self.k_tile_elems + self.v_tile_elems, self.o_tile_elems),
+            space=cutlass.AddressSpace.smem,
+            alignment=128,
+        )
+        sK = sKV
+        sV = sKV.subview(self.k_tile_elems)
+        tma_mbar = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8)
+        k_tma_mbar = tma_mbar
+        v_tma_mbar = tma_mbar.subview(1)
+
+        # Initialize the TMA completion barriers before any load or compute warp
+        # can touch the K/V pipeline.
+        if warp == self.load_warp_id:
+            if prims.elect_sync():
+                prims.prefetch_tensormap(tma_k_desc.get_ptr())
+                prims.prefetch_tensormap(tma_v_desc.get_ptr())
+                prims.mbarrier_init(k_tma_mbar, 1)
+                prims.mbarrier_init(v_tma_mbar, 1)
+        prims.fence_mbarrier_init()
+        prims.barrier_cta_sync(0)
+
+        if cutlass.const_expr(self.thd_varlen and THD_PERSISTENT and self.sched_policy == SCHED_NATURAL):
+            # Persistent grid: the launch is sized to the MACHINE, not to the
+            # plan-time envelope, and each CTA pulls units from a device-side
+            # counter until the live total runs out. That total depends on the
+            # real sequence lengths, which never reach the host (issue #552), so
+            # the setup launch computes it into the metadata buffer. Launching
+            # more CTAs than there is work is harmless — the extra ones fail the
+            # loop test and retire.
+            _meta = cutlass.make_array_view(seq_kv_lens)
+            _nb = (seq_kv_lens.shape[0] - 4) // 4
+            _live = cutlass.Int32(_meta[4 * _nb + 2])
+            _slot = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem, alignment=16)
+            _bidx, _, _ = cute.arch.block_idx()
+            _uid = cutlass.Int32(_bidx)
+            # K/V mbarrier parity is a CTA-lifetime count under a persistent
+            # grid: the barriers are armed once above, so each unit picks the
+            # phase up where the previous one left it.
+            _phase = cutlass.Int32(0)
+            while _uid < _live:
+                _qt, _b, _h = thd_decode_unit(
+                    _meta,
+                    cutlass.Int32(_nb),
+                    _uid,
+                    cutlass.Int32(q.shape[2]),
+                    cutlass.Int32(self.q_tile),
+                    self.is_causal,
+                )
+                _phase = _phase + self._run_unit(
+                    q,
+                    k,
+                    v,
+                    o,
+                    lse,
+                    sinks,
+                    seq_q_lens,
+                    seq_kv_lens,
+                    tma_k_desc,
+                    tma_v_desc,
+                    softmax_scale_log2,
+                    sKV,
+                    sK,
+                    sV,
+                    k_tma_mbar,
+                    v_tma_mbar,
+                    lane,
+                    warp,
+                    _phase,
+                    cutlass.Int32(0),
+                    _b,
+                    _qt,
+                    _b,
+                    _h,
+                )
+                _uid = thd_claim_next(seq_kv_lens, cutlass.Int32(4 * _nb + 3), _slot, cutlass.Int32(tidx))
+        else:
+            q_tile_idx, batch_idx, head_idx = cute.arch.block_idx()
+            # KV split rides the BATCH axis: grid.y = batch + split*B.  Q/K/V and
+            # the per-batch seqlens must use the REAL batch; only the O/LSE
+            # partial slot uses the composite.  Folds away at split_kv == 1.
+            split_idx = cutlass.Int32(0)
+            o_batch_idx = batch_idx
+            if cutlass.const_expr(self.split_kv > 1):
+                n_batch_real = cutlass.Int32(q.shape[0])
+                split_idx = batch_idx // n_batch_real
+                batch_idx = batch_idx % n_batch_real
+                o_batch_idx = split_idx * n_batch_real + batch_idx
+            if cutlass.const_expr(self.sched_policy != SCHED_NATURAL):
+                _n_qh = cutlass.Int32((q.shape[2] // self.qh_per_kh if self.pack_gqa else q.shape[2]))
+                _n_batch = cutlass.Int32(self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0])
+                if cutlass.const_expr(self.sched_policy == SCHED_LPT_L2):
+                    q_tile_idx, head_idx, batch_idx = lpt_l2_tile_coords(
+                        q_tile_idx,
+                        _n_qh,
+                        _n_batch,
+                        n_q_tiles,
+                        _n_qh // cutlass.Int32(k.shape[2]),
+                        cutlass.Int32(k.shape[1]),
+                        (self.head_tile_qk + self.head_tile_v) * self.in_dtype.width // 8,
+                        _SCHED_L2_BUDGET_BYTES,
+                    )
+                else:
+                    q_tile_idx, head_idx, batch_idx = lpt_tile_coords(q_tile_idx, _n_qh, _n_batch, n_q_tiles)
+                o_batch_idx = batch_idx
+            elif cutlass.const_expr(self.is_causal):
+                # Diagonal-bounded work grows with the Q tile. Launch long tiles
+                # first to avoid leaving a few expensive CTAs in the final
+                # scheduler waves (right-band graphs share the causal shape).
+                grid_q, _, _ = cute.arch.grid_dim()
+                q_tile_idx = grid_q - q_tile_idx - 1
+            if cutlass.const_expr(self.thd_varlen):
+                # blockIdx.y is a RANK, not a batch: the remap puts the longest
+                # sequences on the low CTA ids so they dispatch first, leaving
+                # the ragged tail of the launch to short sequences whose
+                # over-length tiles retire without work.
+                _m = cutlass.make_array_view(seq_kv_lens)
+                _n = (seq_kv_lens.shape[0] - 4) // 4
+                batch_idx = cutlass.Int32(_m[3 * _n + 2 + batch_idx])
+                o_batch_idx = batch_idx
+            self._run_unit(
+                q,
+                k,
+                v,
+                o,
+                lse,
+                sinks,
+                seq_q_lens,
+                seq_kv_lens,
+                tma_k_desc,
+                tma_v_desc,
+                softmax_scale_log2,
+                sKV,
+                sK,
+                sV,
+                k_tma_mbar,
+                v_tma_mbar,
+                lane,
+                warp,
+                cutlass.Int32(0),
+                split_idx,
+                o_batch_idx,
+                q_tile_idx,
+                batch_idx,
+                head_idx,
+            )
+
     @cute.jit
     def __call__(
         self,
@@ -1392,6 +1553,7 @@ def __call__(
         thd_q_lens: Optional[cute.Tensor],
         thd_kv_lens: Optional[cute.Tensor],
         thd_lens_form: Optional[cutlass.Int32],
+        thd_n_ctas: cutlass.Int32,
         stream: cuda_driver.CUstream,
     ) -> None:
         """Launch the SM120 cutlass FMHA kernel.
@@ -1489,8 +1651,8 @@ def _static_neq(a, b):
         if cutlass.const_expr(self.thd_varlen):
             if cutlass.const_expr(q.shape[0] != 1):
                 raise ValueError("THD Q/K/V/O must be packed batch-1 views")
-            if cutlass.const_expr(seq_kv_lens.shape != (3 * self.thd_batch + 2,)):
-                raise ValueError("THD seq_kv_lens must be the (3*B+2,) metadata tensor")
+            if cutlass.const_expr(seq_kv_lens.shape != (4 * self.thd_batch + 4,)):
+                raise ValueError("THD seq_kv_lens must be the (4*B+4,) metadata tensor")
 
         # Split D into I contiguous C-element chunks while preserving the
         # per-tensor TMA descriptor over the compact (B, S, H, D) storage.
@@ -1535,16 +1697,19 @@ def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_e
         tma_k_desc = kv_tma_desc(k, head_dim_qk, self.head_tile_qk, self.k_tma_swizzle, self.k_tma_swizzle_chunks, self.k_swizzle_chunk_elems)
         tma_v_desc = kv_tma_desc(v, head_dim_v, self.head_tile_v, self.v_tma_swizzle, self.v_tma_swizzle_chunks, self.v_swizzle_chunk_elems)
         if cutlass.const_expr(self.thd_varlen):
-            # Build the [kv|cu_q|cu_k] metadata buffer DEVICE-side from the
-            # caller's length tensors (no host cumsum, no H2D — issue #552);
-            # the main kernel launched after it on this stream reads it.
+            # Build the [kv|cu_q|cu_k|remap|live|ctr] metadata buffer DEVICE-side
+            # from the caller's length tensors (no host cumsum, no H2D — issue
+            # #552); the main kernel launched after it on this stream reads it.
             _build_thd_meta_kernel(
                 seq_kv_lens,
                 thd_q_lens,
                 thd_kv_lens,
                 thd_lens_form,
                 cutlass.Int32(self.thd_batch),
-            ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+                cutlass.Int32(q.shape[2]),
+                cutlass.Int32(self.q_tile),
+                thd_n_ctas,
+            ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         # Grid geometry. THD: ceil(max_seq_q / q_tile) tiles per sequence over the
         # REAL batch count (the packed view's batch mode is 1); tiles past a shorter
         # sequence's length drain without work. NOTE thd_max_sq is a __call__
@@ -1567,6 +1732,11 @@ def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_e
             # KV split rides the batch axis: y = batch + split*B (config_sm120
             # allows split_kv > 1 only under NATURAL, whose 3-D grid has one).
             grid = (n_q_tiles, n_batch * self.split_kv, n_head)
+        # Persistent THD: a flat, machine-sized grid -- the unit a CTA works on
+        # comes from the claim counter, not from its block index.
+        _persistent = self.thd_varlen and THD_PERSISTENT and self.sched_policy == SCHED_NATURAL
+        if cutlass.const_expr(_persistent):
+            grid = (thd_n_ctas, cutlass.Int32(1), cutlass.Int32(1))
         self.kernel(
             q,
             k,
@@ -1710,7 +1880,7 @@ def _fake_bshd(shape, stride):
     )
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
-        (3 * b + 2,) if PARAMS.thd_varlen else (b,),  # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) ]
+        (4 * b + 4,) if PARAMS.thd_varlen else (b,),  # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) | remap(B) | live | ctr ]
         stride_order=(0,),
         assumed_align=4,
     )
@@ -1741,6 +1911,7 @@ def _fake_bshd(shape, stride):
         fake_thd_q_lens,
         fake_thd_kv_lens,
         fake_thd_lens_form,
+        cutlass.Int32(0),  # thd_n_ctas: persistent THD grid extent (runtime)
         cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
         options="--enable-tvm-ffi",
     )
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py
index a23ce79eb..c1a2fdb09 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py
@@ -73,7 +73,7 @@
     lpt_l2_tile_coords,
 )
 from cudnn.frost.tile_dsl.swizzle import swizzle_xor
-from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_kernel as _build_thd_meta_kernel
+from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_kernel as _build_thd_meta_kernel, THD_SETUP_THREADS
 from cudnn.sdpa.fwd.config_sm120 import (
     FP8_HEAD_TILE_GRANULE,
     SEQ_KV_TILES as _SEQ_KV_TILES,
@@ -1046,11 +1046,17 @@ def kernel(
         q_row_base = cutlass.Int32(0)
         kv_row_base = cutlass.Int32(0)
         if cutlass.const_expr(self.thd_varlen):
-            # THD: seq_kv_lens is the metadata tensor [seq_kv(B) | cu_q(B+1) | cu_k(B+1)].
+            # THD: seq_kv_lens is the metadata tensor
+            # [seq_kv(B) | cu_q(B+1) | cu_k(B+1) | remap(B) | live | ctr].
             # Per-sequence lengths come from the prefix sums; the bases offset
             # every packed (1, T, H, D) access below.
-            n_batch = (seq_kv_lens.shape[0] - 2) // 3
+            n_batch = (seq_kv_lens.shape[0] - 4) // 4
             meta = cutlass.make_array_view(seq_kv_lens)
+            # Longest sequences first: blockIdx.y indexes the remap rather than
+            # the batch directly, so long sequences take the low CTA ids and
+            # dispatch first. The ragged tail of the launch is then short
+            # sequences, whose over-length tiles retire without work.
+            batch_idx = cutlass.Int32(meta[3 * n_batch + 2 + batch_idx])
             q_row_base = cutlass.Int32(meta[n_batch + batch_idx])
             seqlen_q = cutlass.Int32(meta[n_batch + batch_idx + 1]) - q_row_base
             kv_row_base = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx])
@@ -1597,6 +1603,7 @@ def __call__(
         thd_q_lens: Optional[cute.Tensor],
         thd_kv_lens: Optional[cute.Tensor],
         thd_lens_form: Optional[cutlass.Int32],
+        thd_n_ctas: cutlass.Int32,
         stream: cuda_driver.CUstream,
     ) -> None:
         """Launch the SM120 per-tensor FP8 FMHA kernel.
@@ -1693,8 +1700,8 @@ def _static_neq(a, b):
         if cutlass.const_expr(self.thd_varlen):
             if cutlass.const_expr(q.shape[0] != 1):
                 raise ValueError("THD Q/K/V/O must be packed batch-1 views")
-            if cutlass.const_expr(seq_kv_lens.shape != (3 * self.thd_batch + 2,)):
-                raise ValueError("THD seq_kv_lens must be the (3*B+2,) metadata tensor")
+            if cutlass.const_expr(seq_kv_lens.shape != (4 * self.thd_batch + 4,)):
+                raise ValueError("THD seq_kv_lens must be the (4*B+4,) metadata tensor")
 
         # Exact head dims: split D into I contiguous C-element chunks while
         # preserving the compact (B, S, H, D) global-memory address
@@ -1727,16 +1734,19 @@ def kv_tma_desc(t, head_dim, swizzle, swizzle_chunks, swizzle_chunk_elems, envel
         tma_k_desc = kv_tma_desc(k, head_dim_qk, self.k_tma_swizzle, self.k_tma_swizzle_chunks, self.k_swizzle_chunk_elems, head_dim_qk != self.head_tile_qk)
         tma_v_desc = kv_tma_desc(v, head_dim_v, self.v_tma_swizzle, self.v_tma_swizzle_chunks, self.v_swizzle_chunk_elems, head_dim_v != self.head_tile_v)
         if cutlass.const_expr(self.thd_varlen):
-            # Build the [kv|cu_q|cu_k] metadata buffer DEVICE-side from the
-            # caller's length tensors (no host cumsum, no H2D — issue #552);
-            # the main kernel launched after it on this stream reads it.
+            # Build the [kv|cu_q|cu_k|remap|live|ctr] metadata buffer DEVICE-side
+            # from the caller's length tensors (no host cumsum, no H2D — issue
+            # #552); the main kernel launched after it on this stream reads it.
             _build_thd_meta_kernel(
                 seq_kv_lens,
                 thd_q_lens,
                 thd_kv_lens,
                 thd_lens_form,
                 cutlass.Int32(self.thd_batch),
-            ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+                cutlass.Int32(q.shape[2]),
+                cutlass.Int32(self.q_tile),
+                thd_n_ctas,
+            ).launch(grid=(1, 1, 1), block=(THD_SETUP_THREADS, 1, 1), stream=stream)
         # Grid geometry. THD: ceil(max_seq_q / q_tile) tiles per sequence over the
         # REAL batch count (the packed view's batch mode is 1); tiles past a shorter
         # sequence's length drain without work. NOTE thd_max_sq is a __call__
@@ -1901,7 +1911,7 @@ def compile(  # noqa: A001
     )
     fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
-        (3 * b + 2,) if PARAMS.thd_varlen else (b,),  # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) ]
+        (4 * b + 4,) if PARAMS.thd_varlen else (b,),  # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) | remap(B) | live | ctr ]
         stride_order=(0,),
         assumed_align=4,
     )
@@ -1955,6 +1965,7 @@ def _fake_scale():
         fake_thd_q_lens,
         fake_thd_kv_lens,
         fake_thd_lens_form,
+        cutlass.Int32(0),  # thd_n_ctas: unused on this path (no persistent grid)
         cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
         options="--enable-tvm-ffi",
     )
diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
index 106d621c2..085025355 100644
--- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
@@ -8,10 +8,54 @@
 
 import cutlass
 import cutlass.cute as cute
+from cutlass._mlir.dialects import arith
 import cuda.bindings.driver as _cuda_driver
 
 TENSOR_MAP_QWORDS = 128 // 8
 
+# THD metadata layout, in int32 words:
+#   [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) | batch_remap(B) ]
+# The trailing batch_remap is a permutation of [0, B) ordered by DESCENDING
+# Q length, so the tile scheduler walks the longest sequences first (longest
+# processing time): the tail of a THD launch is then made of short sequences,
+# which is what bounds the ragged last wave.
+THD_META_WORDS = lambda b: 4 * b + 4  # noqa: E731
+THD_REMAP_OFF = lambda b: 3 * b + 2  # noqa: E731
+THD_LIVE_OFF = lambda b: 4 * b + 2  # noqa: E731   live unit total (device-computed)
+THD_CTR_OFF = lambda b: 4 * b + 3  # noqa: E731    persistent-scheduler claim counter
+
+# Threads for the THD setup launch. The metadata write itself is one elected
+# thread; the batch-remap ranking that follows is parallel over batches, so the
+# block is sized for that (B > THD_SETUP_THREADS just loops).
+THD_SETUP_THREADS = 256
+
+
+@cute.jit
+def write_thd_batch_remap(meta, n_batch: cutlass.Int32, tid: cutlass.Int32, nthreads: cutlass.Int32) -> None:
+    """Fill batch_remap with [0, B) sorted by descending Q length.
+
+    Rank-by-counting rather than a sort network: each thread owns a batch and
+    counts how many sequences outrank it, which is O(B^2) comparisons but fully
+    parallel, branch-free and trivially deterministic. Ties break on the
+    original index, so the permutation is stable and reproducible run to run.
+
+    Must be called AFTER write_thd_meta (it reads the cu_seqlens_q it wrote)
+    with a barrier in between.
+    """
+    cuq0 = n_batch
+    remap0 = cutlass.Int32(3) * n_batch + cutlass.Int32(2)
+    i = tid
+    while i < n_batch:
+        len_i = cutlass.Int32(meta[cuq0 + i + cutlass.Int32(1)]) - cutlass.Int32(meta[cuq0 + i])
+        rank = cutlass.Int32(0)
+        for j in cutlass.range(0, n_batch, 1, unroll=1):
+            len_j = cutlass.Int32(meta[cuq0 + j + cutlass.Int32(1)]) - cutlass.Int32(meta[cuq0 + j])
+            # Descending by length; ties resolved by the lower original index.
+            outranks = (len_j > len_i) | ((len_j == len_i) & (j < i))
+            rank = rank + cutlass.Int32(arith.select(outranks.ir_value(), cutlass.Int32(1).ir_value(), cutlass.Int32(0).ir_value()))
+        meta[remap0 + rank] = i
+        i = i + nthreads
+
 
 @cute.jit
 def write_thd_meta(meta, ql, kl, lens_form: cutlass.Int32, n_batch: cutlass.Int32) -> None:
@@ -57,6 +101,71 @@ def write_thd_meta(meta, ql, kl, lens_form: cutlass.Int32, n_batch: cutlass.Int3
             meta[cuk0 + b + cutlass.Int32(1)] = acc_k
 
 
+@cute.jit
+def thd_decode_unit(
+    meta,
+    n_batch: cutlass.Int32,
+    uid: cutlass.Int32,
+    n_qh: cutlass.Int32,
+    q_tile: cutlass.Int32,
+    reverse_rows: bool,
+) -> tuple:
+    """Map a linear unit id to ``(q_tile_idx, batch, head)`` through batch_remap.
+
+    A unit is ``q_tile`` rows of one head of one sequence. Sequences are walked
+    LONGEST FIRST (the remap), and the head is the major axis within a sequence
+    so consecutive units sweep the Q tiles of a single head — those share a K/V
+    head, which is what keeps the claim order L2-friendly. ``reverse_rows``
+    walks a sequence's tiles from the diagonal back, putting the causal-heavy
+    tiles first.
+
+    A uid past the live total keeps ``batch == n_batch``; the caller is expected
+    to bound uid against the live count instead of relying on that sentinel.
+    """
+    cuq0 = n_batch
+    remap0 = cutlass.Int32(3) * n_batch + cutlass.Int32(2)
+    f_batch = n_batch
+    f_head = cutlass.Int32(0)
+    f_qt = cutlass.Int32(0)
+    done = cutlass.Int32(0)
+    acc = cutlass.Int32(0)
+    for i in cutlass.range(0, n_batch, 1, unroll=1):
+        b = cutlass.Int32(meta[remap0 + i])
+        s_i = cutlass.Int32(meta[cuq0 + b + cutlass.Int32(1)]) - cutlass.Int32(meta[cuq0 + b])
+        tb = (s_i + q_tile - cutlass.Int32(1)) // q_tile
+        # A zero-length sequence contributes no unit; keep the divisor legal
+        # anyway, since both quotients below are evaluated before the select.
+        tb_nz = cute.math.max(tb, cutlass.Int32(1))
+        units_b = tb * n_qh
+        in_rng = (done == cutlass.Int32(0)) & (uid < acc + units_b)
+        local = uid - acc
+        qt = local % tb_nz
+        if cutlass.const_expr(reverse_rows):
+            qt = tb - cutlass.Int32(1) - qt
+        f_batch = cutlass.Int32(arith.select(in_rng.ir_value(), b.ir_value(), f_batch.ir_value()))
+        f_head = cutlass.Int32(arith.select(in_rng.ir_value(), (local // tb_nz).ir_value(), f_head.ir_value()))
+        f_qt = cutlass.Int32(arith.select(in_rng.ir_value(), qt.ir_value(), f_qt.ir_value()))
+        done = cutlass.Int32(arith.select(in_rng.ir_value(), cutlass.Int32(1).ir_value(), done.ir_value()))
+        acc = acc + units_b
+    return f_qt, f_batch, f_head
+
+
+@cute.jit
+def thd_claim_next(meta_t: cute.Tensor, ctr_off: cutlass.Int32, slot, tidx: cutlass.Int32) -> cutlass.Int32:
+    """Take the next unit from the device-side claim counter.
+
+    One atomic for the whole CTA, broadcast through a single SMEM word. The
+    leading barrier also separates the previous unit's use of the shared K/V
+    staging from the next unit's, so the caller does not need its own.
+    """
+    ctr_ptr = Pointer(meta_t.iterator.raw_ptr(), dtype=cutlass.Int32) + ctr_off
+    nvvm.barrier_cta_sync(0)
+    if tidx == cutlass.Int32(0):
+        slot[0] = cutlass.Int32(nvvm.atomicrmw(nvvm.AtomicOp.ADD, ctr_ptr, cutlass.Int32(1)))
+    nvvm.barrier_cta_sync(0)
+    return cutlass.Int32(slot[0])
+
+
 @cute.kernel
 def build_thd_meta_kernel(
     meta_t: cute.Tensor,
@@ -64,19 +173,40 @@ def build_thd_meta_kernel(
     kv_lens_t: cute.Tensor,
     lens_form: cutlass.Int32,
     n_batch: cutlass.Int32,
+    n_qh: cutlass.Int32,
+    q_tile: cutlass.Int32,
+    n_ctas: cutlass.Int32,
 ) -> None:
     """Meta-only THD setup (SM120: no per-batch O TMA descriptors — O stores
-    are raw pointer writes predicated per row). One elected thread; the main
-    kernel launched after it on the same stream sees the writes by kernel
-    boundary ordering."""
+    are raw pointer writes predicated per row). The metadata write is one
+    elected thread; the batch remap and the live-unit total that follow are
+    whole-block. The main kernel launched after it on the same stream sees the
+    writes by kernel boundary ordering."""
+    meta = cutlass.make_array_view(meta_t)
     if nvvm.elect_sync():
         write_thd_meta(
-            cutlass.make_array_view(meta_t),
+            meta,
             cutlass.make_array_view(q_lens_t),
             cutlass.make_array_view(kv_lens_t),
             lens_form,
             n_batch,
         )
+    # Barrier first: the ranking reads the cu_seqlens_q written above.
+    cute.arch.barrier()
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    write_thd_batch_remap(meta, n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
+    # Live unit total + claim counter, as on SM100 — a SM120 unit is q_tile
+    # rows of one head, so the same count applies with cga_tile_m := q_tile.
+    cute.arch.barrier()
+    if nvvm.elect_sync() and tidx == cutlass.Int32(0):
+        live = cutlass.Int32(0)
+        cuq0 = n_batch
+        for b in cutlass.range(0, n_batch, 1, unroll=1):
+            s_b = cutlass.Int32(meta[cuq0 + b + cutlass.Int32(1)]) - cutlass.Int32(meta[cuq0 + b])
+            live = live + ((s_b + q_tile - cutlass.Int32(1)) // q_tile) * n_qh
+        meta[cutlass.Int32(4) * n_batch + cutlass.Int32(2)] = live
+        meta[cutlass.Int32(4) * n_batch + cutlass.Int32(3)] = n_ctas
 
 
 @cute.kernel
@@ -151,6 +281,14 @@ def build_thd_meta_o_kv_descs_kernel(
             from_proxy=nvvm.Proxy.GENERIC,
             to_proxy=nvvm.Proxy.TENSORMAP,
         )
+    # Outside the elect: every thread helps rank the batches. The barrier makes
+    # the cu_seqlens_q written above visible to the whole block first. The
+    # decode walks this permutation on every THD flavor, so a setup that skips
+    # it leaves the region uninitialized and units decode garbage batches.
+    cute.arch.barrier()
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    write_thd_batch_remap(cutlass.make_array_view(meta_t), n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
 
 
 @cute.kernel
@@ -165,6 +303,8 @@ def build_thd_meta_o_descs_kernel(
     n_qh: cutlass.Int32,
     n_batch: cutlass.Int32,
     o_row_stride: cutlass.Int32,
+    cga_tile_m: cutlass.Int32,
+    n_clusters: cutlass.Int32,
 ) -> None:
     """Per-execute THD setup, one elected thread (issue #552, D2H removal):
     build the [seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1)] metadata
@@ -213,3 +353,23 @@ def build_thd_meta_o_descs_kernel(
             from_proxy=nvvm.Proxy.GENERIC,
             to_proxy=nvvm.Proxy.TENSORMAP,
         )
+    # Outside the elect: every thread helps rank the batches. The barrier makes
+    # the cu_seqlens_q written above visible to the whole block first.
+    cute.arch.barrier()
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    write_thd_batch_remap(cutlass.make_array_view(meta_t), n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
+    # Live unit total + claim counter for the persistent scheduler. The host
+    # cannot know Sigma_b ceil(s_b/tile)*QH without a D2H (issue #552), so the
+    # kernel reads its own bound from here. The counter starts at n_clusters:
+    # cluster c takes unit c from its blockIdx, then pulls from the counter.
+    cute.arch.barrier()
+    if nvvm.elect_sync() and tidx == cutlass.Int32(0):
+        meta_w = cutlass.make_array_view(meta_t)
+        cuq0 = n_batch
+        live = cutlass.Int32(0)
+        for b in cutlass.range(0, n_batch, 1, unroll=1):
+            s_b = cutlass.Int32(meta_w[cuq0 + b + cutlass.Int32(1)]) - cutlass.Int32(meta_w[cuq0 + b])
+            live = live + ((s_b + cga_tile_m - cutlass.Int32(1)) // cga_tile_m) * n_qh
+        meta_w[cutlass.Int32(4) * n_batch + cutlass.Int32(2)] = live
+        meta_w[cutlass.Int32(4) * n_batch + cutlass.Int32(3)] = n_clusters
diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
index c5241fdaf..bc7d1778f 100644
--- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
+++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
@@ -899,6 +899,30 @@ def test_dsl_sm120_thd():
     _run_thd_case(seq_q_lens=[200, 150], seq_kv_lens=[200, 150], is_causal=True)
 
 
+@pytest.mark.L0
+@torch_fork_set_rng(seed=27)
+def test_dsl_sm120_thd_multi_unit_per_cta():
+    """THD with more live units than the machine has CTAs.
+
+    The other THD cases are small enough that every CTA is handed at most one
+    unit, so they never exercise re-entering the K/V pipeline for a second one
+    -- the regime where an unmatched consumer arrival on bar_k/v_consumed used
+    to desynchronise the next unit's producer handshake. These lengths give
+    O(100) units against a grid sized to the SM count, so CTAs claim repeatedly,
+    and each unit spans several K/V tiles.
+    """
+
+    _run_thd_case(
+        seq_q_lens=[1024, 768, 512, 256],
+        seq_kv_lens=[1024, 768, 512, 256],
+        h_q=8,
+        h_kv=2,
+        head_dim=128,
+        is_causal=True,
+        check_stats=True,
+    )
+
+
 @pytest.mark.L0
 @torch_fork_set_rng(seed=23)
 def test_dsl_sm120_thd_cross():

From 8cf7db6fddf825c7a03e296a1c0a303afe563a1b Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Thu, 20 Aug 2026 15:45:12 -0700
Subject: [PATCH 2/9] frost(sdpa): keep THD on its own tile scheduler

The causal auto-upgrade promoted any NATURAL graph to LPT/LPT_L2 with no THD
exclusion, so a ragged batch silently ran a decode that assumes a dense
rectangular tile space. THD carries its own scheduler, which walks the live
units through batch_remap, so exclude it on both arches and drop the policy
term from the SM120 persistent gate.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/cudnn/sdpa/fwd/api_dsl.py                   | 10 ++++++++--
 python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py |  4 ++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 7cfc39921..d6f39b73c 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -1052,7 +1052,10 @@ def compile(self) -> None:
         sched_policy = self.sched_policy
         if sched_policy is None:
             sched_policy = SCHED_NATURAL
-            if self.window_right is not None:
+            # THD is excluded: the LPT decodes assume a dense rectangular
+            # tile space, while a ragged batch carries its own scheduler,
+            # which walks the live units through batch_remap.
+            if self.window_right is not None and not self.thd:
                 # Causal: balance the triangular load; pick the LPT variant by working set.
                 _, _, s_kv_sched, _ = self.k_desc.shape
                 _, _, _, d_qk_sched = self.q_desc.shape
@@ -2584,7 +2587,10 @@ def compile(self) -> None:
         sched_policy = self.sched_policy
         if sched_policy is None:
             sched_policy = SCHED_NATURAL
-            if self.window_right is not None:
+            # THD is excluded: the LPT decodes assume a dense rectangular
+            # tile space, while a ragged batch carries its own scheduler,
+            # which walks the live units through batch_remap.
+            if self.window_right is not None and not self.thd:
                 # Causal: balance the triangular load; pick the LPT variant by working set.
                 _, _, s_kv_sched, _ = self.k_desc.shape
                 _, _, _, d_qk_sched = self.q_desc.shape
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
index 5060fd9d1..96e47cad9 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
@@ -1412,7 +1412,7 @@ def kernel(
         prims.fence_mbarrier_init()
         prims.barrier_cta_sync(0)
 
-        if cutlass.const_expr(self.thd_varlen and THD_PERSISTENT and self.sched_policy == SCHED_NATURAL):
+        if cutlass.const_expr(self.thd_varlen and THD_PERSISTENT):
             # Persistent grid: the launch is sized to the MACHINE, not to the
             # plan-time envelope, and each CTA pulls units from a device-side
             # counter until the live total runs out. That total depends on the
@@ -1734,7 +1734,7 @@ def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_e
             grid = (n_q_tiles, n_batch * self.split_kv, n_head)
         # Persistent THD: a flat, machine-sized grid -- the unit a CTA works on
         # comes from the claim counter, not from its block index.
-        _persistent = self.thd_varlen and THD_PERSISTENT and self.sched_policy == SCHED_NATURAL
+        _persistent = self.thd_varlen and THD_PERSISTENT
         if cutlass.const_expr(_persistent):
             grid = (thd_n_ctas, cutlass.Int32(1), cutlass.Int32(1))
         self.kernel(

From c8f8c88559c368815e569501ce913bcb263a7986 Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 13:12:41 -0700
Subject: [PATCH 3/9] frost(sdpa): resolve the persistent-THD CTA count once
 per device

The SM count was queried through torch.cuda.get_device_properties on every
execute. Measured on a host-bound THD shape (B=1, h=1, d=128, s=256), the
execute drops from 94.1 to 89.2 us/call, level with develop. A device
property cannot change under a live process.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/cudnn/sdpa/fwd/api_dsl.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index d6f39b73c..0f1371f49 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -148,6 +148,11 @@ def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch
 # reverse-row LPT.
 _SCHED_L2_BUDGET_BYTES = 50 * 1024 * 1024
 
+# SM count per device for the persistent THD grid. A device property cannot
+# change under a live process, and the query costs ~7 us on the execute hot
+# path, so resolve it once per device.
+_THD_CTAS_CACHE: dict = {}
+
 
 def _causal_sched_policy(s_kv: int, d_qk: int, d_v: int, elem_bytes: int) -> int:
     """SCHED_LPT_L2 vs SCHED_LPT for a causal graph (see _SCHED_L2_BUDGET_BYTES)."""
@@ -3212,12 +3217,18 @@ def _thd_persistent_ctas(self, device) -> int:
         under-launching only costs parallelism. One CTA per SM matches the
         kernel's ``min_blocks_per_mp``.
         """
-        forced = int(os.environ.get("FROST_THD_CTAS", "0"))
-        if forced > 0:
-            return forced
-        sms = torch.cuda.get_device_properties(device).multi_processor_count
-        per_sm = int(os.environ.get("FROST_THD_CTAS_PER_SM", "1"))
-        return max(1, sms * max(1, per_sm))
+        key = getattr(device, "index", device)
+        n = _THD_CTAS_CACHE.get(key)
+        if n is None:
+            forced = int(os.environ.get("FROST_THD_CTAS", "0"))
+            if forced > 0:
+                n = forced
+            else:
+                sms = torch.cuda.get_device_properties(device).multi_processor_count
+                per_sm = int(os.environ.get("FROST_THD_CTAS_PER_SM", "1"))
+                n = max(1, sms * max(1, per_sm))
+            _THD_CTAS_CACHE[key] = n
+        return n
 
     def scratch_workspace_bytes(self) -> int:
         if self.thd:

From c574630d62a9ddaca5e2eea6dbd861f512f11d8a Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 13:12:51 -0700
Subject: [PATCH 4/9] frost(sdpa): skip the stream switch when already on the
 launch stream

_torch_stream_context built two torch Stream objects (~3.4 us each) and
entered a stream context on every use, several times per execute, even when
the launch stream is the one torch is already on -- where the switch is a
no-op. Compare the raw handle instead and return early.

Measured on a host-bound THD shape (B=1, h=1, d=128, s=256, release cuDNN):
86.7 -> 59.2 us/call. Not THD-specific; this context is used by every SDPA
DSL execute path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/cudnn/sdpa/fwd/api_dsl.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 0f1371f49..46baac3e1 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -130,6 +130,14 @@ def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch
         # ordered against the default stream here, so run in place.
         yield
         return
+    # Fast path: the launch stream is almost always the one torch is already on,
+    # and entering a context for the current stream is a no-op. Building the two
+    # Stream objects below costs ~3.4 us each and this runs several times per
+    # execute; the raw handle getter is ~0.07 us.
+    _idx = device.index if device.index is not None else torch.cuda.current_device()
+    if handle == torch._C._cuda_getCurrentRawStream(_idx):
+        yield
+        return
     torch_current = torch.cuda.current_stream(device)
     torch_default = torch.cuda.default_stream(device)
     if handle == torch_current.cuda_stream:

From a8f8e5bb8783f0b623837a552bb5fa46bfbd743f Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 12:05:56 -0700
Subject: [PATCH 5/9] frost(sdpa): pin THD to NATURAL in the fwd heuristic

develop's #692 promoted the causal LPT/LPT_L2 choice onto the graph path via
_sched_points, which branches on facts.causal alone. The adapters' THD
exclusion therefore covered only the standalone-wrapper tier, and a ragged
graph was ranked straight back onto an LPT decode built for a dense
rectangular tile space. Exclude it there too, alongside the split-KV rule
that already special-cases THD.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/cudnn/sdpa/fwd/heuristics.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py
index ffab44ab8..19fe1b815 100644
--- a/python/cudnn/sdpa/fwd/heuristics.py
+++ b/python/cudnn/sdpa/fwd/heuristics.py
@@ -232,6 +232,14 @@ def _sched_points(caps: Capabilities, facts) -> List[Optional[int]]:
     domain = caps.sched_policies
     if len(domain) <= 1:
         return [_sole(domain)]
+    if facts.thd and SCHED_NATURAL in domain:
+        # A ragged batch carries its own scheduler: it walks the LIVE units
+        # through batch_remap over a machine-sized grid. The LPT decodes map a
+        # linear tile id onto a dense rectangular tile space, so ranking them
+        # here would hand THD a decode built for a geometry it does not have --
+        # and spend autotune slots on it. Same exclusion the adapters apply to
+        # their standalone-wrapper derivation.
+        return [SCHED_NATURAL]
     causal_ish = facts.causal or facts.right_band_widening
     if caps.sm_hi == 80:
         # SM80's measured choices (see the adapter's flavor table): causal

From caebb8f8ae092a943ead7ab7b15db20029a27197 Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 16:23:22 -0700
Subject: [PATCH 6/9] frost(sdpa): narrow the THD setup elect to warp 0,
 decouple the handout stride

Review follow-ups on the persistent THD scheduler.

elect_sync elects one thread PER WARP. The setup kernels ran their
single-thread body under a bare elect_sync, which was safe while they
launched one warp -- this series widened them to THD_SETUP_THREADS so the
batch ranking could run in parallel, leaving 8 elected threads racing.
The metadata writes are idempotent, but the O descriptor build is not: a
lagging warp's base-qword copy can land after another's tensormap_replace
and revert the patched address. Guard all three setup kernels on warp 0.

The persistent handout strided the unit id by the cluster size while the
consumer decodes it back out as linear // CGA_M. Equal while CGA_N == 1,
which is every config today; take both so the two stop being coupled.

Also clamp the decode's per-sequence tile count before dividing (a
zero-length sequence gives cb == 0, and arith.select evaluates both arms),
matching thd_decode_unit's tb_nz, and correct the stale grid comment: the
adapter hands down an occupancy-capped cluster count, not the envelope.

SM100 631 passed, SM120 499 passed, no regressions.
---
 python/cudnn/frost/tile_dsl/scheduler.py      | 10 ++++++-
 .../cudnn/sdpa/fwd/kernels/_common_sm100.py   |  9 ++++--
 .../fwd/kernels/prefill_d128_f16_sm100.py     | 14 +++++----
 python/cudnn/sdpa/fwd/kernels/thd_sm100.py    | 29 +++++++++++++------
 4 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/python/cudnn/frost/tile_dsl/scheduler.py b/python/cudnn/frost/tile_dsl/scheduler.py
index 26fbf3fb7..f3da2f04d 100644
--- a/python/cudnn/frost/tile_dsl/scheduler.py
+++ b/python/cudnn/frost/tile_dsl/scheduler.py
@@ -119,6 +119,7 @@ def scheduler_warp_loop_persistent(
     ctr_off,
     live_off,
     cga_size: int,
+    cga_m: int,
 ):
     """Persistent tile scheduler over a LIVE-ONLY unit range (THD).
 
@@ -132,6 +133,11 @@ def scheduler_warp_loop_persistent(
     peer's scheduler mbarrier -- the same shape as ``read_tile_id_arrive``.
     Multicast is not available here: it is a clusterlaunchcontrol facility, so
     dynamic claiming needs an explicit peer write.
+
+    ``cga_size`` is the DSMEM broadcast fan-out (CTAs per cluster); ``cga_m`` is
+    the stride the consumer's decode divides back out. They are equal while
+    CGA_N == 1, which is every config today -- taking both keeps the handout
+    correct if that ever stops holding.
     """
     meta = cutlass.make_array_view(meta_t)
     ctr_ptr = Pointer(meta_t.iterator.raw_ptr(), dtype=cutlass.Int32) + ctr_off
@@ -151,7 +157,9 @@ def scheduler_warp_loop_persistent(
             uid = cutlass.Int32(nvvm.atomicrmw(nvvm.AtomicOp.ADD, ctr_ptr, cutlass.Int32(1)))
             live = cutlass.Int32(meta[live_off])
             valid = cutlass.Int32(arith.select((uid < live).ir_value(), cutlass.Int32(1).ir_value(), cutlass.Int32(0).ir_value()))
-            linear = uid * cutlass.Int32(cga_size)
+            # Stride by CGA_M, not the cluster size: the consumer decodes the
+            # unit id back out as linear // CGA_M.
+            linear = uid * cutlass.Int32(cga_m)
             # store_async_dsmem wants CuTe pointers; the smem arrays hand out
             # base-DSL ones, so convert through the raw addresses.
             _tile_ptr = cute.make_ptr(
diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
index 52e4a5d66..23b150468 100644
--- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py
@@ -765,12 +765,17 @@ def _thd_decode(linear_cta, seq_kv_lens_t, n_batch, n_qh, cta_in_pair):
             s_i = cutlass.Int32(cu[cuq0 + b + cutlass.Int32(1)]) - cutlass.Int32(cu[cuq0 + b])
             cb = (s_i + cutlass.Int32(cga_tile_m - 1)) // cutlass.Int32(cga_tile_m)
             units_b = cb * n_qh
+            # A zero-length sequence gives cb == 0, and units_b == 0 with it, so
+            # in_rng is false and the quotient is discarded — but arith.select
+            # evaluates BOTH arms, so divide by a clamped copy to keep the dead
+            # arm defined. units_b keeps the true cb (thd_decode_unit's tb_nz).
+            cb_nz = cute.math.max(cb, cutlass.Int32(1))
             in_rng = (done == cutlass.Int32(0)) & (u < acc + units_b)
             local = u - acc
             # Natural order within a sequence (head-major, ascending rows).
             f_batch = cutlass.Int32(arith.select(in_rng.ir_value(), b.ir_value(), f_batch.ir_value()))
-            f_head = cutlass.Int32(arith.select(in_rng.ir_value(), (local // cb).ir_value(), f_head.ir_value()))
-            f_qc = cutlass.Int32(arith.select(in_rng.ir_value(), (local % cb).ir_value(), f_qc.ir_value()))
+            f_head = cutlass.Int32(arith.select(in_rng.ir_value(), (local // cb_nz).ir_value(), f_head.ir_value()))
+            f_qc = cutlass.Int32(arith.select(in_rng.ir_value(), (local % cb_nz).ir_value(), f_qc.ir_value()))
             done = cutlass.Int32(arith.select(in_rng.ir_value(), cutlass.Int32(1).ir_value(), done.ir_value()))
             acc = acc + units_b
         q_super = f_qc * cutlass.Int32(CFG.CTA_MMA) + cta_in_pair
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 27a988f04..9301682df 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
@@ -646,6 +646,7 @@ def _kernel(
                 cutlass.Int32(4) * n_batch + cutlass.Int32(3),
                 cutlass.Int32(4) * n_batch + cutlass.Int32(2),
                 CGA_SIZE,
+                CFG.CGA_M,
             )
         else:
             scheduler_warp_loop(sched, CFG.SCHEDULER_STAGES, is_cga_first_cta)
@@ -2185,11 +2186,14 @@ def _tma_swz(byte_w: int):
         # DEVICE-side from the caller's length tensors (no host cumsum, no
         # H2D — issue #552), then the per-batch O descriptor array (reuse
         # tma_o_desc over the packed [1,T,QH,D_v] O as base). Main grid: the
-        # PLAN-TIME ENVELOPE (n_thd_units = B * ceil(S_q_decl/CGA_TILE_M) * QH,
-        # from the DECLARED S_q — no runtime length reaches the host); units
-        # past a sequence's live tiles decode the batch == n_batch sentinel
-        # and drain without loads or stores. grid_x = n_thd_units * CGA_M.
-        # Works at cga1 (CGA_M=1).
+        # PERSISTENT cluster count — the adapter hands down n_thd_units already
+        # capped to what the device holds resident, min(plan-time envelope,
+        # SMs / CTA_MMA), NOT the envelope itself. It doubles as the claim
+        # counter's seed: cluster c runs unit c off its blockIdx, then pulls
+        # from the counter, so the grid and the seed must be the same number.
+        # Dispatching past the live total stays safe either way — such a unit
+        # decodes the batch == n_batch sentinel and drains without loads or
+        # stores. grid_x = n_thd_units * CGA_M. Works at cga1 (CGA_M=1).
         # ENVELOPE: the packed-O row stride is QH * ACTUAL d_v (o_tensor's
         # static inner extent), not QH * TILE_O — the per-batch descriptor
         # bases must step in real rows or every batch >= 1 lands OOB.
diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
index 085025355..cc17155cc 100644
--- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
@@ -183,7 +183,13 @@ def build_thd_meta_kernel(
     whole-block. The main kernel launched after it on the same stream sees the
     writes by kernel boundary ordering."""
     meta = cutlass.make_array_view(meta_t)
-    if nvvm.elect_sync():
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    # elect_sync elects one thread PER WARP, and this block is THD_SETUP_THREADS
+    # wide so the ranking below can run in parallel — narrow the single-thread
+    # body to warp 0's leader. Every warp still evaluates elect_sync (it is warp
+    # -uniform); only the added predicate is what excludes warps 1..N.
+    if nvvm.elect_sync() and tidx < cutlass.Int32(32):
         write_thd_meta(
             meta,
             cutlass.make_array_view(q_lens_t),
@@ -193,8 +199,6 @@ def build_thd_meta_kernel(
         )
     # Barrier first: the ranking reads the cu_seqlens_q written above.
     cute.arch.barrier()
-    tidx, _, _ = cute.arch.thread_idx()
-    nthreads, _, _ = cute.arch.block_dim()
     write_thd_batch_remap(meta, n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
     # Live unit total + claim counter, as on SM100 — a SM120 unit is q_tile
     # rows of one head, so the same count applies with cga_tile_m := q_tile.
@@ -239,7 +243,12 @@ def build_thd_meta_o_kv_descs_kernel(
     TMA-OOB and land as EXACT ZEROS in SMEM (zero V nulls the masked P·V
     terms; zero K keeps the pre-mask row-max finite). Slot ``n_batch`` stays
     the never-built dead-unit pad slot."""
-    if nvvm.elect_sync():
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    # Warp 0's leader only. This flavor launches one warp, so elect_sync alone
+    # would do; the predicate keeps all three setup kernels safe under any block
+    # width, since elect_sync elects one thread PER WARP.
+    if nvvm.elect_sync() and tidx < cutlass.Int32(32):
         meta = cutlass.make_array_view(meta_t)
         write_thd_meta(meta, cutlass.make_array_view(q_lens_t), cutlass.make_array_view(kv_lens_t), lens_form, n_batch)
         cuq0 = n_batch
@@ -286,8 +295,6 @@ def build_thd_meta_o_kv_descs_kernel(
     # decode walks this permutation on every THD flavor, so a setup that skips
     # it leaves the region uninitialized and units decode garbage batches.
     cute.arch.barrier()
-    tidx, _, _ = cute.arch.thread_idx()
-    nthreads, _, _ = cute.arch.block_dim()
     write_thd_batch_remap(cutlass.make_array_view(meta_t), n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
 
 
@@ -320,7 +327,13 @@ def build_thd_meta_o_descs_kernel(
     O TMA descriptors from the cu_q values just written (same thread, program
     order). Replaces the host tolist → cumsum → H2D round-trip with work
     inside the setup launch that already existed for the descriptors."""
-    if nvvm.elect_sync():
+    tidx, _, _ = cute.arch.thread_idx()
+    nthreads, _, _ = cute.arch.block_dim()
+    # elect_sync elects one thread PER WARP, and this block is THD_SETUP_THREADS
+    # wide so the ranking below can run in parallel — narrow the single-thread
+    # body to warp 0's leader. Without this, one warp's descriptor base-copy can
+    # land after another's tensormap_replace and revert the patched address.
+    if nvvm.elect_sync() and tidx < cutlass.Int32(32):
         meta = cutlass.make_array_view(meta_t)
         write_thd_meta(meta, cutlass.make_array_view(q_lens_t), cutlass.make_array_view(kv_lens_t), lens_form, n_batch)
         cuq0 = n_batch
@@ -356,8 +369,6 @@ def build_thd_meta_o_descs_kernel(
     # Outside the elect: every thread helps rank the batches. The barrier makes
     # the cu_seqlens_q written above visible to the whole block first.
     cute.arch.barrier()
-    tidx, _, _ = cute.arch.thread_idx()
-    nthreads, _, _ = cute.arch.block_dim()
     write_thd_batch_remap(cutlass.make_array_view(meta_t), n_batch, cutlass.Int32(tidx), cutlass.Int32(nthreads))
     # Live unit total + claim counter for the persistent scheduler. The host
     # cannot know Sigma_b ceil(s_b/tile)*QH without a D2H (issue #552), so the

From f9edcf0cbf977bd4c6b69b01eab1a2caf4f850b9 Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 17:34:57 -0700
Subject: [PATCH 7/9] frost(sdpa): cache the SM count off the THD execute path

The SM100 THD pack read multi_processor_count on every execute, including
on the CLC path that never looks at it. Resolve it through a per-device
cache, inside the persistent branch that actually needs it.

Also fix the cache key both this and the CTA-count cache used:
torch.device("cuda") carries index None and means the CURRENT device, so
keying on it directly would hand every device on a multi-GPU host whichever
entry landed first.

SM120 499 passed, no regressions; SM100 suite pending a node.
---
 python/cudnn/sdpa/fwd/api_dsl.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 46baac3e1..579ef766e 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -160,6 +160,27 @@ def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch
 # change under a live process, and the query costs ~7 us on the execute hot
 # path, so resolve it once per device.
 _THD_CTAS_CACHE: dict = {}
+# Raw multi_processor_count, cached separately: _THD_CTAS_CACHE holds an
+# already-scaled CTA count, so the two cannot share a key.
+_THD_SMS_CACHE: dict = {}
+
+
+def _thd_cache_key(device):
+    """Cache key for a device. ``torch.device("cuda")`` carries index None and
+    means the CURRENT device, so resolve it — keying on None would hand every
+    device on a multi-GPU host whichever entry landed first."""
+    key = getattr(device, "index", None)
+    return torch.cuda.current_device() if key is None else key
+
+
+def _device_sm_count(device) -> int:
+    """``multi_processor_count``, resolved once per device (see above)."""
+    key = _thd_cache_key(device)
+    n = _THD_SMS_CACHE.get(key)
+    if n is None:
+        n = torch.cuda.get_device_properties(device).multi_processor_count
+        _THD_SMS_CACHE[key] = n
+    return n
 
 
 def _causal_sched_policy(s_kv: int, d_qk: int, d_v: int, elem_bytes: int) -> int:
@@ -1566,10 +1587,11 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
         # the grid no longer has to cover the work list -- which is what made
         # 38-84% of clusters dead.
         _env = self._thd_unit_envelope()
-        _cta_mma = int(getattr(self._k_mod, "CTA_MMA", 1))
-        _sms = torch.cuda.get_device_properties(q_buf.device).multi_processor_count
         if getattr(self._k_mod, "THD_PERSISTENT", False):
-            units = min(_env, max(1, _sms // max(1, _cta_mma)))
+            # Resolved here, not above: the CLC path below never reads it, and
+            # this runs per execute.
+            _cta_mma = int(getattr(self._k_mod, "CTA_MMA", 1))
+            units = min(_env, max(1, _device_sm_count(q_buf.device) // max(1, _cta_mma)))
             _dbg = int(os.environ.get("FROST_THD_CLUSTERS", "0"))  # debug override
             if _dbg > 0:
                 units = min(_env, _dbg)
@@ -3225,7 +3247,7 @@ def _thd_persistent_ctas(self, device) -> int:
         under-launching only costs parallelism. One CTA per SM matches the
         kernel's ``min_blocks_per_mp``.
         """
-        key = getattr(device, "index", device)
+        key = _thd_cache_key(device)
         n = _THD_CTAS_CACHE.get(key)
         if n is None:
             forced = int(os.environ.get("FROST_THD_CTAS", "0"))

From ecaaf3dbf0a532096260bdedb13fceb5cf577c19 Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 18:49:29 -0700
Subject: [PATCH 8/9] frost(sdpa): guard the raw-stream shortcut, match the
 setup head count

_cuda_getCurrentRawStream is a private torch symbol and only shortcuts the
stream context, so resolve it once via getattr and fall through to the
public API on a build that lacks it rather than failing every execute.

The THD setup launch passed QH as the head count while the main launch
passes QH // HEADS_PER_TILE, so the live-unit total and the decode would
disagree the moment the two differ. They cannot today -- PackGQA is
dense-only and THD runs unpacked, making HEADS_PER_TILE 1 -- so this is
identical in every reachable config; align them so it stays that way.
MXFP8 is left alone: it has no HEADS_PER_TILE and its main launch passes
QH, so the two already agree.

SM120 499 passed, no regressions; SM100 d128/d256/d512/d192 f16 and d128
fp8 trace-compile clean for sm_100a. Full SM100 suite pending a node.
---
 python/cudnn/sdpa/fwd/api_dsl.py                    | 13 +++++++++----
 .../sdpa/fwd/kernels/prefill_d128_f16_sm100.py      |  2 +-
 .../sdpa/fwd/kernels/prefill_d128_fp8_sm100.py      |  2 +-
 .../sdpa/fwd/kernels/prefill_d128_fp8_sm107.py      |  2 +-
 .../sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py |  2 +-
 .../sdpa/fwd/kernels/prefill_d256_f16_sm100.py      |  2 +-
 .../sdpa/fwd/kernels/prefill_d512_f16_sm100.py      |  2 +-
 7 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py
index 579ef766e..c3c2f7720 100644
--- a/python/cudnn/sdpa/fwd/api_dsl.py
+++ b/python/cudnn/sdpa/fwd/api_dsl.py
@@ -114,6 +114,10 @@ def _sm100_fp8_shapes(pertensor: bool, device_cc: tuple[int, int]) -> frozenset[
 # 512 B aligned, so 128 B-multiple offsets stay 128 B aligned absolutely.
 _WS_ALIGN = 128
 
+# Private torch symbol, resolved once. It only shortcuts the stream context, so
+# a build without it must fall through to the public API rather than fail.
+_CUDA_RAW_STREAM = getattr(torch._C, "_cuda_getCurrentRawStream", None)
+
 
 @contextmanager
 def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]:
@@ -134,10 +138,11 @@ def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch
     # and entering a context for the current stream is a no-op. Building the two
     # Stream objects below costs ~3.4 us each and this runs several times per
     # execute; the raw handle getter is ~0.07 us.
-    _idx = device.index if device.index is not None else torch.cuda.current_device()
-    if handle == torch._C._cuda_getCurrentRawStream(_idx):
-        yield
-        return
+    if _CUDA_RAW_STREAM is not None:
+        _idx = device.index if device.index is not None else torch.cuda.current_device()
+        if handle == _CUDA_RAW_STREAM(_idx):
+            yield
+            return
     torch_current = torch.cuda.current_stream(device)
     torch_default = torch.cuda.default_stream(device)
     if handle == torch_current.cuda_stream:
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 9301682df..c1d662570 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
@@ -2205,7 +2205,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
             cutlass.Int32(CGA_TILE_M),
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
index a190613ed..23aac02f1 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
@@ -2181,7 +2181,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
         ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
index bfaa6bbbd..f1c87a1a6 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
@@ -2273,7 +2273,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
         ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
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 dae127c62..8b38aaa68 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
@@ -2272,7 +2272,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
             cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),
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 e0f6ecc84..696bcdc3a 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
@@ -1776,7 +1776,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
             cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),
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 4483bfa42..adb48790e 100644
--- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
@@ -1981,7 +1981,7 @@ def _tma_swz(byte_w: int):
             thd_q_lens_tensor,
             thd_kv_lens_tensor,
             thd_lens_form,
-            cutlass.Int32(QH),
+            cutlass.Int32(QH // HEADS_PER_TILE),
             cutlass.Int32(B),
             cutlass.Int32(o_tensor.stride[1]),
             cutlass.Int32(CFG.TILES_Q * CFG.TILE_M * CFG.CTA_MMA),

From 05269f4a5be655b2f0b4cced3bc75e6572b9496b Mon Sep 17 00:00:00 2001
From: yanzhuoc <yanzhuoc@nvidia.com>
Date: Mon, 24 Aug 2026 19:27:03 -0700
Subject: [PATCH 9/9] frost(sdpa): drop elect_sync from the single-thread THD
 setup guards

elect.sync elects an implementation-defined lane, so conjoining it with
tidx == 0 can select NO thread at all and leave the live-unit total and
the claim-counter seed unwritten. Thread 0 alone is the guard that was
meant, and it matches thd_claim_next.

The warp-0 guards keep their elect_sync: those only need some one thread
of warp 0, not a specific lane.

SM120 499 passed, no regressions; SM100 f16 d128/d256/d512/d192 and d128
fp8 trace-compile clean for sm_100a.
---
 python/cudnn/sdpa/fwd/kernels/thd_sm100.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
index cc17155cc..556c58369 100644
--- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
+++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py
@@ -203,7 +203,10 @@ def build_thd_meta_kernel(
     # Live unit total + claim counter, as on SM100 — a SM120 unit is q_tile
     # rows of one head, so the same count applies with cga_tile_m := q_tile.
     cute.arch.barrier()
-    if nvvm.elect_sync() and tidx == cutlass.Int32(0):
+    # Thread 0 alone, WITHOUT elect_sync: elect.sync picks an
+    # implementation-defined lane, so conjoining it with tidx == 0 can
+    # select no thread at all and leave these words unwritten.
+    if tidx == cutlass.Int32(0):
         live = cutlass.Int32(0)
         cuq0 = n_batch
         for b in cutlass.range(0, n_batch, 1, unroll=1):
@@ -375,7 +378,10 @@ def build_thd_meta_o_descs_kernel(
     # kernel reads its own bound from here. The counter starts at n_clusters:
     # cluster c takes unit c from its blockIdx, then pulls from the counter.
     cute.arch.barrier()
-    if nvvm.elect_sync() and tidx == cutlass.Int32(0):
+    # Thread 0 alone, WITHOUT elect_sync: elect.sync picks an
+    # implementation-defined lane, so conjoining it with tidx == 0 can
+    # select no thread at all and leave these words unwritten.
+    if tidx == cutlass.Int32(0):
         meta_w = cutlass.make_array_view(meta_t)
         cuq0 = n_batch
         live = cutlass.Int32(0)