diff --git a/README.md b/README.md index 2b715f540..774d597f3 100644 --- a/README.md +++ b/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). diff --git a/benchmark/e2e/Qwen-Image/run_model.py b/benchmark/e2e/Qwen-Image/run_model.py index 751395db5..5fdf9b6bd 100644 --- a/benchmark/e2e/Qwen-Image/run_model.py +++ b/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") diff --git a/benchmark/e2e/Qwen3.8/run_matrix.py b/benchmark/e2e/Qwen3.8/run_matrix.py index 63f8b2e8d..a22b4e55b 100644 --- a/benchmark/e2e/Qwen3.8/run_matrix.py +++ b/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 diff --git a/benchmark/e2e/Qwen3.8/run_model.py b/benchmark/e2e/Qwen3.8/run_model.py index f422e7d4c..0649492b1 100644 --- a/benchmark/e2e/Qwen3.8/run_model.py +++ b/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. diff --git a/docs/adding_torch_custom_ops.md b/docs/adding_torch_custom_ops.md index 207ea059f..2f75259ea 100644 --- a/docs/adding_torch_custom_ops.md +++ b/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 diff --git a/docs/operations/Attention.md b/docs/operations/Attention.md index bcabe6133..093cb874f 100644 --- a/docs/operations/Attention.md +++ b/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 diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 22ffbf7a8..68ca74e51 100644 --- a/python/cudnn/__init__.py +++ b/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"), diff --git a/python/cudnn/experimental/ops/__init__.py b/python/cudnn/experimental/ops/__init__.py index a39f8ba92..803d48569 100644 --- a/python/cudnn/experimental/ops/__init__.py +++ b/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", ] diff --git a/python/cudnn/experimental/ops/sdpa.py b/python/cudnn/experimental/ops/sdpa.py deleted file mode 100644 index faea0e475..000000000 --- a/python/cudnn/experimental/ops/sdpa.py +++ /dev/null @@ -1,911 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -PyTorch custom operator wrapping cuDNN's Scaled Dot-Product Attention (SDPA). - -Provides ``scaled_dot_product_attention`` as the public entry point, closely -matching the signature of ``torch.nn.functional.scaled_dot_product_attention``. - -**Layout convention**: tensors are expected in **BHSD** layout -``(batch, num_heads, seq_len, head_dim)`` — matching both the cuDNN convention -and PyTorch's ``torch.nn.functional.scaled_dot_product_attention`` layout. - -Graph caching ensures cuDNN graphs are built once per unique configuration -and reused across calls. -""" - -import logging -import math -from typing import Optional, Tuple, Dict -from enum import IntEnum - -import torch -import cudnn - -_logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Module-level state -# --------------------------------------------------------------------------- - -_cudnn_handles = {} -_fprop_cache: Dict[tuple, tuple] = {} -_bprop_cache: Dict[tuple, tuple] = {} - -# Dtype mapping (module-level constant) -_TORCH_DTYPE_TO_CUDNN = { - torch.float16: cudnn.data_type.HALF, - torch.bfloat16: cudnn.data_type.BFLOAT16, - torch.float32: cudnn.data_type.FLOAT, - torch.int32: cudnn.data_type.INT32, - torch.int64: cudnn.data_type.INT64, -} - - -# --------------------------------------------------------------------------- -# UID enum — explicit tensor UIDs for graph caching -# --------------------------------------------------------------------------- - - -class _UIDs(IntEnum): - Q = 1 - K = 2 - V = 3 - O = 100 - STATS = 101 - DO = 200 - DQ = 201 - DK = 202 - DV = 203 - SEQ_LEN_Q = 300 - SEQ_LEN_KV = 301 - CUM_SEQ_LEN_Q = 302 - CUM_SEQ_LEN_KV = 303 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _get_handle(device: torch.device): - """Return a lazily-initialised cuDNN handle with the current CUDA stream.""" - if device not in _cudnn_handles: - _cudnn_handles[device] = cudnn.create_handle() - stream = _get_current_stream(device) - cudnn.set_stream(handle=_cudnn_handles[device], stream=stream) - return _cudnn_handles[device] - - -def _get_current_stream(device: torch.device): - """Return the caller's active CUDA stream for the given device.""" - return torch.cuda.current_stream(device).cuda_stream - - -def _torch_dtype_to_cudnn(dtype: torch.dtype): - """Map a PyTorch dtype to a cuDNN data_type enum.""" - return _TORCH_DTYPE_TO_CUDNN[dtype] - - -def _diagonal_alignment_enum(val: int): - """Convert int sentinel to cudnn.diagonal_alignment enum.""" - if val == 0: - return cudnn.diagonal_alignment.TOP_LEFT - return cudnn.diagonal_alignment.BOTTOM_RIGHT - - -def _make_fprop_cache_key( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_scale: float, - is_causal: bool, - diagonal_alignment: int, - left_bound: int, - right_bound: int, - has_seq_len_q: bool, - has_seq_len_kv: bool, - has_cum_q: bool, - has_cum_kv: bool, -): - q_shape, q_stride = tuple(q.shape), tuple(q.stride()) - k_shape, k_stride = tuple(k.shape), tuple(k.stride()) - v_shape, v_stride = tuple(v.shape), tuple(v.stride()) - return ( - "fprop", - q_shape, - q_stride, - q.dtype, - k_shape, - k_stride, - k.dtype, - v_shape, - v_stride, - v.dtype, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - has_seq_len_q, - has_seq_len_kv, - has_cum_q, - has_cum_kv, - q.device, - ) - - -def _make_bprop_cache_key( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_scale: float, - is_causal: bool, - diagonal_alignment: int, - left_bound: int, - right_bound: int, - has_seq_len_q: bool, - has_seq_len_kv: bool, - has_cum_q: bool, - has_cum_kv: bool, - is_deterministic: bool, -): - q_shape, q_stride = tuple(q.shape), tuple(q.stride()) - k_shape, k_stride = tuple(k.shape), tuple(k.stride()) - v_shape, v_stride = tuple(v.shape), tuple(v.stride()) - return ( - "bprop", - q_shape, - q_stride, - q.dtype, - k_shape, - k_stride, - k.dtype, - v_shape, - v_stride, - v.dtype, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - has_seq_len_q, - has_seq_len_kv, - has_cum_q, - has_cum_kv, - is_deterministic, - q.device, - ) - - -# --------------------------------------------------------------------------- -# Forward graph builder -# --------------------------------------------------------------------------- - - -def _build_fprop_graph( - handle, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_scale: float, - is_causal: bool, - diagonal_alignment: int, - left_bound: int, - right_bound: int, - seq_len_q: Optional[torch.Tensor], - seq_len_kv: Optional[torch.Tensor], - cumulative_seq_len_q: Optional[torch.Tensor], - cumulative_seq_len_kv: Optional[torch.Tensor], -): - """Build, validate, and compile a forward SDPA cuDNN graph.""" - - io_dtype = _torch_dtype_to_cudnn(q.dtype) - - _logger.debug(f"Building forward graph for q: {q.shape}, k: {k.shape}, v: {v.shape}") - - q_shape, q_stride = tuple(q.shape), tuple(q.stride()) - k_shape, k_stride = tuple(k.shape), tuple(k.stride()) - v_shape, v_stride = tuple(v.shape), tuple(v.stride()) - - B, H_q, S_q, D_qk = q_shape - _, H_v, S_kv, D_v = v_shape - - graph = cudnn.pygraph( - handle=handle, - io_data_type=io_dtype, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - - # -- Input tensors -- - q_t = graph.tensor(name="q", dim=list(q_shape), stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q) - k_t = graph.tensor(name="k", dim=list(k_shape), stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K) - v_t = graph.tensor(name="v", dim=list(v_shape), stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V) - - # -- Optional tensors -- - seq_len_q_t = None - seq_len_kv_t = None - cum_q_t = None - cum_kv_t = None - - if seq_len_q is not None: - seq_len_q_t = graph.tensor( - name="seq_len_q", dim=list(seq_len_q.shape), stride=list(seq_len_q.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q - ) - if seq_len_kv is not None: - seq_len_kv_t = graph.tensor( - name="seq_len_kv", dim=list(seq_len_kv.shape), stride=list(seq_len_kv.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV - ) - if cumulative_seq_len_q is not None: - cum_q_t = graph.tensor( - name="cum_seq_len_q", - dim=list(cumulative_seq_len_q.shape), - stride=list(cumulative_seq_len_q.stride()), - data_type=cudnn.data_type.INT32, - uid=_UIDs.CUM_SEQ_LEN_Q, - ) - if cumulative_seq_len_kv is not None: - cum_kv_t = graph.tensor( - name="cum_seq_len_kv", - dim=list(cumulative_seq_len_kv.shape), - stride=list(cumulative_seq_len_kv.stride()), - data_type=cudnn.data_type.INT32, - uid=_UIDs.CUM_SEQ_LEN_KV, - ) - - # -- Ragged offsets -- - if cum_q_t is not None: - q_t.set_ragged_offset(cum_q_t) - if cum_kv_t is not None: - k_t.set_ragged_offset(cum_kv_t) - v_t.set_ragged_offset(cum_kv_t) - - # -- Mask configuration -- - use_padding = seq_len_q is not None or seq_len_kv is not None - lb = left_bound if left_bound >= 0 else None - rb = right_bound if right_bound >= 0 else None - if is_causal and rb is None: - rb = 0 - - # -- SDPA forward -- - o_t, stats_t = graph.sdpa( - name="sdpa", - q=q_t, - k=k_t, - v=v_t, - generate_stats=True, - attn_scale=attn_scale, - use_padding_mask=use_padding, - seq_len_q=seq_len_q_t, - seq_len_kv=seq_len_kv_t, - diagonal_alignment=_diagonal_alignment_enum(diagonal_alignment), - diagonal_band_left_bound=lb, - diagonal_band_right_bound=rb, - compute_data_type=cudnn.data_type.FLOAT, - ) - - # -- Output shapes (BHSD contiguous) -- - o_shape = (B, H_q, S_q, D_v) - o_stride = (H_q * S_q * D_v, S_q * D_v, D_v, 1) - o_t.set_uid(_UIDs.O).set_output(True).set_dim(list(o_shape)).set_stride(list(o_stride)) - o_t.set_data_type(io_dtype) - - if cum_q_t is not None: - o_t.set_ragged_offset(cum_q_t) - - stats_t.set_uid(_UIDs.STATS).set_output(True).set_data_type(cudnn.data_type.FLOAT) - if cum_q_t is not None: - stats_t.set_ragged_offset(cum_q_t) - - # -- Build -- - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - graph.build_plans() - - workspace_size = graph.get_workspace_size() - - return graph, workspace_size - - -# --------------------------------------------------------------------------- -# Backward graph builder -# --------------------------------------------------------------------------- - - -def _build_bprop_graph( - handle, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - dO: torch.Tensor, - stats: torch.Tensor, - attn_scale: float, - is_causal: bool, - diagonal_alignment: int, - left_bound: int, - right_bound: int, - seq_len_q: Optional[torch.Tensor], - seq_len_kv: Optional[torch.Tensor], - cumulative_seq_len_q: Optional[torch.Tensor], - cumulative_seq_len_kv: Optional[torch.Tensor], - is_deterministic: bool, -): - """Build, validate, and compile a backward SDPA cuDNN graph.""" - - io_dtype = _torch_dtype_to_cudnn(q.dtype) - - _logger.debug(f"Building backward graph for q: {q.shape}, k: {k.shape}, v: {v.shape}, o: {o.shape}, dO: {dO.shape}") - - q_shape, q_stride = tuple(q.shape), tuple(q.stride()) - k_shape, k_stride = tuple(k.shape), tuple(k.stride()) - v_shape, v_stride = tuple(v.shape), tuple(v.stride()) - o_shape, o_stride = tuple(o.shape), tuple(o.stride()) - dO_shape, dO_stride = tuple(dO.shape), tuple(dO.stride()) - - B, H_q, S_q, D_qk = q_shape - _, H_k, S_kv, _ = k_shape - _, H_v, _, D_v = v_shape - - graph = cudnn.pygraph( - handle=handle, - io_data_type=io_dtype, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - - # -- Input tensors -- - q_t = graph.tensor(name="q", dim=list(q_shape), stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q) - k_t = graph.tensor(name="k", dim=list(k_shape), stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K) - v_t = graph.tensor(name="v", dim=list(v_shape), stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V) - o_t = graph.tensor(name="o", dim=list(o_shape), stride=list(o_stride), data_type=io_dtype, uid=_UIDs.O) - dO_t = graph.tensor(name="dO", dim=list(dO_shape), stride=list(dO_stride), data_type=io_dtype, uid=_UIDs.DO) - stats_t = graph.tensor(name="stats", dim=list(stats.shape), stride=list(stats.stride()), data_type=cudnn.data_type.FLOAT, uid=_UIDs.STATS) - - # -- Optional tensors -- - seq_len_q_t = None - seq_len_kv_t = None - cum_q_t = None - cum_kv_t = None - - if seq_len_q is not None: - seq_len_q_t = graph.tensor( - name="seq_len_q", dim=list(seq_len_q.shape), stride=list(seq_len_q.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q - ) - if seq_len_kv is not None: - seq_len_kv_t = graph.tensor( - name="seq_len_kv", dim=list(seq_len_kv.shape), stride=list(seq_len_kv.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV - ) - if cumulative_seq_len_q is not None: - cum_q_t = graph.tensor( - name="cum_seq_len_q", - dim=list(cumulative_seq_len_q.shape), - stride=list(cumulative_seq_len_q.stride()), - data_type=cudnn.data_type.INT32, - uid=_UIDs.CUM_SEQ_LEN_Q, - ) - if cumulative_seq_len_kv is not None: - cum_kv_t = graph.tensor( - name="cum_seq_len_kv", - dim=list(cumulative_seq_len_kv.shape), - stride=list(cumulative_seq_len_kv.stride()), - data_type=cudnn.data_type.INT32, - uid=_UIDs.CUM_SEQ_LEN_KV, - ) - - # -- Ragged offsets -- - if cum_q_t is not None: - q_t.set_ragged_offset(cum_q_t) - o_t.set_ragged_offset(cum_q_t) - dO_t.set_ragged_offset(cum_q_t) - if cum_kv_t is not None: - k_t.set_ragged_offset(cum_kv_t) - v_t.set_ragged_offset(cum_kv_t) - - # -- Mask configuration -- - use_padding = seq_len_q is not None or seq_len_kv is not None - lb = left_bound if left_bound >= 0 else None - rb = right_bound if right_bound >= 0 else None - if is_causal and rb is None: - rb = 0 - - # Compute max_total_seq_len for ragged backward - max_total_seq_len_q = None - max_total_seq_len_kv = None - if cumulative_seq_len_q is not None and seq_len_q is not None: - total = torch.sum(seq_len_q).item() - max_total_seq_len_q = ((total + 63) // 64) * 64 - if cumulative_seq_len_kv is not None and seq_len_kv is not None: - total = torch.sum(seq_len_kv).item() - max_total_seq_len_kv = ((total + 63) // 64) * 64 - - # -- SDPA backward -- - dQ_t, dK_t, dV_t = graph.sdpa_backward( - name="sdpa_backward", - q=q_t, - k=k_t, - v=v_t, - o=o_t, - dO=dO_t, - stats=stats_t, - attn_scale=attn_scale, - use_padding_mask=use_padding, - seq_len_q=seq_len_q_t, - seq_len_kv=seq_len_kv_t, - max_total_seq_len_q=max_total_seq_len_q, - max_total_seq_len_kv=max_total_seq_len_kv, - diagonal_alignment=_diagonal_alignment_enum(diagonal_alignment), - diagonal_band_left_bound=lb, - diagonal_band_right_bound=rb, - use_deterministic_algorithm=is_deterministic, - ) - - # -- Output shapes (BHSD order) -- - dq_shape, dq_stride = q_shape, q_stride - dk_shape, dk_stride = k_shape, k_stride - dv_shape, dv_stride = v_shape, v_stride - - dQ_t.set_uid(_UIDs.DQ).set_output(True).set_dim(list(dq_shape)).set_stride(list(dq_stride)) - dK_t.set_uid(_UIDs.DK).set_output(True).set_dim(list(dk_shape)).set_stride(list(dk_stride)) - dV_t.set_uid(_UIDs.DV).set_output(True).set_dim(list(dv_shape)).set_stride(list(dv_stride)) - - if cum_q_t is not None: - dQ_t.set_ragged_offset(cum_q_t) - if cum_kv_t is not None: - dK_t.set_ragged_offset(cum_kv_t) - dV_t.set_ragged_offset(cum_kv_t) - - # -- Build -- - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - graph.build_plans() - - workspace_size = graph.get_workspace_size() - - return graph, workspace_size - - -# --------------------------------------------------------------------------- -# Forward custom op -# --------------------------------------------------------------------------- - - -_lib = torch.library.Library("cudnn", "DEF") - -_lib.define( - "sdpa(Tensor q, Tensor k, Tensor v, 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, " - "Tensor? cumulative_seq_len_q=None, Tensor? cumulative_seq_len_kv=None" - ") -> (Tensor, Tensor)" -) - -_lib.define( - "sdpa_bwd(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, " - "Tensor? cumulative_seq_len_q=None, Tensor? cumulative_seq_len_kv=None, " - "bool is_deterministic=False" - ") -> (Tensor, Tensor, Tensor)" -) - - -def _sdpa_impl( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_scale: float, - is_causal: bool = False, - diagonal_alignment: int = 0, - left_bound: int = -1, - right_bound: int = -1, - seq_len_q: Optional[torch.Tensor] = None, - seq_len_kv: Optional[torch.Tensor] = None, - cumulative_seq_len_q: Optional[torch.Tensor] = None, - cumulative_seq_len_kv: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - cuDNN SDPA forward (internal). BHSD layout. - - Args: - q: Query tensor (B, H_q, S_q, D_qk) - k: Key tensor (B, H_k, S_kv, D_qk) - v: Value tensor (B, H_v, S_kv, D_v) - attn_scale: Attention scale factor (typically 1/sqrt(D_qk)) - is_causal: If True, apply causal mask (right_bound defaults to 0) - diagonal_alignment: 0 = TOP_LEFT, 1 = BOTTOM_RIGHT - left_bound: Left sliding window bound (>= 0 to enable, -1 = disabled) - right_bound: Right sliding window bound (>= 0 to enable, -1 = disabled) - seq_len_q: Actual query sequence lengths (B, 1, 1, 1) INT32 - seq_len_kv: Actual key/value sequence lengths (B, 1, 1, 1) INT32 - cumulative_seq_len_q: Ragged offset for Q (B+1, 1, 1, 1) INT32 - cumulative_seq_len_kv: Ragged offset for KV (B+1, 1, 1, 1) INT32 - - Returns: - (O, Stats): Output tensor (B, H_q, S_q, D_v) and softmax stats (B, H_q, S_q, 1) - """ - - handle = _get_handle(q.device) - - cache_key = _make_fprop_cache_key( - q, - k, - v, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - seq_len_q is not None, - seq_len_kv is not None, - cumulative_seq_len_q is not None, - cumulative_seq_len_kv is not None, - ) - - if cache_key not in _fprop_cache: - graph, workspace_size = _build_fprop_graph( - handle, - q, - k, - v, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - seq_len_q, - seq_len_kv, - cumulative_seq_len_q, - cumulative_seq_len_kv, - ) - _fprop_cache[cache_key] = (graph, workspace_size) - - graph, workspace_size = _fprop_cache[cache_key] - - # Allocate outputs and workspace (BHSD layout) - # Workspace is per-call — PyTorch's caching allocator recycles the allocation. - B, H_q, S_q, D_qk = q.shape - _, H_v, S_kv, D_v = v.shape - o_gpu = torch.empty(B, H_q, S_q, D_v, dtype=q.dtype, device=q.device) - stats_gpu = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device) - workspace = torch.empty(max(workspace_size, 1), device=q.device, dtype=torch.uint8) - - # UID → tensor map for sorted pointer extraction - uid_to_tensor = { - int(_UIDs.Q): q, - int(_UIDs.K): k, - int(_UIDs.V): v, - int(_UIDs.O): o_gpu, - int(_UIDs.STATS): stats_gpu, - } - if seq_len_q is not None: - uid_to_tensor[int(_UIDs.SEQ_LEN_Q)] = seq_len_q - if seq_len_kv is not None: - uid_to_tensor[int(_UIDs.SEQ_LEN_KV)] = seq_len_kv - if cumulative_seq_len_q is not None: - uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_Q)] = cumulative_seq_len_q - if cumulative_seq_len_kv is not None: - uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv - - graph.execute(uid_to_tensor, workspace, handle=handle) - - return o_gpu, stats_gpu - - -_lib.impl("sdpa", _sdpa_impl, "CUDA") - - -@torch.library.register_fake("cudnn::sdpa") -def _sdpa_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_scale: float, - is_causal: bool = False, - diagonal_alignment: int = 0, - left_bound: int = -1, - right_bound: int = -1, - seq_len_q: Optional[torch.Tensor] = None, - seq_len_kv: Optional[torch.Tensor] = None, - cumulative_seq_len_q: Optional[torch.Tensor] = None, - cumulative_seq_len_kv: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - B, H_q, S_q, D_qk = q.shape - _, H_v, S_kv, D_v = v.shape - O = torch.empty(B, H_q, S_q, D_v, dtype=q.dtype, device=q.device) - Stats = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device) - return O, Stats - - -def _sdpa_bwd_impl( - dO: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - stats: torch.Tensor, - attn_scale: float, - is_causal: bool = False, - diagonal_alignment: int = 0, - left_bound: int = -1, - right_bound: int = -1, - seq_len_q: Optional[torch.Tensor] = None, - seq_len_kv: Optional[torch.Tensor] = None, - cumulative_seq_len_q: Optional[torch.Tensor] = None, - cumulative_seq_len_kv: Optional[torch.Tensor] = None, - is_deterministic: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - cuDNN SDPA backward (graph-based). - - Returns: - (dQ, dK, dV) in BHSD layout matching Q, K, V shapes. - """ - handle = _get_handle(dO.device) - - cache_key = _make_bprop_cache_key( - q, - k, - v, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - seq_len_q is not None, - seq_len_kv is not None, - cumulative_seq_len_q is not None, - cumulative_seq_len_kv is not None, - is_deterministic, - ) - - if cache_key not in _bprop_cache: - graph, workspace_size = _build_bprop_graph( - handle, - q, - k, - v, - o, - dO, - stats, - attn_scale, - is_causal, - diagonal_alignment, - left_bound, - right_bound, - seq_len_q, - seq_len_kv, - cumulative_seq_len_q, - cumulative_seq_len_kv, - is_deterministic, - ) - _bprop_cache[cache_key] = (graph, workspace_size) - - graph, workspace_size = _bprop_cache[cache_key] - - # Allocate gradient outputs and workspace (same shapes as Q, K, V) - dQ_gpu = torch.empty_like(q) - dK_gpu = torch.empty_like(k) - dV_gpu = torch.empty_like(v) - workspace = torch.empty(max(workspace_size, 1), device=dO.device, dtype=torch.uint8) - - # UID → tensor map for sorted pointer extraction - uid_to_tensor = { - int(_UIDs.Q): q, - int(_UIDs.K): k, - int(_UIDs.V): v, - int(_UIDs.O): o, - int(_UIDs.DO): dO, - int(_UIDs.STATS): stats, - int(_UIDs.DQ): dQ_gpu, - int(_UIDs.DK): dK_gpu, - int(_UIDs.DV): dV_gpu, - } - if seq_len_q is not None: - uid_to_tensor[int(_UIDs.SEQ_LEN_Q)] = seq_len_q - if seq_len_kv is not None: - uid_to_tensor[int(_UIDs.SEQ_LEN_KV)] = seq_len_kv - if cumulative_seq_len_q is not None: - uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_Q)] = cumulative_seq_len_q - if cumulative_seq_len_kv is not None: - uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv - - graph.execute(uid_to_tensor, workspace, handle=handle) - - return dQ_gpu, dK_gpu, dV_gpu - - -_lib.impl("sdpa_bwd", _sdpa_bwd_impl, "CUDA") - - -@torch.library.register_fake("cudnn::sdpa_bwd") -def _sdpa_bwd_fake( - dO: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - stats: torch.Tensor, - attn_scale: float, - is_causal: bool = False, - diagonal_alignment: int = 0, - left_bound: int = -1, - right_bound: int = -1, - seq_len_q: Optional[torch.Tensor] = None, - seq_len_kv: Optional[torch.Tensor] = None, - cumulative_seq_len_q: Optional[torch.Tensor] = None, - cumulative_seq_len_kv: Optional[torch.Tensor] = None, - is_deterministic: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) - - -# --------------------------------------------------------------------------- -# Autograd registration -# --------------------------------------------------------------------------- - - -def _sdpa_setup_context(ctx, inputs, output): - q, k, v, attn_scale, is_causal, diagonal_alignment, left_bound, right_bound, seq_len_q, seq_len_kv, cumulative_seq_len_q, cumulative_seq_len_kv = inputs - o, stats = output - - tensors_to_save = [q, k, v, o, stats] - ctx.has_seq_len_q = seq_len_q is not None - ctx.has_seq_len_kv = seq_len_kv is not None - ctx.has_cum_q = cumulative_seq_len_q is not None - ctx.has_cum_kv = cumulative_seq_len_kv is not None - - if ctx.has_seq_len_q: - tensors_to_save.append(seq_len_q) - if ctx.has_seq_len_kv: - tensors_to_save.append(seq_len_kv) - if ctx.has_cum_q: - tensors_to_save.append(cumulative_seq_len_q) - if ctx.has_cum_kv: - tensors_to_save.append(cumulative_seq_len_kv) - - ctx.save_for_backward(*tensors_to_save) - - ctx.attn_scale = attn_scale - ctx.is_causal = is_causal - ctx.diagonal_alignment = diagonal_alignment - ctx.left_bound = left_bound - ctx.right_bound = right_bound - - -def _sdpa_backward(ctx, dO, dStats): - # dO from autograd may have zero strides (e.g. from o.sum().backward()), - # which cuDNN cannot handle. Make it contiguous. - dO = dO.contiguous() - - saved = list(ctx.saved_tensors) - q, k, v, o, stats = saved[:5] - idx = 5 - seq_len_q = saved[idx] if ctx.has_seq_len_q else None - if ctx.has_seq_len_q: - idx += 1 - seq_len_kv = saved[idx] if ctx.has_seq_len_kv else None - if ctx.has_seq_len_kv: - idx += 1 - cum_q = saved[idx] if ctx.has_cum_q else None - if ctx.has_cum_q: - idx += 1 - cum_kv = saved[idx] if ctx.has_cum_kv else None - - dQ, dK, dV = torch.ops.cudnn.sdpa_bwd( - dO, - q, - k, - v, - o, - stats, - ctx.attn_scale, - ctx.is_causal, - ctx.diagonal_alignment, - ctx.left_bound, - ctx.right_bound, - seq_len_q, - seq_len_kv, - cum_q, - cum_kv, - ) - - # Return gradients for: q, k, v, attn_scale, is_causal, diagonal_alignment, - # left_bound, right_bound, seq_len_q, seq_len_kv, cum_q, cum_kv - return dQ, dK, dV, None, None, None, None, None, None, None, None, None - - -torch.library.register_autograd("cudnn::sdpa", _sdpa_backward, setup_context=_sdpa_setup_context) - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - - -def scaled_dot_product_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, - dropout_p: float = 0.0, - is_causal: bool = False, - scale: Optional[float] = None, - enable_gqa: bool = False, - *, - diagonal_alignment: int = 0, - left_bound: int = -1, - right_bound: int = -1, - seq_len_q: Optional[torch.Tensor] = None, - seq_len_kv: Optional[torch.Tensor] = None, - cumulative_seq_len_q: Optional[torch.Tensor] = None, - cumulative_seq_len_kv: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """cuDNN-accelerated Scaled Dot-Product Attention. - - API closely mirrors ``torch.nn.functional.scaled_dot_product_attention``. - - **Layout**: tensors use **BHSD** layout ``(batch, num_heads, seq_len, head_dim)``, - matching PyTorch's convention. - - Args: - query: Query tensor ``(B, H_q, S_q, D)``. - key: Key tensor ``(B, H_k, S_kv, D)``. - value: Value tensor ``(B, H_v, S_kv, D_v)``. - attn_mask: Not yet supported. Must be ``None``. - dropout_p: Not yet supported. Must be ``0.0``. - is_causal: If ``True``, applies a causal (upper-triangular) mask. - scale: Attention scale factor. Defaults to ``1 / sqrt(D)`` when ``None``. - enable_gqa: When ``False``, raises ``ValueError`` if ``H_q != H_k``. - When ``True``, grouped-query attention is enabled (cuDNN handles - this automatically via head dimension broadcast). - - diagonal_alignment: cuDNN extension. ``0`` = TOP_LEFT, ``1`` = BOTTOM_RIGHT. - left_bound: cuDNN extension. Left sliding-window bound (``-1`` = disabled). - right_bound: cuDNN extension. Right sliding-window bound (``-1`` = disabled). - seq_len_q: cuDNN extension. Actual query sequence lengths ``(B, 1, 1, 1)`` INT32. - seq_len_kv: cuDNN extension. Actual key/value sequence lengths ``(B, 1, 1, 1)`` INT32. - cumulative_seq_len_q: cuDNN extension. Ragged offset for Q ``(B+1, 1, 1, 1)`` INT32. - cumulative_seq_len_kv: cuDNN extension. Ragged offset for KV ``(B+1, 1, 1, 1)`` INT32. - - Note: - For head dimension ``256``, the specialized backward path currently supports - only plain BHSD tensors and does not support ``seq_len_q``, ``seq_len_kv``, - ``cumulative_seq_len_q``, or ``cumulative_seq_len_kv``. - - Returns: - Output tensor ``(B, H_q, S_q, D_v)``. - """ - if attn_mask is not None: - raise NotImplementedError("attn_mask is not yet supported by cuDNN SDPA") - if dropout_p != 0.0: - raise NotImplementedError("dropout is not yet supported by cuDNN SDPA") - if not enable_gqa and query.shape[1] != key.shape[1]: - raise ValueError(f"query has {query.shape[1]} heads but key has {key.shape[1]} heads. " f"Set enable_gqa=True for grouped-query attention.") - - d = query.shape[-1] - attn_scale = scale if scale is not None else (1.0 / math.sqrt(d)) - - o, _stats = torch.ops.cudnn.sdpa( - query, - key, - value, - attn_scale, - is_causal=is_causal, - diagonal_alignment=diagonal_alignment, - left_bound=left_bound, - right_bound=right_bound, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - cumulative_seq_len_q=cumulative_seq_len_q, - cumulative_seq_len_kv=cumulative_seq_len_kv, - ) - return o diff --git a/python/cudnn/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py new file mode 100644 index 000000000..805b3fab0 --- /dev/null +++ b/python/cudnn/sdpa/fwd/torch_op.py @@ -0,0 +1,973 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PyTorch custom ops exposing the FULL cuDNN SDPA feature surface. + +``torch.ops.cudnn.sdpa_fwd`` / ``sdpa_bwd`` — the family-local torch contract +for features ``torch.nn.functional.scaled_dot_product_attention`` cannot +express: + +- **attention sinks** (per-Q-head logits folded into the softmax denominator) +- **sliding window** (``window_left``) +- **bottom-right causal alignment** (inference-style diagonals) +- **padded batches** (per-batch actual sequence lengths) +- **THD / varlen packing** (FlashAttention-style ``(T, H, D)`` + ``cu_seqlens``) + +Layout: dense tensors are BHSD ``(B, H, S, D)``; varlen tensors are packed +``(T, H, D)`` with ``cu_seqlens_q/kv`` (``(B+1,)`` int32 token prefix sums, as +in FA's ``flash_attn_varlen_func``). + +The ops build cuDNN pygraph ``sdpa`` / ``sdpa_backward`` nodes; the engine +Router then picks the best serving plan (FROST OSS kernels or cuDNN-backend +engines) per config. + +Backward contract: ``sdpa_bwd`` serves the THD/varlen path (dense backward and +sink backward are follow-ups and raise ``NotImplementedError``). It consumes a +PADDED ``(B, H, max_seqlen_q, 1)`` fp32 LSE — a backend restriction (bprop THD +rejects ragged LSE on SM8X/SM12X). ``sdpa_fwd`` is differentiable on the +varlen path via ``torch.library.register_autograd`` when called with +``return_lse=True``; the autograd glue converts the packed TH1 stats to the +padded layout the backward needs. + +Public entry point: ``cudnn.sdpa_torch`` (lazy export — accessing it imports +this module, which registers the ops; ``torch`` is not imported before then). +``import cudnn.sdpa.fwd.torch_op`` works too. +""" + +import logging +import threading +from enum import IntEnum +from typing import Dict, Optional, Tuple + +import torch + +import cudnn + +_logger = logging.getLogger(__name__) + +_TORCH_DTYPE_TO_CUDNN = { + torch.float16: cudnn.data_type.HALF, + torch.bfloat16: cudnn.data_type.BFLOAT16, +} + +# cuDNN handles are NOT thread-safe (simultaneous use of one handle from two +# threads is undefined) — keep them thread-local. Graph builds are serialized +# by a lock; built plans are immutable, so cached-graph EXECUTION stays +# lock-free (each thread executes with its own handle). +_tls = threading.local() +_graph_cache: Dict[tuple, tuple] = {} +_graph_cache_lock = threading.Lock() +# Bounded: one entry per distinct (shape, stride, flags) config. FIFO eviction +# keeps pathological shape-churn workloads from accumulating plans without +# bound (each holds device workspace-size metadata and backend plans). +_GRAPH_CACHE_MAX = 128 + + +class _UIDs(IntEnum): + Q = 1 + K = 2 + V = 3 + SINKS = 4 + SEQ_LEN_Q = 5 + SEQ_LEN_KV = 6 + RAGGED_Q = 7 + RAGGED_KV = 8 + RAGGED_O = 9 + RAGGED_STATS = 10 + RAGGED_V = 11 + RAGGED_DQ = 12 + RAGGED_DK = 13 + RAGGED_DV = 14 + O = 100 # noqa: E741 — matches the SDPA output tensor name + STATS = 101 + DO = 200 + DQ = 201 + DK = 202 + DV = 203 + + +def _get_handle(device: torch.device): + """This thread's cuDNN handle for ``device``, bound to torch's current stream.""" + handles = getattr(_tls, "handles", None) + if handles is None: + handles = _tls.handles = {} + if device not in handles: + # create_handle() binds to the CURRENT device — pin it explicitly so a + # tensor on cuda:1 never gets a handle created against cuda:0. + with torch.cuda.device(device): + handles[device] = cudnn.create_handle() + cudnn.set_stream(handle=handles[device], stream=torch.cuda.current_stream(device).cuda_stream) + return handles[device] + + +def _cached_graph(key, build): + """Graph-cache lookup with serialized builds and bounded (FIFO) size.""" + hit = _graph_cache.get(key) + if hit is not None: + return hit + with _graph_cache_lock: + hit = _graph_cache.get(key) # racing builder may have won + if hit is None: + hit = _graph_cache[key] = build() + while len(_graph_cache) > _GRAPH_CACHE_MAX: + _graph_cache.pop(next(iter(_graph_cache))) + return hit + + +def _check_io_dtypes(name: str, **tensors: torch.Tensor) -> None: + """One io dtype for the whole graph: reject mixes and unsupported dtypes + loudly (cuDNN would otherwise write q.dtype bits into buffers torch + believes hold another dtype).""" + ref = next(iter(tensors.values())).dtype + if ref not in _TORCH_DTYPE_TO_CUDNN: + raise ValueError(f"{name}: unsupported dtype {ref}; supported: {sorted(str(d) for d in _TORCH_DTYPE_TO_CUDNN)}") + mismatched = {n: str(t.dtype) for n, t in tensors.items() if t.dtype != ref} + if mismatched: + raise ValueError(f"{name}: all io tensors must share one dtype ({ref}); got {mismatched}") + + +def _stride_order(t: torch.Tensor) -> Tuple[int, ...]: + return tuple(sorted(range(t.ndim), key=lambda dim: t.stride()[dim])) + + +def _like_layout_stride(shape: Tuple[int, ...], like: torch.Tensor) -> Tuple[int, ...]: + """Compact strides for ``shape`` in ``like``'s dim-permutation — O adopts + Q's layout (all B/H/S permutations; D innermost stays D innermost). + Broadcast (stride-0) inputs have no meaningful order: fall back to + contiguous.""" + if 0 in like.stride(): + return tuple(torch.empty(shape, device="meta").stride()) + stride = [0] * len(shape) + acc = 1 + for dim in _stride_order(like): # innermost outward + stride[dim] = acc + acc *= shape[dim] + return tuple(stride) + + +def _packed_bhsd_stride(b: int, h: int, s: int, d: int) -> Tuple[int, int, int, int]: + """Token-major stride for THD descriptors (ragged offsets address batches).""" + return (s * h * d, d, h * d, 1) + + +def _thd_desc_stride(t: torch.Tensor, s_max: int) -> Tuple[int, int, int, int]: + """(B,H,S,D) descriptor stride for a packed (T,H,D) tensor, honoring the + tensor's ACTUAL strides — a (T,H,D) view of a kv-packed (T,2,H,D) buffer + has token stride 2*H*D, not H*D. The batch stride is a placeholder (the + ragged offset supplies per-batch bases).""" + s_t, s_h, s_d = t.stride() + return (s_max * s_t, s_h, s_t, s_d) + + +def _normalize_thd(t: torch.Tensor, name: str) -> torch.Tensor: + """Innermost dim must be dense and the base pointer 16B-aligned for the + cuDNN descriptors (an odd-element storage offset has equal strides but + faults the kernels with a misaligned address). clone(), NOT contiguous(): + contiguous() returns ``self`` unchanged for an already-contiguous tensor, + whatever its storage offset, so it cannot repair a misaligned base.""" + if t.stride(-1) != 1 or t.data_ptr() % 16: + _logger.warning("sdpa_fwd: copying %s to normalize layout/alignment (slow path)", name) + t = t.clone(memory_format=torch.contiguous_format) + return t + + +def _int32_col(t: torch.Tensor) -> torch.Tensor: + """View a 1-D int tensor as the (N, 1, 1, 1) INT32 column cuDNN expects.""" + return t.to(torch.int32).reshape(-1, 1, 1, 1) + + +def _int64_col(t: torch.Tensor) -> torch.Tensor: + """(N, 1, 1, 1) INT64 column: ragged offsets are int64 so element offsets + (token prefix sums x token stride) cannot overflow.""" + return t.to(torch.int64).reshape(-1, 1, 1, 1) + + +def _round64(n: int) -> int: + return ((n + 63) // 64) * 64 + + +def _check_same_device(q: torch.Tensor, **tensors) -> None: + """Every operand is bound into the variant pack as a DEVICE pointer — a + CPU (or other-device) tensor would hand cuDNN a foreign address and fault + as an illegal memory access instead of a clear error. None entries are + skipped.""" + for name, t in tensors.items(): + if t is not None and t.device != q.device: + raise ValueError(f"{name} must be on {q.device} (bound as a device pointer); got {t.device}") + + +# --------------------------------------------------------------------------- +# Graph builder +# --------------------------------------------------------------------------- + + +def _build_graph( + handle, + *, + dtype: torch.dtype, + B: int, + H_q: int, + H_k: int, + H_v: int, + S_q: int, + S_kv: int, + D_qk: int, + D_v: int, + q_stride, + k_stride, + v_stride, + o_stride, + attn_scale: float, + is_causal: bool, + causal_bottom_right: bool, + window_left: int, + has_sinks: bool, + has_seq_lens: bool, + is_thd: bool, + return_lse: bool, + stats_stride, +): + io_dtype = _TORCH_DTYPE_TO_CUDNN[dtype] + g = cudnn.pygraph( + handle=handle, + io_data_type=io_dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q_t = g.tensor(name="q", dim=[B, H_q, S_q, D_qk], stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q) + k_t = g.tensor(name="k", dim=[B, H_k, S_kv, D_qk], stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K) + v_t = g.tensor(name="v", dim=[B, H_v, S_kv, D_v], stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V) + + sinks_t = None + if has_sinks: + sinks_t = g.tensor(name="sinks", dim=[1, H_q, 1, 1], stride=[H_q, 1, 1, 1], data_type=cudnn.data_type.FLOAT, uid=_UIDs.SINKS) + + seq_q_t = seq_kv_t = None + if has_seq_lens or is_thd: + seq_q_t = g.tensor(name="seq_len_q", dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q) + seq_kv_t = g.tensor(name="seq_len_kv", dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV) + + if is_thd: + rq = g.tensor(name="ragged_q", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_Q) + rk = g.tensor(name="ragged_k", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_KV) + rv = g.tensor(name="ragged_v", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_V) + ro = g.tensor(name="ragged_o", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_O) + q_t.set_ragged_offset(rq) + k_t.set_ragged_offset(rk) + v_t.set_ragged_offset(rv) + + rb = 0 if is_causal else None + lb = window_left if window_left >= 0 else None + alignment = cudnn.diagonal_alignment.BOTTOM_RIGHT if causal_bottom_right else cudnn.diagonal_alignment.TOP_LEFT + + o_t, stats_t = g.sdpa( + name="sdpa_fwd", + q=q_t, + k=k_t, + v=v_t, + generate_stats=return_lse, + attn_scale=attn_scale, + use_padding_mask=has_seq_lens or is_thd, + seq_len_q=seq_q_t, + seq_len_kv=seq_kv_t, + diagonal_alignment=alignment, + diagonal_band_left_bound=lb, + diagonal_band_right_bound=rb, + sink_token=sinks_t, + ) + + o_t.set_uid(_UIDs.O).set_output(True).set_dim([B, H_q, S_q, D_v]).set_stride(list(o_stride)).set_data_type(io_dtype) + if is_thd: + o_t.set_ragged_offset(ro) + + if return_lse: + stats_t.set_uid(_UIDs.STATS).set_output(True).set_dim([B, H_q, S_q, 1]).set_stride(list(stats_stride)).set_data_type(cudnn.data_type.FLOAT) + if is_thd: + rs = g.tensor(name="ragged_stats", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_STATS) + stats_t.set_ragged_offset(rs) + + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + g.check_support() + g.build_plans() + return g, g.get_workspace_size() + + +# --------------------------------------------------------------------------- +# Custom op +# --------------------------------------------------------------------------- + +# FRAGMENT: extend the "cudnn" namespace (import-order independent with any +# other registrant into it). +_lib = torch.library.Library("cudnn", "FRAGMENT") + +_lib.define( + "sdpa_fwd(Tensor q, Tensor k, Tensor v, float attn_scale, " + "bool is_causal=False, bool causal_bottom_right=False, int window_left=-1, " + "Tensor? sinks=None, " + "Tensor? seq_len_q=None, Tensor? seq_len_kv=None, " + "Tensor? cu_seqlens_q=None, Tensor? cu_seqlens_kv=None, " + "int max_seqlen_q=0, int max_seqlen_kv=0, " + "bool return_lse=True) -> (Tensor, Tensor)" +) + + +def _sdpa_fwd_impl( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_scale: float, + is_causal: bool = False, + causal_bottom_right: bool = False, + window_left: int = -1, + sinks: Optional[torch.Tensor] = None, + seq_len_q: Optional[torch.Tensor] = None, + seq_len_kv: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: int = 0, + max_seqlen_kv: int = 0, + return_lse: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + _check_io_dtypes("sdpa_fwd", q=q, k=k, v=v) + if causal_bottom_right and not (is_causal or window_left >= 0): + raise ValueError("causal_bottom_right only re-anchors an active diagonal band; set is_causal or window_left too") + is_thd = cu_seqlens_q is not None + if is_thd: + if cu_seqlens_kv is None or max_seqlen_q <= 0 or max_seqlen_kv <= 0: + raise ValueError("varlen path needs cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv") + if q.ndim != 3: + raise ValueError(f"varlen path expects packed (T, H, D) tensors, got q.ndim={q.ndim}") + if seq_len_q is not None or seq_len_kv is not None: + raise ValueError("varlen path derives seq lens from cu_seqlens; do not pass seq_len_q/kv") + B = cu_seqlens_q.numel() - 1 + q = _normalize_thd(q, "q") + k = _normalize_thd(k, "k") + v = _normalize_thd(v, "v") + T_q, H_q, D_qk = q.shape + T_kv, H_v, D_v = v.shape + H_k = k.shape[1] + # cuDNN supports h_k != h_v (each must divide h_q). + if k.shape != (T_kv, H_k, D_qk): + raise ValueError(f"k shape {tuple(k.shape)} must be (T_kv={T_kv}, H_k, D_qk={D_qk}) to match q and v") + if H_q % H_k or H_q % H_v: + raise ValueError(f"GQA head counts must divide H_q={H_q}; got H_k={H_k}, H_v={H_v}") + S_q, S_kv = max_seqlen_q, max_seqlen_kv + _check_same_device(q, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv) + # Descriptors honor ACTUAL strides (kv-packed views etc.); O is ours, + # allocated packed. + q_stride = _thd_desc_stride(q, S_q) + k_stride = _thd_desc_stride(k, S_kv) + v_stride = _thd_desc_stride(v, S_kv) + o_stride = _packed_bhsd_stride(B, H_q, S_q, D_v) + stats_stride = (S_q * H_q, 1, H_q, 1) # TH1 token-major + else: + if q.ndim != 4: + raise ValueError(f"dense path expects BHSD tensors, got q.ndim={q.ndim}") + B, H_q, S_q, D_qk = q.shape + _, H_v, S_kv, D_v = v.shape + H_k = k.shape[1] + # cuDNN supports h_k != h_v (each must divide h_q). + if k.shape != (B, H_k, S_kv, D_qk) or v.shape[0] != B: + raise ValueError(f"k shape {tuple(k.shape)} must be (B={B}, H_k, S_kv={S_kv}, D_qk={D_qk}) to match q and v") + if H_q % H_k or H_q % H_v: + raise ValueError(f"GQA head counts must divide H_q={H_q}; got H_k={H_k}, H_v={H_v}") + q_stride, k_stride, v_stride = q.stride(), k.stride(), v.stride() + o_stride = _like_layout_stride((B, H_q, S_q, D_v), q) # O adopts Q's layout + stats_stride = (H_q * S_q, S_q, 1, 1) + + _check_same_device(q, k=k, v=v, sinks=sinks, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv) + has_sinks = sinks is not None + has_seq_lens = seq_len_q is not None or seq_len_kv is not None + if has_seq_lens and (seq_len_q is None or seq_len_kv is None): + raise ValueError("padded path needs both seq_len_q and seq_len_kv") + + key = ( + "sdpa_fwd", + q.dtype, + B, + H_q, + H_k, + H_v, + S_q, + S_kv, + D_qk, + D_v, + tuple(q.stride()), + tuple(k.stride()), + tuple(v.stride()), + attn_scale, + is_causal, + causal_bottom_right, + window_left, + has_sinks, + has_seq_lens, + is_thd, + return_lse, + q.device, + ) + + handle = _get_handle(q.device) + g, ws = _cached_graph( + key, + lambda: _build_graph( + handle, + dtype=q.dtype, + B=B, + H_q=H_q, + H_k=H_k, + H_v=H_v, + S_q=S_q, + S_kv=S_kv, + D_qk=D_qk, + D_v=D_v, + q_stride=q_stride, + k_stride=k_stride, + v_stride=v_stride, + o_stride=o_stride, + attn_scale=attn_scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_left=window_left, + has_sinks=has_sinks, + has_seq_lens=has_seq_lens, + is_thd=is_thd, + return_lse=return_lse, + stats_stride=stats_stride, + ), + ) + + # Outputs + workspace (workspace per call: the torch allocator recycles it). + # The no-LSE placeholder is a 0-elem CUDA fp32 tensor (a CUDA op must not + # hand back a CPU tensor; the fake kernel mirrors this). + if is_thd: + o = torch.empty(T_q, H_q, D_v, dtype=q.dtype, device=q.device) + stats = torch.empty(T_q, H_q, 1, dtype=torch.float32, device=q.device) if return_lse else torch.empty(0, dtype=torch.float32, device=q.device) + else: + o = torch.empty_strided((B, H_q, S_q, D_v), o_stride, dtype=q.dtype, device=q.device) + stats = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device) if return_lse else torch.empty(0, dtype=torch.float32, device=q.device) + workspace = torch.empty(max(ws, 1), dtype=torch.uint8, device=q.device) + + variant = {int(_UIDs.Q): q, int(_UIDs.K): k, int(_UIDs.V): v, int(_UIDs.O): o} + if return_lse: + variant[int(_UIDs.STATS)] = stats + if has_sinks: + variant[int(_UIDs.SINKS)] = sinks.to(torch.float32).reshape(1, H_q, 1, 1) + if is_thd: + # cuDNN ragged offsets are int64 ELEMENT offsets per tensor, so each + # scales its token prefix sums by that tensor's OWN token stride — + # widened BEFORE the multiply (an int32 product would wrap before + # _int64_col ever sees it). Small on-stream int ops — CUDA-graph- + # capture safe. + cu_q64 = cu_seqlens_q.to(torch.int64) + cu_kv64 = cu_seqlens_kv.to(torch.int64) + variant[int(_UIDs.RAGGED_Q)] = _int64_col(cu_q64 * q.stride(0)) + variant[int(_UIDs.RAGGED_KV)] = _int64_col(cu_kv64 * k.stride(0)) + variant[int(_UIDs.RAGGED_V)] = _int64_col(cu_kv64 * v.stride(0)) + variant[int(_UIDs.RAGGED_O)] = _int64_col(cu_q64 * (H_q * D_v)) + variant[int(_UIDs.SEQ_LEN_Q)] = _int32_col(cu_seqlens_q[1:] - cu_seqlens_q[:-1]) + variant[int(_UIDs.SEQ_LEN_KV)] = _int32_col(cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]) + if return_lse: + variant[int(_UIDs.RAGGED_STATS)] = _int64_col(cu_q64 * H_q) + elif has_seq_lens: + variant[int(_UIDs.SEQ_LEN_Q)] = _int32_col(seq_len_q) + variant[int(_UIDs.SEQ_LEN_KV)] = _int32_col(seq_len_kv) + + g.execute(variant, workspace, handle=handle) + return o, stats + + +_lib.impl("sdpa_fwd", _sdpa_fwd_impl, "CUDA") + + +@torch.library.register_fake("cudnn::sdpa_fwd") +def _sdpa_fwd_fake( + q, + k, + v, + attn_scale, + is_causal=False, + causal_bottom_right=False, + window_left=-1, + sinks=None, + seq_len_q=None, + seq_len_kv=None, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=0, + max_seqlen_kv=0, + return_lse=True, +): + # Mirrors the REAL kernel's output metadata exactly — torch.compile plans + # downstream layouts from these strides. + D_v = v.shape[-1] + if cu_seqlens_q is not None: # THD: (T, H, D) packed-contiguous + T_q, H_q = q.shape[0], q.shape[1] + o = torch.empty(T_q, H_q, D_v, dtype=q.dtype, device=q.device) + stats = torch.empty(T_q, H_q, 1, dtype=torch.float32, device=q.device) if return_lse else torch.empty(0, dtype=torch.float32, device=q.device) + else: + B, H_q, S_q = q.shape[0], q.shape[1], q.shape[2] + o_stride = _like_layout_stride((B, H_q, S_q, D_v), q) # O adopts Q's layout + o = torch.empty_strided((B, H_q, S_q, D_v), o_stride, dtype=q.dtype, device=q.device) + stats = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device) if return_lse else torch.empty(0, dtype=torch.float32, device=q.device) + return o, stats + + +# --------------------------------------------------------------------------- +# Backward (THD/varlen-capable). Prototype home: will move next to the bwd +# engine family (sdpa/bwd/) when this graduates. +# --------------------------------------------------------------------------- + + +def _build_bwd_graph( + handle, + *, + dtype: torch.dtype, + B: int, + H_q: int, + H_k: int, + H_v: int, + S_q: int, + S_kv: int, + D_qk: int, + D_v: int, + total_q: int, + total_kv: int, + attn_scale: float, + is_causal: bool, + causal_bottom_right: bool, + window_left: int, + q_stride, + k_stride, + v_stride, + o_stride, + stats_stride, + is_deterministic: bool, +): + """THD/varlen backward graph (the only path the bwd op serves today).""" + io_dtype = _TORCH_DTYPE_TO_CUDNN[dtype] + g = cudnn.pygraph( + handle=handle, + io_data_type=io_dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q_t = g.tensor(name="q", dim=[B, H_q, S_q, D_qk], stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q) + k_t = g.tensor(name="k", dim=[B, H_k, S_kv, D_qk], stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K) + v_t = g.tensor(name="v", dim=[B, H_v, S_kv, D_v], stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V) + o_t = g.tensor(name="o", dim=[B, H_q, S_q, D_v], stride=list(o_stride), data_type=io_dtype, uid=_UIDs.O) + do_t = g.tensor(name="dO", dim=[B, H_q, S_q, D_v], stride=list(o_stride), data_type=io_dtype, uid=_UIDs.DO) + # Stats stay PADDED dense even in THD: the backend rejects ragged LSE for + # bprop THD on SM8X/SM12X ("Packed/ragged LSE is not supported"). + stats_t = g.tensor(name="stats", dim=[B, H_q, S_q, 1], stride=list(stats_stride), data_type=cudnn.data_type.FLOAT, uid=_UIDs.STATS) + + seq_q_t = g.tensor(name="seq_len_q", dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q) + seq_kv_t = g.tensor(name="seq_len_kv", dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV) + rq = g.tensor(name="ragged_q", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_Q) + rk = g.tensor(name="ragged_k", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_KV) + rv = g.tensor(name="ragged_v", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_V) + ro = g.tensor(name="ragged_o", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_O) + rdq = g.tensor(name="ragged_dq", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_DQ) + rdk = g.tensor(name="ragged_dk", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_DK) + rdv = g.tensor(name="ragged_dv", dim=[B + 1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT64, uid=_UIDs.RAGGED_DV) + q_t.set_ragged_offset(rq) + k_t.set_ragged_offset(rk) + v_t.set_ragged_offset(rv) + o_t.set_ragged_offset(ro) + do_t.set_ragged_offset(ro) + + rb = 0 if is_causal else None + lb = window_left if window_left >= 0 else None + alignment = cudnn.diagonal_alignment.BOTTOM_RIGHT if causal_bottom_right else cudnn.diagonal_alignment.TOP_LEFT + + dq_t, dk_t, dv_t = g.sdpa_backward( + name="sdpa_bwd", + q=q_t, + k=k_t, + v=v_t, + o=o_t, + dO=do_t, + stats=stats_t, + attn_scale=attn_scale, + use_padding_mask=True, + seq_len_q=seq_q_t, + seq_len_kv=seq_kv_t, + # Actual packed token totals (rounded up to the backend's 64-token + # accumulator granularity) — sizes the dq accumulator. + max_total_seq_len_q=_round64(total_q), + max_total_seq_len_kv=_round64(total_kv), + diagonal_alignment=alignment, + diagonal_band_left_bound=lb, + diagonal_band_right_bound=rb, + use_deterministic_algorithm=is_deterministic, + ) + + # Gradients are OURS: always packed-contiguous, independent of the input views. + dq_stride = _packed_bhsd_stride(B, H_q, S_q, D_qk) + dk_stride = _packed_bhsd_stride(B, H_k, S_kv, D_qk) + dv_stride = _packed_bhsd_stride(B, H_v, S_kv, D_v) + dq_t.set_uid(_UIDs.DQ).set_output(True).set_dim([B, H_q, S_q, D_qk]).set_stride(list(dq_stride)).set_data_type(io_dtype) + dk_t.set_uid(_UIDs.DK).set_output(True).set_dim([B, H_k, S_kv, D_qk]).set_stride(list(dk_stride)).set_data_type(io_dtype) + dv_t.set_uid(_UIDs.DV).set_output(True).set_dim([B, H_v, S_kv, D_v]).set_stride(list(dv_stride)).set_data_type(io_dtype) + dq_t.set_ragged_offset(rdq) + dk_t.set_ragged_offset(rdk) + dv_t.set_ragged_offset(rdv) + + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + g.check_support() + g.build_plans() + return g, g.get_workspace_size() + + +_lib.define( + "sdpa_bwd(Tensor grad_out, Tensor q, Tensor k, Tensor v, Tensor o, Tensor lse, float attn_scale, " + "bool is_causal=False, bool causal_bottom_right=False, int window_left=-1, " + "Tensor? sinks=None, " + "Tensor? cu_seqlens_q=None, Tensor? cu_seqlens_kv=None, " + "int max_seqlen_q=0, int max_seqlen_kv=0, " + "bool is_deterministic=False) -> (Tensor, Tensor, Tensor)" +) + + +def _sdpa_bwd_impl( + grad_out: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + attn_scale: float, + is_causal: bool = False, + causal_bottom_right: bool = False, + window_left: int = -1, + sinks: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: int = 0, + max_seqlen_kv: int = 0, + is_deterministic: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if sinks is not None: + # A sink forward folds the sink logits into the softmax denominator; + # this backward has no dSink support yet, and silently ignoring the + # sink term would produce numerically wrong dq/dk/dv. + raise NotImplementedError("cudnn::sdpa_bwd does not support attention sinks yet (dSink is a follow-up); gradients would be wrong") + if cu_seqlens_q is None: + raise NotImplementedError("cudnn::sdpa_bwd currently serves the THD/varlen path; dense backward is a follow-up") + if cu_seqlens_kv is None or max_seqlen_q <= 0 or max_seqlen_kv <= 0: + raise ValueError("varlen path needs cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv") + if q.ndim != 3: + raise ValueError(f"varlen path expects packed (T, H, D) tensors, got q.ndim={q.ndim}") + _check_io_dtypes("sdpa_bwd", grad_out=grad_out, q=q, k=k, v=v, o=o) + + B = cu_seqlens_q.numel() - 1 + q = _normalize_thd(q, "q") + k = _normalize_thd(k, "k") + v = _normalize_thd(v, "v") + o = _normalize_thd(o, "o") + T_q, H_q, D_qk = q.shape + T_kv, H_v, D_v = v.shape + H_k = k.shape[1] + if k.shape != (T_kv, H_k, D_qk): + raise ValueError(f"k shape {tuple(k.shape)} must be (T_kv={T_kv}, H_k, D_qk={D_qk}) to match q and v") + if H_q % H_k or H_q % H_v: + raise ValueError(f"GQA head counts must divide H_q={H_q}; got H_k={H_k}, H_v={H_v}") + if o.shape != (T_q, H_q, D_v) or grad_out.shape != o.shape: + raise ValueError(f"o {tuple(o.shape)} / grad_out {tuple(grad_out.shape)} must be (T_q={T_q}, H_q={H_q}, D_v={D_v})") + S_q, S_kv = max_seqlen_q, max_seqlen_kv + _check_same_device(q, k=k, v=v, o=o, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, lse=lse, grad_out=grad_out) + + # lse arrives PADDED (B, H, max_seqlen_q) or (B, H, max_seqlen_q, 1) fp32 + # — a backend restriction (bprop THD rejects ragged LSE on SM8X/SM12X). + # Rows past each sequence's length are ignored. Normalize BEFORE reshape + # (reshape of a non-contiguous tensor silently copies or raises), and also + # on base-pointer misalignment. + if lse.dtype != torch.float32: + raise ValueError(f"lse must be float32, got {lse.dtype}") + if not lse.is_contiguous() or lse.data_ptr() % 16: + lse = lse.clone(memory_format=torch.contiguous_format) + lse = lse.reshape(B, H_q, S_q, 1) + # Normalize dO to O's layout; ALSO on base-pointer misalignment — equal + # strides with an odd storage offset would fault the kernels. + if grad_out.stride() != o.stride() or grad_out.data_ptr() % 16: + grad_out = torch.empty_strided(o.shape, o.stride(), dtype=grad_out.dtype, device=grad_out.device).copy_(grad_out) + + q_stride = _thd_desc_stride(q, S_q) + k_stride = _thd_desc_stride(k, S_kv) + v_stride = _thd_desc_stride(v, S_kv) + o_stride = _thd_desc_stride(o, S_q) + stats_stride = (H_q * S_q, S_q, 1, 1) # padded dense BHS1 + + key = ( + "sdpa_bwd", + q.dtype, + B, + H_q, + H_k, + H_v, + # The packed token totals are BAKED into the graph (they size the dq + # accumulator via max_total_seq_len_*): a plan built for smaller + # totals must not serve a call with larger ones. Rounded to the same + # 64-token granularity the graph uses, so the cache still hits across + # calls that share an accumulator size. + _round64(T_q), + _round64(T_kv), + S_q, + S_kv, + D_qk, + D_v, + tuple(q.stride()), + tuple(k.stride()), + tuple(v.stride()), + tuple(o.stride()), + attn_scale, + is_causal, + causal_bottom_right, + window_left, + is_deterministic, + q.device, + ) + + handle = _get_handle(q.device) + g, ws = _cached_graph( + key, + lambda: _build_bwd_graph( + handle, + dtype=q.dtype, + B=B, + H_q=H_q, + H_k=H_k, + H_v=H_v, + S_q=S_q, + S_kv=S_kv, + D_qk=D_qk, + D_v=D_v, + total_q=T_q, + total_kv=T_kv, + attn_scale=attn_scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_left=window_left, + q_stride=q_stride, + k_stride=k_stride, + v_stride=v_stride, + o_stride=o_stride, + stats_stride=stats_stride, + is_deterministic=is_deterministic, + ), + ) + + # Gradients are ours: packed-contiguous, one io dtype (validated above). + dq = torch.empty(T_q, H_q, D_qk, dtype=q.dtype, device=q.device) + dk = torch.empty(T_kv, H_k, D_qk, dtype=q.dtype, device=q.device) + dv = torch.empty(T_kv, H_v, D_v, dtype=q.dtype, device=q.device) + workspace = torch.empty(max(ws, 1), dtype=torch.uint8, device=q.device) + + variant = { + int(_UIDs.Q): q, + int(_UIDs.K): k, + int(_UIDs.V): v, + int(_UIDs.O): o, + int(_UIDs.DO): grad_out, + int(_UIDs.STATS): lse, + int(_UIDs.DQ): dq, + int(_UIDs.DK): dk, + int(_UIDs.DV): dv, + # Widened BEFORE the multiply: int32 products wrap before _int64_col. + int(_UIDs.RAGGED_Q): _int64_col(cu_seqlens_q.to(torch.int64) * q.stride(0)), + int(_UIDs.RAGGED_KV): _int64_col(cu_seqlens_kv.to(torch.int64) * k.stride(0)), + int(_UIDs.RAGGED_V): _int64_col(cu_seqlens_kv.to(torch.int64) * v.stride(0)), + int(_UIDs.RAGGED_O): _int64_col(cu_seqlens_q.to(torch.int64) * o.stride(0)), + int(_UIDs.RAGGED_DQ): _int64_col(cu_seqlens_q.to(torch.int64) * (H_q * D_qk)), + int(_UIDs.RAGGED_DK): _int64_col(cu_seqlens_kv.to(torch.int64) * (H_k * D_qk)), + int(_UIDs.RAGGED_DV): _int64_col(cu_seqlens_kv.to(torch.int64) * (H_v * D_v)), + int(_UIDs.SEQ_LEN_Q): _int32_col(cu_seqlens_q[1:] - cu_seqlens_q[:-1]), + int(_UIDs.SEQ_LEN_KV): _int32_col(cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]), + } + + g.execute(variant, workspace, handle=handle) + return dq, dk, dv + + +_lib.impl("sdpa_bwd", _sdpa_bwd_impl, "CUDA") + + +@torch.library.register_fake("cudnn::sdpa_bwd") +def _sdpa_bwd_fake( + grad_out, + q, + k, + v, + o, + lse, + attn_scale, + is_causal=False, + causal_bottom_right=False, + window_left=-1, + sinks=None, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=0, + max_seqlen_kv=0, + is_deterministic=False, +): + # The real kernel returns FRESH packed-contiguous gradients in q.dtype — + # not views of the inputs (q/k/v may be non-contiguous kv-interleaved + # views), so empty_like would report strides that never materialize. + dq = torch.empty(q.shape, dtype=q.dtype, device=q.device) + dk = torch.empty(k.shape, dtype=q.dtype, device=q.device) + dv = torch.empty(v.shape, dtype=q.dtype, device=q.device) + return dq, dk, dv + + +# --------------------------------------------------------------------------- +# Autograd: sdpa_fwd is differentiable on the varlen path (dense/sink +# backward raise until their engine contracts land). The glue converts the +# forward's packed TH1 stats to the padded (B, H, max_seqlen_q, 1) layout the +# backward requires. +# --------------------------------------------------------------------------- + + +def _sdpa_setup_context(ctx, inputs, output): + ( + q, + k, + v, + attn_scale, + is_causal, + causal_bottom_right, + window_left, + sinks, + seq_len_q, + seq_len_kv, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + return_lse, + ) = inputs + o, stats = output + ctx.save_for_backward(q, k, v, o, stats, cu_seqlens_q, cu_seqlens_kv) + ctx.attn_scale = attn_scale + ctx.is_causal = is_causal + ctx.causal_bottom_right = causal_bottom_right + ctx.window_left = window_left + ctx.has_sinks = sinks is not None + ctx.has_seq_lens = seq_len_q is not None or seq_len_kv is not None + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_kv = max_seqlen_kv + ctx.return_lse = return_lse + # The returned stats/lse is NOT differentiable through this backward + # (cuDNN's sdpa_backward consumes lse, it does not produce dLSE): + # mark it so autograd refuses a caller's lse-gradient with a clear error + # instead of this backward silently dropping it. Unused-o grads stay + # None (no zero materialization) and short-circuit below. + ctx.set_materialize_grads(False) + ctx.mark_non_differentiable(output[1]) + + +def _sdpa_backward(ctx, grad_o, _grad_stats): # stats marked non-differentiable + q, k, v, o, stats, cu_q, cu_kv = ctx.saved_tensors + if grad_o is None: # o unused in the loss; stats is non-differentiable + return (None,) * 15 + if cu_q is None: + raise NotImplementedError("cudnn::sdpa_fwd autograd serves the THD/varlen path; dense backward is a follow-up") + if ctx.has_sinks: + raise NotImplementedError("cudnn::sdpa_fwd autograd does not support attention sinks yet (dSink is a follow-up)") + if ctx.has_seq_lens: + raise NotImplementedError("cudnn::sdpa_fwd autograd does not support the padded dense path yet") + if not ctx.return_lse: + raise RuntimeError("cudnn::sdpa_fwd autograd requires return_lse=True (the backward consumes the forward stats)") + + # Packed TH1 (T, H, 1) -> padded (B, H, max_seqlen_q, 1): the backend + # rejects ragged LSE for bprop THD on SM8X/SM12X. Entirely device-side + # (no host reads of cu values): traceable under dynamic-shape AOT + # dispatch, and no D2H sync on the backward hot path. + B = cu_q.numel() - 1 + H_q = q.shape[1] + T_q = stats.shape[0] + token = torch.arange(T_q, device=stats.device) + seq_of_token = torch.searchsorted(cu_q[1:].long(), token, right=True) # token t in [cu[i], cu[i+1]) -> i + pos_in_seq = token - cu_q.long()[seq_of_token] + lse_padded = torch.zeros(B, H_q, ctx.max_seqlen_q, 1, dtype=torch.float32, device=stats.device) + lse_padded[seq_of_token, :, pos_in_seq, 0] = stats[:, :, 0] + + dq, dk, dv = torch.ops.cudnn.sdpa_bwd( + grad_o, + q, + k, + v, + o, + lse_padded, + ctx.attn_scale, + is_causal=ctx.is_causal, + causal_bottom_right=ctx.causal_bottom_right, + window_left=ctx.window_left, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_kv=ctx.max_seqlen_kv, + is_deterministic=torch.are_deterministic_algorithms_enabled(), + ) + # One grad slot per op input: (q, k, v, attn_scale, is_causal, + # causal_bottom_right, window_left, sinks, seq_len_q, seq_len_kv, + # cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, return_lse). + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None + + +torch.library.register_autograd("cudnn::sdpa_fwd", _sdpa_backward, setup_context=_sdpa_setup_context) + + +# --------------------------------------------------------------------------- +# Public wrapper +# --------------------------------------------------------------------------- + + +def sdpa( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + scale: Optional[float] = None, + is_causal: bool = False, + causal_bottom_right: bool = False, + window_left: int = -1, + sinks: Optional[torch.Tensor] = None, + seq_len_q: Optional[torch.Tensor] = None, + seq_len_kv: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: int = 0, + max_seqlen_kv: int = 0, + return_lse: bool = False, +): + """cuDNN SDPA forward with the extended feature surface (see module docstring). + + Returns ``o`` or ``(o, lse)`` when ``return_lse=True``. + """ + import math + + attn_scale = scale if scale is not None else 1.0 / math.sqrt(query.shape[-1]) + o, lse = torch.ops.cudnn.sdpa_fwd( + query, + key, + value, + attn_scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_left=window_left, + sinks=sinks, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + return_lse=return_lse, + ) + return (o, lse) if return_lse else o diff --git a/test/python/sdpa/test_torch_ops.py b/test/python/sdpa/test_torch_ops.py new file mode 100644 index 000000000..ab0c19e63 --- /dev/null +++ b/test/python/sdpa/test_torch_ops.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the extended SDPA torch custom ops (cudnn.sdpa.fwd.torch_op). + +``torch.ops.cudnn.sdpa_fwd`` / ``sdpa_bwd`` expose the cuDNN feature +surface that ``torch.nn.functional.scaled_dot_product_attention``'s aten +contract cannot: attention sinks, sliding windows, bottom-right causal +diagonals, padded batches, and THD/varlen packing (FlashAttention-style +``(T, H, D)`` + ``cu_seqlens``). Each case checks numerics against a pure +fp32 PyTorch reference. The engine Router picks the serving plan (FROST OSS +kernels or cuDNN-backend engines) per configuration — these tests pass on +either route. +""" + +import math + +import pytest +import torch + +import cudnn + +if not torch.cuda.is_available(): + pytest.skip("CUDA device required", allow_module_level=True) +if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("cuDNN SDPA requires sm80+", allow_module_level=True) +if cudnn.backend_version() < 90600: + pytest.skip("requires cuDNN >= 9.6 (THD token-major stats)", allow_module_level=True) + +from cudnn.sdpa.fwd import torch_op # noqa: E402 (registers torch.ops.cudnn.sdpa_fwd / sdpa_bwd) + +# SDPA rejects sink_token below 9.13 (scaled_dot_product_flash_attention.h: +# "SDPA with sink_token is not supported before 9.13."), while the module +# gate above only needs 9.6 for token-major THD stats. +_SINKS_UNSUPPORTED = pytest.mark.skipif(cudnn.backend_version() < 91300, reason="sink_token requires cuDNN >= 9.13") + +TOL = 2.5e-2 # bf16 rounding at these magnitudes + + +def ref_attention(q, k, v, scale, is_causal=False, bottom_right=False, window_left=-1, sinks=None, return_lse=False): + """fp32 reference in BHSD. window_left counts VISIBLE tokens including self + (the cuDNN diagonal_band_left_bound convention). sinks: (H,) extra softmax + logit per query head, contributing no value (but part of the softmax + denominator, so part of the LSE too).""" + q, k, v = q.float(), k.float(), v.float() + B, Hq, Sq, _ = q.shape + Hkv, Skv = k.shape[1], k.shape[2] + if Hq != Hkv: + k = k.repeat_interleave(Hq // Hkv, dim=1) + v = v.repeat_interleave(Hq // Hkv, dim=1) + s = torch.einsum("bhqd,bhkd->bhqk", q, k) * scale + + i = torch.arange(Sq, device=q.device).view(-1, 1) + j = torch.arange(Skv, device=q.device).view(1, -1) + off = (Skv - Sq) if bottom_right else 0 + mask = torch.zeros(Sq, Skv, dtype=torch.bool, device=q.device) + if is_causal: + mask |= j > (i + off) + if window_left >= 0: + mask |= j <= (i + off - window_left) + s = s.masked_fill(mask, float("-inf")) + + if sinks is not None: + sink_col = sinks.float().view(1, Hq, 1, 1).expand(B, Hq, Sq, 1) + s = torch.cat([s, sink_col], dim=-1) + p = torch.softmax(s, dim=-1)[..., :-1] + else: + p = torch.softmax(s, dim=-1) + o = torch.einsum("bhqk,bhkd->bhqd", p, v) + if return_lse: + return o, torch.logsumexp(s, dim=-1, keepdim=True) # (B, Hq, Sq, 1) + return o + + +def bshd(B, H, S, D, dtype=torch.bfloat16, requires_grad=False): + t = torch.randn(B, S, H, D, dtype=dtype, device="cuda").transpose(1, 2) + return t.requires_grad_(True) if requires_grad else t + + +class TestSdpaFwdDense: + @pytest.mark.L0 + @_SINKS_UNSUPPORTED + @pytest.mark.parametrize("is_causal", [False, True]) + def test_sinks(self, is_causal): + torch.manual_seed(0) + B, H, S, D = 2, 8, 512, 128 + q, k, v = bshd(B, H, S, D), bshd(B, H, S, D), bshd(B, H, S, D) + sinks = torch.randn(H, device="cuda", dtype=torch.float32) + scale = D**-0.5 + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=is_causal, sinks=sinks, return_lse=True) + ref, ref_lse = ref_attention(q, k, v, scale, is_causal=is_causal, sinks=sinks, return_lse=True) + assert (o.float() - ref).abs().max().item() < TOL + # LSE values (not just metadata): the sink logit is part of the denominator. + assert lse.shape == (B, H, S, 1) and lse.dtype == torch.float32 + assert (lse - ref_lse).abs().max().item() < TOL + + @pytest.mark.L0 + def test_padded_seq_lens(self): + """Dense padded batches: per-batch actual lengths via seq_len_q/kv. + Only rows/cols inside each batch's actual lengths are compared (rows + past seq_len_q are dead by contract).""" + torch.manual_seed(0) + B, H, S, D = 2, 8, 256, 128 + q, k, v = bshd(B, H, S, D), bshd(B, H, S, D), bshd(B, H, S, D) + len_q = torch.tensor([200, 96], device="cuda", dtype=torch.int32) + len_kv = torch.tensor([128, 256], device="cuda", dtype=torch.int32) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, seq_len_q=len_q, seq_len_kv=len_kv, return_lse=False) + for b in range(B): + lq, lkv = int(len_q[b]), int(len_kv[b]) + ref = ref_attention(q[b : b + 1, :, :lq], k[b : b + 1, :, :lkv], v[b : b + 1, :, :lkv], scale) + assert (o[b, :, :lq].float() - ref[0]).abs().max().item() < TOL, f"batch {b}" + + @pytest.mark.L0 + @pytest.mark.parametrize("window_left", [64, 128]) + def test_sliding_window_causal(self, window_left): + torch.manual_seed(0) + B, H, S, D = 2, 8, 512, 128 + q, k, v = bshd(B, H, S, D), bshd(B, H, S, D), bshd(B, H, S, D) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, window_left=window_left, return_lse=False) + ref = ref_attention(q, k, v, scale, is_causal=True, window_left=window_left) + assert (o.float() - ref).abs().max().item() < TOL + + @pytest.mark.L0 + def test_bottom_right_causal_cross_seqlen(self): + torch.manual_seed(0) + B, H, Sq, Skv, D = 2, 8, 128, 512, 128 + q, k, v = bshd(B, H, Sq, D), bshd(B, H, Skv, D), bshd(B, H, Skv, D) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, causal_bottom_right=True, return_lse=False) + ref = ref_attention(q, k, v, scale, is_causal=True, bottom_right=True) + assert (o.float() - ref).abs().max().item() < TOL + + @pytest.mark.L0 + def test_gqa_hk_ne_hv(self): + """cuDNN supports h_k != h_v (each dividing h_q) — K and V carry + independent head counts.""" + torch.manual_seed(0) + B, S, D = 2, 256, 128 + q, k, v = bshd(B, 32, S, D), bshd(B, 8, S, D), bshd(B, 4, S, D) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=False) + kx = k.repeat_interleave(4, dim=1) # expand both to H_q for the reference + vx = v.repeat_interleave(8, dim=1) + ref = ref_attention(q, kx, vx, scale, is_causal=True) + assert (o.float() - ref).abs().max().item() < TOL + + @pytest.mark.L0 + @pytest.mark.parametrize("permute", [(0, 1, 2, 3), (0, 2, 1, 3), (1, 2, 0, 3), (2, 1, 0, 3)]) + def test_output_adopts_query_layout(self, permute): + """O is allocated in Q's dim-permutation (any B/H/S order, D innermost) + — the aten contract test_cudnn_attention_preserves_query_layout relies + on this.""" + torch.manual_seed(0) + BHSD = (2, 8, 256, 64) + shape = tuple(BHSD[i] for i in permute) + reverse = [permute.index(i) for i in range(4)] + q = torch.randn(*shape, dtype=torch.bfloat16, device="cuda").permute(reverse) + k = torch.randn(*shape, dtype=torch.bfloat16, device="cuda").permute(reverse) + v = torch.randn(*shape, dtype=torch.bfloat16, device="cuda").permute(reverse) + scale = BHSD[3] ** -0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=False) + assert o.permute(permute).is_contiguous(), f"O stride {o.stride()} does not follow Q layout {q.stride()}" + ref = ref_attention(q, k, v, scale, is_causal=True) + assert (o.float() - ref).abs().max().item() < TOL + + @pytest.mark.L0 + @_SINKS_UNSUPPORTED + def test_sinks_with_window(self): + torch.manual_seed(0) + B, H, S, D = 2, 8, 512, 128 + q, k, v = bshd(B, H, S, D), bshd(B, H, S, D), bshd(B, H, S, D) + sinks = torch.randn(H, device="cuda", dtype=torch.float32) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, window_left=256, sinks=sinks, return_lse=False) + ref = ref_attention(q, k, v, scale, is_causal=True, window_left=256, sinks=sinks) + assert (o.float() - ref).abs().max().item() < TOL + + +class TestOpContract: + @pytest.mark.L0 + def test_opcheck(self): + """torch.library.opcheck: fake-vs-real metadata agreement, schema + round-trip, and autograd registration — including dynamic-shape AOT + dispatch (the torch.compile contract).""" + torch.manual_seed(0) + q = torch.randn(2, 128, 8, 64, dtype=torch.bfloat16, device="cuda").transpose(1, 2) + torch.library.opcheck(torch.ops.cudnn.sdpa_fwd, (q, q.clone(), q.clone(), 0.125), dict(is_causal=True, return_lse=True)) + + lens = torch.tensor([100, 156], device="cuda") + cu = torch.nn.functional.pad(lens.cumsum(0), (1, 0)).to(torch.int32) + T, mx = int(cu[-1]), int(lens.max()) + q, k, v = (torch.randn(T, 8, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) for _ in range(3)) + torch.library.opcheck( + torch.ops.cudnn.sdpa_fwd, + (q, k, v, 0.125), + dict(is_causal=True, cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=True), + ) + + +class TestSdpaVarlen: + """THD (packed varlen) forward + backward through the ops directly.""" + + def _make(self, lens, Hq, Hkv, D, dtype=torch.bfloat16): + lens_t = torch.tensor(lens, device="cuda") + cu = torch.nn.functional.pad(lens_t.cumsum(0), (1, 0)).to(torch.int32) + T, mx = int(cu[-1]), int(lens_t.max()) + q = torch.randn(T, Hq, D, dtype=dtype, device="cuda") + k = torch.randn(T, Hkv, D, dtype=dtype, device="cuda") + v = torch.randn(T, Hkv, D, dtype=dtype, device="cuda") + return q, k, v, cu, T, mx + + def _ref(self, q, k, v, cu, is_causal, grad=None): + """Per-sequence fp32 reference; returns (out, lse, dq, dk, dv) — the + packed (T, Hq) log-sum-exp always, grads None unless grad given.""" + qr, kr, vr = (t.detach().float().requires_grad_(grad is not None) for t in (q, k, v)) + Hq, Hkv = q.shape[1], k.shape[1] + outs, lses = [], [] + for i in range(cu.numel() - 1): + a, b = int(cu[i]), int(cu[i + 1]) + qi, ki, vi = (t[a:b].transpose(0, 1).unsqueeze(0) for t in (qr, kr, vr)) + if Hq != Hkv: + ki = ki.repeat_interleave(Hq // Hkv, dim=1) + vi = vi.repeat_interleave(Hq // Hkv, dim=1) + s = torch.einsum("bhqd,bhkd->bhqk", qi, ki) * q.shape[-1] ** -0.5 + if is_causal: + Sq = qi.shape[2] + m = torch.ones(Sq, Sq, dtype=torch.bool, device=q.device).triu(1) + s = s.masked_fill(m, float("-inf")) + outs.append(torch.einsum("bhqk,bhkd->bhqd", torch.softmax(s, dim=-1), vi)[0].transpose(0, 1)) + lses.append(torch.logsumexp(s.detach(), dim=-1)[0].transpose(0, 1)) # (b-a, Hq) + out = torch.cat(outs) + lse = torch.cat(lses) + if grad is None: + return out, lse, None, None, None + out.backward(grad.float()) + return out, lse, qr.grad, kr.grad, vr.grad + + @pytest.mark.L0 + @pytest.mark.parametrize("lens", [[333, 128, 512, 47], [256, 384]]) + def test_thd_forward(self, lens): + torch.manual_seed(0) + H, D = 8, 128 + q, k, v, cu, T, mx = self._make(lens, H, H, D) + o, lse = torch.ops.cudnn.sdpa_fwd( + q, k, v, D**-0.5, is_causal=True, + cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=True, + ) # fmt: skip + ref, ref_lse, _, _, _ = self._ref(q, k, v, cu, is_causal=True) + assert (o.float() - ref).abs().max().item() < TOL + assert lse.shape == (T, H, 1) and lse.dtype == torch.float32 + assert (lse[:, :, 0] - ref_lse).abs().max().item() < TOL + + @pytest.mark.L0 + @pytest.mark.parametrize("gqa", [False, True]) + def test_thd_forward_backward(self, gqa): + torch.manual_seed(0) + lens, D = [200, 312, 96], 128 + Hq, Hkv = (16, 4) if gqa else (8, 8) + q, k, v, cu, _, mx = self._make(lens, Hq, Hkv, D) + scale = D**-0.5 + 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, + ) # fmt: skip + grad = torch.randn_like(o) + + # bwd takes PADDED (B, H, maxS, 1) LSE (backend restriction: bprop THD + # rejects ragged LSE on SM8X/SM12X) — scatter the packed TH1 stats. + B = cu.numel() - 1 + lse_padded = torch.zeros(B, Hq, mx, 1, dtype=torch.float32, device="cuda") + for i in range(B): + a, b = int(cu[i]), int(cu[i + 1]) + lse_padded[i, :, : b - a, 0] = lse[a:b, :, 0].transpose(0, 1) + + dq, dk, dv = torch.ops.cudnn.sdpa_bwd( + grad, q, k, v, o, lse_padded, scale, is_causal=True, + cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, + ) # fmt: skip + + ref, _, rdq, rdk, rdv = self._ref(q, k, v, cu, is_causal=True, grad=grad) + group = Hq // Hkv + assert (o.float() - ref).abs().max().item() < TOL + assert (dq.float() - rdq).abs().max().item() < TOL + # dk/dv accumulate GQA groups in bf16 — error grows ~sqrt(group) + assert (dk.float() - rdk).abs().max().item() < TOL * group**0.5 + assert (dv.float() - rdv).abs().max().item() < TOL * group**0.5 + + @pytest.mark.L0 + def test_thd_autograd(self): + """sdpa_fwd is differentiable end to end on the varlen path: the + registered autograd glue converts the packed TH1 stats to the padded + LSE layout and routes grads through cudnn::sdpa_bwd.""" + torch.manual_seed(0) + lens, H, D = [200, 312, 96], 8, 128 + q, k, v, cu, _, mx = self._make(lens, H, H, D) + q.requires_grad_(True) + k.requires_grad_(True) + v.requires_grad_(True) + o, _ = torch.ops.cudnn.sdpa_fwd( + q, k, v, D**-0.5, is_causal=True, + cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=True, + ) # fmt: skip + assert o.grad_fn is not None, "sdpa_fwd output is detached from autograd" + grad = torch.randn_like(o) + o.backward(grad) + + ref, _, rdq, rdk, rdv = self._ref(q, k, v, cu, is_causal=True, grad=grad) + assert (o.float() - ref).abs().max().item() < TOL + assert (q.grad.float() - rdq).abs().max().item() < TOL + assert (k.grad.float() - rdk).abs().max().item() < TOL + assert (v.grad.float() - rdv).abs().max().item() < TOL + + @pytest.mark.L0 + def test_thd_kv_packed_views(self): + """K/V as views of a kv-interleaved [T, 2, H, D] buffer (token stride + 2*H*D) — the layout torch.nn.attention.varlen users produce. Must be + served correctly (by whichever engine the router picks) or declined.""" + torch.manual_seed(0) + H, D = 8, 128 + q, _, _, cu, T, mx = self._make([333, 128, 512, 47], H, H, D) + kv = torch.randn(T, 2, H, D, dtype=torch.bfloat16, device="cuda") + k, v = kv[:, 0], kv[:, 1] + o, _ = torch.ops.cudnn.sdpa_fwd( + q, k, v, D**-0.5, is_causal=True, + cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=False, + ) # fmt: skip + ref, _, _, _, _ = self._ref(q, k, v, cu, is_causal=True) + assert (o.float() - ref).abs().max().item() < TOL diff --git a/test/python/test_cudnn_sdpa_op.py b/test/python/test_cudnn_sdpa_op.py deleted file mode 100644 index 138946225..000000000 --- a/test/python/test_cudnn_sdpa_op.py +++ /dev/null @@ -1,593 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Tests for the cuDNN SDPA PyTorch custom operator (cudnn.experimental.ops.sdpa). - -Each test runs both forward and backward, verifying output shapes, dtypes, -numerical correctness against a PyTorch reference for both O and dQ/dK/dV. - -All tensors use BHSD layout (batch, num_heads, seq_len, head_dim). -""" - -import math - -import pytest -import torch - -import cudnn -from cudnn.experimental.ops.sdpa import ( - _fprop_cache, - scaled_dot_product_attention, -) - -# --------------------------------------------------------------------------- -# PyTorch reference implementation (differentiable) -# --------------------------------------------------------------------------- - - -def sdpa_reference_fwd_bwd( - q, - k, - v, - attn_scale=None, - is_causal=False, - diagonal_alignment=0, - left_bound=-1, - right_bound=-1, - seq_len_q=None, - seq_len_kv=None, -): - """ - Pure-PyTorch differentiable SDPA reference in BHSD layout. - - Runs forward and backward (via .sum().backward()) and returns (o, dq, dk, dv). - All computation in float32 for numerical stability. - - Args: - q: (B, H_q, S_q, D_qk) — will be cloned with requires_grad - k: (B, H_k, S_kv, D_qk) - v: (B, H_v, S_kv, D_v) - Returns: - (o, dq, dk, dv) all in the original dtype - """ - if attn_scale is None: - attn_scale = 1.0 / math.sqrt(q.shape[-1]) - - # Clone to float32 with grad tracking - q_ref = q.detach().float().requires_grad_(True) - k_ref = k.detach().float().requires_grad_(True) - v_ref = v.detach().float().requires_grad_(True) - - B, H_q, S_q, D_qk = q_ref.shape - _, H_k, S_kv, _ = k_ref.shape - _, H_v, _, D_v = v_ref.shape - - q_t = q_ref - k_t = k_ref - v_t = v_ref - - # Expand for GQA/MQA - if H_q != H_k: - assert H_q % H_k == 0 - k_t = k_t.unsqueeze(2).expand(-1, -1, H_q // H_k, -1, -1).reshape(B, H_q, S_kv, D_qk) - if H_q != H_v: - assert H_q % H_v == 0 - v_t = v_t.unsqueeze(2).expand(-1, -1, H_q // H_v, -1, -1).reshape(B, H_q, S_kv, D_v) - - # Attention scores - s = torch.einsum("bhqd,bhkd->bhqk", q_t, k_t) * attn_scale - - # Causal / sliding window mask - rb = right_bound if right_bound >= 0 else None - lb = left_bound if left_bound >= 0 else None - if is_causal and rb is None: - rb = 0 - - if rb is not None: - if diagonal_alignment == 0: # TOP_LEFT - causal_mask = torch.ones(S_q, S_kv, dtype=torch.bool, device=q.device) - causal_mask.triu_(diagonal=1 + rb) - else: # BOTTOM_RIGHT - if seq_len_q is not None and seq_len_kv is not None: - causal_mask = torch.ones(B, 1, S_q, S_kv, dtype=torch.bool, device=q.device) - sl_q = seq_len_q.flatten() - sl_kv = seq_len_kv.flatten() - for i in range(B): - causal_mask[i, :, :, :].triu_(diagonal=int(sl_kv[i]) - int(sl_q[i]) + 1 + rb) - else: - causal_mask = torch.ones(S_q, S_kv, dtype=torch.bool, device=q.device) - causal_mask.triu_(diagonal=S_kv - S_q + 1 + rb) - s = s.masked_fill(causal_mask, float("-inf")) - - if lb is not None: - if diagonal_alignment == 0: # TOP_LEFT - swa_mask = torch.ones(S_q, S_kv, dtype=torch.bool, device=q.device) - swa_mask.tril_(diagonal=-1 * lb) - else: # BOTTOM_RIGHT - if seq_len_q is not None and seq_len_kv is not None: - swa_mask = torch.ones(B, 1, S_q, S_kv, dtype=torch.bool, device=q.device) - sl_q = seq_len_q.flatten() - sl_kv = seq_len_kv.flatten() - for i in range(B): - swa_mask[i, :, :, :].tril_(diagonal=int(sl_kv[i]) - int(sl_q[i]) - lb) - else: - swa_mask = torch.ones(S_q, S_kv, dtype=torch.bool, device=q.device) - swa_mask.tril_(diagonal=-1 * lb + (S_kv - S_q)) - s = s.masked_fill(swa_mask, float("-inf")) - - # Padding mask on scores - if seq_len_kv is not None: - sl_kv = seq_len_kv.flatten() - s_mask = torch.zeros(B, 1, S_q, S_kv, dtype=torch.bool, device=q.device) - for i in range(B): - s_mask[i, :, :, sl_kv[i] :] = True - s = s.masked_fill(s_mask, float("-inf")) - - p = torch.softmax(s, dim=-1) - - # Padding mask on probabilities - if seq_len_q is not None: - sl_q = seq_len_q.flatten() - p_mask = torch.zeros(B, 1, S_q, S_kv, dtype=torch.bool, device=q.device) - for i in range(B): - p_mask[i, :, sl_q[i] :, :] = True - p = p.masked_fill(p_mask, 0.0) - - o = torch.einsum("bhqk,bhkd->bhqd", p, v_t) - - # Backward - o.sum().backward() - - return ( - o.detach().to(q.dtype), - q_ref.grad.detach().to(q.dtype), - k_ref.grad.detach().to(k.dtype), - v_ref.grad.detach().to(v.dtype), - ) - - -# cuDNN 9.23.0 added d=256 SDPA fprop and bprop support in the backend; the -# op no longer carries an OSS fallback for older backends. -_D256_BACKEND_MIN_VERSION = 92300 - - -def _skip_if_unsupported_d256(D): - if D != 256: - return - if cudnn.backend_version() < _D256_BACKEND_MIN_VERSION: - pytest.skip(f"d=256 SDPA requires cuDNN backend >= {_D256_BACKEND_MIN_VERSION}") - major, _ = torch.cuda.get_device_capability() - if major < 9: - pytest.skip("d=256 backward path requires SM90+") - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -class TestCudnnSdpa: - """Combined forward + backward tests with numerical gradient verification.""" - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128, 256]) - def test_basic(self, D): - """Basic forward + backward, no masking.""" - _skip_if_unsupported_d256(D) - B, H, S = 2, 8, 128 - - 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) - - o = scaled_dot_product_attention(q, k, v) - assert o.shape == (B, H, S, D) - assert o.dtype == torch.float16 - - loss = o.sum() - loss.backward() - - # Reference - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd(q, k, v) - - torch.testing.assert_close(o.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(q.grad.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_basic_bf16(self, D): - """BFloat16 forward + backward.""" - B, H, S = 2, 4, 64 - - q = torch.randn(B, H, S, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) - k = torch.randn(B, H, S, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) - v = torch.randn(B, H, S, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) - - o = scaled_dot_product_attention(q, k, v) - assert o.shape == (B, H, S, D) - assert o.dtype == torch.bfloat16 - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd(q, k, v) - - torch.testing.assert_close(o.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(q.grad.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_causal_top_left(self, D): - """Causal mask with TOP_LEFT alignment, forward + backward.""" - B, H, S = 2, 4, 128 - - 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) - - o = scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True) - assert o.shape == (B, H, S, D) - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd(q, k, v, is_causal=True, diagonal_alignment=0) - - torch.testing.assert_close(o.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(q.grad.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_causal_bottom_right_with_padding(self, D): - """BOTTOM_RIGHT causal with variable sequence lengths, forward + backward.""" - B, H, S = 2, 4, 128 - - 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) - - seq_len_q = torch.tensor([[64], [96]], dtype=torch.int32, device="cuda").reshape(B, 1, 1, 1) - seq_len_kv = torch.tensor([[80], [128]], dtype=torch.int32, device="cuda").reshape(B, 1, 1, 1) - - o = scaled_dot_product_attention( - q, - k, - v, - is_causal=True, - enable_gqa=True, - diagonal_alignment=1, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - ) - assert o.shape == (B, H, S, D) - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd( - q, - k, - v, - is_causal=True, - diagonal_alignment=1, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - ) - - # Zero out padded regions for comparison (seq dim is index 2 in BHSD) - o_cmp = o.detach().clone() - dq_cmp = q.grad.detach().clone() - for i in range(B): - m = seq_len_q[i].item() - o_cmp[i, :, m:, :] = 0 - o_ref[i, :, m:, :] = 0 - dq_cmp[i, :, m:, :] = 0 - dq_ref[i, :, m:, :] = 0 - - torch.testing.assert_close(o_cmp.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(dq_cmp.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_sliding_window(self, D): - """Sliding window attention, forward + backward.""" - B, H, S = 2, 4, 256 - - 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) - - o = scaled_dot_product_attention( - q, - k, - v, - is_causal=True, - enable_gqa=True, - diagonal_alignment=0, - left_bound=32, - right_bound=0, - ) - assert o.shape == (B, H, S, D) - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd( - q, - k, - v, - diagonal_alignment=0, - left_bound=32, - right_bound=0, - ) - - torch.testing.assert_close(o.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(q.grad.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_variable_sequence_lengths(self, D): - """Padding mask with actual sequence lengths, forward + backward.""" - B, H, S = 2, 4, 128 - - 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) - - seq_len_q = torch.tensor([[64], [100]], dtype=torch.int32, device="cuda").reshape(B, 1, 1, 1) - seq_len_kv = torch.tensor([[80], [128]], dtype=torch.int32, device="cuda").reshape(B, 1, 1, 1) - - o = scaled_dot_product_attention( - q, - k, - v, - enable_gqa=True, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - ) - assert o.shape == (B, H, S, D) - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd( - q, - k, - v, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - ) - - # Zero out padded regions for comparison (seq dim is index 2 in BHSD) - o_cmp = o.detach().clone() - dq_cmp = q.grad.detach().clone() - for i in range(B): - m = seq_len_q[i].item() - o_cmp[i, :, m:, :] = 0 - o_ref[i, :, m:, :] = 0 - dq_cmp[i, :, m:, :] = 0 - dq_ref[i, :, m:, :] = 0 - - torch.testing.assert_close(o_cmp.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(dq_cmp.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - -class TestCudnnSdpaGQA: - """Grouped Query Attention tests.""" - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128]) - def test_gqa(self, D): - """GQA forward + backward with H_q > H_k = H_v.""" - B, S = 2, 64 - H_q, H_k, H_v = 8, 2, 2 - - q = torch.randn(B, H_q, S, D, dtype=torch.float16, device="cuda", requires_grad=True) - k = torch.randn(B, H_k, S, D, dtype=torch.float16, device="cuda", requires_grad=True) - v = torch.randn(B, H_v, S, D, dtype=torch.float16, device="cuda", requires_grad=True) - - o = scaled_dot_product_attention(q, k, v, enable_gqa=True) - assert o.shape == (B, H_q, S, D) - - loss = o.sum() - loss.backward() - - o_ref, dq_ref, dk_ref, dv_ref = sdpa_reference_fwd_bwd(q, k, v) - - torch.testing.assert_close(o.float(), o_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(q.grad.float(), dq_ref.float(), atol=2e-2, rtol=2e-2) - # GQA: dK/dV reference grads are summed across the GQA groups by autograd, - # so shapes match k/v directly - torch.testing.assert_close(k.grad.float(), dk_ref.float(), atol=2e-2, rtol=2e-2) - torch.testing.assert_close(v.grad.float(), dv_ref.float(), atol=2e-2, rtol=2e-2) - - -class TestCudnnSdpaCaching: - """Graph caching tests.""" - - @pytest.mark.L0 - def test_fprop_cache_reuse(self): - """Same config should reuse the cached graph.""" - B, H, S, D = 2, 4, 64, 64 - - q = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - k = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - v = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - - initial_cache_size = len(_fprop_cache) - - scaled_dot_product_attention(q, k, v) - after_first = len(_fprop_cache) - assert after_first == initial_cache_size + 1 - - q2 = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - k2 = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - v2 = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - scaled_dot_product_attention(q2, k2, v2) - after_second = len(_fprop_cache) - assert after_second == after_first, "Cache should be reused for same config" - - @pytest.mark.L0 - def test_different_shapes_create_new_entry(self): - """Different shapes should create a new cache entry.""" - D = 64 - - q1 = torch.randn(1, 4, 32, D, dtype=torch.float16, device="cuda") - k1 = torch.randn(1, 4, 32, D, dtype=torch.float16, device="cuda") - v1 = torch.randn(1, 4, 32, D, dtype=torch.float16, device="cuda") - - q2 = torch.randn(1, 4, 64, D, dtype=torch.float16, device="cuda") - k2 = torch.randn(1, 4, 64, D, dtype=torch.float16, device="cuda") - v2 = torch.randn(1, 4, 64, D, dtype=torch.float16, device="cuda") - - initial = len(_fprop_cache) - scaled_dot_product_attention(q1, k1, v1) - scaled_dot_product_attention(q2, k2, v2) - assert len(_fprop_cache) == initial + 2, "Different shapes should create exactly 2 cache entries" - - -class TestCudnnSdpaAPIValidation: - """API validation tests.""" - - @pytest.mark.L0 - def test_attn_mask_not_supported(self): - """attn_mask should raise NotImplementedError.""" - q = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - k = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - v = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - mask = torch.ones(1, 1, 16, 16, dtype=torch.float16, device="cuda") - - with pytest.raises(NotImplementedError, match="attn_mask"): - scaled_dot_product_attention(q, k, v, attn_mask=mask) - - @pytest.mark.L0 - def test_dropout_not_supported(self): - """dropout_p > 0 should raise NotImplementedError.""" - q = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - k = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - v = torch.randn(1, 4, 16, 64, dtype=torch.float16, device="cuda") - - with pytest.raises(NotImplementedError, match="dropout"): - scaled_dot_product_attention(q, k, v, dropout_p=0.1) - - @pytest.mark.L0 - def test_enable_gqa_validation(self): - """enable_gqa=False with mismatched heads should raise ValueError.""" - q = torch.randn(1, 8, 16, 64, dtype=torch.float16, device="cuda") - k = torch.randn(1, 2, 16, 64, dtype=torch.float16, device="cuda") - v = torch.randn(1, 2, 16, 64, dtype=torch.float16, device="cuda") - - with pytest.raises(ValueError, match="enable_gqa"): - scaled_dot_product_attention(q, k, v, enable_gqa=False) - - # Should work with enable_gqa=True - o = scaled_dot_product_attention(q, k, v, enable_gqa=True) - assert o.shape == (1, 8, 16, 64) - - @pytest.mark.L0 - def test_default_scale(self): - """Default scale should be 1/sqrt(D).""" - B, H, S, D = 1, 2, 32, 64 - - q = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - k = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - v = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - - # Default scale - o1 = scaled_dot_product_attention(q, k, v) - # Explicit scale = 1/sqrt(D) - o2 = scaled_dot_product_attention(q, k, v, scale=1.0 / math.sqrt(D)) - - torch.testing.assert_close(o1, o2) - - -class TestCudnnSdpaTorchCompile: - """Tests for torch.compile compatibility.""" - - @pytest.mark.L0 - def test_torch_compile_forward(self): - """torch.compile should work for forward pass.""" - B, H, S, D = 2, 4, 64, 128 - - compiled_sdpa = torch.compile(scaled_dot_product_attention, fullgraph=True) - - q = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - k = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - v = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda") - - # Eager - o_eager = scaled_dot_product_attention(q, k, v) - - # Compiled - o_compiled = compiled_sdpa(q, k, v) - - torch.testing.assert_close(o_eager, o_compiled) - - @pytest.mark.L0 - @pytest.mark.parametrize("D", [128, 256]) - def test_torch_compile_backward(self, D): - """torch.compile should work for forward + backward pass.""" - _skip_if_unsupported_d256(D) - B, H, S = 2, 4, 64 - - compiled_sdpa = torch.compile(scaled_dot_product_attention, fullgraph=True) - - 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) - - # Eager forward + backward - o_eager = scaled_dot_product_attention(q, k, v) - o_eager.sum().backward() - dq_eager = q.grad.clone() - dk_eager = k.grad.clone() - dv_eager = v.grad.clone() - - q.grad, k.grad, v.grad = None, None, None - - # Compiled forward + backward - o_compiled = compiled_sdpa(q, k, v) - o_compiled.sum().backward() - - torch.testing.assert_close(o_eager, o_compiled) - torch.testing.assert_close(dq_eager, q.grad) - torch.testing.assert_close(dk_eager, k.grad) - torch.testing.assert_close(dv_eager, v.grad) - - @pytest.mark.L0 - def test_torch_compile_causal(self): - """torch.compile with causal masking.""" - B, H, S, D = 2, 4, 128, 128 - - compiled_sdpa = torch.compile(scaled_dot_product_attention, fullgraph=True) - - 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) - - # Eager - o_eager = scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True) - o_eager.sum().backward() - dq_eager = q.grad.clone() - - q.grad, k.grad, v.grad = None, None, None - - # Compiled - o_compiled = compiled_sdpa(q, k, v, is_causal=True, enable_gqa=True) - o_compiled.sum().backward() - - torch.testing.assert_close(o_eager, o_compiled) - torch.testing.assert_close(dq_eager, q.grad)