frost(sdpa): native THD declared-stride support in the SM100/SM120 f16 fwd kernels

python/cudnn/AGENTS.md

@@ -39,6 +39,34 @@ Numbered so reviews can cite them; the list grows — append, never renumber.
   extent would be zero, bind a never-dereferenced dummy view over storage
   the contract already guarantees.
 
+**Rule 2 — `execute()` launches exactly the kernels the plan promised:
+serve the declared layout natively, or decline — never adapt.**
+
+Rule 1 bans implicit conversions and allocations; this rule bans the loophole
+that survives its letter: "helpful" adapter-side work that makes an
+unsupported input runnable.
+
+- **No hidden kernel launches.** A gather/scatter "normalization" copy, a
+  `.contiguous()`, a layout repack, a scatter-back after the launch — each is
+  an extra kernel that silently changes the measured perf profile per
+  configuration. **Carving the copy's scratch from the caller's workspace
+  does NOT make it acceptable**: Rule 1's workspace-carve exemption covers
+  metadata buffers and dead-slot dummies, never data-tensor copies.
+- **Can't address the declared layout natively? Decline in
+  `check_support()`** (`NotImplementedError` naming the offending tensor and
+  its strides) so the Router picks an engine that honors the declaration.
+  Silent wrong results are the worst failure mode; a silent slow path is the
+  second worst — both hide behind a green test. See
+  `_thd_check_strides_native` in `sdpa/fwd/api_dsl.py`.
+- **Precedent is not a license.** The SM100 dense path's compact-BSHD
+  normalization (`dense_layout_ok`: "one gather/scatter copy otherwise")
+  predates this rule and is grandfathered — do not cite it to justify a new
+  copy path, and treat migrating it to serve-or-decline as open cleanup.
+- The flip side of declining: whatever `check_support()` ACCEPTS, the kernel
+  must address natively (layout-driven offset math, strides encoded in TMA
+  descriptors) — acceptance is a promise about the execute path, not about
+  what the adapter can patch up.
+
 ## Frontend-only kernel package layout
 
 ```

python/cudnn/sdpa/fwd/api_dsl.py

@@ -360,6 +360,82 @@ def _to_bshd_writable(tensor: torch.Tensor):
         scratch = torch.empty_like(view, memory_format=torch.contiguous_format)
         return view, True, scratch
 
+    # -- THD declared-stride binding ------------------------------------------
+    # A THD tensor may DECLARE a wider token stride than the packed h*d — e.g.
+    # a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d),
+    # the layout torch.nn.attention.varlen users produce by slicing a fused KV
+    # projection. The f16 kernels address declared strides NATIVELY
+    # (layout-driven offset math + TMA-encoded strides); declarations the
+    # hardware cannot express are REJECTED in check_support — no
+    # normalization-copy fallback (AGENTS.md Hard Rule 2) — so the router
+    # picks an engine that honors them instead.
+
+    @staticmethod
+    def _thd_declared(desc: TensorDesc):
+        """(token, head, elem) strides a THD tensor declares, and whether they
+        are the packed contract (h*d, d, 1)."""
+        h, d = desc.shape[1], desc.shape[3]
+        st = (int(desc.stride[2]), int(desc.stride[1]), int(desc.stride[3]))
+        return st, st == (h * d, d, 1)
+
+    def _thd_check_strides_native(self) -> None:
+        """Reject THD stride declarations the kernels cannot address
+        natively: TMA's 16-byte global-stride rule — the head dim must be
+        innermost-contiguous (elem stride 1) and the token/head strides
+        multiples of ``16 // itemsize`` elements (which also keeps every
+        per-sequence ragged base 16-byte aligned). Whole-token gaps always
+        qualify for supported head dims; sub-token gaps only in 16-byte
+        multiples. The strides must also COVER the tensor (head >= d,
+        token >= h*head): an overlapping declaration would alias distinct O
+        rows onto the same storage (a write race) and is outside the
+        kernels' addressing contract."""
+        for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc):
+            (ts, hs, es), _ = self._thd_declared(desc)
+            h, d = desc.shape[1], desc.shape[3]
+            # The 16-byte TMA rule in this tensor's OWN element units: 8 at
+            # 2 B/elem (f16/bf16), 16 at 1 B/elem (fp8), 4 at 4 B/elem.
+            quantum = 16 // desc.dtype.itemsize
+            self._not_implemented_error_if(
+                es != 1 or ts % quantum != 0 or hs % quantum != 0 or hs < d or ts < h * hs,
+                f"{desc.name} THD strides {tuple(desc.stride)} are not TMA-expressible "
+                f"(head dim must be innermost-contiguous, token/head strides 16-byte — "
+                f"{quantum}-element — multiples, and non-overlapping: head stride >= {d}, "
+                f"token stride >= heads * head stride)",
+            )
+
+    def _thd_check_strides_packed(self) -> None:
+        """FP8 THD serves only the packed contract for now (its kernel and
+        harness are not audited for declared strides) — decline anything else
+        rather than adapt (AGENTS.md Hard Rule 2)."""
+        for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc):
+            _, packed = self._thd_declared(desc)
+            self._not_implemented_error_if(
+                not packed,
+                f"{desc.name}: non-packed THD strides {tuple(desc.stride)} are not supported by the FP8 path yet",
+            )
+
+    def _thd_view(self, buf: torch.Tensor, desc: TensorDesc, tokens: int) -> torch.Tensor:
+        """The declared-stride ``(1, T, H, D)`` view over a THD buffer's storage.
+
+        Validates the RUNTIME buffer against the declaration before
+        reinterpreting its storage: the dtype/device must match what was
+        declared, and the base address must be 16-byte aligned — the kernels
+        are compiled with ``assumed_align=16`` and TMA requires it of the
+        descriptor's global address, so a misaligned slice would fault (or
+        worse) instead of erroring here. ``as_strided`` itself rejects views
+        that extend past the underlying storage."""
+        h, d = desc.shape[1], desc.shape[3]
+        (ts, hs, es), _ = self._thd_declared(desc)
+        self._value_error_if(
+            buf.dtype != desc.dtype or buf.device != desc.device,
+            f"{desc.name}: runtime buffer ({buf.dtype}, {buf.device}) does not match its declaration ({desc.dtype}, {desc.device})",
+        )
+        self._value_error_if(
+            buf.data_ptr() % 16 != 0,
+            f"{desc.name}: runtime buffer base address must be 16-byte aligned (TMA global-address rule); got data_ptr() % 16 == {buf.data_ptr() % 16}",
+        )
+        return buf.as_strided((1, tokens, h, d), (max(tokens, 1) * ts, ts, hs, es), buf.storage_offset())
+
     def _amax_slot(self, tensor, name: str, device: torch.device) -> torch.Tensor:
         """The caller's 1-element amax storage, or a cached dummy.
 
@@ -570,6 +646,9 @@ def check_support(self) -> bool:
                     f"strides allowed); got stride {_stride} shape {_shape}",
                 )
 
+        if self.thd:
+            self._thd_check_strides_native()
+
         b, h_qo, s_qo, d_qk = self.q_desc.shape
         _, h_kv, s_kv, _ = self.k_desc.shape
         _, _, _, d_v = self.v_desc.shape
@@ -1117,11 +1196,12 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se
         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)
 
-        def _packed(buf, t, h, d):
-            return buf.as_strided((1, t, h, d), (t * h * d, h * d, d, 1), buf.storage_offset())
-
-        Q = _packed(q_buf, t_q, qh, d_qk)
-        O = _packed(o_buf, t_q, qh, d_v)
+        # 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)
         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
@@ -1138,8 +1218,8 @@ def _packed(buf, t, h, d):
             K = q_buf.as_strided((1, 1, kh, d_qk), (kh * d_qk, kh * d_qk, d_qk, 1), q_buf.storage_offset())
             V = o_buf.as_strided((1, 1, kh, d_v), (kh * d_v, kh * d_v, d_v, 1), o_buf.storage_offset())
         else:
-            K = _packed(k_buf, t_kv, kh, d_qk)
-            V = _packed(v_buf, t_kv, kh, d_v)
+            K = self._thd_view(k_buf, self.k_desc, t_kv)
+            V = self._thd_view(v_buf, self.v_desc, t_kv)
         # LSE binding: the caller's ragged Stats buffer in its declared layout
         # when a Stats output exists; None otherwise — the kernel compiles the
         # LSE store out (has_lse=False), so no dummy buffer exists at all.
@@ -1167,6 +1247,13 @@ def _packed(buf, t, h, d):
             has_lse=lse is not None,
             lse_head_major=lse is not None and self.thd_stats_head_major,
             lse_head_stride=(self.thd_stats_head_stride if (lse is not None and self.thd_stats_head_major) else 0),
+            # Declared strides of the bound views (cache-key): compact views
+            # reproduce the packed specialization; native non-packed views
+            # compile their strides into the kernel's addressing.
+            q_stride=tuple(Q.stride()),
+            k_stride=tuple(K.stride()),
+            v_stride=tuple(V.stride()),
+            o_stride=tuple(O.stride()),
         )
         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)
         self._logger.debug("execute (THD) completed")
@@ -1653,6 +1740,11 @@ def check_support(self) -> bool:
                     f"non-broadcast, non-overlapping strides (any B/H/S order, padded "
                     f"strides allowed); got stride {desc.stride} shape {desc.shape}",
                 )
+        if self.thd:
+            if self._pertensor:
+                self._thd_check_strides_packed()
+            else:
+                self._thd_check_strides_native()
 
         b, h_q, s_q, d_q = self.q_desc.shape
         _, h_kv, s_kv, _ = self.k_desc.shape
@@ -2138,7 +2230,7 @@ def _scalar(t, default=1.0):
                 amax_o_buf.div_(max(so, 1e-30))
         self._logger.debug("execute (SM120 FP8 per-tensor) completed")
 
-    def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label):
+    def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label, declared_views=False):
         """Shared THD (ragged) packing: cu_seqlens metadata + ``(1, T, H, D)`` views.
 
         Serves the same fully-packed contract as the SM100 THD path
@@ -2183,6 +2275,13 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa
         def _packed(buf, tokens, heads, d):
             return buf.as_strided((1, tokens, heads, d), (tokens * heads * d, heads * d, d, 1), buf.storage_offset())
 
+        def _view(buf, desc, tokens, heads, d):
+            # declared_views: the f16 kernel addresses declared strides
+            # natively (check_support rejected inexpressible ones); the FP8
+            # path keeps the packed contract (check_support declined
+            # anything else).
+            return self._thd_view(buf, desc, tokens) if declared_views else _packed(buf, tokens, heads, d)
+
         if t_kv == 0:
             # Every query row is dead (all-zero seq_kv_lens): served by the
             # KERNEL's own dead-row path (row_sum <= 0 -> O := 0 and
@@ -2199,18 +2298,18 @@ def _packed(buf, tokens, heads, d):
             K = q_buf.as_strided((1, 1, kh, d_qk), (kh * d_qk, kh * d_qk, d_qk, 1), q_buf.storage_offset())
             V = o_buf.as_strided((1, 1, kh, d_v), (kh * d_v, kh * d_v, d_v, 1), o_buf.storage_offset())
         else:
-            K = _packed(k_buf, t_kv, kh, d_qk)
-            V = _packed(v_buf, t_kv, kh, d_v)
+            K = _view(k_buf, self.k_desc, t_kv, kh, d_qk)
+            V = _view(v_buf, self.v_desc, t_kv, kh, d_v)
 
         return SimpleNamespace(
             meta=meta,
             t_q=t_q,
             t_kv=t_kv,
             max_sq=max_sq,
-            Q=_packed(q_buf, t_q, qh, d_qk),
+            Q=_view(q_buf, self.q_desc, t_q, qh, d_qk),
             K=K,
             V=V,
-            O=_packed(o_buf, t_q, qh, d_v),
+            O=_view(o_buf, self.o_desc, t_q, qh, d_v),
             seq_q_dummy=self._dummy("seq_q_lens", dev, lambda: torch.zeros(b, dtype=torch.int32, device=dev)),
         )
 
@@ -2225,7 +2324,7 @@ def _execute_thd(
         within each head row.
         """
 
-        pack = self._thd_pack(q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)")
+        pack = self._thd_pack(q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)", declared_views=True)
         if pack is None:
             return
 
@@ -2263,6 +2362,14 @@ def _execute_thd(
             has_lse=self.lse_desc is not None,
             lse_head_major=self.thd_stats_head_major,
             lse_head_stride=self.thd_stats_head_stride,
+            # Declared strides of the bound views (cache-key): compact views
+            # reproduce the packed specialization; native non-packed views
+            # compile their strides into the Q/O offset math and K/V TMA
+            # descriptors.
+            q_stride=tuple(pack.Q.stride()),
+            k_stride=tuple(pack.K.stride()),
+            v_stride=tuple(pack.V.stride()),
+            o_stride=tuple(pack.O.stride()),
         )
         fn(
             pack.Q,

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

@@ -2007,7 +2007,7 @@ def _tma_swz(byte_w: int):
             seq_kv_lens_tensor,
             cutlass.Int32(QH),
             cutlass.Int32(B),
-            cutlass.Int32(o_tensor.shape[3]),
+            cutlass.Int32(o_tensor.stride[1]),
         ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
@@ -2050,6 +2050,10 @@ def compile(  # noqa: A001
     has_lse: bool = True,
     lse_head_major: bool = False,
     lse_head_stride: int = 0,
+    q_stride: Optional[tuple] = None,
+    k_stride: Optional[tuple] = None,
+    v_stride: Optional[tuple] = None,
+    o_stride: Optional[tuple] = None,
 ) -> Callable:
     """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time.
 
@@ -2077,30 +2081,21 @@ def compile(  # noqa: A001
     if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0:
         raise ValueError(f"d128 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}")
     _fake_batch = 1 if CFG.THD_VARLEN else b
-    fake_q = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, sq, qh, d_qk),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_k = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, skv, kh, d_qk),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_v = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, skv, kh, d_v),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_o = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, sq, qh, d_v),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
+
+    def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
+        if stride is None:
+            return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16)
+        if stride[3] != 1:
+            raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)")
+        for axis in (1, 2):  # seq/head global strides feed TMA: 16-byte rule
+            if (stride[axis] * bpe) % 16 != 0:
+                raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)")
+        return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16)
+
+    fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride)
+    fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride)
+    fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride)
+    fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE)
     if not has_lse:
         # No Stats output: the LSE argument is None-specialized and the store
         # is compiled out entirely — no dummy buffer exists at any level.

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

@@ -2102,7 +2102,7 @@ def _tma_swz(byte_w: int):
             seq_kv_lens_tensor,
             cutlass.Int32(QH),
             cutlass.Int32(B),
-            cutlass.Int32(o_tensor.shape[3]),
+            cutlass.Int32(o_tensor.stride[1]),
         ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
         grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
     else:
@@ -2145,6 +2145,10 @@ def compile(  # noqa: A001
     has_lse: bool = True,
     lse_head_major: bool = False,
     lse_head_stride: int = 0,
+    q_stride: Optional[tuple] = None,
+    k_stride: Optional[tuple] = None,
+    v_stride: Optional[tuple] = None,
+    o_stride: Optional[tuple] = None,
 ) -> Callable:
     """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time.
 
@@ -2164,30 +2168,21 @@ def compile(  # noqa: A001
     if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0:
         raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}")
     _fake_batch = 1 if CFG.THD_VARLEN else b
-    fake_q = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, sq, qh, d_qk),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_k = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, skv, kh, d_qk),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_v = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, skv, kh, d_v),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
-    fake_o = cute.runtime.make_fake_compact_tensor(
-        STORAGE_DTYPE,
-        (_fake_batch, sq, qh, d_v),
-        stride_order=(3, 2, 1, 0),
-        assumed_align=16,
-    )
+
+    def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE):
+        if stride is None:
+            return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16)
+        if stride[3] != 1:
+            raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)")
+        for axis in (1, 2):  # seq/head global strides feed TMA: 16-byte rule
+            if (stride[axis] * bpe) % 16 != 0:
+                raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)")
+        return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16)
+
+    fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride)
+    fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride)
+    fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride)
+    fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE)
     if not has_lse:
         # No Stats output: the LSE argument is None-specialized and the store
         # is compiled out entirely — no dummy buffer exists at any level.