Make the FROST SDPA execute path async where it can be, and stop it re-deriving build-time facts

python/cudnn/AGENTS.md

@@ -67,6 +67,61 @@ unsupported input runnable.
   descriptors) — acceptance is a promise about the execute path, not about
   what the adapter can patch up.
 
+**Rule 3 — `execute()` never reads device memory to the host.**
+
+Rules 1 and 2 both cite CUDA-graph capture as the reason for what they ban, but
+neither names the thing that breaks it most directly: a device-to-host read.
+
+- **No `.item()` / `.tolist()` / `.cpu()` / `.to("cpu")` / `.numpy()` /
+  `float(tensor)` / `int(tensor)` / `torch.is_nonzero`**, and no branch or
+  f-string that forces one, on an execute argument or anything derived from one.
+  A D2H read makes `execute()` synchronous — the whole point of an async launch
+  API is gone. **Nor may it block**: no `torch.cuda.synchronize()`, no
+  stream/event `synchronize()`. A sync reads nothing but costs the same.
+- **It is a functional gap, not a slow path.** A blocking D2H during stream
+  capture is illegal, so a path that does one **cannot be CUDA-graph captured
+  at all** — which is how every inference stack runs decode.
+- **Its cost is the queue, not the transfer.** Measured on SM100: one
+  `.tolist()` costs 11 µs against a drained queue, 2.6 ms behind 16 queued
+  matmuls. Any figure you measure in a microbenchmark is the floor.
+- **If a device value must shape the launch**, pass its pointer and dereference
+  in-kernel, or compile on an envelope and let the kernel read the real extent
+  from device metadata (the f16 prefill kernels already do this for head dims).
+- **A validation that needs a device read is not a validation.** Decline the
+  declaration in `check_support()` — per Rule 2, the graph says what it will
+  hand you — or assert in-kernel. Reading lengths back to decide whether to
+  raise buys nothing: the Router had to choose an engine before any buffer
+  existed.
+
+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). `t_q`/`t_kv` reach the host only because `T` is a
+  compile-time constant. Compile on the `b * s_q_max` envelope and pass `T` as a
+  runtime argument, as the f16 kernels already do for head dims;
+  `sdpa_fwd_wrapper_sm80` shows the other half — it requires `max_s_q` from the
+  caller 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
+  lifted to `check_support()`: `use_padding_mask=True` requires a `seq_len_q`
+  tensor even when only KV is padded, so no static rule separates "declares
+  per-batch Q lengths" from "the lengths are actually short" — declining the
+  declaration would drop the KV-only-padding population these kernels serve
+  correctly. It goes away when the FP8 kernels get the epilogue trim; until
+  then the read is what keeps a short length from being silently ignored.
+- The ragged cache-key `max()` in `sdpa/{fwd,bwd}/api.py`.
+- `cu_seqlens_{q,k}.to(dtype=..., device="cpu")` in the SM80 packed-THD backward
+  (`sdpa/bwd/kernels/bprop_f16_sm80.py`). Reachable only through the standalone
+  wrapper: the registered `sdpa_bwd_sm80` spec declares `thd=False`, so
+  `graph.execute()` does not route here. Still a violation, and it is the one to
+  fix first if that spec ever gains THD.
+
+When auditing this list, grep for the ARGUMENT, not the call shape:
+`device="cpu"` finds `to(dtype=..., device="cpu")`, which `to(device="cpu")`
+misses.
+
 ## Frontend-only kernel package layout
 
 ```

python/cudnn/_pygraph.py

@@ -433,12 +433,16 @@ def _rename_tensor(self, t: Tensor, name: str) -> None:
         becomes ambiguous and leaves the unique-name index."""
         if name == t.name:
             return
-        # NOT freeze-guarded: names are labels (classic allows renaming after
-        # build — the lowered graph already carries the old label, and labels
-        # have no execution semantics).
+        # Freeze-guarded like every other setter. A name is a label, but it is
+        # also a variant-pack key: a compiled plan may be holding the name it
+        # was built with, and the lowered graph keeps the old one, so a rename
+        # after planning leaves two answers to "which tensor is 'q'" -- and
+        # swapping two names would silently rebind buffers. Nothing needs to
+        # rename a planned graph; build another one.
+        self._check_mutable("rename a tensor")
         if self._tensors.get(t.name) is t:
             del self._tensors[t.name]
-        object.__setattr__(t, "name", name)  # label write is exempt from the freeze
+        object.__setattr__(t, "name", name)
         if name in self._tensors or name in self._ambiguous_names:
             self._tensors.pop(name, None)
             self._ambiguous_names.add(name)

python/cudnn/api_base.py

@@ -41,6 +41,21 @@ def _torch():
     return sys.modules.get("torch")
 
 
+_RAW_STREAM = None
+
+
+def _raw_stream(torch):
+    """``device_index -> raw CUstream int``, resolved once.
+
+    ``torch._C._cuda_getCurrentRawStream`` is private, so fall back to the
+    documented accessor if a torch build ever drops it.
+    """
+    global _RAW_STREAM
+    if _RAW_STREAM is None:
+        _RAW_STREAM = getattr(torch._C, "_cuda_getCurrentRawStream", None) or (lambda _dev: torch.cuda.current_stream().cuda_stream)
+    return _RAW_STREAM
+
+
 def _is_framework_tensor(obj: Any) -> bool:
     """True for framework tensors (torch/jax/numpy/...) as opposed to dtypes or shape/stride tuples."""
     return hasattr(obj, "__dlpack__")
@@ -566,14 +581,15 @@ def _get_default_stream(self, stream: Optional[cuda.CUstream]) -> cuda.CUstream:
             ...     current_stream = self._get_default_stream(current_stream)
             ...     # Now current_stream is guaranteed to be a valid stream
         """
-        if stream is None:
-            torch = _torch()
-            if torch is not None:
-                self._logger.debug(f"{self.__class__.__name__}: No CUDA stream provided, using torch current stream")
-                return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-            self._logger.debug(f"{self.__class__.__name__}: No CUDA stream provided and torch not imported, using CUDA legacy default stream")
+        if stream is not None:
+            return stream
+        torch = _torch()
+        if torch is None:
             return cuda.CUstream(0)
-        return stream
+        # _cuda_getCurrentRawStream is the same value torch.cuda.current_stream()
+        # reports, without building the Stream wrapper: 0.1 us against 4.3 us,
+        # and execute() calls this once per launch.
+        return cuda.CUstream(_raw_stream(torch)(torch.cuda.current_device()))
 
     def _pad_tensor_to_ndim(
         self,

python/cudnn/sdpa/bwd/engine.py

@@ -45,24 +45,25 @@ def __init__(self, name: str, compiled: Any):
         # graph API hands us covers every IO tensor of the graph, so key the
         # kernel's own operands out of it by uid (uids are eager and unique).
         self._tensors = list(compiled.binding.bound_tensors())
+        # A bound tensor's uid is fixed once the graph is frozen, so read them
+        # here rather than re-walking the list on every execute.
+        self._uids = [t.get_uid() for t in self._tensors]
+        self._workspace_bytes = int(getattr(compiled, "workspace_bytes", 0) or 0)
 
     def get_workspace_size(self) -> int:
-        return int(getattr(self._compiled, "workspace_bytes", 0) or 0)
+        return self._workspace_bytes
 
     def execute(self, graph: "pygraph", uid_to_data, ctx: ExecutionContext) -> None:
         # Keyed by IR tensor object: that is the binding's own identity, and the
         # only key resolve_variant_pack() accepts for an auto-assigned uid.
         pack = {}
-        missing = []
-        for t in self._tensors:
-            buf = uid_to_data.get(t.get_uid())
+        for t, uid in zip(self._tensors, self._uids):
+            buf = uid_to_data.get(uid)
             if buf is None:
-                missing.append(t.get_name() or t.get_uid())
-            else:
-                pack[t] = buf
-        if missing:
-            raise ValueError(f"{self._name}: the variant pack is missing buffers for {missing}")
-        required = self.get_workspace_size()
+                missing = [t.get_name() or uid for t, uid in zip(self._tensors, self._uids) if uid_to_data.get(uid) is None]
+                raise ValueError(f"{self._name}: the variant pack is missing buffers for {missing}")
+            pack[t] = buf
+        required = self._workspace_bytes
         if required:
             _check_workspace(ctx.workspace, required, self._name)
             self._compiled(pack, ctx.workspace, stream=ctx.stream)

python/cudnn/sdpa/fwd/engine.py

@@ -46,24 +46,25 @@ def __init__(self, name: str, compiled: Any):
         # graph API hands us covers every IO tensor of the graph, so key the
         # kernel's own operands out of it by uid (uids are eager and unique).
         self._tensors = list(compiled.binding.bound_tensors())
+        # A bound tensor's uid is fixed once the graph is frozen, so read them
+        # here rather than re-walking the list on every execute.
+        self._uids = [t.get_uid() for t in self._tensors]
+        self._workspace_bytes = int(getattr(compiled, "workspace_bytes", 0) or 0)
 
     def get_workspace_size(self) -> int:
-        return int(getattr(self._compiled, "workspace_bytes", 0) or 0)
+        return self._workspace_bytes
 
     def execute(self, graph: "pygraph", uid_to_data, ctx: ExecutionContext) -> None:
         # Keyed by IR tensor object: that is the binding's own identity, and the
         # only key resolve_variant_pack() accepts for an auto-assigned uid.
         pack = {}
-        missing = []
-        for t in self._tensors:
-            buf = uid_to_data.get(t.get_uid())
+        for t, uid in zip(self._tensors, self._uids):
+            buf = uid_to_data.get(uid)
             if buf is None:
-                missing.append(t.get_name() or t.get_uid())
-            else:
-                pack[t] = buf
-        if missing:
-            raise ValueError(f"{self._name}: the variant pack is missing buffers for {missing}")
-        required = self.get_workspace_size()
+                missing = [t.get_name() or uid for t, uid in zip(self._tensors, self._uids) if uid_to_data.get(uid) is None]
+                raise ValueError(f"{self._name}: the variant pack is missing buffers for {missing}")
+            pack[t] = buf
+        required = self._workspace_bytes
         if required:
             _check_workspace(ctx.workspace, required, self._name)
             self._compiled(pack, ctx.workspace, stream=ctx.stream)

python/cudnn/sdpa/graph_analyzer.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Optional
 
 import cudnn
@@ -665,8 +665,23 @@ class SdpaBinding:
     dbias: Any = None
     dsink: Any = None
 
-    def bound_tensors(self) -> list:
-        return [
+    # Built once on first use and reused. Rebuilding it per execute cost ~1.3 us
+    # per bound operand: three passes over the bound list and five dict
+    # constructions, not any one expensive getter.
+    #
+    # What makes the cache safe is the graph, not this class: a binding is
+    # constructed by the engine's lowering AFTER the graph is frozen, and a
+    # frozen graph can no longer re-uid or rename a tensor, so the names and
+    # uids indexed here cannot move. The binding itself is still an ordinary
+    # mutable dataclass -- reassigning a field after the first index() would go
+    # unnoticed. Nothing does; an ordered-slot binding would remove the question.
+    # init=False so a replace()d binding rebuilds rather than inheriting a
+    # cache for the operands it no longer has; compare/repr excluded so the
+    # cache cannot change how a binding prints or compares.
+    _index: Optional[tuple] = field(default=None, init=False, repr=False, compare=False)
+
+    def _build_index(self) -> tuple:
+        bound = [
             t
             for t in (
                 self.q,
@@ -701,6 +716,33 @@ def bound_tensors(self) -> list:
             )
             if t is not None
         ]
+        name_counts: dict = {}
+        uid_counts: dict = {}
+        names, uids = [], []
+        for t in bound:
+            nm = _safe_name(t)
+            names.append(nm)
+            if nm is not None:
+                name_counts[nm] = name_counts.get(nm, 0) + 1
+            uid = _safe_uid(t)
+            uids.append(uid)
+            if uid is not None:
+                uid_counts[uid] = uid_counts.get(uid, 0) + 1
+        # A name or uid carried by two bound tensors identifies neither.
+        by_obj = {id(t): t for t in bound}
+        by_name = {nm: t for nm, t in zip(names, bound) if nm is not None and name_counts[nm] == 1}
+        by_uid = {uid: t for uid, t in zip(uids, bound) if uid is not None and uid_counts[uid] == 1}
+        self._index = (tuple(bound), by_obj, by_uid, by_name)
+        return self._index
+
+    def index(self) -> tuple:
+        """``(bound, by_obj, by_uid, by_name)`` — the resolution tables."""
+        return self._index or self._build_index()
+
+    def bound_tensors(self) -> tuple:
+        """The bound tensors, in slot order. A tuple: this is the binding's own
+        record, not a working list for a caller to edit."""
+        return self.index()[0]
 
 
 def _safe_name(t: Any) -> Optional[str]:
@@ -730,22 +772,7 @@ def resolve_variant_pack(variant_pack: dict, binding: SdpaBinding) -> dict:
         raise TypeError(
             f"cudnn.sdpa: compiled plans are called with a variant-pack dict {{cudnn_tensor | uid | name: buffer}}; got {type(variant_pack).__name__}"
         )
-    bound = binding.bound_tensors()
-    by_obj = {id(t): t for t in bound}
-
-    name_counts: dict = {}
-    uid_counts: dict = {}
-    for t in bound:
-        nm = _safe_name(t)
-        if nm is not None:
-            name_counts[nm] = name_counts.get(nm, 0) + 1
-        uid = _safe_uid(t)
-        if uid is not None:
-            uid_counts[uid] = uid_counts.get(uid, 0) + 1
-    by_name = {_safe_name(t): t for t in bound if name_counts.get(_safe_name(t)) == 1}
-    by_uid = {_safe_uid(t): t for t in bound if uid_counts.get(_safe_uid(t)) == 1}
-    by_name.pop(None, None)
-    by_uid.pop(None, None)
+    _bound, by_obj, by_uid, by_name = binding.index()
 
     resolved: dict = {}
     for key, buf in variant_pack.items():

samples/python/24_layernorm_zero_centered_gamma_forward_training_and_backward.ipynb

@@ -510,7 +510,7 @@
     "d_out = bwd_graph.tensor(name=\"d_out\", dim=x_gpu.size(), stride=x_gpu.stride(), data_type=x_gpu.dtype)\n",
     "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n",
     "gamma_bwd = bwd_graph.tensor_like(gamma, name=\"gamma\")\n",
-    "one_bwd = graph.tensor_like(one_cpu).set_name(\"one\")\n",
+    "one_bwd = bwd_graph.tensor_like(one_cpu, name=\"one\")\n",
     "mean_bwd = bwd_graph.tensor_like(mean, name=\"mean\")\n",
     "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n",
     "\n",