frost(sdpa): strict lse/sink/seq-lens execute contract, no torch.empty in execute

python/cudnn/AGENTS.md

@@ -8,6 +8,31 @@ The `cudnn` Python package: pybind11-backed graph API plus pure-Python **fronten
 - Never add an eager `import torch` / `import cutlass` to `__init__.py` or anything it imports transitively. `api_base.py` itself imports them at top level, which is why kernel classes must only be reachable through the lazy table.
 - Reuse the existing `[cutedsl]` extra (`pyproject.toml` optional-dependencies) unless a kernel truly needs a new package.
 
+## Hard rules
+
+Numbered so reviews can cite them; the list grows — append, never renumber.
+
+**Rule 1 — `execute()` is a zero-surprise hot path: validate, never convert, never allocate.**
+
+- **No implicit conversions.** Never `.to(dtype)`, and never a `reshape()` that can
+  copy, on an execute argument: both silently allocate and launch a kernel per
+  call, and the fresh pointer breaks CUDA-graph capture. Worse, for an *output*
+  tensor a reshape copy swallows the kernel's write. Validate dtype / shape /
+  contiguity and bind a true view (`.view()` or a checked `reshape`), raising
+  `ValueError` otherwise — see `_checked_lse_view` / `_checked_sinks_1d` /
+  `_checked_seq_lens` in `sdpa/fwd/api_dsl.py`.
+- **No per-execute allocations.** No `torch.empty`/`torch.zeros` inside
+  `execute()`: scratch is carved from the caller's workspace
+  (`scratch_workspace_bytes()` contract), and a dead ABI slot may use a
+  one-time cached dummy (`_dummy`) at most. Prefer compiling the unused
+  operand out entirely (CuTeDSL specializes on `None` via
+  `cutlass.const_expr` — see the SM120 SDPA kernel's optional lse/sinks).
+- **Init-time flags are compile-time specializations; `execute()` must match
+  them exactly, in both directions.** A required-but-missing tensor must
+  raise, never fall back to a zeros dummy (zeros sinks change the softmax
+  denominator; zeros seq lens mask every row — silently wrong output). A
+  provided-but-uncompiled tensor must also raise, never be silently ignored.
+
 ## Frontend-only kernel package layout
 
 ```

python/cudnn/frost/README.md

@@ -100,8 +100,10 @@ Reading that sequence line by line:
   still forwarded to the lowered C++ graph.
 - **`get_workspace_size()` is honest.** For a python plan it returns
   `CompiledPlan.get_workspace_size()`, the executor's real requirement (for
-  the graph above: the dummy-LSE scratch `b*h*s*4 = 16384` bytes, because an
-  inference graph has no Stats output). `execute()` forwards the caller's
+  the graph above on an SM100 engine: the dummy-LSE scratch `b*h*s*4 = 16384`
+  bytes, because an inference graph has no Stats output and the SM100 kernels
+  always write an LSE; `lse_optional` adapters like SM120 compile the LSE
+  store out instead and report 0). `execute()` forwards the caller's
   buffer through `ExecutionContext.workspace` and the executor carves its
   scratch out of it in 128-byte-aligned chunks, never touching bytes at or
   beyond the reported size -- no hidden per-execute allocation, stable

python/cudnn/sdpa/fwd/engines.py

@@ -115,6 +115,11 @@ class Capabilities:
     padded: bool = False
     sink: bool = False
     stats: bool = False
+    # The adapter accepts lse_tensor=None (its kernel None-specializes the LSE
+    # store), so a stats-less graph needs no dummy-LSE workspace chunk. Rows
+    # that keep False (the SM100 flavors) always write an LSE and get a carved
+    # dummy from lower_dsl_prefill when the graph has no Stats output.
+    lse_optional: bool = False
     thd: bool = False
     # True = the kernel anchors the THD bottom-right diagonal at each sequence's
     # own (seq_len_q[b], seq_len_kv[b]). Rows that keep False (the SM100 flavors)
@@ -426,6 +431,7 @@ def _sm120_spec() -> EngineSpec:
             padded=True,
             sink=True,
             stats=True,
+            lse_optional=True,
             padded_stats=True,
             thd=True,
             thd_bottom_right=True,
@@ -495,6 +501,13 @@ def lower_dsl_prefill(
 
     seq_q_t = facts.seq_q_t if facts.padded else None
     seq_kv_t = facts.seq_kv_t if facts.padded else None
+    # Mirrors the seq_q_lens_present constructor argument below. Execute
+    # forwards seq_q only when the compiled specialization consumes it (or THD,
+    # which sources cu_seqlens from it) — the adapter rejects mismatches, so a
+    # buffer the FP8/MXFP8 kernels can't honor (dense padded-Q trim is not
+    # plumbed there — known gap) is dropped here rather than erroring at
+    # execute.
+    seq_q_lens_present = facts.padded and not facts.thd and facts.seq_q_t is not None and not (facts.is_mxfp8 or facts.is_fp8)
     api = api_type(
         sample_q=ga.tensor_desc_from_ir(facts.q_t, name="q"),
         sample_k=ga.tensor_desc_from_ir(facts.k_t, name="k"),
@@ -510,7 +523,7 @@ def lower_dsl_prefill(
         # enabled whenever a dense padded graph carries per-batch Q lengths.
         # THD carries Q lengths via cu_seqlens; the FP8/MXFP8 kernels are not
         # plumbed (their specs also keep padded_stats=False).
-        seq_q_lens_present=(facts.padded and not facts.thd and facts.seq_q_t is not None and not (facts.is_mxfp8 or facts.is_fp8)),
+        seq_q_lens_present=seq_q_lens_present,
         has_sink=facts.has_sink,
         thd=facts.thd,
         dtype_o=facts.dtype_o if (facts.is_mxfp8 or facts.is_fp8) else None,
@@ -527,14 +540,20 @@ def lower_dsl_prefill(
     # buffer is carved from the CALLER's workspace, so its size is fixed here at
     # build time and recorded on the executor as ``workspace_bytes`` — that
     # number is what the plan's CompiledPlan.get_workspace_size() reports.
-    #   - dummy LSE (dense, stats absent): the kernel always writes an LSE;
-    #     without a Stats output it lands in b*h_q*s_q fp32 scratch.
-    #     (THD needs no engine-level LSE chunk — the packed THD LSE is part
-    #     of the api-level scratch below.)
+    #   - dummy LSE (dense, stats absent, non-lse_optional adapters): the
+    #     SM100 kernels always write an LSE; without a Stats output it lands
+    #     in b*h_q*s_q fp32 scratch. lse_optional adapters (SM120) compile the
+    #     LSE store out instead and bind no buffer. (THD needs no engine-level
+    #     LSE chunk — the packed THD LSE is part of the api-level scratch
+    #     below.)
     #   - synthesized seq_len_kv (skv_tail_via_padding rows): b int32.
     #   - api-level scratch (api.scratch_workspace_bytes()): the dense padded
     #     [seq_kv|seq_q] combine and the THD metadata/LSE buffers.
-    dummy_lse_bytes = 0 if (not spec.capabilities.stats or facts.stats_t is not None or facts.thd) else ws_align(facts.b * facts.h_q * facts.s_q * 4)
+    dummy_lse_bytes = (
+        0
+        if (not spec.capabilities.stats or spec.capabilities.lse_optional or facts.stats_t is not None or facts.thd)
+        else ws_align(facts.b * facts.h_q * facts.s_q * 4)
+    )
     synth_kv_bytes = ws_align(facts.b * 4) if synth_kv_padding else 0
     api_scratch_bytes = api.scratch_workspace_bytes()
     total_workspace_bytes = dummy_lse_bytes + synth_kv_bytes + api_scratch_bytes
@@ -593,9 +612,10 @@ def _execute(variant_pack, workspace=None, stream=None):
         # re-validates so a direct call cannot silently corrupt memory.
         carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None
         lse_buf = resolved.get(id(binding.stats)) if binding.stats is not None else None
-        if lse_buf is None and spec.capabilities.stats and not facts.thd:
+        if lse_buf is None and spec.capabilities.stats and not spec.capabilities.lse_optional and not facts.thd:
             # Dummy LSE for stats-less dense graphs — carved, not allocated
-            # (uninitialized is fine: the kernel writes every row).
+            # (uninitialized is fine: the kernel writes every row). lse_optional
+            # adapters take lse_tensor=None instead.
             lse_buf = carver.take(facts.b * facts.h_q * facts.s_q, torch.float32)
         sinks_buf = resolved.get(id(binding.sink_token)) if binding.sink_token is not None else None
         seq_kv_buf = resolved.get(id(binding.seq_len_kv)) if binding.seq_len_kv is not None else None
@@ -623,7 +643,7 @@ def _execute(variant_pack, workspace=None, stream=None):
             scale_softmax=facts.scale,
             sinks=sinks_buf,
             seq_kv_lens=seq_kv_buf,
-            seq_q_lens=seq_q_buf,
+            seq_q_lens=seq_q_buf if (seq_q_lens_present or facts.thd) else None,
             # Stream from the execute-time handle (raw CUstream int, the
             # ExecutionContext's stream); None keeps the default stream.
             current_stream=_cuda_driver.CUstream(stream) if stream is not None else None,

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

@@ -40,7 +40,7 @@
 
 from functools import lru_cache, partial
 from types import SimpleNamespace
-from typing import Callable, Type
+from typing import Callable, Optional, Type
 
 import cuda.bindings.driver as cuda_driver
 import cutlass
@@ -772,8 +772,8 @@ def kernel(
         k: cute.Tensor,
         v: cute.Tensor,
         o: cute.Tensor,
-        lse: cute.Tensor,
-        sinks: 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],
@@ -786,9 +786,10 @@ def kernel(
         :param k: Key tensor.
         :param v: Value tensor.
         :param o: Output tensor.
-        :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output.
-        :param sinks: ``(H,)`` fp32 per-Q-head sink logits, or an unused
-            dummy tensor when the kernel is configured without ``has_sink``.
+        :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, 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.
@@ -1140,25 +1141,26 @@ def kernel(
                         lse_val = -cutlass.Float32.inf
                     row_lse[row_half] = lse_val
 
-            if lane % 4 == 0:
-                lse_arr = cutlass.make_array_view(lse)
-                for row_half in cutlass.range_constexpr(2):
-                    lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8
-                    lse_out = cutlass.Float32(row_lse[row_half])
-                    if cutlass.const_expr(self.thd_varlen):
-                        # Packed (1, H, T) LSE: rows past this sequence's Q
-                        # length belong to the NEXT sequence — never written,
-                        # and there is no padded region to trim.
-                        if lse_q_idx < seqlen_q:
-                            lse_row = lse_arr[0, head_idx, :]
-                            lse_row[q_row_base + lse_q_idx] = lse_out
-                    else:
-                        # Rows at/past this batch's Q length trim to -inf.
-                        if lse_q_idx >= seqlen_q:
-                            lse_out = -cutlass.Float32.inf
-                        if lse_q_idx < q.shape[1]:
-                            lse_row = lse_arr[batch_idx, head_idx, :]
-                            lse_row[lse_q_idx] = lse_out
+            if cutlass.const_expr(lse is not None):
+                if lane % 4 == 0:
+                    lse_arr = cutlass.make_array_view(lse)
+                    for row_half in cutlass.range_constexpr(2):
+                        lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8
+                        lse_out = cutlass.Float32(row_lse[row_half])
+                        if cutlass.const_expr(self.thd_varlen):
+                            # Packed (1, H, T) LSE: rows past this sequence's Q
+                            # length belong to the NEXT sequence — never written,
+                            # and there is no padded region to trim.
+                            if lse_q_idx < seqlen_q:
+                                lse_row = lse_arr[0, head_idx, :]
+                                lse_row[q_row_base + lse_q_idx] = lse_out
+                        else:
+                            # Rows at/past this batch's Q length trim to -inf.
+                            if lse_q_idx >= seqlen_q:
+                                lse_out = -cutlass.Float32.inf
+                            if lse_q_idx < q.shape[1]:
+                                lse_row = lse_arr[batch_idx, head_idx, :]
+                                lse_row[lse_q_idx] = lse_out
 
             prims.barrier_cta_sync(self.bar_compute_sync, thread_count=self.threads_compute)
 
@@ -1237,8 +1239,8 @@ def __call__(
         k: cute.Tensor,
         v: cute.Tensor,
         o: cute.Tensor,
-        lse: cute.Tensor,
-        sinks: cute.Tensor,
+        lse: Optional[cute.Tensor],
+        sinks: Optional[cute.Tensor],
         seq_q_lens: cute.Tensor,
         seq_kv_lens: cute.Tensor,
         softmax_scale_log2: cutlass.Float32,
@@ -1250,9 +1252,10 @@ def __call__(
         :param k: Key tensor with shape ``(B, Sk, H, D)``.
         :param v: Value tensor with shape ``(B, Sk, H, D)``.
         :param o: Output tensor with shape ``(B, Sq, H, D)``.
-        :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output.
-        :param sinks: ``(H,)`` fp32 per-Q-head sink logits, or an unused
-            dummy tensor when the kernel is configured without ``has_sink``.
+        :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to
+            compile the LSE store out entirely (no dummy buffer needed).
+        :param sinks: ``(H,)`` fp32 per-Q-head sink logits; must be ``None``
+            exactly when 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 softmax_scale_log2: ``softmax_scale * log2(e)``.
@@ -1273,11 +1276,14 @@ def __call__(
         for name, tensor in (("Q", q), ("K", k), ("V", v), ("O", o)):
             if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)):
                 raise ValueError(f"{name} must use compact BSHD storage")
-        if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])):
-            raise ValueError("LSE must have shape (B, H, Sq)")
-        if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)):
-            raise ValueError("LSE must be compact row-major")
-        if cutlass.const_expr(self.has_sink and sinks.shape != (q.shape[2],)):
+        if cutlass.const_expr(lse is not None):
+            if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])):
+                raise ValueError("LSE must have shape (B, H, Sq)")
+            if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)):
+                raise ValueError("LSE must be compact row-major")
+        if cutlass.const_expr(self.has_sink != (sinks is not None)):
+            raise ValueError("sinks must be provided exactly when the kernel is configured with has_sink")
+        if cutlass.const_expr(sinks is not None and sinks.shape != (q.shape[2],)):
             raise ValueError("sinks must have shape (H,)")
         if cutlass.const_expr(self.thd_varlen):
             if cutlass.const_expr(q.shape[0] != 1):
@@ -1360,12 +1366,17 @@ def compile(  # noqa: A001
     skv: int = 128,
     d: int = 128,
     max_sq: int = 0,
+    has_lse: bool = True,
 ) -> Callable:
     """Compile and cache one architecture-specific compact BSHD shape.
 
     THD specializations pack the batch: ``b`` is the real sequence count,
     ``sq``/``skv`` are the packed token totals, and ``max_sq`` (the longest
     sequence's Q length) sizes the per-sequence grid.
+
+    ``has_lse=False`` compiles the LSE store out (the kernel specializes on a
+    ``None`` LSE argument) — callers that don't want stats pass no LSE buffer
+    at all instead of a dummy.
     """
 
     kernel = SM120FusedMultiHeadAttentionForward(
@@ -1409,17 +1420,25 @@ def compile(  # noqa: A001
         stride_order=(3, 2, 1, 0),
         assumed_align=16,
     )
-    fake_lse = cute.runtime.make_fake_compact_tensor(
-        cutlass.Float32,
-        (fake_batch, qh, sq),
-        stride_order=(2, 1, 0),
-        assumed_align=4,
+    fake_lse = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32,
+            (fake_batch, qh, sq),
+            stride_order=(2, 1, 0),
+            assumed_align=4,
+        )
+        if has_lse
+        else None
     )
-    fake_sinks = cute.runtime.make_fake_compact_tensor(
-        cutlass.Float32,
-        (qh,),
-        stride_order=(0,),
-        assumed_align=4,
+    fake_sinks = (
+        cute.runtime.make_fake_compact_tensor(
+            cutlass.Float32,
+            (qh,),
+            stride_order=(0,),
+            assumed_align=4,
+        )
+        if PARAMS.has_sink
+        else None
     )
     fake_seq_q_lens = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,

python/cudnn/sdpa/graph_analyzer.py

@@ -350,11 +350,26 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts:
         seq_q_trim = seq_len_q is not None and not use_padding_mask
         padded = use_padding_mask and seq_len_kv is not None
 
+    # The kernels consume per-batch lengths as int32 directly; there is no
+    # implicit conversion anywhere on the execute path (it would allocate and
+    # launch a cast kernel).
+    for name, t in (("seq_len_q", seq_len_q), ("seq_len_kv", seq_len_kv)):
+        if t is not None:
+            t_dtype = _DTYPE_FROM_CUDNN.get(t.get_data_type())
+            if t_dtype != torch.int32:
+                return _invalid(f"{name} must be int32; got {t_dtype}")
+
     sink_token = rec.get("sink_token")
     if sink_token is not None:
         sink_dim = tuple(sink_token.get_dim())
         if sink_dim != (1, h_q, 1, 1):
             return _invalid(f"sink_token must be (1, H_q, 1, 1); got {sink_dim}")
+        # The kernels consume fp32 sink logits directly; there is no implicit
+        # conversion anywhere on the execute path (it would allocate and
+        # launch a cast kernel).
+        sink_dtype = _DTYPE_FROM_CUDNN.get(sink_token.get_data_type())
+        if sink_dtype != torch.float32:
+            return _invalid(f"sink_token must be float32; got {sink_dtype}")
 
     generate_stats = rec.get("generate_stats")
     is_inference = rec.get("is_inference")