From dcea8c82f6619636d848404f0b243c5602bdcec4 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 15:07:24 -0700 Subject: [PATCH 1/3] sdpa: add torch custom ops cudnn::sdpa_fwd / cudnn::sdpa_bwd Family-local torch contract for the features torch.nn.functional.scaled_dot_product_attention cannot express: attention sinks, sliding window, bottom-right causal, padded batches, and THD/varlen packing (FA-style (T,H,D) + cu_seqlens). The ops build pygraph sdpa/sdpa_backward nodes; the Router picks the serving plan (FROST OSS kernels or backend engines) per config. Contract highlights: - register_fake meta kernels mirror the real kernels' output strides; torch.library.opcheck passes on both paths, including dynamic-shape AOT dispatch (torch.compile contract), and is locked in by a test. - sdpa_fwd is differentiable on the varlen path via register_autograd; the glue converts packed TH1 stats to the padded LSE layout device-side (no host reads, capture/tracing-safe). Dense and sink backward raise NotImplementedError until their engine contracts land. - Thread-safe: thread-local cuDNN handles (a handle must not be used from two threads), serialized graph builds, bounded (FIFO) graph cache. - Validation: one io dtype per call, k/o/grad_out shape checks, int32 ragged-offset overflow guards, inert-flag rejection (causal_bottom_right without an active band), clone() not contiguous() for base-pointer realignment (contiguous() cannot fix a misaligned base). cudnn::sdpa_fwd / cudnn::sdpa_bwd are the canonical names; the experimental dense module's backward is renamed cudnn::sdpa_bwd_legacy so both modules coexist in one process until it is removed. Tests (14, L0): sinks/window/bottom-right/padded dense with LSE value checks against an fp32 reference; THD fwd/bwd incl. GQA, kv-interleaved views, end-to-end autograd; opcheck. Docs: docs/fe-oss-apis/sdpa-torch-ops.md. Co-Authored-By: Claude Fable 5 --- docs/fe-oss-apis/sdpa-torch-ops.md | 69 ++ python/cudnn/__init__.py | 1 + python/cudnn/experimental/ops/sdpa.py | 15 +- python/cudnn/sdpa/fwd/torch_op.py | 973 +++++++++++++++++++++++ test/python/test_cudnn_sdpa_torch_ops.py | 324 ++++++++ 5 files changed, 1376 insertions(+), 6 deletions(-) create mode 100644 docs/fe-oss-apis/sdpa-torch-ops.md create mode 100644 python/cudnn/sdpa/fwd/torch_op.py create mode 100644 test/python/test_cudnn_sdpa_torch_ops.py diff --git a/docs/fe-oss-apis/sdpa-torch-ops.md b/docs/fe-oss-apis/sdpa-torch-ops.md new file mode 100644 index 000000000..96ac8d181 --- /dev/null +++ b/docs/fe-oss-apis/sdpa-torch-ops.md @@ -0,0 +1,69 @@ +# SDPA torch custom ops: `cudnn::sdpa_fwd` / `cudnn::sdpa_bwd` + +PyTorch custom ops (`torch.library`) exposing the full cuDNN SDPA feature +surface — the features `torch.nn.functional.scaled_dot_product_attention`'s +aten contract cannot express: + +- **attention sinks** — per-Q-head logits folded into the softmax denominator +- **sliding window** — `window_left` (cuDNN convention: visible tokens + *including* self; FA2's `(w, 0)` maps to `window_left = w + 1`) +- **bottom-right causal alignment** — inference-style diagonals +- **padded batches** — per-batch actual lengths via `seq_len_q` / `seq_len_kv` +- **THD / varlen packing** — FlashAttention-style `(T, H, D)` + `cu_seqlens` + +The ops build cuDNN pygraph `sdpa` / `sdpa_backward` nodes; the engine Router +picks the best serving plan (FROST OSS kernels or cuDNN-backend engines) per +configuration. Graphs are cached per configuration (bounded, thread-safe; +cuDNN handles are thread-local). + +## Usage + +```python +import torch +import cudnn + +_ = cudnn.sdpa_torch # lazy public export: importing registers cudnn::sdpa_fwd / cudnn::sdpa_bwd + +# Dense BHSD with sinks + sliding window +o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, + window_left=128, sinks=sinks, return_lse=True) + +# THD / varlen (FA-style packed (T, H, D) + cu_seqlens), differentiable: +q, k, v = (t.requires_grad_(True) for t in (q_thd, k_thd, v_thd)) +o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, + cu_seqlens_q=cu, cu_seqlens_kv=cu, + max_seqlen_q=mx, max_seqlen_kv=mx, + return_lse=True) +o.backward(grad) # routes through cudnn::sdpa_bwd via register_autograd + +# Or through the python wrapper (same op underneath): +o = cudnn.sdpa_torch(q, k, v, is_causal=True, cu_seqlens_q=cu, cu_seqlens_kv=cu, + max_seqlen_q=mx, max_seqlen_kv=mx) +``` + +## Contracts and limits + +- Dense tensors are BHSD `(B, H, S, D)` (any strides; the graph declares the + actual layout). Varlen tensors are packed `(T, H, D)`; non-contiguous views + (e.g. K/V slices of a fused `(T, 2, H, D)` KV projection) are declared with + their true strides. On the varlen path, a non-dense innermost dim or a + misaligned base pointer is repaired by one copy (warned as slow path); the + dense path declares the given strides as-is. +- One io dtype per call (`fp16` or `bf16`); mixed-dtype inputs are rejected. +- `sdpa_bwd` serves the **THD/varlen** path. Dense backward and sink backward + (dSink) are follow-ups and raise `NotImplementedError`. It consumes a + **padded** `(B, H, max_seqlen_q, 1)` fp32 LSE (backend restriction: bprop + THD rejects ragged LSE on SM8X/SM12X). +- Autograd (`register_autograd`) requires `return_lse=True` on the forward; + the glue converts the packed TH1 stats to the padded layout device-side. +- Both ops ship `register_fake` meta kernels. `cudnn::sdpa_fwd` passes + `torch.library.opcheck` on the dense and varlen paths, including + dynamic-shape AOT dispatch (`torch.compile`-ready); the opcheck autograd + case exercises `cudnn::sdpa_bwd` through the registered backward. + +## Requirements + +- `nvidia-cudnn-frontend[cutedsl]`, cuDNN backend ≥ 9.6 (THD token-major + stats), sm80+. + +Tests: `test/python/test_cudnn_sdpa_torch_ops.py`. diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 43bf10692..63dd74c88 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/sdpa.py b/python/cudnn/experimental/ops/sdpa.py index faea0e475..57fae9598 100644 --- a/python/cudnn/experimental/ops/sdpa.py +++ b/python/cudnn/experimental/ops/sdpa.py @@ -488,7 +488,10 @@ def _build_bprop_graph( ) _lib.define( - "sdpa_bwd(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, " + # Renamed from cudnn::sdpa_bwd: the canonical name now belongs to the + # consolidated op family in cudnn.sdpa.fwd.torch_op (this experimental + # module is slated to fold into it). + "sdpa_bwd_legacy(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, " "float attn_scale, bool is_causal=False, int diagonal_alignment=0, " "int left_bound=-1, int right_bound=-1, " "Tensor? seq_len_q=None, Tensor? seq_len_kv=None, " @@ -625,7 +628,7 @@ def _sdpa_fake( return O, Stats -def _sdpa_bwd_impl( +def _sdpa_bwd_legacy_impl( dO: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -723,11 +726,11 @@ def _sdpa_bwd_impl( return dQ_gpu, dK_gpu, dV_gpu -_lib.impl("sdpa_bwd", _sdpa_bwd_impl, "CUDA") +_lib.impl("sdpa_bwd_legacy", _sdpa_bwd_legacy_impl, "CUDA") -@torch.library.register_fake("cudnn::sdpa_bwd") -def _sdpa_bwd_fake( +@torch.library.register_fake("cudnn::sdpa_bwd_legacy") +def _sdpa_bwd_legacy_fake( dO: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -800,7 +803,7 @@ def _sdpa_backward(ctx, dO, dStats): idx += 1 cum_kv = saved[idx] if ctx.has_cum_kv else None - dQ, dK, dV = torch.ops.cudnn.sdpa_bwd( + dQ, dK, dV = torch.ops.cudnn.sdpa_bwd_legacy( dO, q, k, diff --git a/python/cudnn/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py new file mode 100644 index 000000000..5e86e7813 --- /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, 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(q.shape[0], H_q, D_v, dtype=q.dtype, device=q.device) + stats = torch.empty(q.shape[0], 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, 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/test_cudnn_sdpa_torch_ops.py b/test/python/test_cudnn_sdpa_torch_ops.py new file mode 100644 index 000000000..81351c150 --- /dev/null +++ b/test/python/test_cudnn_sdpa_torch_ops.py @@ -0,0 +1,324 @@ +# 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) + +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 + @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 + 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 + @pytest.mark.xfail(reason="#613: the zero-host-read THD extents (#606) under-claim non-packed views on the FROST route", strict=False) + 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 From 4af94b048cc74f1d372afd257dcd0816d7e8b209 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 25 Aug 2026 23:17:24 -0700 Subject: [PATCH 2/3] test(sdpa): drop the #613 xfail from test_thd_kv_packed_views The span-derived THD capacity landed on develop in 3631ecb44, so K/V bound as views of a kv-interleaved [T, 2, H, D] buffer are served correctly rather than silently truncated. The test XPASSes; unmark it. Co-Authored-By: Claude Fable 5 --- test/python/test_cudnn_sdpa_torch_ops.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/python/test_cudnn_sdpa_torch_ops.py b/test/python/test_cudnn_sdpa_torch_ops.py index 81351c150..0d7f1f48f 100644 --- a/test/python/test_cudnn_sdpa_torch_ops.py +++ b/test/python/test_cudnn_sdpa_torch_ops.py @@ -306,7 +306,6 @@ def test_thd_autograd(self): assert (v.grad.float() - rdv).abs().max().item() < TOL @pytest.mark.L0 - @pytest.mark.xfail(reason="#613: the zero-host-read THD extents (#606) under-claim non-packed views on the FROST route", strict=False) 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 From 18b446e125c13044ec486370da7eebab896a8d8f Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 11 Aug 2026 10:29:24 -0700 Subject: [PATCH 3/3] torch: add the "CUDNN" torch.nn.attention provider (cudnn.torch) Graduates the PyTorch-integration bridge into the wheel. Importing cudnn.torch registers the "CUDNN" provider with torch.nn.attention's flash-attention implementation registry (PyTorch 2.13+, the same mechanism FA3/FA4 use); activation stays explicit: import cudnn.torch torch.nn.attention.activate_flash_attention_impl("CUDNN") After activation, F.scaled_dot_product_attention's cuDNN backend and torch.nn.attention.varlen.varlen_attn run on the cuDNN Python API via the cudnn::sdpa_fwd / sdpa_bwd custom ops, with hybrid fallback to the existing implementations for what the python path does not serve yet: dense bias/dropout forwards, every dense backward (until dense lands in cudnn::sdpa_bwd), and paged/asymmetric-window varlen (stock flash kernels). Dense forwards adopt the query's layout permutation. Unlike the in-tree cuDNN varlen branch, the provider serves GQA and causal sliding windows. restore_flash_attention_impl() reverts; torch < 2.13 uses install(). No dependency on the experimental ops module. The [cutedsl] import boundary is preserved (import cudnn stays torch-free; cudnn.torch is explicit opt-in). Stacked on the cudnn::sdpa_fwd / cudnn::sdpa_bwd ops PR. Co-Authored-By: Claude Fable 5 --- python/cudnn/__init__.py | 1 + python/cudnn/torch/__init__.py | 25 ++ python/cudnn/torch/sdpa_provider.py | 298 +++++++++++++++++++++++ test/python/test_cudnn_torch_provider.py | 230 +++++++++++++++++ 4 files changed, 554 insertions(+) create mode 100644 python/cudnn/torch/__init__.py create mode 100644 python/cudnn/torch/sdpa_provider.py create mode 100644 test/python/test_cudnn_torch_provider.py diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 63dd74c88..f4471d3ef 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -310,6 +310,7 @@ def _dlopen_cudnn(): _LAZY_OPTIONAL_IMPORTS = { "gnn": (".gnn", None), "sdpa_torch": (".sdpa.fwd.torch_op", "sdpa"), + "torch": (".torch", None), "BSA": (".block_sparse_attention", "BSA"), "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"), diff --git a/python/cudnn/torch/__init__.py b/python/cudnn/torch/__init__.py new file mode 100644 index 000000000..2f5871dec --- /dev/null +++ b/python/cudnn/torch/__init__.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PyTorch integration for the cuDNN frontend Python API. + +Importing this package registers the ``"CUDNN"`` provider with +``torch.nn.attention``'s flash-attention implementation registry +(PyTorch 2.13+, the same mechanism FA3/FA4 use). Registration is passive — +activation stays explicit: + + import cudnn.torch + torch.nn.attention.activate_flash_attention_impl("CUDNN") + +After activation, ``F.scaled_dot_product_attention`` under +``sdpa_kernel([SDPBackend.CUDNN_ATTENTION])`` and +``torch.nn.attention.varlen.varlen_attn`` run on the cuDNN *Python* API +(pygraph + engine Router: FROST OSS kernels or cuDNN-backend engines), with +hybrid fallback to the existing implementations for configurations the +python path does not serve yet. ``restore_flash_attention_impl()`` reverts. + +On torch < 2.13 (no registry), ``cudnn.torch.install()`` applies the +``F.scaled_dot_product_attention`` overrides directly. +""" + +from cudnn.torch.sdpa_provider import calls, install, served_plan_names # noqa: F401 diff --git a/python/cudnn/torch/sdpa_provider.py b/python/cudnn/torch/sdpa_provider.py new file mode 100644 index 000000000..3c2ab3b70 --- /dev/null +++ b/python/cudnn/torch/sdpa_provider.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The "CUDNN" torch.nn.attention provider: torch.sdpa on the cuDNN Python API. + +Routes PyTorch's cuDNN SDPA backend through the cudnn-frontend Python API +instead of the vendored C++ frontend. Overrides the CUDA dispatch-key kernels +of + + aten::_scaled_dot_product_cudnn_attention + aten::_scaled_dot_product_cudnn_attention_backward + +with Python implementations that call the cudnn-frontend Python API custom ops +(``torch.ops.cudnn.sdpa_fwd`` / ``sdpa_bwd`` from ``cudnn.sdpa.fwd.torch_op``). +The native Autograd wrapper of the aten op is untouched: it saves our forward's +outputs and routes grad through the (also overridden) aten backward, so vanilla + + with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): + F.scaled_dot_product_attention(q, k, v, is_causal=True) + +transparently runs on the Python API after ``install()``. + +Conveniently, aten's logsumexp convention for this op is (B, H, S, 1) float32 +(keepdim) — bit-identical in layout to cuDNN's Stats tensor, so tensors cross +the boundary with no reshape or copy. + +Hybrid fallback to the C++ worker ops (bit-exact with the shadowed native +kernel): attn_bias, dropout_p > 0, and — until dense backward lands in +cudnn::sdpa_bwd — every dense backward. The forward runs on the python API +either way, so training still exercises the python fwd path. +""" + +import math +from typing import Optional + +import torch + +# Importing this module registers torch.ops.cudnn.sdpa_fwd / sdpa_bwd. +import cudnn.sdpa.fwd.torch_op as _cudnn_ops # noqa: F401 + +_lib: Optional[torch.library.Library] = None + +# Observability for tests: how many aten calls the bridge served on the +# python API vs fell back (cpp = C++ worker ops; fa2 = flash varlen kernels). +calls = {"fwd": 0, "bwd": 0, "fwd_cpp": 0, "bwd_cpp": 0, "fwd_fa2": 0, "bwd_fa2": 0} + + +def _fwd( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_bias: Optional[torch.Tensor], + compute_log_sumexp: bool, + dropout_p: float = 0.0, + is_causal: bool = False, + return_debug_mask: bool = False, + *, + scale: Optional[float] = None, +): + if attn_bias is not None or dropout_p != 0.0 or return_debug_mask: + # Not wired in the python path yet — fall back to the C++ implementation + # through the (un-shadowed) worker op. Bit-exact with the native kernel. + calls["fwd_cpp"] += 1 + return torch.ops.aten._cudnn_attention_forward( + query, key, value, attn_bias, None, None, + query.size(-2), key.size(-2), compute_log_sumexp, + dropout_p, is_causal, return_debug_mask, scale=scale, + ) # fmt: skip + + calls["fwd"] += 1 + attn_scale = scale if scale is not None else 1.0 / math.sqrt(query.size(-1)) + + # Below-autograd call: runs the raw CUDA impl (graph-cached cuDNN execute). + o, stats = torch.ops.cudnn.sdpa_fwd(query, key, value, attn_scale, is_causal=is_causal, return_lse=compute_log_sumexp) + + # aten contract: (output, logsumexp(B,H,S,1) f32, cum_seq_q, cum_seq_k, + # max_q, max_k, philox_seed, philox_offset, debug_attn_mask) + philox_seed = torch.zeros((), dtype=torch.long, device=query.device) + philox_offset = torch.zeros((), dtype=torch.long, device=query.device) + return (o, stats, None, None, query.size(-2), key.size(-2), philox_seed, philox_offset, None) + + +def _bwd( + grad_out: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + out: torch.Tensor, + logsumexp: torch.Tensor, + philox_seed: torch.Tensor, + philox_offset: torch.Tensor, + attn_bias: Optional[torch.Tensor], + cum_seq_q: Optional[torch.Tensor], + cum_seq_k: Optional[torch.Tensor], + max_q: int, + max_k: int, + dropout_p: float, + is_causal: bool, + *, + scale: Optional[float] = None, +): + # Dense backward is not wired in cudnn::sdpa_bwd yet — route to the C++ + # worker op (bit-exact with the shadowed native kernel). The forward + # already ran on the python API, whose (B,H,S,1) fp32 logsumexp is the + # exact layout this worker consumes. + calls["bwd_cpp"] += 1 + return torch.ops.aten._cudnn_attention_backward( + grad_out, query, key, value, out, logsumexp, + philox_seed, philox_offset, attn_bias, cum_seq_q, cum_seq_k, + max_q, max_k, dropout_p, is_causal, scale=scale, + ) # fmt: skip + + +def install() -> None: + """Register the overrides (idempotent per process: last registration wins).""" + global _lib + if _lib is None: + _lib = torch.library.Library("aten", "IMPL") + _lib.impl("_scaled_dot_product_cudnn_attention", _fwd, "CUDA") + _lib.impl("_scaled_dot_product_cudnn_attention_backward", _bwd, "CUDA") + + +# --------------------------------------------------------------------------- +# varlen_attn (THD) via the cuDNN python API +# +# torch.nn.attention.varlen.varlen_attn routes to flash kernels in 2.13 (its +# in-tree cuDNN branch is dead: `_should_use_cudnn` is hardcoded False, and +# its `_cudnn_attention_backward` call predates the 2.13 schema). We hook one +# level up instead: override the `torch_attn::_varlen_attn{,_backward}` +# custom ops at the CUDA key. Their autograd wiring is untouched; unlike the +# dead branch we also serve GQA and causal sliding windows. +# --------------------------------------------------------------------------- + + +def _norm_window(window_size): + ws = list(window_size) if window_size is not None else [-1, -1] + if len(ws) != 2: + raise ValueError(f"window_size must have length 2, got {len(ws)}") + return ws + + +def _fa_window_left_to_cudnn(w: int) -> int: + """FA2 window_size=(w, 0) attends to [i-w, i] — w tokens back PLUS self. + cuDNN's diagonal_band_left_bound=lb masks j <= i-lb, i.e. lb visible + tokens including self. So lb = w + 1.""" + return w + 1 if w >= 0 else -1 + + +def _varlen_supported(ws, seqused_k=None, block_table=None, num_splits=None) -> bool: + """Configs the cudnn python varlen path serves today; everything else falls + back to the flash kernels (exactly what the stock op body runs).""" + if seqused_k is not None or block_table is not None: # paged KV not wired yet + return False + if num_splits is not None and num_splits != 1: + return False + # left-window + causal only; asymmetric/right bounds pending window_right in sdpa_*_ex + return ws[1] in (-1, 0) and not (ws[0] >= 0 and ws[1] != 0) + + +def _varlen_fwd_flash(query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, ws, seqused_k, block_table, num_splits): + calls["fwd_fa2"] += 1 + output, softmax_lse, _rng, _, _ = torch.ops.aten._flash_attention_forward( + query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, 0.0, is_causal, + return_debug_mask=False, scale=scale, + window_size_left=ws[0], window_size_right=ws[1], + seqused_k=seqused_k, block_table=block_table, num_splits=num_splits, + ) # fmt: skip + rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device) + return output, softmax_lse, rng_state + + +def _varlen_fwd(query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, is_causal=False, scale=None, window_size=None, enable_gqa=False, seqused_k=None, block_table=None, num_splits=None,): # fmt: skip + ws = _norm_window(window_size) + if not _varlen_supported(ws, seqused_k, block_table, num_splits): + return _varlen_fwd_flash(query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, ws, seqused_k, block_table, num_splits) + is_causal = is_causal or ws[1] == 0 + + calls["fwd"] += 1 + attn_scale = scale if scale is not None else query.shape[-1] ** -0.5 + o, stats = torch.ops.cudnn.sdpa_fwd( + query, key, value, attn_scale, + is_causal=is_causal, window_left=_fa_window_left_to_cudnn(ws[0]), + cu_seqlens_q=cu_seq_q, cu_seqlens_kv=cu_seq_k, + max_seqlen_q=max_q, max_seqlen_kv=max_k, return_lse=True, + ) # fmt: skip + lse = stats.squeeze(-1).transpose(0, 1).contiguous() # (T,H,1) -> (H,T) flash convention + rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device) + return o, lse, rng_state + + +def _varlen_fwd_out(out, query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, is_causal=False, scale=None, window_size=None, enable_gqa=False, seqused_k=None, block_table=None, num_splits=None,): # fmt: skip + """torch_attn::_varlen_attn_out — same as fwd but writes into `out`; returns lse.""" + ws = _norm_window(window_size) + if not _varlen_supported(ws, seqused_k, block_table, num_splits): + calls["fwd_fa2"] += 1 + return torch.ops.aten._flash_attention_forward_no_dropout_inplace( + out, query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, 0.0, is_causal, + False, scale=scale, window_size_left=ws[0], window_size_right=ws[1], + seqused_k=seqused_k, block_table=block_table, num_splits=num_splits, + ) # fmt: skip + o, lse, _rng = _varlen_fwd( + query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, + is_causal=is_causal, scale=scale, window_size=window_size, enable_gqa=enable_gqa, + seqused_k=seqused_k, block_table=block_table, num_splits=num_splits, + ) # fmt: skip + out.copy_(o) + return lse + + +def _varlen_bwd(grad_out, query, key, value, out, lse, cu_seq_q, cu_seq_k, max_q, max_k, is_causal, rng_state, scale=None, window_size=None,): # fmt: skip + ws = _norm_window(window_size) + if not _varlen_supported(ws): + calls["bwd_fa2"] += 1 # fwd for this config ran flash too (same predicate) + unused = torch.empty(0, device=query.device) + dq, dk, dv = torch.ops.aten._flash_attention_backward( + grad_out, query, key, value, out, lse, cu_seq_q, cu_seq_k, max_q, max_k, + 0.0, is_causal, rng_state, unused, scale=scale, + window_size_left=ws[0], window_size_right=ws[1], + ) # fmt: skip + return dq, dk, dv + is_causal = is_causal or ws[1] == 0 + + calls["bwd"] += 1 + attn_scale = scale if scale is not None else query.shape[-1] ** -0.5 + # (H, T) packed -> (B, H, max_q, 1) padded: the backend rejects ragged LSE + # for bprop THD on SM8X/SM12X, so the bwd op takes the padded layout. + B = cu_seq_q.numel() - 1 + H = query.shape[1] + lse_padded = torch.zeros(B, H, max_q, 1, dtype=torch.float32, device=lse.device) + for i in range(B): + a, b = int(cu_seq_q[i]), int(cu_seq_q[i + 1]) + lse_padded[i, :, : b - a, 0] = lse[:, a:b] + dq, dk, dv = torch.ops.cudnn.sdpa_bwd( + grad_out, query, key, value, out, lse_padded, attn_scale, + is_causal=is_causal, window_left=_fa_window_left_to_cudnn(ws[0]), + cu_seqlens_q=cu_seq_q, cu_seqlens_kv=cu_seq_k, + max_seqlen_q=max_q, max_seqlen_kv=max_k, + is_deterministic=torch.are_deterministic_algorithms_enabled(), + ) # fmt: skip + return dq, dk, dv + + +# --------------------------------------------------------------------------- +# torch.nn.attention flash-impl registry integration (PyTorch 2.13+) +# +# The same mechanism FA3/FA4 use: activation registers python overrides of +# existing CUDA kernels; restore drops the Library handles to deregister. +# +# import cudnn.torch # registers "CUDNN" (no activation) +# torch.nn.attention.activate_flash_attention_impl("CUDNN") +# --------------------------------------------------------------------------- + + +class _RegistryHandle: + def __init__(self, *libs: torch.library.Library): + self._libs = list(libs) + + def remove(self) -> None: + for lib in self._libs: + lib._destroy() + self._libs = [] + + +def _registry_register() -> _RegistryHandle: + import cudnn.sdpa.fwd.torch_op # noqa: F401 — registers cudnn::sdpa_fwd / sdpa_bwd + + lib = torch.library.Library("aten", "IMPL") + lib.impl("_scaled_dot_product_cudnn_attention", _fwd, "CUDA") + lib.impl("_scaled_dot_product_cudnn_attention_backward", _bwd, "CUDA") + vlib = torch.library.Library("torch_attn", "IMPL") + from torch.nn.attention import varlen as _varlen_mod # noqa: F401 — ensure torch_attn ops are defined + + vlib.impl("_varlen_attn", _varlen_fwd, "CUDA") + vlib.impl("_varlen_attn_out", _varlen_fwd_out, "CUDA") + vlib.impl("_varlen_attn_backward", _varlen_bwd, "CUDA") + return _RegistryHandle(lib, vlib) + + +def _register_with_torch() -> None: + try: + from torch.nn.attention import register_flash_attention_impl + except ImportError: + return # torch < 2.13: use install() directly + register_flash_attention_impl("CUDNN", register_fn=_registry_register) + + +_register_with_torch() + + +def served_plan_names() -> list: + """Which execution plan served each cached graph (debug/reporting).""" + names = [] + for graph, _ws in _cudnn_ops._graph_cache.values(): + try: + names.append(graph.get_plan_name_at_index(graph._plan_index)) + except Exception as e: # noqa: BLE001 + names.append(f"") + return names diff --git a/test/python/test_cudnn_torch_provider.py b/test/python/test_cudnn_torch_provider.py new file mode 100644 index 000000000..c4ae00efe --- /dev/null +++ b/test/python/test_cudnn_torch_provider.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the "CUDNN" torch.nn.attention provider (cudnn.torch). + +Two surfaces, both served by the cuDNN *Python* API after activation: + +- vanilla ``F.scaled_dot_product_attention`` under + ``sdpa_kernel([SDPBackend.CUDNN_ATTENTION])`` — fwd + autograd bwd, checked + against the fp32 math backend with the stock flash backend's error on the + same inputs as the rounding yardstick (same-precision kernels differ only + in accumulation order, so cuDNN passes within 3x of flash's error); +- ``torch.nn.attention.varlen.varlen_attn`` — fwd + bwd against a + per-sequence fp32 dense reference, including GQA, causal sliding windows + (which the in-tree cuDNN varlen branch rejects), and non-contiguous + kv-interleaved K/V views (the layout users produce by slicing a fused KV + projection). + +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 # noqa: F401 + +if not torch.cuda.is_available(): + pytest.skip("CUDA device required", allow_module_level=True) + +try: + from torch.nn.attention import SDPBackend, activate_flash_attention_impl, restore_flash_attention_impl, sdpa_kernel + from torch.nn.attention.varlen import AuxRequest, varlen_attn +except ImportError: + pytest.skip("torch.nn.attention flash-impl registry required (torch >= 2.13)", allow_module_level=True) + +import torch.nn.functional as F # noqa: E402 + +import cudnn.torch as provider # noqa: E402 (registers the "CUDNN" provider) + + +@pytest.fixture(autouse=True) +def _activate_provider(): + activate_flash_attention_impl("CUDNN") + yield + restore_flash_attention_impl() + + +def math_ref(q, k, v, is_causal, scale, enable_gqa=False): + """fp32 math-backend reference, differentiable.""" + q_, k_, v_ = (t.detach().float().requires_grad_(True) for t in (q, k, v)) + with sdpa_kernel([SDPBackend.MATH]): + o = F.scaled_dot_product_attention(q_, k_, v_, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa) + return o, q_, k_, v_ + + +DENSE_CASES = [ + # B, Hq, Hkv, Sq, Skv, D, dtype, is_causal, scale, enable_gqa, bshd + pytest.param(2, 8, 8, 512, 512, 128, torch.bfloat16, False, None, False, False, id="bf16-plain"), + pytest.param(2, 8, 8, 512, 512, 128, torch.bfloat16, True, None, False, False, id="bf16-causal"), + pytest.param(2, 8, 8, 512, 512, 128, torch.float16, True, None, False, False, id="fp16-causal"), + pytest.param(2, 16, 4, 512, 512, 128, torch.bfloat16, True, None, True, False, id="gqa"), + pytest.param(1, 8, 8, 1024, 2048, 64, torch.bfloat16, True, None, False, False, id="cross-seqlen-d64"), + pytest.param(2, 8, 8, 512, 512, 128, torch.float16, True, 0.05, False, False, id="custom-scale"), + pytest.param(2, 8, 8, 512, 512, 128, torch.bfloat16, True, None, False, True, id="bshd-projection"), + pytest.param(2, 16, 4, 1024, 1024, 128, torch.bfloat16, True, None, True, True, id="bshd-gqa"), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("B,Hq,Hkv,Sq,Skv,D,dtype,is_causal,scale,enable_gqa,bshd", DENSE_CASES) +def test_sdpa_dense_parity(B, Hq, Hkv, Sq, Skv, D, dtype, is_causal, scale, enable_gqa, bshd): + torch.manual_seed(0) + if bshd: # realistic transformer layout: (B,S,H,D) projections viewed as BHSD + q = torch.randn(B, Sq, Hq, D, dtype=dtype, device="cuda").transpose(1, 2).requires_grad_(True) + k = torch.randn(B, Skv, Hkv, D, dtype=dtype, device="cuda").transpose(1, 2).requires_grad_(True) + v = torch.randn(B, Skv, Hkv, D, dtype=dtype, device="cuda").transpose(1, 2).requires_grad_(True) + else: + q = torch.randn(B, Hq, Sq, D, dtype=dtype, device="cuda", requires_grad=True) + k = torch.randn(B, Hkv, Skv, D, dtype=dtype, device="cuda", requires_grad=True) + v = torch.randn(B, Hkv, Skv, D, dtype=dtype, device="cuda", requires_grad=True) + + # Dense: fwd runs on the python API; bwd routes to the C++ worker op + # (bit-exact hybrid) until dense backward lands in cudnn::sdpa_bwd. + fwd_before, bwd_before = provider.calls["fwd"], provider.calls["bwd_cpp"] + with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): + o = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa) + grad_o = torch.randn_like(o) + o.backward(grad_o) + assert provider.calls["fwd"] == fwd_before + 1, "provider fwd did not intercept" + assert provider.calls["bwd_cpp"] == bwd_before + 1, "provider bwd did not intercept" + assert o.grad_fn.__class__.__name__.startswith("ScaledDotProductCudnnAttention"), o.grad_fn + + o_ref, q_ref, k_ref, v_ref = math_ref(q, k, v, is_causal, scale, enable_gqa) + o_ref.backward(grad_o.float()) + + # Rounding yardstick: the stock flash backend's error on identical inputs. + qf, kf, vf = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) + with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]): + o_fa = F.scaled_dot_product_attention(qf, kf, vf, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa) + o_fa.backward(grad_o) + + def err(a, b): + return (a.float() - b).abs().max().item() + + floor = {torch.bfloat16: 1e-2, torch.float16: 2e-3}[dtype] + for name, ours, flash in ( + ("o", err(o, o_ref), err(o_fa, o_ref)), + ("dq", err(q.grad, q_ref.grad), err(qf.grad, q_ref.grad)), + ("dk", err(k.grad, k_ref.grad), err(kf.grad, k_ref.grad)), + ("dv", err(v.grad, v_ref.grad), err(vf.grad, v_ref.grad)), + ): + assert ours <= max(3 * flash, floor), f"{name}: cudnn err {ours:.4f} vs flash {flash:.4f}" + + +def ref_varlen(q, k, v, cu_q, cu_kv, is_causal, window_left=-1): + """Per-sequence fp32 dense reference; returns (out, q_ref, k_ref, v_ref).""" + qr, kr, vr = (t.detach().float().requires_grad_(True) for t in (q, k, v)) + Hq, Hkv = q.shape[1], k.shape[1] + outs = [] + for i in range(cu_q.numel() - 1): + aq, bq = int(cu_q[i]), int(cu_q[i + 1]) + ak, bk = int(cu_kv[i]), int(cu_kv[i + 1]) + qi = qr[aq:bq].transpose(0, 1).unsqueeze(0) + ki = kr[ak:bk].transpose(0, 1).unsqueeze(0) + vi = vr[ak:bk].transpose(0, 1).unsqueeze(0) + 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 + Sq, Skv = qi.shape[2], ki.shape[2] + ii = torch.arange(Sq, device=q.device).view(-1, 1) + jj = torch.arange(Skv, device=q.device).view(1, -1) + mask = torch.zeros(Sq, Skv, dtype=torch.bool, device=q.device) + if is_causal: + mask |= jj > ii + if window_left >= 0: + mask |= jj < (ii - window_left) # FA2: window (w, 0) attends [i-w, i] + s = s.masked_fill(mask, float("-inf")) + outs.append(torch.einsum("bhqk,bhkd->bhqd", torch.softmax(s, dim=-1), vi)[0].transpose(0, 1)) + out = torch.cat(outs) + return out, qr, kr, vr + + +VARLEN_CASES = [ + # Hq, Hkv, D, lens, window, enable_gqa, kv_packed + pytest.param(8, 8, 128, [333, 128, 512, 47], (-1, 0), False, False, id="causal"), + pytest.param(8, 8, 128, [256, 384], (-1, -1), False, False, id="non-causal"), + pytest.param(16, 4, 128, [200, 312, 96], (-1, 0), True, False, id="gqa"), + pytest.param(8, 8, 128, [400, 288], (128, 0), False, False, id="window-128"), + pytest.param(8, 8, 128, [400, 288], (4, 0), False, False, id="window-4-tight"), + pytest.param(8, 8, 64, [512, 512], (-1, 0), False, False, id="d64"), + pytest.param( + 8, + 8, + 128, + [333, 128, 512, 47], + (-1, 0), + False, + True, + id="kv-interleaved", + ), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("Hq,Hkv,D,lens,window,enable_gqa,kv_packed", VARLEN_CASES) +def test_varlen_attn(Hq, Hkv, D, lens, window, enable_gqa, kv_packed): + torch.manual_seed(0) + 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=torch.bfloat16, device="cuda", requires_grad=True) + if kv_packed: # non-contiguous k/v views of one buffer (token stride 2*H*D) + kv = torch.randn(T, 2, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) + k, v = kv[:, 0], kv[:, 1] + else: + k = torch.randn(T, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) + v = torch.randn(T, Hkv, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + fwd0, bwd0 = provider.calls["fwd"], provider.calls["bwd"] + out, lse = varlen_attn(q, k, v, cu, cu, mx, mx, window_size=window, enable_gqa=enable_gqa, return_aux=AuxRequest(lse=True)) + grad = torch.randn_like(out) + out.backward(grad) + assert provider.calls["fwd"] == fwd0 + 1 and provider.calls["bwd"] == bwd0 + 1, "provider did not intercept" + + is_causal = window[1] == 0 + ref, qr, kr, vr = ref_varlen(q, k, v, cu, cu, is_causal, window[0]) + ref.backward(grad.float()) + if kv_packed: + kv_grad = torch.stack([kr.grad, vr.grad], dim=1) + dk_err = dv_err = (kv.grad.float() - kv_grad).abs().max().item() + else: + dk_err = (k.grad.float() - kr.grad).abs().max().item() + dv_err = (v.grad.float() - vr.grad).abs().max().item() + + # dk/dv accumulate Hq/Hkv gradient groups in bf16 — error grows ~sqrt(group) + # (stock flash shows the same inflation on GQA). + group = Hq // Hkv + tol = {"o": 2.5e-2, "dq": 2.5e-2, "dk": 2.5e-2 * group**0.5, "dv": 2.5e-2 * group**0.5} + errs = { + "o": (out.float() - ref).abs().max().item(), + "dq": (q.grad.float() - qr.grad).abs().max().item(), + "dk": dk_err, + "dv": dv_err, + } + for name, e in errs.items(): + assert e < tol[name], f"{name}: err {e:.4f} tol {tol[name]:.4f}" + + +@pytest.mark.L0 +def test_d256_direct_aten_op(): + """d=256: torch's C++ fused_sdp_choice still gates cuDNN to head_dim<=128, + so F.sdpa cannot reach it — but the python path serves it through the aten + op directly (what a fixed selection gate would dispatch to).""" + torch.manual_seed(0) + q = torch.randn(2, 4, 384, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + k, v = torch.randn_like(q, requires_grad=True), torch.randn_like(q, requires_grad=True) + try: + out = torch.ops.aten._scaled_dot_product_cudnn_attention(q, k, v, None, True, 0.0, False) + except RuntimeError as e: + pytest.skip(f"no engine serves d=256 on this arch: {str(e).splitlines()[0][:80]}") + o = out[0] + o.backward(torch.ones_like(o)) + o_ref, q_ref, _, _ = math_ref(q, k, v, False, None) + o_ref.backward(torch.ones_like(o_ref)) + assert (o.float() - o_ref).abs().max().item() < 0.05 + assert (q.grad.float() - q_ref.grad).abs().max().item() < 0.5