cudnn-frontend / cudnn

# python/cudnn — Agent Guide

The `cudnn` Python package: pybind11-backed graph API plus pure-Python **frontend-only OSS kernels** (CuTeDSL). See `README.md` in this directory for the package inventory and [../../AGENTS.md](../../AGENTS.md) for build/test commands.

## Import-time rules (the most common way to break this package)

- `import cudnn` must work **without** torch/cutlass/cuda-python installed. Everything that needs them is exported lazily via `_LAZY_OPTIONAL_IMPORTS` in `__init__.py` — a module-level `__getattr__` imports the submodule on first attribute access and re-raises failures as `ImportError` that names the missing framework module (torch, jax, cuda-python) alongside the base `pip install nvidia-cudnn-frontend` hint — or, for a CuTe DSL below `CUTEDSL_MIN_VERSION`, points at the DSL upgrade (Rule 7).
- A function and its implementation module can share a name. Python installs an imported
  submodule directly on its parent, bypassing module `__getattr__`. Check direct-submodule,
  sibling-symbol, and public-symbol import orders in fresh interpreters; the detector is
  `test_ops_callable_exports_survive_import_order` in `test/python/core/test_import_boundaries.py`.
- 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 required CuTeDSL dependencies (`pyproject.toml` `[project] dependencies`) unless a kernel truly needs a new package. The `[cutedsl]` extra now holds only `cuda-python`.

## 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.
- **No degenerate-path fixups.** Runtime-degenerate inputs (e.g. all-zero
  THD ``seq_kv_lens``) go through the kernel's own dead-row path — never
  re-implemented adapter-side with `fill_`/`copy_` writes (surprise kernel
  launches, and a second copy of the semantics that can drift). If a packed
  extent would be zero, bind a never-dereferenced dummy view over storage
  the contract already guarantees.
- **Overlapping optional declarations are validated as a set, not one by
  one.** When two mechanisms can declare the same thing (ragged offsets vs
  `cu_seqlen` vs plain `seq_len` tensors), each combination is either
  defined or explicitly rejected — an unhandled overlap is an untested code
  path with unspecified semantics, and "both supplied" is exactly the case
  no per-argument check catches (raised in review on PR #266).
- **An execute-time shape/stride override must reach the executor or raise
  before launch.** A raw uid-map plan cannot consume it; do not silently drop
  the override triple. Filter override-enabled graphs per plan's binding capability,
  preserving prepared plans while declining tensor-only ones. Detectors:
  `test_uid_map_plan_rejects_runtime_overrides_before_execute` and
  `test_override_filter_preserves_compatible_split_candidates`.
- **Shape overrides do not enlarge the producer's storage.** Validate metadata
  inputs against their observed span as well as their effective shape, and check
  pointer alignment for the element type. The host-only detector is
  `test_dense_metadata_rejects_short_observed_storage_and_misalignment`.

**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.

**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.
- **Prove it; do not grep for it.** The list above is a reminder, not a
  detector — the spellings are many (`int(cu[i])` on a CUDA tensor is a
  blocking copy that a search for `.item()` will not find) and a reviewer who
  greps a subset concludes "clean". Assert the property instead:

  ```python
  torch.cuda.set_sync_debug_mode("error")   # any blocking D2H now raises
  try:
      out.backward(grad)                    # or graph.execute(...)
  finally:
      torch.cuda.set_sync_debug_mode("default")
  ```

  Put that in a test (see `test_varlen_backward_does_not_sync`), and check the
  test is RED against the old code before trusting it — a sync test that was
  never seen to fail is asserting nothing.
- **Suspect duplicated logic first.** Every violation found so far has been a
  *second* copy of a conversion that was already device-side somewhere else:
  the packed-to-padded LSE repad existed in both `sdpa/fwd/torch_op.py` (with
  `searchsorted`, device-side) and `torch/sdpa_provider.py` (a `for i in
  range(B): int(cu[i])` loop). Extract the correct one and call it from both
  rather than writing the obvious loop again.

The SM80 packed backward wrapper now shares the prepared graph chain. When
`max_s_kv` is absent, the packed capacity bounds the grid; B+1 prefixes stay
on device. `test_wrapper_without_length_hints_does_not_sync` detects the old
CPU copy after warmup (verified RED on the preceding wrapper). Capacity,
Stats head pitch, launch bounds and deterministic-counter size are distinct:
a caller may reserve more storage than `B * max_sequence_length`. Preserve
that capacity while honoring the caller's valid launch bounds; cover poisoned
slack and changed device prefixes under replay.

When auditing this list, grep for the ARGUMENT, not the call shape:
`device="cpu"` finds `to(dtype=..., device="cpu")`, which `to(device="cpu")`
misses.

**Rule 4 — compile keys are PLAN-TIME-ONLY: never key a kernel compile on
runtime data values.**

`cute.compile` takes seconds. Anything an execute path feeds into a
compile-cache key (an `lru_cache`d `compile()` wrapper, a template parameter,
a fake-tensor extent) must be derivable from the graph declaration alone —
tensor dtypes, declared strides, head counts, head dims, flags. Values read
out of runtime tensors (THD packed token totals, max sequence lengths, batch
contents) change every step under continuous batching, so a key that includes
them degenerates into a fresh multi-second compile per `execute()` — a
pathology that no correctness test catches (issue #552 is the case study:
`sq=t_q, skv=t_kv` in the THD compile key). Rule 3 bans the read that feeds
such a key; this rule bans the key itself — a runtime value that arrives
legally (a caller-passed host scalar, an `int(tensor.shape[...])`) still must
not become a compile key.

- **Runtime extents compile DYNAMIC.** Use `cute.sym_int()` in the fake
  tensors (one symbol per ragged group) so one compiled artifact re-binds any
  total; runtime scalars the launch needs (grid extents like THD `max_sq`)
  are `cutlass.Int32` call arguments, never compile parameters.
- **Derived values count.** A stride tuple whose batch stride is
  `t_q * token_stride` smuggles the runtime total into the key just as
  surely as `sq=t_q` — normalize it out (zero the never-stepped batch
  stride, rebuild it symbolically kernel-side).
- **Compile at plan time, re-bind at execute.** With a plan-time-only key
  there is no reason to defer: `compile()` builds the artifact once and the
  execute path's cached call must be a guaranteed hit. Guard it with a
  cache-miss regression test (see
  `test_dsl_sm100_thd_compile_key_plan_time_only`), not by inspection.
- **Prove artifact reuse across plans.** A stable key does not guarantee that
  `compile_cached` exported an artifact: its wrapper serializer declines
  dataclass compile arguments, including `Constexpr` ones. Read an immutable
  template module's `PARAMS` inside the host instead of passing the same
  dataclass again. Build a second plan with `cute.compile` forbidden, assert
  a real cache hit, and check the reloaded artifact's outputs and graph replay;
  `test_replan_reloads_prepared_artifact` is the SM80 detector.
- **Issue #604 is closed**: SM80 THD compiles use symbolic packed extents.
  The prepared backward host takes Int64 capacities and launch bounds at
  runtime, including the compact Stats head pitch and deterministic-counter
  size; the retained tensor compiles use `cute.sym_int` and key on
  `b = 1, sq = skv = 0` plus the plan-time sequence count. The regression tests are
  `test_sm80_bwd_thd_compile_key_plan_time_only` (wrapper) and
  `test_graph_thd_compile_key_is_plan_time_only` (graph path). Copy that
  pattern, not a shape-keyed one.
- **Key on exactly the contract-relevant set — no more, no less.** Both
  failure modes shipped on PR #553 and were caught in review: *under-keying*
  (the cache keyed only `x.shape`/`w.shape` while `check_support()`
  validated weight, RoPE, and scale descriptors — a hit can return an
  artifact compiled for a different contract, i.e. wrong results) and
  *over-keying* (`alpha` passed at launch, `m`/`n`/`k` on a shape-generic
  kernel — every miss is a spurious multi-second recompile). Enumerate what
  `check_support()` validates and what the kernel specializes on; the key
  is that set.

**Rule 5 — every torch operation on the execute path is ordered on the
LAUNCH stream, never implicitly on torch's current stream.**

The kernel launches on the stream carried by the execute-time handle
(`ExecutionContext.stream`), but torch enqueues work — H2D metadata uploads,
buffer resets (`zero_()`), post-kernel reductions (`div_()`, `copy_()`),
and the caching allocator's stream-tagging of fresh blocks — on
`torch.cuda.current_stream()`. When the two differ, the prep and the kernel
race (PR #543 is the case study: the THD `[seq_kv | cu_q | cu_k]` upload vs
the kernel that reads it).

- **Resolve the launch stream FIRST**, before any torch work in the execute
  path, and run every torch op (including allocator calls: workspace-less
  fallback allocations, cached-dummy first use) inside
  `_torch_stream_context(current_stream, device)` — see the fp8/mxfp8 amax
  resets and both `_execute_thd` paths in `sdpa/fwd/api_dsl.py`.
- **Consumers too, not just producers**: anything reading what the kernel
  wrote (`amax_o.div_()`, an O scratch copy-back) belongs on the launch
  stream for the same reason.
- The PyTorch-integration path launches on torch's current stream, where the
  context is a no-op — the race only bites direct graph-API users with an
  explicit handle stream, which is exactly why tests miss it. Order the work
  by construction rather than relying on the common case.
- **The device is implicit state exactly like the stream.** A `torch.empty`
  (or any allocator call) without a device context silently allocates on the
  *current* GPU, not the input tensor's — wrap execute-path allocations in
  the right device context as well as the stream context. And a raw pointer
  argument is a contract: validate device-residency and dtype (a CUDA int64
  tensor, not a host tensor) before handing its address to a kernel — both
  flagged in review on PR #517.
- **A raw stream handle never goes straight into `torch.cuda.ExternalStream`.**
  Every eager caller on torch's default stream hands us a default-stream
  sentinel (`0`, `cudaStreamLegacy` = 1, `cudaStreamPerThread` = 2), and torch
  before PR pytorch/pytorch#183258 (in v2.13.0; NGC 26.06 and torch <= 2.12
  lack it) returns a fresh NON-BLOCKING pool stream for `ExternalStream(0)`.
  Torch work issued in that context is unordered with a kernel launched on
  `CUstream(0)`: on an idle GPU the copies win the race and every isolated
  test passes; under xdist load the kernel reads stale conversion buffers and
  a staged output is copied back before it is written (the qa sm90
  `hopper_cuda` reds, PR #1165 — the same trap FROST SDPA hit in #682/#717/#860).
  Map the sentinels and torch's own default stream to
  `torch.cuda.default_stream(device)`, the current stream to itself, and only a
  genuine side stream to `ExternalStream(handle, device=device)`. The one
  implementation is `cudnn._torch_stream` (`as_torch_stream`, `stream_context`,
  with the raw-handle fast path); every engine calls it, none writes its own
  wrapper. Detector:
  monkeypatch `torch.cuda.ExternalStream` to raise and drive the execute path
  with handle 0 (`test_hopper_marshal_stream.py`).

SDPA-specific hard rules (cited as Rule S1, S2, ...) live in
[sdpa/AGENTS.md](sdpa/AGENTS.md) — read it before touching anything under
`python/cudnn/sdpa/`.

**Rule 6 — every Frost-generated kernel has a cuDNN-attributable symbol.**

- Immediately after every ``@cute.kernel`` definition, including auxiliary and
  generated-template kernels, call the public naming API:

  ```python
  kernel.set_name_prefix("cudnn", remove_cutlass_symbol=True)
  ```

- Use the decorated function's actual name, keep the default
  ``keep_mangled_name=True``, and do not use compiler flags or symbol rewriting
  instead.
- Verify with ``(cd test/python && pytest -q test_frost_kernel_name_prefix.py)``.
- This call runs at module import, and DSL APIs used this way can be newer than
  the `pyproject.toml` floor admits. It is legal only because Rule 7's gate runs
  before the kernel module is imported — do not add an import path that skips
  it.

**Rule 7 — gate the CuTe DSL version at runtime; never assume the installed
DSL satisfies your kernel.**

- The `pyproject.toml` floor on `nvidia-cutlass-dsl` (`>=4.6.2`) is the
  **downstream** floor, not ours: vLLM and SGLang inherit quack-kernels'
  `==4.6.2`, and a higher floor would make this package uninstallable next to
  them. The FROST-derived kernels need more (`CUTEDSL_MIN_VERSION`, 4.7.0). So
  an installed DSL that satisfies pip can still be below what a kernel needs,
  and every backend/kernel must cope with that at runtime.
- Before a path imports a DSL-version-specific API, check the installed version
  with `cudnn.frost.buffers.cutedsl_state()` / `cutedsl_too_old()` (floor:
  `CUTEDSL_MIN_VERSION`) and **decline, or raise an error that names the
  version** — `cutedsl_requirement_error(what)` builds it. Never let the failure
  surface as an `AttributeError` / `TypeError` / `ModuleNotFoundError` from
  inside the DSL, and never let it read as a missing-dependency install hint:
  the package is installed, and that `pip install` changes nothing.
- The gate lives at the entry the caller hits, before the kernel module is
  imported: the semantic op's route check (`_can_route_causal_conv1d_bulk` in
  `ops/causal_conv1d.py`, `_validated_native_update` in
  `ops/_causal_conv1d_update.py`), an engine's `check_support`, or the family
  `__init__`'s lazy import. Module-scope code in kernel files may assume the
  floor only because that gate ran first.
- A version floor does not guarantee target-architecture support. For example,
  public 4.7.0 lacks `sm_107a`; check the actual target capability before
  admitting SM107 plans. Test graph and standalone declines with the real
  unsupported wheel as well as a controlled missing-target probe.
- Known floors — extend this list when you take a dependency on a newer API,
  and say so in the PR body if it raises the floor of a user-facing op:
  `cutlass.experimental.*` (primitives, `cuda.tensor_map`; everything under
  `cudnn/frost/tile_dsl` inherits it) → 4.7.0.
- Tests that import a kernel module directly `pytest.skip` on a too-old DSL —
  they do not fail. CI runs the `oss:` lanes across the supported DSL versions
  (`ci/stages/oss_tests/jobs.yml` in internal CI); a lane below your floor
  must show skips, not errors.
- **MLIR argument widths are part of a DSL API migration.**
  `FastDivmodDivisorV2` carries two MLIR values (encoded divisor and scalar),
  while its predecessor carried one. Custom `__new_from_mlir_values__` methods
  must slice by the lengths recorded during extraction, rather than assuming
  one value per field. Static coordinates can contribute zero values, too.
  `test_cutlass_schedulers.py` checks these round trips on real MLIR values.
- Why: PR #799's `causal_conv1d_update` imported `frost.tile_dsl` from a route
  with no version check and broke the 4.6.2 lane — the version vLLM and SGLang
  ship — with a bare `ModuleNotFoundError: cutlass.experimental`; the bulk
  route next to it had the check and declined cleanly. Earlier, PR #854's
  module-scope `set_name_prefix(..., remove_cutlass_symbol=True)` failed the
  same way on a since-dropped 4.5.x lane, reported as "install optional
  dependencies".

**Rule 8 — the graph API owns no device memory and never blocks the host: at
build AND at execute.**

Rules 1 and 3 say this for `execute()` in torch vocabulary. This rule closes
the two gaps that produced the #1151 sm103 red: plan build, and the driver-API
spellings the torch vocabulary does not name.

- **No plan- or engine-owned device allocation, ever.** Not `cuMemAlloc` /
  `cudaMalloc`, not `torch.empty` / `torch.zeros`, not `frost.buffers.DeviceBuffer`,
  at build or at execute, cached or not. Execute scratch — per-batch metadata,
  on-device TMA descriptors, an output the kernel always writes but the graph
  did not request — is carved from the caller's workspace and declared through
  `get_workspace_size()` (`prepared.py`: `meta_ptr = workspace_ptr`,
  `o_desc_ptr = workspace_ptr + off_o_desc`). A path with no workspace contract
  (the direct `jit_from_cudnn_graph` MoE call) gets one; it does not get an
  allocation. Owned device memory has a GC-timed release, and a cyclic
  collection inside someone else's `torch.cuda.graph` window turned three
  per-plan dummies into `cuMemFree -> 900` and an invalidated capture (#1151 on
  sm103). #1152's relaxed-mode guard on the two remaining finalizers is defence
  in depth, not a licence.
- **A dead ABI slot is compiled out, or bound to `0`, or borrowed — never
  allocated.** Prefer `cutlass.const_expr` on `None` so the operand does not
  exist. Otherwise `0` is a legitimate address for a slot the kernel never
  dereferences: `cuTensorMapEncodeTiled` accepts a NULL global address (only
  misalignment is rejected), the DSL rejects only negative pointer addresses
  and ships `cute.runtime.nullptr`, and the tvm-ffi positional entry carries
  `cute.Pointer` parameters as plain integers — 12 prepared THD launches with
  `sinks_ptr = 0` pass bit-exact. A dead slot that turns out to be live then
  faults loudly instead of reading garbage. Borrowing an aligned address the
  contract already guarantees (`sinks_ptr = q.ptr` when `HAS_SINK` is off,
  descriptor-only THD rows aliasing live Q/O storage) is the other acceptable
  form. A cached `torch.zeros` dummy is not: Rule 1's `_dummy` exemption is
  grandfathered for `sdpa/fwd/api_dsl.py` and closed for new code.
- **No host-blocking call, build or execute.** `cuStreamSynchronize`,
  `cuCtxSynchronize`, `cudaDeviceSynchronize`, `cuEventSynchronize`, the
  synchronous `cuMemcpy*` / `cuMemsetD*` forms, `torch.cuda.synchronize()`,
  `.item()`. Build is lazy — it runs on the first execute of a shape, which in
  a serving stack is inside a stream capture (a FlashInfer graph-cache miss) —
  so build is held to the execute standard: uploads via `cuMemcpyHtoDAsync` on
  the launch stream, fills via `cuMemsetD32Async`, constants baked into the
  kernel image or written by a setup kernel. `cuMemAlloc` + `cuMemsetD32` +
  `cuStreamSynchronize(0)` at build was the #1151 anti-pattern.
- **An engine that cannot be async declines; it does not sync.** The CAKE KDA
  route plans its work items on the host (`cuMemcpyDtoHAsync` +
  `cuStreamSynchronize`) and therefore checks `cuStreamIsCapturing` and raises
  under capture (`linear_attention/cake/compiler.py::check_not_capturing`).
  That is the only legal shape of an exception: declared in the engine, loud,
  never on a captured stream.
- **Detectors.** `test_execute_allocates_nothing_and_never_synchronizes`
  (`torch.cuda.set_sync_debug_mode("error")` plus allocator accounting around
  a prepared execute) — run the same assertion around the BUILD;
  `test_collect_unrelated_resources_during_capture` (GC inside a global-mode
  capture window, then replay and a native cuDNN launch); and a capture test
  whose first execute of a shape happens inside `torch.cuda.graph`, so the lazy
  build runs under capture.


### Rule 8 / Rule 5 recipes — do exactly this, do not reinvent

One canonical answer per situation. Every engine copies the recipe; a reviewer
cites the recipe name. If a recipe does not fit, say so in the PR and extend the
recipe here — do not write a local variant (the audit behind #1165/#1167/this
section found 18 hand-rolled stream wrappers, 5 of them wrong).

**R1 — a raw stream handle becomes a torch stream.**
```python
from cudnn._torch_stream import as_torch_stream, stream_context

with stream_context(ctx.stream, device):          # torch work on the launch stream
    buf.copy_(src)
tensor.record_stream(as_torch_stream(ctx.stream, device))
```
Never call `torch.cuda.ExternalStream` / `get_stream_from_external` directly.
`stream_context(None)` is a no-op; a handle equal to torch's current stream is
a no-op via the raw-handle fast path; `0`/`1`/`2` and torch's default stream
resolve to `torch.cuda.default_stream(device)`.

**R2 — execute needs scratch (metadata, on-device descriptors, an output the
kernel always writes but the graph did not request, staging for a dead-but-
required tensor slot).** Declare it, carve it, never allocate it:
```python
def get_workspace_size(self) -> int:                 # BaseEngine / CompiledPlan
    return ws_align(meta_bytes) + ws_align(desc_bytes) + ...
def execute(self, graph, variant_pack, ctx):
    ws = Workspace.over(variant_pack, self.get_workspace_size(), type(self).__name__)  # frost/workspace.py
    meta = ws.take(4 * b + 4, "int32"); desc = ws.view(off, "int64", (slots * 16,))
```
(APIBase adapters: `scratch_workspace_bytes()` + `WorkspaceCarver(workspace, bytes, label).take(numel, dtype)`
in `sdpa/fwd/api_dsl.py`.) `Workspace(None, ...)` already raises
`"<owner> requires a N-byte workspace but execute() received none; allocate
graph.get_workspace_size() bytes and pass the buffer to execute()"` — reuse
that error, never fall back to `torch.empty` / `DeviceBuffer` when the caller
passed nothing. A path with no workspace contract gets one. Where a wrapper
allocates per call on the caller's behalf, it allocates under
`stream_context(<launch stream>)` (R1): the caching allocator orders a block's
reuse only against the stream it was allocated on, so scratch allocated on
torch's ambient stream for a handle re-streamed to a side stream is a
use-after-free waiting for load. For eager JAX wrappers, allocation readiness
(`block_until_ready`) does not extend storage lifetime through a foreign CUDA
consumer. Pair caller-layer stream-ordered allocation/free around the launch
(`grouped.backend_utils.wrapper_workspace`), or use an XLA custom call that owns
scratch. A cached plan's latest reference cannot cover overlapping calls. Test
multiple pending consumers on independent streams with an intercepted bounded
byte copy; never run a real kernel on an intentionally recycled scratch pointer.
For direct API workspace, validate CUDA device type and the operand's ordinal
before launch (`Workspace(..., device=...)`); byte size/alignment alone also
accept host memory.

**R3 — a dead ABI slot (the compiled kernel never dereferences it).** In order
of preference: (1) compile it out — an `Optional`/`None`-typed kernel parameter
read only under `cutlass.const_expr(flag)`, with `flag` in the compile key, and
`None` passed at BOTH compile and launch (DSA sm90 `mTopkIdx`, `mTopkLength`);
(2) pointer ABI: bind `0` (`prepared.py` `sinks_ptr`, dense `o_desc_ptr` /
`meta_ptr`); (3) tensor ABI with a required `cute.Tensor` parameter: borrow the
bytes from the workspace via R2 (`sdpa/bwd/api_dsl.py` dense `seq_kv` /
`desc_words`). Never a cached `torch.zeros` dummy, never `q.ptr` borrowing
(a live read then reads Q bytes silently instead of faulting). `0` is legal at
every layer: `cuTensorMapEncodeTiled` accepts NULL, the DSL rejects only negative
addresses, the tvm-ffi positional entry carries pointers as ints.

**R4 — a live per-batch table the caller did not give you (lengths, offsets,
scale scalars).** R2 (carve) + fill on the launch stream with one async op
(`cuMemsetD32Async`, `cuMemcpyHtoDAsync`, `buffers.memset_zero_async`), or make
it a scalar kernel argument. Never `torch.tensor(values, device=...)` per execute
(pageable H2D + implicit sync), never build it at `compile()` into a plan-owned
tensor.

**R5 — the input is not in the layout/dtype the kernel takes.** Decline in
`check_support()` with a `NotImplementedError` naming the tensor and its
strides/dtype (`_thd_check_strides_native` in `sdpa/fwd/api_dsl.py`), so the
Router picks another engine. Not `.contiguous()`, not `.to(dtype)`, not a
repack/copy-back — even into the workspace (Rule 2). If the engine is meant to
serve that input, the kernel reads it natively.

**R6 — something must block the host (host-planned work items, a D2H read
of lengths).** The engine is non-capturable: check `cuStreamIsCapturing` and
raise before doing it (`linear_attention/cake/compiler.py::check_not_capturing`),
and say so in its docstring. Never a silent `cuStreamSynchronize` /
`torch.cuda.synchronize()` / `.item()` on a build or execute path.

**R7 — you need a cuDNN handle and the caller gave none.** Graph API lowering:
`_pygraph._backend_handle_for_lowering` (process default, one per thread and
device, stream 0, destroyed at exit) — the graph never owns one. torch-op
layers: the per-device cached handle re-streamed to torch's current stream
before every call (`linear_attention/ops/common.py::get_handle`,
`ops/norm/_common.py`). Never `cudnn.create_handle()` inside a plan, engine or
C++ graph object.

**R8 — you own a CUDA resource whose release can be GC-timed (only tests and
the C++ PyGraph may).** Release inside `cuThreadExchangeStreamCaptureMode(RELAXED)`
and restore in `finally` (`frost/buffers.py DeviceBuffer.__del__`,
`pygraph.h CaptureModeGuard`). Production plans and engines reach this recipe
only if R2/R7 were skipped — fix that instead.

**R9 — proving it.** Around a warm `build()` + `execute()`:
`torch.cuda.set_sync_debug_mode("error")` (no sync), and
`torch.cuda.memory_stats()["allocation.all.allocated"]` unchanged across three
executes (no allocation) — `test_sdpa_prepared_thd.py::test_execute_allocates_nothing_and_never_synchronizes`,
`test_sdpa_bwd_thd_sm80.py::test_graph_thd_execute_does_not_allocate`. For a
capture-safety claim, `test_cuda_capture_lifetime.py` (GC inside a global-mode
window, then replay and a native launch). For R1, monkeypatch
`torch.cuda.ExternalStream` to raise and drive the path with handle 0
(`test_torch_stream.py`).

For a host-overhead migration, measure warm enqueue and captured GPU replay
separately, and check changed-input outputs after timing. A tiny async memset
can introduce cross-engine waits between graph kernels even when their kernel
durations improve. If replay regresses, inspect profiler timestamps between
nodes as well as individual kernel durations; the prepared FP8 scalar amax
reset uses an SM kernel to avoid that wait. Keep the empty-input reduction
identity correct when no attention host is launched
(`test_prepared_fp8_empty_thd_resets_amax_without_attention`).


**Rule 9 — backend and FROST share one FE Python graph contract; no special
treatment at the caller boundary.**

- **Same API and semantics.** For workloads supported by both engines, graph
  preparation and execution must accept the same arguments and preserve the
  same meaning of bindings, shape/stride overrides, outputs, workspace, streams,
  capture/replay, ownership, and invalidation. Users and integrations such as
  FlashInfer must not branch on the selected engine to invoke a different API.
  A private method consumed outside FE is still a caller boundary for this rule.
- **Preparation belongs inside FE.** Prepare plan-derived binding state during
  build/compile. Geometry first supplied through execute-time overrides can be
  prepared and reused internally when it becomes known; it does not require an
  extra caller-visible preparation step. Keep current buffer observations and
  required checks per call. Any optional lower-level interface needs a demonstrated
  benefit beyond internal caching and must serve both engines under one contract.
- **Optimize the shared path.** Reuse provider-independent preparation,
  UID ordering, metadata validation, binding, and caching work so host-overhead
  improvements benefit both engines. A backend-only prepared entry plus the old
  FROST path is not a unified optimization; neither is merely renaming that entry.
  Engine-specific launch lowering stays internal, without forcing either engine
  through a slower common implementation. This does not require factoring kernel
  pipeline templates into a common implementation.
- **Capabilities may differ; the contract must not.** Express real differences
  through support checks and plan selection. Decline unsupported requests under
  the common contract; never silently ignore an argument, weaken validation, or
  require a caller-side workaround for one engine. Any internal fallback must
  preserve the same execution and lifetime semantics.
- **Verify both routes explicitly.** Changes to this shared contract need tests
  using supported plans from both engines with the same caller code, covering
  fresh bindings, overrides, replan/invalidation, workspace/stream changes, and
  capture/replay as applicable. Check outputs and invalid-input behavior, not
  heuristic ranking or plan-list order. For shared host optimizations, measure
  both routes and report their scope separately; one engine's result is not proof
  for the other.

## Frontend-only kernel package layout

```
python/cudnn/<operation>/            # or sdpa/<direction>/, gemm/cutedsl/<layout>/<fusion>/
├── __init__.py                      # exports API class + wrapper via __all__
├── api.py                           # APIBase subclass + <operation>_wrapper() function
└── <kernel_module>.py               # CuTeDSL kernel implementation(s); some families use csrc/ per-arch trees
```

All GEMM fusions live under `gemm/`, grouped by how the operands are laid out:

```
python/cudnn/gemm/
├── cutedsl/
│   ├── dense/<fusion>/              # amax, dsrelu, proj_rope_mxfp8, srelu, swiglu
│   ├── grouped/<fusion>/            # dglu, dsrelu, dswiglu, glu, glu_hadamard,
│   │                                #   quant, srelu, swiglu, unfused, wgrad
│   └── discrete_grouped/<fusion>/   # dswiglu, swiglu (per-expert weight pointers)
├── frost/                           # the FROST GEMM engine (JIT fused matmul chains from cuDNN graphs)
│   ├── sm100/, sm120/               #   one tree per arch family: compiler.py + epilogue_codegen.py + kernel_templates/
│   ├── compiler.py, epilogue_codegen.py  # facades: become the active family's module (arch_family.py)
│   └── kernel_templates/            #   template code SHARED by both trees (split-K reduction)
├── ops/                             # backend-independent torch custom-op contracts
└── reference/                       # pure-PyTorch MATMUL/POINTWISE correctness engine
```

Shared helpers (schedulers, metadata utils, e.g. `gemm/cutedsl/grouped/moe_*.py`) stay internal to the family package — never exported through `cudnn`.

## CuTeDSL kernel bodies

**Do not factor code out of a `@cute.kernel` body into a plain Python helper.**
The DSL AST-transforms only the decorated function's own source: `for` becomes
an `ir_loop`, `if` becomes an `scf` region. A helper called from the kernel is
not transformed, so the ops it emits can land outside the enclosing region.

Hoisting an 11-line block that ran correctly inline into a
`write_clamped_kv_descs(...)` helper — called from inside
`if nvvm.elect_sync() and tidx < 32:` — turned 212 passing forward tests into
31 failures (`Error building ...`, traceback through `ir_loop` →
`scf_execute_dynamic`). Unrolling the helper's own loop did not help; the
helper *call* was the problem. Duplicating the block across flavors is the
correct trade here. Factor only host-side code, or code you can mark
`@cute.jit`.

Related: inside a kernel body, `for x in (a, b)` over a Python tuple is
rewritten into a dynamic `ir_loop` and cannot iterate heterogeneous objects
(e.g. `GridConstant[TensorMap]`). Unroll it, or use `cutlass.range_constexpr`.

**Detector.** These break at `compile()`, not at import — `python -c "import ..."`
and `pytest --collect-only` both stay green. After any refactor of a kernel
body, run that flavor's own tests.

## The APIBase contract (`api_base.py`)

Every OSS kernel API extends `APIBase` and implements:

- `check_support() -> bool` — validate dtype/shape/stride/arch/config via the `_check_tensor_*` / `_value_error_if` helpers; must set `self._is_supported`. Works on `TensorDesc` (metadata-only tensors), so it runs without GPU storage.
- `compile()` — calls `self._ensure_support_checked()`, builds and `cute.compile`s the kernel, caches in `self._compiled_kernel`.
- `execute(..., current_stream=None)` — runs the cached kernel.

`__call__` = compile-if-needed + execute. High-level wrappers (`<op>_wrapper_sm100(...)`) allocate outputs and return a **`TupleDict`** (dict that also unpacks as a tuple) with stable, documented key order. FP4x2 packing: use `_tensor_shape`/`_tensor_stride`, which double the innermost dim when `interpret_uint8_as_fp4x2` is set.

## Adding a new frontend-only API — required checklist

1. Kernel package under the closest existing family (layout above).
2. `APIBase` subclass + wrapper in `api.py`.
3. Exports: family `__init__.py` `__all__` **and** `_LAZY_OPTIONAL_IMPORTS` in `python/cudnn/__init__.py`; register any new package dir in `pyproject.toml` packages list.
4. Docs: page under `docs/fe-oss-apis/` (family subdir) + link it from `docs/fe-oss-apis/overview.md`.
5. Tests: `test/python/<op>/cutedsl/test_<op>.py` (+ `_utils.py`/reference), covering check_support pass/fail and numerical reference comparison.
6. DSL version gate (Rule 7): the route/`check_support` declines with a version-naming error below `CUTEDSL_MIN_VERSION`, and the tests skip there instead of failing.

The `cutedsl-kernel-integration` skill (`skills/cutedsl-kernel-integration/`) documents this workflow in detail, including how to classify a kernel into a family — follow it for any kernel integration.

## Other notes

- `wrapper.py` `Graph` context manager (the pythonic graph builder) requires cuDNN backend ≥ 9.12 (`backend_version() >= 91200`) and builds plans on `__exit__`.
- Torch custom-op implementations live with their owning operation family and may be re-exported from `experimental/ops/` while maturing (pattern doc: `docs/utilities/adding_torch_custom_ops.md`); they cache built graphs per config and use stable `_UIDs` enums.
- dtype conversions go through `datatypes.py`, which probes torch/cutlass availability lazily — keep it that way.
- Formatting: black, line length 160.

CUDA-owning objects, GC-timed release and stream capture: Rule 8.


## DSA training composition

Metadata for backward and targets has two different ordering contracts. Training
may compact valid active indices only if forward and backward share that exact
metadata. Score targets must preserve original slots, mask invalid/inactive
entries, and normalize over retained slots after summing heads. Never reuse
compacted training indices as the caller-visible target order. Detector:
`deepseek_sparse_attention/cutedsl/test_DSA_training.py::test_native_training_and_original_score_slots`
checks holes, duplicates, bounded lengths, all-masked rows and target alignment.