sdpa: replace the experimental torch op with cudnn::sdpa_fwd / cudnn::sdpa_bwd

README.md

@@ -43,7 +43,6 @@ We are now shipping **OSS kernels**, allowing you to inspect, modify, and contri
 *   **[SDPA Backward: SM100, D=256](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/sdpa):** SDPA Backward pass for D=256 on SM100.
 *   **[cudnn SDPA Fprop](https://github.com/NVIDIA/cudnn-frontend/tree/main/include/cudnn_frontend/generated/sdpa):** Open sourcing the Hopper and Blackwell fprop kernels with stats.
 *   **[Fused RMSNorm + SiLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/include/cudnn_frontend/generated/rms_norm_silu):** Implementation of a fused kernel of RMS normalization followed by SiLU (Swish) activation.
-*   **[SDPA PyTorch Op](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/experimental/ops):** PyTorch custom operator for cuDNN-accelerated Scaled Dot-Product Attention with autograd and `torch.compile` support.
 *   **[DSA](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/deepseek_sparse_attention):** DSA/CSA kernels for DSv4 and DSv3.2 for fprop and bprop.
 
 Contributor credits for these OSS CuTe DSL kernels are listed in [Acknowledgements](ACKNOWLEDGEMENTS.md).

benchmark/e2e/Qwen-Image/run_model.py

@@ -101,13 +101,14 @@ def is_right_padded(mask):
 
 def load_runtime():
     import cudnn
-    import cudnn.experimental.ops.sdpa as cudnn_sdpa_module
+
+    _ = cudnn.sdpa_torch  # registers cudnn::sdpa_fwd / cudnn::sdpa_bwd
     import diffusers
     import diffusers.models.transformers.transformer_qwenimage as qwen_module
     import torch
     import torch.nn.functional as F
 
-    return torch, F, cudnn, cudnn_sdpa_module, diffusers, qwen_module
+    return torch, F, cudnn, diffusers, qwen_module
 
 
 def install_joint_attention_dispatch(qwen_module, *, text_tokens, counters=None, torch_probe=None):
@@ -117,7 +118,7 @@ def install_joint_attention_dispatch(qwen_module, *, text_tokens, counters=None,
     projections around this function.  The treatment therefore changes only
     the joint SDPA core (plus the exact padding-layout adapter when needed).
     """
-    import cudnn.experimental.ops.sdpa as cudnn_sdpa_module
+    import cudnn
     import torch
     import torch.nn.functional as F
 
@@ -128,7 +129,7 @@ def install_joint_attention_dispatch(qwen_module, *, text_tokens, counters=None,
     if torch_probe is None:
         torch_probe = {}
     original = qwen_module.dispatch_attention_fn
-    cudnn_sdpa = cudnn_sdpa_module.scaled_dot_product_attention
+    _ = cudnn.sdpa_torch  # registers cudnn::sdpa_fwd / cudnn::sdpa_bwd
 
     def _validate(q, k, v, attn_mask, dropout_p, is_causal, parallel_config):
         if q.ndim != 4 or k.shape != q.shape or v.shape != q.shape:
@@ -238,16 +239,22 @@ def cudnn_dispatch(
                 seq_len_kv = (image_tokens + text_valid.sum(dim=-1, dtype=torch.int32)).reshape(batch, 1, 1, 1)
                 reordered = True
         qt, kt, vt = (tensor.transpose(1, 2) for tensor in (q, k, v))
-        out = cudnn_sdpa(
+        if dropout_p:
+            raise NotImplementedError("cudnn::sdpa_fwd does not serve dropout")
+        # torch's SDPA APIs treat scale=None as "use the default"; the op's
+        # schema takes a required float, so resolve it here.
+        attn_scale = scale if scale is not None else qt.shape[-1] ** -0.5
+        out, _ = torch.ops.cudnn.sdpa_fwd(
             qt,
             kt,
             vt,
-            dropout_p=dropout_p,
+            attn_scale,
             is_causal=is_causal,
-            scale=scale,
             seq_len_q=seq_len_q,
             seq_len_kv=seq_len_kv,
-        ).transpose(1, 2)
+            return_lse=False,
+        )
+        out = out.transpose(1, 2)
         if reordered:
             out = torch.cat([out[:, image_tokens:], out[:, :image_tokens]], dim=1)
         return out
@@ -327,7 +334,7 @@ def _main():
     if args.inspect:
         print(json.dumps({"shape": shape, "model": OFFICIAL_MODEL, "diffusers": DIFFUSERS_ANCHOR, "recipe": NUMERICAL_RECIPE}, indent=2))
         return
-    torch, _, cudnn, _, _, qwen_module = load_runtime()
+    torch, _, cudnn, _, qwen_module = load_runtime()
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
     device = torch.device("cuda")

benchmark/e2e/Qwen3.8/run_matrix.py

@@ -385,7 +385,7 @@ def _pick_device(mode):
 def _run_experiment(args, qwen, device, properties, orders, started_utc):
     import cudnn
     from cudnn import _env as cudnn_env
-    import cudnn.experimental.ops.sdpa as sdpamod
+    import cudnn.sdpa.fwd.torch_op as sdpamod
     import cudnn.fla as cfla
     import fla
     import fla.layers.attn as fla_attn
@@ -396,8 +396,6 @@ def _run_experiment(args, qwen, device, properties, orders, started_utc):
     backend_floor = 92300
     if cudnn.backend_version() < backend_floor:
         raise RuntimeError("d256 FE arm requires cuDNN backend " f">= {backend_floor}; got {cudnn.backend_version()}")
-    if any(hasattr(sdpamod, name) for name in ("sdpa_fwd_d256", "sdpa_bwd_d256")):
-        raise RuntimeError("loaded FE SDPA module predates #682 and still exposes the legacy standalone d256 stacks")
 
     from cudnn.gemm.ops import swiglu_mlp as public_swiglu_mlp
 

benchmark/e2e/Qwen3.8/run_model.py

@@ -51,7 +51,8 @@ def _wire_sdpa_attention():
     only the SDPA core changes across this axis.
     """
     import cudnn
-    import cudnn.experimental.ops.sdpa as cudnn_sdpa_module
+
+    _ = cudnn.sdpa_torch  # registers cudnn::sdpa_fwd / cudnn::sdpa_bwd
     import fla.layers.attn as fla_attn
     import torch.nn.functional as F
 
@@ -71,7 +72,6 @@ def _wire_sdpa_attention():
     backend_version = cudnn.backend_version()
     if backend_version < backend_floor:
         raise RuntimeError(f"d256 SDPA requires cuDNN backend >= {backend_floor}; got {backend_version}")
-    cudnn_sdpa = cudnn_sdpa_module.scaled_dot_product_attention
 
     def _prepare(q, k, v, window_size):
         qt, kt, vt = (x.transpose(1, 2) for x in (q, k, v))  # [B,L,H,D] -> [B,H,L,D]
@@ -116,16 +116,21 @@ def _cudnn_sdpa_flash(
         **kw,
     ):
         qt, kt, vt = _prepare(q, k, v, window_size)
-        o = cudnn_sdpa(
+        if dropout_p:
+            raise NotImplementedError("cudnn::sdpa_fwd does not serve dropout")
+        if window_size[1] not in (-1, 0):
+            raise NotImplementedError(f"cudnn::sdpa_fwd has no right window bound; got {window_size[1]}")
+        # torch's SDPA APIs treat scale=None as "use the default"; the op's
+        # schema takes a required float, so resolve it here.
+        attn_scale = softmax_scale if softmax_scale is not None else qt.shape[-1] ** -0.5
+        o, _ = torch.ops.cudnn.sdpa_fwd(
             qt,
             kt,
             vt,
+            attn_scale,
             is_causal=causal,
-            scale=softmax_scale,
-            dropout_p=dropout_p,
-            enable_gqa=qt.shape[1] != kt.shape[1],
-            left_bound=window_size[0],
-            right_bound=window_size[1],
+            window_left=window_size[0],
+            return_lse=False,
         )
         # flash-attn's adapter contract is packed [B,L,H,D].  The direct cuDNN
         # wrapper returns packed BHSD, so normalize once before FLA flattens H*D.

docs/adding_torch_custom_ops.md

@@ -5,7 +5,7 @@ Best practices for wrapping cuDNN graph ops as PyTorch custom ops with minimal C
 ## File location
 
 Custom ops live in `python/cudnn/experimental/ops/`. Each op gets its own file
-(e.g., `sdpa.py`, `rmsnorm.py`, `layernorm.py`, `moe.py`). Export from
+(e.g., `rmsnorm.py`, `layernorm.py`, `moe.py`). Export from
 `python/cudnn/experimental/ops/__init__.py`.
 
 ## Registration: use torch.Library, NOT @torch.library.custom_op

docs/operations/Attention.md

@@ -720,94 +720,77 @@ forward and backward automatically build private block metadata on the active
 CUDA stream without adding public API parameters; D256 backward builds both
 Q-to-K and K-to-Q views from one coarse classification.
 
-(scaled-dot-product-attention-pytorch-op)=
-### SDPA PyTorch Custom Op (Experimental)
+(scaled-dot-product-attention-torch-ops)=
+### SDPA PyTorch Custom Ops (`cudnn::sdpa_fwd` / `cudnn::sdpa_bwd`)
 
-A high-level PyTorch custom operator that wraps the cuDNN SDPA forward and backward graphs into a single, autograd-compatible function. This provides a drop-in replacement for `torch.nn.functional.scaled_dot_product_attention` that routes computation through cuDNN.
+PyTorch custom ops (`torch.library`) exposing the full cuDNN SDPA feature
+surface — the features `torch.nn.functional.scaled_dot_product_attention`'s
+aten contract cannot express:
 
-**Key features:**
-- Full autograd support (forward + backward)
-- `torch.compile` compatible via FakeTensor/meta registration
-- Graph caching for efficient repeated execution
-- Supports FP16, BF16 datatypes
-- Supports causal masking, sliding window, padding mask, GQA/MQA, and ragged tensors
+- **attention sinks** — per-Q-head logits folded into the softmax denominator
+- **sliding window** — `window_left` (cuDNN convention: visible tokens
+  *including* self; FA2's `(w, 0)` maps to `window_left = w + 1`)
+- **bottom-right causal alignment** — inference-style diagonals
+- **padded batches** — per-batch actual lengths via `seq_len_q` / `seq_len_kv`
+- **THD / varlen packing** — FlashAttention-style `(T, H, D)` + `cu_seqlens`
 
-**Limitations:**
-- `attn_mask` and `dropout` are not yet supported
-- FP8 is not supported (use the Graph API directly)
-- For head dimension `256`, the specialized backward path currently supports only plain BHSD inputs. `seq_len_q`, `seq_len_kv`, `cumulative_seq_len_q`, and `cumulative_seq_len_kv` are not supported on that backward path.
+The ops build cuDNN pygraph `sdpa` / `sdpa_backward` nodes; the engine Router
+picks the best serving plan (FROST OSS kernels or cuDNN-backend engines) per
+configuration. Graphs are cached per configuration (bounded, thread-safe;
+cuDNN handles are thread-local).
 
-#### Python API
-
-```python
-from cudnn.experimental.ops import scaled_dot_product_attention
-
-output = scaled_dot_product_attention(
-    query,                        # (B, H_q, S_q, D) — FP16 or BF16
-    key,                          # (B, H_k, S_kv, D)
-    value,                        # (B, H_v, S_kv, D_v)
-    attn_mask=None,               # Not yet supported, must be None
-    dropout_p=0.0,                # Not yet supported, must be 0.0
-    is_causal=False,              # Apply causal (upper-triangular) mask
-    scale=None,                   # Attention scale, defaults to 1/sqrt(D)
-    enable_gqa=False,             # Enable grouped-query attention (H_q > H_k)
-    *,
-    diagonal_alignment=0,         # 0 = TOP_LEFT, 1 = BOTTOM_RIGHT
-    left_bound=-1,                # Sliding window left bound (-1 = disabled)
-    right_bound=-1,               # Sliding window right bound (-1 = disabled)
-    seq_len_q=None,               # Actual query seq lengths (B, 1, 1, 1) INT32
-    seq_len_kv=None,              # Actual key/value seq lengths (B, 1, 1, 1) INT32
-    cumulative_seq_len_q=None,    # Ragged offset for Q (B+1, 1, 1, 1) INT32
-    cumulative_seq_len_kv=None,   # Ragged offset for KV (B+1, 1, 1, 1) INT32
-)
-```
-
-**Args:**
-- `query` (torch.Tensor): Query tensor in BHSD layout `(B, H_q, S_q, D)`.
-- `key` (torch.Tensor): Key tensor in BHSD layout `(B, H_k, S_kv, D)`.
-- `value` (torch.Tensor): Value tensor in BHSD layout `(B, H_v, S_kv, D_v)`.
-- `attn_mask` (Optional[torch.Tensor]): Not yet supported. Must be `None`.
-- `dropout_p` (float): Not yet supported. Must be `0.0`.
-- `is_causal` (bool): If `True`, applies a causal mask (sets `right_bound=0`).
-- `scale` (Optional[float]): Attention scale factor. Defaults to `1/sqrt(D)`.
-- `enable_gqa` (bool): When `False`, raises `ValueError` if `H_q != H_k`. Set to `True` for grouped-query or multi-query attention.
-- `diagonal_alignment` (int): `0` for TOP_LEFT, `1` for BOTTOM_RIGHT alignment.
-- `left_bound` (int): Left sliding-window bound. `-1` disables.
-- `right_bound` (int): Right sliding-window bound. `-1` disables. `0` for causal.
-- `seq_len_q` (Optional[torch.Tensor]): Per-batch query sequence lengths `(B, 1, 1, 1)` INT32.
-- `seq_len_kv` (Optional[torch.Tensor]): Per-batch key/value sequence lengths `(B, 1, 1, 1)` INT32.
-- `cumulative_seq_len_q` (Optional[torch.Tensor]): Ragged offset for Q `(B+1, 1, 1, 1)` INT32.
-- `cumulative_seq_len_kv` (Optional[torch.Tensor]): Ragged offset for KV `(B+1, 1, 1, 1)` INT32.
-
-For head dimension `256`, backward support is narrower than the general SDPA op contract: the specialized `d=256` backward path requires plain BHSD tensors and does not support `seq_len_q`, `seq_len_kv`, `cumulative_seq_len_q`, or `cumulative_seq_len_kv`.
-
-**Returns:**
-- `output` (torch.Tensor): Attention output `(B, H_q, S_q, D_v)`.
-
-#### Example Usage
+#### Usage
 
 ```python
 import torch
-from cudnn.experimental.ops import scaled_dot_product_attention
+import cudnn
 
-B, H, S, D = 2, 8, 1024, 128
+_ = cudnn.sdpa_torch  # lazy public export: importing registers cudnn::sdpa_fwd / cudnn::sdpa_bwd
 
-q = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)
-k = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)
-v = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)
+# Dense BHSD with sinks + sliding window
+o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True,
+                                  window_left=128, sinks=sinks, return_lse=True)
 
-# Forward
-output = scaled_dot_product_attention(q, k, v, is_causal=True)
+# THD / varlen (FA-style packed (T, H, D) + cu_seqlens), differentiable:
+q, k, v = (t.requires_grad_(True) for t in (q_thd, k_thd, v_thd))
+o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True,
+                                  cu_seqlens_q=cu, cu_seqlens_kv=cu,
+                                  max_seqlen_q=mx, max_seqlen_kv=mx,
+                                  return_lse=True)
+o.backward(grad)  # routes through cudnn::sdpa_bwd via register_autograd
 
-# Backward (autograd handles this automatically)
-loss = output.sum()
-loss.backward()
-# q.grad, k.grad, v.grad are now populated
+# Or through the python wrapper (same op underneath). It defaults to
+# return_lse=False; autograd needs the stats, so ask for them explicitly:
+o, lse = cudnn.sdpa_torch(q, k, v, is_causal=True, cu_seqlens_q=cu, cu_seqlens_kv=cu,
+                          max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=True)
 ```
 
-#### Tests
-
-- Python tests: [test/python/test_cudnn_sdpa_op.py](https://github.com/NVIDIA/cudnn-frontend/blob/main/test/python/test_cudnn_sdpa_op.py)
+#### Contracts and limits
+
+- Dense tensors are BHSD `(B, H, S, D)` (any strides; the graph declares the
+  actual layout). Varlen tensors are packed `(T, H, D)`; non-contiguous views
+  (e.g. K/V slices of a fused `(T, 2, H, D)` KV projection) are declared with
+  their true strides. On the varlen path, a non-dense innermost dim or a
+  misaligned base pointer is repaired by one copy (warned as slow path); the
+  dense path declares the given strides as-is.
+- One io dtype per call (`fp16` or `bf16`); mixed-dtype inputs are rejected.
+- `sdpa_bwd` serves the **THD/varlen** path. Dense backward and sink backward
+  (dSink) are follow-ups and raise `NotImplementedError`. It consumes a
+  **padded** `(B, H, max_seqlen_q, 1)` fp32 LSE (backend restriction: bprop
+  THD rejects ragged LSE on SM8X/SM12X).
+- Autograd (`register_autograd`) requires `return_lse=True` on the forward;
+  the glue converts the packed TH1 stats to the padded layout device-side.
+- Both ops ship `register_fake` meta kernels. `cudnn::sdpa_fwd` passes
+  `torch.library.opcheck` on the dense and varlen paths, including
+  dynamic-shape AOT dispatch (`torch.compile`-ready); the opcheck autograd
+  case exercises `cudnn::sdpa_bwd` through the registered backward.
+
+#### Requirements
+
+- `nvidia-cudnn-frontend[cutedsl]`, cuDNN backend ≥ 9.6 (THD token-major
+  stats), sm80+.
+
+Tests: [test/python/sdpa/test_torch_ops.py](https://github.com/NVIDIA/cudnn-frontend/blob/main/test/python/sdpa/test_torch_ops.py).
 
 (scaled-dot-product-attention-fp8-forward)=
 ### SDPA FP8 Forward

python/cudnn/__init__.py

@@ -309,6 +309,7 @@ def _dlopen_cudnn():
 
 _LAZY_OPTIONAL_IMPORTS = {
     "gnn": (".gnn", None),
+    "sdpa_torch": (".sdpa.fwd.torch_op", "sdpa"),
     "BSA": (".block_sparse_attention", "BSA"),
     "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"),
     "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"),

python/cudnn/experimental/ops/__init__.py

@@ -5,8 +5,6 @@
 import sys
 from typing import Any
 
-from .sdpa import scaled_dot_product_attention
-
 # moe_grouped_matmul / swiglu_mlp live with the rest of the GEMM family in
 # cudnn.gemm.ops (their modules import torch). Expose them here lazily so that
 # importing this package does not eagerly pull in those kernel modules; the
@@ -31,7 +29,6 @@ def __getattr__(name: str) -> Any:
 
 
 __all__ = [
-    "scaled_dot_product_attention",
     "moe_grouped_matmul",
     "swiglu_mlp",
 ]