frost(sdpa): THD execute with zero host reads on SM100 — device-built metadata, plan-time envelope grid, CUDA-graph capturable (issue #552)

python/cudnn/AGENTS.md

@@ -97,14 +97,21 @@ Known violations, all pre-existing and each needing a kernel-side change, so
 none is precedent:
 
 - THD `cu_seqlens` host cumsum (`sdpa/fwd/api_dsl.py`, `_execute_thd` on both
-  SM100 and SM120). The compile-side half is DONE — the THD kernels compile
-  with dynamic token extents (Rule 4), so `T` is no longer a compile-time
-  constant and no compile is keyed on it. `t_q`/`t_kv` still reach the host
-  for the metadata upload, the ragged views' extents, and the launch
-  grid; removing that needs the plan-time-max (`b * s_q_max`) grid with
-  in-kernel dead-tile exit and the device-side metadata read (issue #552).
-  `sdpa_fwd_wrapper_sm80` shows the other half — it requires `max_s_q` from
-  the caller rather than deriving it.
+  SM100 and SM120). RESOLVED on all SM100 f16 families
+  (issue #552) — the reference for the SM120 port: the kernels compile
+  with dynamic token extents (Rule 4), the setup kernel builds the metadata
+  buffer device-side from the caller's length tensors, every ragged view
+  binds its buffer's capacity, and the grid is the plan-time envelope
+  (`b * ceil(s_q_declared / CGA_TILE_M) * qh`) with dead units exiting via
+  `_thd_decode`'s `batch == n_batch` sentinel — zero host reads, pinned by
+  the sync-debug and CUDA-graph capture tests. Still open: SM120 (exact
+  host-computed grid + host-built metadata; its own engine class and grid
+  mechanism, so its own port), and
+  the envelope's dead-tile tax under far-oversized declarations (measured
+  ~145 ns/dead unit exposed, <0.2% at realistic declarations; a capped
+  persistent grid reading a device-side live-unit count would bound it by
+  resident clusters). `sdpa_fwd_wrapper_sm80` shows the caller-provided
+  alternative — it requires `max_s_q` rather than deriving it.
 - Per-tensor FP8 descale readback (`_scalar` in the same file): fold on device,
   passing the pointers, as the backend FP8 sdpa does.
 - The FP8/MXFP8 `seq_len_q` guard in `sdpa/fwd/engines.py`. This one cannot be

python/cudnn/sdpa/fwd/api_dsl.py

@@ -1236,6 +1236,24 @@ def _key(desc):
             o_stride=_key(self.o_desc),
         )
 
+    def _thd_unit_envelope(self) -> int:
+        """PLAN-TIME upper bound on live THD units:
+        ``B * ceil(S_q_declared / CGA_TILE_M) * QH``.
+
+        Every sequence's length is bounded by the declared S_q (the padding
+        contract), so this covers ``Σ_b ceil(s_b / tile) * QH``. Units past
+        the live total are DEAD by kernel contract (the decode's
+        ``batch == n_batch`` sentinel): no loads, no O/LSE writes, one
+        empty-mainloop barrier dance each. The grid is host-known at PLAN
+        time — execute reads nothing from the lengths — which removes the
+        last THD D2H sync and unblocks CUDA-graph capture (issue #552). The
+        dead-tile tax mirrors the C++ backend's THD grid strategy; callers
+        declaring S_q far above their live totals pay it. The over-launch
+        test pads this to pin the dead-unit contract."""
+        cga_tile_m = int(self._k_mod.CGA_TILE_M)
+        s_q_decl = int(self.q_desc.shape[2])
+        return self.batch_size * ((s_q_decl + cga_tile_m - 1) // cga_tile_m) * self.h_q
+
     def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_len_kv, seq_q_lens, lse_tensor=None, workspace=None, current_stream=None):
         """THD / varlen execute: reconstruct the kernel's packed [1, T, H, D] views and metadata buffer from the cuDNN ragged buffers, then launch.
 
@@ -1248,54 +1266,65 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se
         in the first ``T*H`` elements, or head-major ``(H, head_stride)``
         with tokens contiguous within each head row; when ``None`` the kernel
         compiles the LSE store out (has_lse=False) and no scratch exists.
-        The host round-trip for the runtime totals (t_q / t_kv / unit count)
-        feeds the metadata upload, the ragged views' extents and the exact
-        launch grid, and costs one D2H sync per length tensor, no device
-        allocation. It no longer keys any compile: the kernels compile with
-        DYNAMIC token extents (issue #552), so a new packed total re-binds
-        the same artifact; removing the sync itself needs the plan-time-max
-        grid redesign tracked there."""
+        Host round-trips (issue #552): NONE — the lengths never reach the
+        host (the setup kernel builds the metadata buffer device-side),
+        every ragged view binds its buffer's capacity, and the launch grid
+        is the plan-time envelope (dead units exit by kernel contract) —
+        the execute is fully async and CUDA-graph capturable. No compile is
+        keyed on runtime data: the kernels compile with DYNAMIC token
+        extents, so a new packed total re-binds the same artifact."""
         import cutlass
 
         dev = q_buf.device
         b = self.batch_size
         carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm100 (THD)") if workspace is not None else None
         # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ],
-        # built HOST-side from the (inherent) tolist round-trip and uploaded in
-        # ONE H2D copy: a device-side cumsum would allocate its scan-temp
-        # storage and launch kernels on the execute hot path. Either length
-        # form feeds it — per-batch (B,) lengths or the (B+1,) cu_seq_len
-        # prefix sums — at identical cost. The torch work (allocation, D2H
-        # length reads, the H2D upload) runs on the LAUNCH stream so it is
-        # ordered against the kernel that consumes it — the execute-time
-        # handle may carry a stream that is not torch's current.
+        # built DEVICE-side by the setup kernel from the caller's length
+        # tensors (either form) — no host cumsum, no H2D, and the lengths
+        # NEVER reach the host (issue #552); the prefix-sum invariants are
+        # caller contract (a validation that needs a device read is not a
+        # validation — AGENTS.md Rule 3). The torch allocations run on the
+        # LAUNCH stream so they are ordered against the kernels that consume
+        # them — the execute-time handle may carry a stream that is not
+        # torch's current.
         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)
-            slq_host, cu_q_host = self._thd_host_lens(seq_q_lens, "cu_seq_len_q" if self.cu_seq_q_lens else "seq_q_lens", self.cu_seq_q_lens)
-            slk_host, cu_k_host = self._thd_host_lens(seq_len_kv, "cu_seq_len_kv" if self.cu_seq_kv_lens else "seq_kv_lens", self.cu_seq_kv_lens)
-            meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32))
-        t_q = cu_q_host[-1]
-        t_kv = cu_k_host[-1]
+        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_len_kv, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_len_kv, "seq_kv_lens")
 
         qh, kh = self.h_q, self.h_kv
         d_qk, d_v = self.head_dim_qk, self.head_dim_v
 
-        # Degenerate total (runtime value, invisible to the plan-time probe):
-        # t_q == 0 means no query token exists anywhere, so the packed O/LSE
-        # have zero rows — nothing to compute or write. (t_kv == 0 launches
-        # normally through the kernel's dead-row path; see the K/V binding
-        # below.)
+        # Q/O token extent = buffer CAPACITY (issue #552: the packed Q total
+        # lives on device only). Q, O and a token-major (or compact
+        # head-major) LSE bind ONE dynamic token symbol, so take the shared
+        # floor. Writes never step past the real per-sequence lengths the
+        # kernel reads from the device metadata (O through the per-batch
+        # descriptors' extents, LSE via the row predicate), so the
+        # over-claim only widens the TMA descriptors' bound.
+        (q_ts, _, _), _ = self._thd_declared(self.q_desc)
+        (o_ts, _, _), _ = self._thd_declared(self.o_desc)
+        t_q = min(q_buf.numel() // q_ts, o_buf.numel() // o_ts)
+        if lse_tensor is not None and not (self.thd_stats_head_major and self.thd_stats_head_stride):
+            t_q = min(t_q, lse_tensor.numel() // qh)
+
+        # Degenerate Q side: zero CAPACITY — no token is addressable, and a
+        # zero-token view cannot back a TMA descriptor — nothing to compute
+        # or write. All-zero LENGTHS over live storage launch normally:
+        # every unit decodes dead and neither O nor LSE is touched.
+        # (t_kv == 0 launches normally through the kernel's dead-row path;
+        # see the K/V binding below.)
         if t_q == 0:
-            self._logger.debug("execute (THD): t_q == 0, nothing to do")
+            self._logger.debug("execute (THD): no addressable Q token, nothing to do")
             return
         lse = None
         if lse_tensor is not None:
             if self.thd_stats_head_major:
+                # head_stride covering the packed total is caller contract
+                # (t_q is a device value — Rule 3).
                 head_stride = self.thd_stats_head_stride
-                self._value_error_if(
-                    head_stride < t_q,
-                    f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})",
-                )
+                if head_stride == 0:
+                    head_stride = t_q  # compact: the token extent itself
                 lse = lse_tensor.as_strided((1, qh, head_stride), (qh * head_stride, head_stride, 1), lse_tensor.storage_offset())
             else:
                 # Token-major (TH1, the default): natural packed rank-2 (T, H)
@@ -1311,25 +1340,35 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se
         # vs the builder pass).
         with _torch_stream_context(current_stream, dev):
             o_desc = carver.take(b * 16 + 16, torch.int64) if carver is not None else torch.empty(b * 16 + 16, dtype=torch.int64, device=dev)
-        # One THD unit per CGA-height slice of each sequence's Q rows.
-        cga_tile_m = int(self._k_mod.CGA_TILE_M)
-        units = qh * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host)
+        # The PLAN-TIME envelope grid — dead units exit by kernel contract.
+        units = self._thd_unit_envelope()
 
         # Declared-stride (1, T, H, D) views, addressed NATIVELY by the kernel
         # (the Q/K/V/O TMA descriptors are built from the tensor views, and
         # the THD O-descriptor builder steps by O's declared seq stride);
         # check_support rejected any declaration TMA cannot express.
         Q = self._thd_view(q_buf, self.q_desc, t_q)
         O = self._thd_view(o_buf, self.o_desc, t_q)
+        # KV token extent = buffer CAPACITY (issue #552: the packed KV total
+        # lives on device only). K and V bind the same dynamic token symbol,
+        # so take the shared floor. Loads never step past the REAL
+        # per-sequence lengths the kernel reads from the device metadata, so
+        # an over-claimed extent only widens the TMA descriptors' bound.
+        # All-zero lengths with a live buffer launch normally through the
+        # kernel's per-sequence dead-row path.
+        (k_ts, _, _), _ = self._thd_declared(self.k_desc)
+        (v_ts, _, _), _ = self._thd_declared(self.v_desc)
+        t_kv = min(k_buf.numel() // k_ts, v_buf.numel() // v_ts)
         if t_kv == 0:
-            # Every query row is dead (all-zero seq_kv_lens): served by the
-            # KERNEL's own dead-row path (total_sum <= 0 -> O := 0 and
-            # LSE := -inf, or the sink alone — its column keeps the softmax
-            # denominator alive), exactly like a live launch's zero-KV
-            # sequences — no adapter-side fills on the execute hot path
-            # (AGENTS.md Rule 1). A zero-token packed K/V view cannot back a
-            # CuTe layout / TMA descriptor, so clamp the packed KV extent to
-            # ONE never-dereferenced token (every tile sees kv_left ==
+            # No KV storage at all:
+            # every query row is dead — served by the KERNEL's own dead-row
+            # path (total_sum <= 0 -> O := 0 and LSE := -inf, or the sink
+            # alone — its column keeps the softmax denominator alive),
+            # exactly like a live launch's zero-KV sequences — no
+            # adapter-side fills on the execute hot path (AGENTS.md Rule 1).
+            # A zero-token packed K/V view cannot back a CuTe layout / TMA
+            # descriptor, so clamp the packed KV extent to ONE
+            # never-dereferenced token (every tile sees kv_left ==
             # kv_right == 0, so no K/V load is ever issued) bound over
             # storage guaranteed large enough: Q backs K (kh*d_qk <=
             # t_q*qh*d_qk) and O backs V (kh*d_v <= t_q*qh*d_v).
@@ -1364,22 +1403,50 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se
         kwargs = self._thd_compile_kwargs()
         kwargs.update(k_stride=(0, *K.stride()[1:]), v_stride=(0, *V.stride()[1:]))
         fn = self._k_mod.compile(**kwargs)
-        fn(Q, K, V, O, LSE, sinks_t, meta, o_desc, (b, qh, kh, t_q, t_kv, 0), cutlass.Float32(scale_softmax_log2), cutlass.Int32(units), stream=current_stream)
+        # The caller's length tensors ride to the setup kernel, which builds
+        # the metadata buffer device-side; the form bitmask is a runtime
+        # value (no compile key grows). problem_size sq/skv slots are 0 by
+        # the THD contract (_host reads the dynamic extents).
+        lens_form = (1 if self.cu_seq_q_lens else 0) | (2 if self.cu_seq_kv_lens else 0)
+        fn(
+            Q,
+            K,
+            V,
+            O,
+            LSE,
+            sinks_t,
+            meta,
+            o_desc,
+            (b, qh, kh, 0, 0, 0),
+            cutlass.Float32(scale_softmax_log2),
+            cutlass.Int32(units),
+            None,
+            q_lens_dev,
+            kv_lens_dev,
+            cutlass.Int32(lens_form),
+            stream=current_stream,
+        )
         self._logger.debug("execute (THD) completed")
 
     @staticmethod
     def _ceil_div(x: int, a: int) -> int:
         return (x + a - 1) // a
 
     def _reshape_sf(self, sf: torch.Tensor, h: int, n_tiles: int, sf_smem_size: int) -> torch.Tensor:
-        """cuDNN F8_128x4 scale-factor tensor (FP8_E8M0, ``[B, H, *, *]``) → the
-        kernel's per-tile int8 view ``[B, H, n_tiles, sf_smem_size]``.
+        """cuDNN F8_128x4 scale-factor tensor (FP8_E8M0) → the kernel's per-tile
+        int8 view ``[B, H, n_tiles, sf_smem_size]``.
 
         cuDNN packs the 128×4 SF atom contiguously (``F8_128x4`` reordering); a Q/K
         tile is 128 rows × d/32 d-blocks and a V tile is 128 rows × 4 s-blocks, so
         each tile is exactly ``sf_smem_size`` E8M0 bytes and this is a pure reshape.
+
+        A reordered tensor is an opaque byte layout: callers legally bind it under
+        any shape with the right byte count (the graph declares logical
+        ``[B, H, s_padded, d/32]`` dims; TE-style producers hand over flat
+        ``[B·H·s_padded, 4]`` swizzle output).  So B comes from the graph facts,
+        never from ``sf.shape[0]``, and only the total size is validated.
         """
-        b = sf.shape[0]
+        b = self.batch_size
         flat = sf.contiguous()
         if flat.dtype != torch.int8:
             flat = flat.view(torch.int8)

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

@@ -449,7 +449,14 @@ def _thd_decode(linear_cta, seq_kv_lens_t, n_batch, n_qh, cta_in_pair):
         cu = cutlass.make_array_view(seq_kv_lens_t)
         cuq0 = n_batch
         acc = cutlass.Int32(0)
-        f_batch = cutlass.Int32(0)
+        # DEAD-unit sentinel (issue #552 over-launch): a unit no sequence
+        # claims (u >= sum of live units) keeps batch == n_batch.  That index
+        # makes every downstream consumer a no-op through IN-BOUNDS metadata
+        # reads: _resolve_seqlen_kv reads meta[n_batch] = cu_q[0] = 0 (empty
+        # KV range in every role), the epilogue's per-sequence Q length
+        # cu[2n+1]-cu[2n] goes negative (LSE predicate never fires), and the
+        # O-store role skips the TMA store explicitly (batch >= n_batch).
+        f_batch = n_batch
         f_head = cutlass.Int32(0)
         f_qc = cutlass.Int32(0)
         done = cutlass.Int32(0)