torch: add the "CUDNN" torch.nn.attention provider (cudnn.torch)

docs/fe-oss-apis/sdpa-torch-ops.md

@@ -0,0 +1,69 @@
+# SDPA torch custom ops: `cudnn::sdpa_fwd` / `cudnn::sdpa_bwd`
+
+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:
+
+- **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`
+
+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).
+
+## Usage
+
+```python
+import torch
+import cudnn
+
+_ = cudnn.sdpa_torch  # lazy public export: importing registers cudnn::sdpa_fwd / cudnn::sdpa_bwd
+
+# 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)
+
+# 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
+
+# Or through the python wrapper (same op underneath):
+o = 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)
+```
+
+## 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/test_cudnn_sdpa_torch_ops.py`.

python/cudnn/__init__.py

@@ -309,6 +309,8 @@ def _dlopen_cudnn():
 
 _LAZY_OPTIONAL_IMPORTS = {
     "gnn": (".gnn", None),
+    "sdpa_torch": (".sdpa.fwd.torch_op", "sdpa"),
+    "torch": (".torch", None),
     "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/sdpa.py

@@ -488,7 +488,10 @@ def _build_bprop_graph(
 )
 
 _lib.define(
-    "sdpa_bwd(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, "
+    # Renamed from cudnn::sdpa_bwd: the canonical name now belongs to the
+    # consolidated op family in cudnn.sdpa.fwd.torch_op (this experimental
+    # module is slated to fold into it).
+    "sdpa_bwd_legacy(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, "
     "float attn_scale, bool is_causal=False, int diagonal_alignment=0, "
     "int left_bound=-1, int right_bound=-1, "
     "Tensor? seq_len_q=None, Tensor? seq_len_kv=None, "
@@ -625,7 +628,7 @@ def _sdpa_fake(
     return O, Stats
 
 
-def _sdpa_bwd_impl(
+def _sdpa_bwd_legacy_impl(
     dO: torch.Tensor,
     q: torch.Tensor,
     k: torch.Tensor,
@@ -723,11 +726,11 @@ def _sdpa_bwd_impl(
     return dQ_gpu, dK_gpu, dV_gpu
 
 
-_lib.impl("sdpa_bwd", _sdpa_bwd_impl, "CUDA")
+_lib.impl("sdpa_bwd_legacy", _sdpa_bwd_legacy_impl, "CUDA")
 
 
-@torch.library.register_fake("cudnn::sdpa_bwd")
-def _sdpa_bwd_fake(
+@torch.library.register_fake("cudnn::sdpa_bwd_legacy")
+def _sdpa_bwd_legacy_fake(
     dO: torch.Tensor,
     q: torch.Tensor,
     k: torch.Tensor,
@@ -800,7 +803,7 @@ def _sdpa_backward(ctx, dO, dStats):
         idx += 1
     cum_kv = saved[idx] if ctx.has_cum_kv else None
 
-    dQ, dK, dV = torch.ops.cudnn.sdpa_bwd(
+    dQ, dK, dV = torch.ops.cudnn.sdpa_bwd_legacy(
         dO,
         q,
         k,