Thd lpt remap

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,89 @@ 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,
+    cga_m: 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.
+
+    ``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
+    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()))
+            # 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(
+                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()

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]:
@@ -130,6 +134,15 @@ 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.
+    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:
@@ -148,6 +161,32 @@ 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 = {}
+# 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:
     """SCHED_LPT_L2 vs SCHED_LPT for a causal graph (see _SCHED_L2_BUDGET_BYTES)."""
@@ -1052,7 +1091,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
@@ -1203,7 +1245,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 +1562,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 +1586,22 @@ 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()
+        if getattr(self._k_mod, "THD_PERSISTENT", False):
+            # 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)
+        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)
@@ -2570,7 +2627,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
@@ -2778,6 +2838,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 +3005,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 +3087,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 +3239,32 @@ 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``.
+        """
+        key = _thd_cache_key(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:
             # [meta(seq_kv, cu_q, cu_k)].
@@ -3192,7 +3277,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.

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

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

@@ -757,15 +757,25 @@ 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
+            # 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