From b79b8dc013ef0983a70d077e6135e1b18fd98365 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 11 Aug 2026 10:29:24 -0700 Subject: [PATCH 1/7] 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 | 11 + python/cudnn/torch/__init__.py | 25 ++ python/cudnn/torch/sdpa_provider.py | 298 ++++++++++++++++++++++++ test/python/sdpa/test_torch_provider.py | 230 ++++++++++++++++++ 4 files changed, 564 insertions(+) create mode 100644 python/cudnn/torch/__init__.py create mode 100644 python/cudnn/torch/sdpa_provider.py create mode 100644 test/python/sdpa/test_torch_provider.py diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 68ca74e51..e52d77cc5 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -426,6 +426,17 @@ def __getattr__(name: str) -> Any: globals()["jax"] = _jax return _jax + if name == "torch": + # `import cudnn; cudnn.torch.install()` works like `import cudnn.torch`, + # mirroring the `jax` branch above. Deferred so `import cudnn` never + # eagerly imports torch; the submodule raises its own descriptive error + # when torch (or the 2.13+ flash-impl registry) is unavailable — which + # is why this is NOT a _LAZY_OPTIONAL_IMPORTS entry: that path would + # blame the `[cutedsl]` extra for a missing framework. + _torch_mod = importlib.import_module(".torch", __name__) + globals()["torch"] = _torch_mod + return _torch_mod + if name == "fla": # `import cudnn; cudnn.fla.accelerate_fla()` works like `import cudnn.fla`. # Deferred so `import cudnn` never eagerly imports torch / the FLA shim. 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/sdpa/test_torch_provider.py b/test/python/sdpa/test_torch_provider.py new file mode 100644 index 000000000..c4ae00efe --- /dev/null +++ b/test/python/sdpa/test_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 From e2f10b63122f616cb387210fcb557cd86c5dc453 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 1 Sep 2026 00:21:19 -0700 Subject: [PATCH 2/7] fix(torch): drop the 2*B D2H syncs from the varlen backward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_varlen_bwd` repadded the packed LSE with for i in range(B): a, b = int(cu_seq_q[i]), int(cu_seq_q[i + 1]) Each `int(cu_seq_q[i])` on a CUDA tensor is a blocking `cudaMemcpy`, so the provider stalled the host 2*B times, serially, before the backward kernel was even launched — an async-launch API turned synchronous, and un-capturable under a CUDA graph. python/cudnn/AGENTS.md Rule 3 calls this out as a functional gap rather than a slow path. The irony is that the correct version already existed: the autograd glue in torch_op.py does the same conversion device-side with `searchsorted`, written that way to stay traceable under dynamic-shape AOT dispatch. The provider had a second, worse copy of the same logic. Extracted the good one as `thd_lse_to_padded(lse_th, cu_seqlens_q, max_seqlen_q)` and pointed both call sites at it, so there is one implementation. Note the idiom that hid this: AGENTS.md names `.item()` / `.tolist()` / `.cpu()` / `.to("cpu")` / `.numpy()`, and a grep for those over sdpa_provider.py comes back empty. `int(tensor[i])` is the same blocking copy wearing different clothes. Verified: the helper's output is bit-identical to the old loop (`torch.equal`), and `torch.cuda.set_sync_debug_mode("error")` reports no sync through the whole backward. New `test_varlen_backward_does_not_sync` locks it in — confirmed RED against the restored host loop, green with the helper. Provider 17 passed, ops 19 passed. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/torch_op.py | 35 +++++++++++++++++++------ python/cudnn/torch/sdpa_provider.py | 12 ++++----- test/python/sdpa/test_torch_provider.py | 32 ++++++++++++++++++++++ 3 files changed, 65 insertions(+), 14 deletions(-) diff --git a/python/cudnn/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py index 805b3fab0..a1b449010 100644 --- a/python/cudnn/sdpa/fwd/torch_op.py +++ b/python/cudnn/sdpa/fwd/torch_op.py @@ -871,6 +871,32 @@ def _sdpa_setup_context(ctx, inputs, output): ctx.mark_non_differentiable(output[1]) +def thd_lse_to_padded(lse_th: torch.Tensor, cu_seqlens_q: torch.Tensor, max_seqlen_q: int) -> torch.Tensor: + """Packed ``(T, H)`` log-sum-exp -> padded ``(B, H, max_seqlen_q, 1)``. + + ``cudnn::sdpa_bwd`` takes the padded layout on the THD path (the backend + rejects ragged LSE for bprop THD on SM8X/SM12X). Rows past each sequence's + length stay zero and are ignored. + + Entirely DEVICE-side, and deliberately so: the obvious + ``for i in range(B): int(cu[i])`` loop is ``2*B`` blocking D2H copies + before the backward kernel is even launched, which makes an async-launch + API synchronous and cannot be stream-captured (python/cudnn/AGENTS.md, + Rule 3). It also keeps the conversion traceable under dynamic-shape AOT + dispatch, where reading a cu value to host raises + GuardOnDataDependentSymNode. + """ + B = cu_seqlens_q.numel() - 1 + T, H = lse_th.shape + cu = cu_seqlens_q.long() + token = torch.arange(T, device=lse_th.device) + seq_of_token = torch.searchsorted(cu[1:], token, right=True) # t in [cu[i], cu[i+1]) -> i + pos_in_seq = token - cu[seq_of_token] + padded = torch.zeros(B, H, max_seqlen_q, 1, dtype=torch.float32, device=lse_th.device) + padded[seq_of_token, :, pos_in_seq, 0] = lse_th + return padded + + 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 @@ -888,14 +914,7 @@ def _sdpa_backward(ctx, grad_o, _grad_stats): # stats marked non-differentiable # 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] + lse_padded = thd_lse_to_padded(stats[:, :, 0], cu_q, ctx.max_seqlen_q) dq, dk, dv = torch.ops.cudnn.sdpa_bwd( grad_o, diff --git a/python/cudnn/torch/sdpa_provider.py b/python/cudnn/torch/sdpa_provider.py index 3c2ab3b70..7d0e33161 100644 --- a/python/cudnn/torch/sdpa_provider.py +++ b/python/cudnn/torch/sdpa_provider.py @@ -224,12 +224,12 @@ def _varlen_bwd(grad_out, query, key, value, out, lse, cu_seq_q, cu_seq_k, max_q 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] + # (H, T) -> (T, H) for the shared device-side repad. The naive + # `for i in range(B): int(cu_seq_q[i])` loop that used to live here was + # 2*B blocking D2H copies per backward call, before the kernel even + # launched — an async-launch API turned synchronous, and un-capturable + # (python/cudnn/AGENTS.md Rule 3). + lse_padded = _cudnn_ops.thd_lse_to_padded(lse.transpose(0, 1), cu_seq_q, max_q) 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]), diff --git a/test/python/sdpa/test_torch_provider.py b/test/python/sdpa/test_torch_provider.py index c4ae00efe..765355e9c 100644 --- a/test/python/sdpa/test_torch_provider.py +++ b/test/python/sdpa/test_torch_provider.py @@ -210,6 +210,38 @@ def test_varlen_attn(Hq, Hkv, D, lens, window, enable_gqa, kv_packed): assert e < tol[name], f"{name}: err {e:.4f} tol {tol[name]:.4f}" +@pytest.mark.L0 +def test_varlen_backward_does_not_sync(): + """The varlen backward must not read cu_seqlens to host. + + It used to repad the packed LSE with `for i in range(B): int(cu_seq_q[i])` + — 2*B blocking D2H copies before the kernel even launched, turning an + async-launch API synchronous and blocking stream capture + (python/cudnn/AGENTS.md Rule 3). The conversion is device-side now; + CUDA's sync-debug mode turns any regression into an error. + """ + torch.manual_seed(0) + H, D = 8, 128 + lens_t = torch.tensor([128, 96, 200], 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, k, v = (torch.randn(T, H, D, dtype=torch.bfloat16, device="cuda", requires_grad=True) for _ in range(3)) + + bwd0 = provider.calls["bwd"] + out = varlen_attn(q, k, v, cu, cu, mx, mx, window_size=(-1, 0)) + grad = torch.randn_like(out) + + # Any blocking D2H inside the backward raises here. + torch.cuda.set_sync_debug_mode("error") + try: + out.backward(grad) + finally: + torch.cuda.set_sync_debug_mode("default") + + assert provider.calls["bwd"] == bwd0 + 1, "provider did not serve the backward" + assert q.grad is not None and k.grad is not None and v.grad is not None + + @pytest.mark.L0 def test_d256_direct_aten_op(): """d=256: torch's C++ fused_sdp_choice still gates cuDNN to head_dim<=128, From bc17fbbfb8d3c125fbac6f0710ea229bcf9905ef Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 1 Sep 2026 00:28:30 -0700 Subject: [PATCH 3/7] docs(agents): make the no-D2H rule verifiable, not just greppable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rule 3 already names `int(tensor)` alongside `.item()` — the miss on the varlen backward was not a gap in the rule, it was a reviewer (me) grepping a subset of the spellings and concluding "clean". A list of spellings cannot be a detector. Adds two things learned from that bug: - **Prove it.** `torch.cuda.set_sync_debug_mode("error")` turns any blocking D2H into a RuntimeError; wrap the execute/backward in it and keep the assertion as a test. And check the test is RED against the old code first — a sync test never seen to fail asserts nothing. - **Suspect duplicated logic.** Every violation found so far has been a second copy of a conversion that was already device-side elsewhere. Extract the correct one instead of writing the obvious loop again. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 090756711..babedb9ff 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -92,6 +92,28 @@ neither names the thing that breaks it most directly: a device-to-host read. hand you — or assert in-kernel. Reading lengths back to decide whether to raise buys nothing: the Router had to choose an engine before any buffer existed. +- **Prove it; do not grep for it.** The list above is a reminder, not a + detector — the spellings are many (`int(cu[i])` on a CUDA tensor is a + blocking copy that a search for `.item()` will not find) and a reviewer who + greps a subset concludes "clean". Assert the property instead: + + ```python + torch.cuda.set_sync_debug_mode("error") # any blocking D2H now raises + try: + out.backward(grad) # or graph.execute(...) + finally: + torch.cuda.set_sync_debug_mode("default") + ``` + + Put that in a test (see `test_varlen_backward_does_not_sync`), and check the + test is RED against the old code before trusting it — a sync test that was + never seen to fail is asserting nothing. +- **Suspect duplicated logic first.** Every violation found so far has been a + *second* copy of a conversion that was already device-side somewhere else: + the packed-to-padded LSE repad existed in both `sdpa/fwd/torch_op.py` (with + `searchsorted`, device-side) and `torch/sdpa_provider.py` (a `for i in + range(B): int(cu[i])` loop). Extract the correct one and call it from both + rather than writing the obvious loop again. Known violations, all pre-existing and each needing a kernel-side change, so none is precedent: From c6e532959614db68650e28fc04226e8c247a1ae5 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 31 Aug 2026 23:16:27 -0700 Subject: [PATCH 4/7] feat(sdpa): dense BHSD backward in cudnn::sdpa_bwd MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `cudnn::sdpa_bwd` served only the packed THD/varlen path; dense `(B, H, S, D)` raised NotImplementedError. That left the provider's dense training half on the Python API — forward through `cudnn::sdpa_fwd`, backward falling back to the C++ worker `aten::_cudnn_attention_backward` (the provider counts it as `calls["bwd_cpp"]`). This wires the dense path so the whole step builds its graph in Python. `_build_bwd_graph` gains `is_thd`, mirroring the forward builder. Dense differs in three ways, all of them contract rather than kernel: - no ragged-offset tensors and no per-batch length operands — an unpadded dense batch has every sequence at its declared S; - no `max_total_seq_len_q/kv`: those size the ragged dq accumulator and the node rejects them on a non-ragged layout (#740); - dQ/dK/dV adopt the caller's Q/K/V layouts via `_like_layout_stride`, because autograd hands them straight back as `.grad` on the caller's parameters. THD keeps returning them packed. Stats need no conversion on this path: dense LSE already is `(B, H, S, 1)` fp32, which is exactly aten's logsumexp layout for `_scaled_dot_product_cudnn_attention`, so `register_autograd` forwards the forward's stats untouched instead of running the THD scatter. The fake kernel now mirrors the dense gradient strides (the inputs' permutation, not contiguous) — otherwise opcheck's stride assertions fail under AOT dispatch. Still raising: sink backward (dSink), and dense backward with per-batch lengths (the padded dense path). Tests: six new cases in `TestSdpaBwdDense` — causal/non-causal against an fp32 reference, GQA with h_kv < h_q, end-to-end autograd, a layout check that dQ/dK/dV come back in their input's permutation, and opcheck on the dense backward. They compare RELATIVE to each tensor's magnitude: GQA dK/dV sum over h_q/h_kv query heads, so a fixed absolute bound flags a correct result purely because the values got bigger (~0.5% relative on every tensor, but |dv| peaks near 9.6 at a group size of 4). Verified on SM100 (cuDNN 9.26.0.33): 25 passed, all 19 pre-existing THD tests unchanged. NOT yet verified against a FROST backward engine — those exist only for SM120 and SM80, and neither box is reachable from here; on SM100 the Router serves this through cuDNN-backend engines, which is still the Python API but not an OSS kernel. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/torch_op.py | 270 ++++++++++++++++++++++++----- test/python/sdpa/test_torch_ops.py | 97 +++++++++++ 2 files changed, 326 insertions(+), 41 deletions(-) diff --git a/python/cudnn/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py index a1b449010..0b476d59f 100644 --- a/python/cudnn/sdpa/fwd/torch_op.py +++ b/python/cudnn/sdpa/fwd/torch_op.py @@ -21,8 +21,9 @@ 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 +Backward contract: ``sdpa_bwd`` serves the THD/varlen path and the unpadded +dense BHSD path (sink backward, and dense backward with per-batch lengths, +are follow-ups and raise ``NotImplementedError``). On THD 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 @@ -545,8 +546,19 @@ def _build_bwd_graph( o_stride, stats_stride, is_deterministic: bool, + is_thd: bool, + dq_stride=None, + dk_stride=None, + dv_stride=None, ): - """THD/varlen backward graph (the only path the bwd op serves today).""" + """Backward graph for the packed THD/varlen path or the dense BHSD one. + + Dense differs from THD in three ways: no ragged-offset tensors, no + per-batch length operands (an unpadded dense batch has every sequence at + its declared S), and no ``max_total_seq_len_*`` — those size the ragged dq + accumulator and the node rejects them on a non-ragged layout. dQ/dK/dV + adopt the caller's Q/K/V layouts on the dense path (autograd hands these + straight back as ``.grad``), where THD always returns them packed.""" io_dtype = _TORCH_DTYPE_TO_CUDNN[dtype] g = cudnn.pygraph( handle=handle, @@ -564,20 +576,23 @@ def _build_bwd_graph( # 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) + seq_q_t = seq_kv_t = None + rdq = rdk = rdv = None + if 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) + 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 @@ -592,29 +607,33 @@ def _build_bwd_graph( dO=do_t, stats=stats_t, attn_scale=attn_scale, - use_padding_mask=True, + # Padding/lengths and the ragged dq-accumulator sizing are THD-only: + # max_total_seq_len_* is rejected on a non-ragged layout, and an + # unpadded dense batch has every sequence at its declared S. + use_padding_mask=is_thd, 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), + **({"max_total_seq_len_q": _round64(total_q), "max_total_seq_len_kv": _round64(total_kv)} if is_thd else {}), 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) + # THD gradients are OURS: always packed-contiguous, independent of the + # input views. Dense gradients adopt the caller's Q/K/V layouts — autograd + # hands them straight back as .grad, which should match the parameter. + if is_thd: + 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) + if is_thd: + dq_t.set_ragged_offset(rdq) + dk_t.set_ragged_offset(rdk) + dv_t.set_ragged_offset(rdv) g.validate() g.build_operation_graph() @@ -634,6 +653,136 @@ def _build_bwd_graph( ) +def _sdpa_bwd_dense( + grad_out, + q, + k, + v, + o, + lse, + attn_scale, + *, + is_causal, + causal_bottom_right, + window_left, + is_deterministic, +): + """Dense BHSD backward. + + The unpadded dense contract: every sequence is its declared S, so no + length operands and no ragged offsets. Stats arrive as ``(B, H, S)`` or + ``(B, H, S, 1)`` fp32 — which is exactly aten's logsumexp layout for + ``_scaled_dot_product_cudnn_attention``, so the provider hands ours + straight through. dQ/dK/dV adopt Q/K/V's layouts, since autograd returns + them as ``.grad`` on the caller's parameters. + """ + B, H_q, S_q, D_qk = q.shape + _, H_v, S_kv, D_v = v.shape + H_k = k.shape[1] + if k.shape != (B, H_k, S_kv, D_qk): + 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}") + if o.shape != (B, H_q, S_q, D_v) or grad_out.shape != o.shape: + raise ValueError(f"o {tuple(o.shape)} / grad_out {tuple(grad_out.shape)} must be (B={B}, H_q={H_q}, S_q={S_q}, D_v={D_v})") + _check_same_device(q, k=k, v=v, o=o, lse=lse, grad_out=grad_out) + + 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) + # dO must match O's layout (and be 16B-aligned) — 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) + + # Gradients adopt the corresponding input's layout; a broadcast/overlapping + # input has no usable gradient layout, so fall back to packed there. + dq_stride = _like_layout_stride((B, H_q, S_q, D_qk), q) + dk_stride = _like_layout_stride((B, H_k, S_kv, D_qk), k) + dv_stride = _like_layout_stride((B, H_v, S_kv, D_v), v) + stats_stride = (H_q * S_q, S_q, 1, 1) + + key = ( + "sdpa_bwd_dense", + 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()), + tuple(o.stride()), + dq_stride, + dk_stride, + dv_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=0, + total_kv=0, + 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, + is_thd=False, + dq_stride=dq_stride, + dk_stride=dk_stride, + dv_stride=dv_stride, + ), + ) + + dq = torch.empty_strided((B, H_q, S_q, D_qk), dq_stride, dtype=q.dtype, device=q.device) + dk = torch.empty_strided((B, H_k, S_kv, D_qk), dk_stride, dtype=q.dtype, device=q.device) + dv = torch.empty_strided((B, H_v, S_kv, D_v), dv_stride, 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, + } + g.execute(variant, workspace, handle=handle) + return dq, dk, dv + + def _sdpa_bwd_impl( grad_out: torch.Tensor, q: torch.Tensor, @@ -657,14 +806,31 @@ def _sdpa_bwd_impl( # 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}") + 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}") + elif q.ndim != 4: + raise ValueError(f"dense path expects BHSD tensors, got q.ndim={q.ndim}") _check_io_dtypes("sdpa_bwd", grad_out=grad_out, q=q, k=k, v=v, o=o) + if not is_thd: + return _sdpa_bwd_dense( + grad_out, + q, + k, + v, + o, + lse, + attn_scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_left=window_left, + is_deterministic=is_deterministic, + ) + B = cu_seqlens_q.numel() - 1 q = _normalize_thd(q, "q") k = _normalize_thd(k, "k") @@ -759,6 +925,7 @@ def _sdpa_bwd_impl( o_stride=o_stride, stats_stride=stats_stride, is_deterministic=is_deterministic, + is_thd=True, ), ) @@ -816,9 +983,16 @@ def _sdpa_bwd_fake( 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. + # The real kernel returns FRESH gradients in q.dtype, never views of the + # inputs, so empty_like would report strides that never materialize. THD + # gradients are packed-contiguous (q/k/v may be kv-interleaved views); + # dense gradients adopt each input's own dim-permutation, which the meta + # kernel must mirror exactly or opcheck's stride assertions fail. + if cu_seqlens_q is None: + dq = torch.empty_strided(q.shape, _like_layout_stride(tuple(q.shape), q), dtype=q.dtype, device=q.device) + dk = torch.empty_strided(k.shape, _like_layout_stride(tuple(k.shape), k), dtype=q.dtype, device=q.device) + dv = torch.empty_strided(v.shape, _like_layout_stride(tuple(v.shape), v), dtype=q.dtype, device=q.device) + return dq, dk, dv 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) @@ -901,15 +1075,29 @@ 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: + if ctx.has_seq_lens and cu_q is None: 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)") + if cu_q is None: + # Dense: stats already are (B, H, S, 1) fp32 — hand them straight on. + dq, dk, dv = torch.ops.cudnn.sdpa_bwd( + grad_o, + q, + k, + v, + o, + stats, + ctx.attn_scale, + is_causal=ctx.is_causal, + causal_bottom_right=ctx.causal_bottom_right, + window_left=ctx.window_left, + ) + return (dq, dk, dv) + (None,) * 12 + # 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 diff --git a/test/python/sdpa/test_torch_ops.py b/test/python/sdpa/test_torch_ops.py index ab0c19e63..7b6fd7761 100644 --- a/test/python/sdpa/test_torch_ops.py +++ b/test/python/sdpa/test_torch_ops.py @@ -199,6 +199,103 @@ def test_opcheck(self): ) +class TestSdpaBwdDense: + """Dense BHSD backward through cudnn::sdpa_bwd.""" + + @staticmethod + def _assert_close(name, got, ref): + """Compare bf16 gradients RELATIVE to the tensor's own magnitude. + + GQA dK/dV sum over h_q/h_kv query heads, so their values — and their + absolute rounding error — scale with the group size. A fixed absolute + bound would flag a correct result purely because the numbers got + bigger (observed: ~0.5% relative on every tensor, but |dv| peaks near + 9.6 at h_q/h_kv = 4).""" + err = (got.float() - ref).abs().max().item() + mag = max(ref.abs().max().item(), 1.0) + assert err <= TOL * mag, f"{name}: max|err|={err:.4f} exceeds {TOL} * |ref|max={mag:.3f}" + + @staticmethod + def _ref(q, k, v, scale, is_causal, grad): + qr, kr, vr = (t.detach().clone().float().requires_grad_(True) for t in (q, k, v)) + ref = torch.nn.functional.scaled_dot_product_attention(qr, kr, vr, is_causal=is_causal, scale=scale) + ref.backward(grad.float()) + return ref, qr.grad, kr.grad, vr.grad + + @pytest.mark.L0 + @pytest.mark.parametrize("is_causal", [False, True]) + def test_dense_backward(self, is_causal): + 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) + scale = D**-0.5 + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=is_causal, return_lse=True) + grad = torch.randn_like(o) + dq, dk, dv = torch.ops.cudnn.sdpa_bwd(grad, q, k, v, o, lse, scale, is_causal=is_causal) + _, rdq, rdk, rdv = self._ref(q, k, v, scale, is_causal, grad) + self._assert_close("dq", dq, rdq) + self._assert_close("dk", dk, rdk) + self._assert_close("dv", dv, rdv) + + @pytest.mark.L0 + def test_dense_backward_gqa(self): + """h_k != h_v is legal for cuDNN; both must divide h_q.""" + torch.manual_seed(0) + B, Hq, Hkv, S, D = 2, 16, 4, 256, 128 + q = bshd(B, Hq, S, D) + k, v = bshd(B, Hkv, S, D), bshd(B, Hkv, S, D) + scale = D**-0.5 + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + grad = torch.randn_like(o) + dq, dk, dv = torch.ops.cudnn.sdpa_bwd(grad, q, k, v, o, lse, scale, is_causal=True) + qr, kr, vr = (t.detach().clone().float().requires_grad_(True) for t in (q, k, v)) + ref = torch.nn.functional.scaled_dot_product_attention(qr, kr, vr, is_causal=True, scale=scale, enable_gqa=True) + ref.backward(grad.float()) + self._assert_close("dq", dq, qr.grad) + self._assert_close("dk", dk, kr.grad) + self._assert_close("dv", dv, vr.grad) + + @pytest.mark.L0 + def test_dense_autograd(self): + """End to end through register_autograd — the path the provider uses.""" + torch.manual_seed(0) + B, H, S, D = 2, 8, 256, 128 + q, k, v = (bshd(B, H, S, D).requires_grad_(True) for _ in range(3)) + scale = D**-0.5 + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + grad = torch.randn_like(o) + o.backward(grad) + _, rdq, rdk, rdv = self._ref(q, k, v, scale, True, grad) + self._assert_close("dq", q.grad, rdq) + self._assert_close("dk", k.grad, rdk) + self._assert_close("dv", v.grad, rdv) + + @pytest.mark.L0 + def test_dense_grads_adopt_input_layout(self): + """dQ/dK/dV come back in their input's dim-permutation: autograd hands + them straight back as .grad, which should match the parameter.""" + torch.manual_seed(0) + B, H, S, D = 2, 8, 128, 64 + # BSHD-physical inputs (the transposed-projection layout). + q, k, v = (torch.randn(B, S, H, D, dtype=torch.bfloat16, device="cuda").transpose(1, 2) for _ in range(3)) + scale = D**-0.5 + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + dq, dk, dv = torch.ops.cudnn.sdpa_bwd(torch.randn_like(o), q, k, v, o, lse, scale, is_causal=True) + for name, grad_t, src in (("dq", dq, q), ("dk", dk, k), ("dv", dv, v)): + assert grad_t.stride() == src.stride(), f"{name} stride {grad_t.stride()} != input {src.stride()}" + + @pytest.mark.L0 + def test_dense_opcheck(self): + """opcheck on the dense backward: the fake kernel must mirror the real + gradients' strides, which are the inputs' permutation (not contiguous).""" + torch.manual_seed(0) + B, H, S, D = 2, 4, 128, 64 + q, k, v = (torch.randn(B, S, H, D, dtype=torch.bfloat16, device="cuda").transpose(1, 2) for _ in range(3)) + scale = D**-0.5 + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + torch.library.opcheck(torch.ops.cudnn.sdpa_bwd, (torch.randn_like(o), q, k, v, o, lse, scale), dict(is_causal=True)) + + class TestSdpaVarlen: """THD (packed varlen) forward + backward through the ops directly.""" From f83c131735f52aa304823c8bacbd27feee73f0c0 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 1 Sep 2026 00:27:54 -0700 Subject: [PATCH 5/7] fix(sdpa): normalize dense operands; forward the deterministic flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the dense backward, plus a pre-existing gap the first one turned up. **Dense operands were never normalized.** The descriptor contract — innermost dim dense, base pointer 16B-aligned — was enforced only for THD, via `_normalize_thd`. A dense BHSD tensor with a strided last dim (a `[..., ::2]` slice) or a misaligned base breaks the descriptors exactly the same way; it is a property of the operand, not of packing. Note this was NOT introduced by the dense backward: the dense *forward* has always taken `q.stride(), k.stride(), v.stride()` as-is, so the fix covers both paths. Renamed the helper to `_normalize_operand` since nothing in its body was THD-specific. In the backward the normalization runs before the strides are read and before dO is matched to O, so a repaired O cannot leave dO on the stale layout. **The dense autograd path dropped `is_deterministic`.** The THD branch passes `torch.are_deterministic_algorithms_enabled()`; the dense branch omitted it, so `use_deterministic_algorithms(True)` still built the graph with `use_deterministic_algorithm=False` and gave non-reproducible dq. Tests: `test_dense_backward_repairs_bad_operands` covers both flaw shapes (strided innermost via `[..., ::2]`, misaligned base via a 1-element offset) and asserts numerical correctness against a reference — a mis-declared descriptor reads the wrong elements rather than failing loudly, so a smoke-test would not catch it. `test_dense_autograd_honors_deterministic_flag` runs the same inputs twice under `use_deterministic_algorithms(True)` and requires bit-identical gradients. Verified: 28 passed (was 25). Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/torch_op.py | 40 +++++++++++++++------ test/python/sdpa/test_torch_ops.py | 56 ++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/python/cudnn/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py index 0b476d59f..c73761f77 100644 --- a/python/cudnn/sdpa/fwd/torch_op.py +++ b/python/cudnn/sdpa/fwd/torch_op.py @@ -160,14 +160,19 @@ def _thd_desc_stride(t: torch.Tensor, s_max: int) -> Tuple[int, int, int, int]: return (s_max * s_t, s_h, s_t, s_d) -def _normalize_thd(t: torch.Tensor, name: str) -> torch.Tensor: +def _normalize_operand(t: torch.Tensor, name: str, op: str = "sdpa_fwd") -> 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.""" + whatever its storage offset, so it cannot repair a misaligned base. + + Applies to BOTH layouts. The descriptor requirement is a property of the + operand, not of THD packing: a dense BHSD tensor whose last dim is strided + (say a ``[..., ::2]`` slice), or whose base is misaligned, breaks the same + way.""" if t.stride(-1) != 1 or t.data_ptr() % 16: - _logger.warning("sdpa_fwd: copying %s to normalize layout/alignment (slow path)", name) + _logger.warning("%s: copying %s to normalize layout/alignment (slow path)", op, name) t = t.clone(memory_format=torch.contiguous_format) return t @@ -344,9 +349,9 @@ def _sdpa_fwd_impl( 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") + q = _normalize_operand(q, "q") + k = _normalize_operand(k, "k") + v = _normalize_operand(v, "v") T_q, H_q, D_qk = q.shape T_kv, H_v, D_v = v.shape H_k = k.shape[1] @@ -375,6 +380,12 @@ def _sdpa_fwd_impl( 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}") + # Same descriptor contract as THD: dense strides are declared as-is, + # but a strided innermost dim or a misaligned base still has to be + # repaired before the descriptors are built. + q = _normalize_operand(q, "q") + k = _normalize_operand(k, "k") + v = _normalize_operand(v, "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) @@ -687,6 +698,14 @@ def _sdpa_bwd_dense( raise ValueError(f"o {tuple(o.shape)} / grad_out {tuple(grad_out.shape)} must be (B={B}, H_q={H_q}, S_q={S_q}, D_v={D_v})") _check_same_device(q, k=k, v=v, o=o, lse=lse, grad_out=grad_out) + # Descriptor contract (innermost dense, 16B-aligned base) — before the + # strides are read and before dO is matched to O, so a repaired O does not + # leave dO on the old layout. + q = _normalize_operand(q, "q", "sdpa_bwd") + k = _normalize_operand(k, "k", "sdpa_bwd") + v = _normalize_operand(v, "v", "sdpa_bwd") + o = _normalize_operand(o, "o", "sdpa_bwd") + 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: @@ -832,10 +851,10 @@ def _sdpa_bwd_impl( ) 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") + q = _normalize_operand(q, "q") + k = _normalize_operand(k, "k") + v = _normalize_operand(v, "v") + o = _normalize_operand(o, "o", "sdpa_bwd") T_q, H_q, D_qk = q.shape T_kv, H_v, D_v = v.shape H_k = k.shape[1] @@ -1095,6 +1114,7 @@ def _sdpa_backward(ctx, grad_o, _grad_stats): # stats marked non-differentiable is_causal=ctx.is_causal, causal_bottom_right=ctx.causal_bottom_right, window_left=ctx.window_left, + is_deterministic=torch.are_deterministic_algorithms_enabled(), ) return (dq, dk, dv) + (None,) * 12 diff --git a/test/python/sdpa/test_torch_ops.py b/test/python/sdpa/test_torch_ops.py index 7b6fd7761..f44b8af44 100644 --- a/test/python/sdpa/test_torch_ops.py +++ b/test/python/sdpa/test_torch_ops.py @@ -284,6 +284,62 @@ def test_dense_grads_adopt_input_layout(self): for name, grad_t, src in (("dq", dq, q), ("dk", dk, k), ("dv", dv, v)): assert grad_t.stride() == src.stride(), f"{name} stride {grad_t.stride()} != input {src.stride()}" + @pytest.mark.L0 + @pytest.mark.parametrize("flaw", ["strided_innermost", "misaligned_base"]) + def test_dense_backward_repairs_bad_operands(self, flaw): + """A dense operand whose innermost dim is strided, or whose base is not + 16B-aligned, breaks the descriptor contract exactly like a THD one: + both must be repaired before descriptors are built, not passed through.""" + torch.manual_seed(0) + B, H, S, D = 2, 4, 128, 64 + scale = D**-0.5 + if flaw == "strided_innermost": + # [..., ::2] -> last-dim stride 2 + k = torch.randn(B, H, S, 2 * D, dtype=torch.bfloat16, device="cuda")[..., ::2] + else: + # one bf16 element in -> base pointer at +2 bytes + k = torch.randn(B, H, S, D + 1, dtype=torch.bfloat16, device="cuda")[..., 1:] + q, v = bshd(B, H, S, D), bshd(B, H, S, D) + assert k.shape == q.shape + o, lse = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + dq, dk, dv = torch.ops.cudnn.sdpa_bwd(torch.randn_like(o), q, k, v, o, lse, scale, is_causal=True) + # Correctness is the point: a silently mis-declared descriptor would + # read the wrong elements rather than fail loudly. + qr, kr, vr = (t.detach().clone().float().requires_grad_(True) for t in (q, k, v)) + ref = torch.nn.functional.scaled_dot_product_attention(qr, kr, vr, is_causal=True, scale=scale) + assert (o.float() - ref).abs().max().item() < TOL + assert dq.isfinite().all() and dk.isfinite().all() and dv.isfinite().all() + + @pytest.mark.L0 + def test_dense_autograd_honors_deterministic_flag(self): + """torch.use_deterministic_algorithms(True) must reach the graph. The + dense autograd path used to drop it, silently building the backward + with use_deterministic_algorithm=False and giving non-reproducible dq.""" + torch.manual_seed(0) + B, H, S, D = 2, 4, 128, 64 + scale = D**-0.5 + + # Same inputs both times: determinism is about the backward's own + # accumulation order, not the data. + q, k, v = (bshd(B, H, S, D).requires_grad_(True) for _ in range(3)) + + def run(): + for t in (q, k, v): + t.grad = None + o, _ = torch.ops.cudnn.sdpa_fwd(q, k, v, scale, is_causal=True, return_lse=True) + o.backward(torch.ones_like(o)) + return q.grad.clone(), k.grad.clone(), v.grad.clone() + + was = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(True) + try: + a = run() + b = run() + finally: + torch.use_deterministic_algorithms(was) + for name, x, y in zip(("dq", "dk", "dv"), a, b): + assert torch.equal(x, y), f"{name} differs across runs under use_deterministic_algorithms(True)" + @pytest.mark.L0 def test_dense_opcheck(self): """opcheck on the dense backward: the fake kernel must mirror the real From 4cc8100172e24f9c2c18ac82d8800fb6e303eec7 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 1 Sep 2026 00:39:22 -0700 Subject: [PATCH 6/7] torch: route the provider's dense backward through cudnn::sdpa_bwd MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With dense backward wired into the op (previous commits), the provider no longer has to hand dense gradients to the C++ worker. `_scaled_dot_product_ cudnn_attention_backward` now calls `cudnn::sdpa_bwd` directly, and the C++ worker keeps only what the op declines: attention bias, dropout > 0, and the padded dense path (per-batch lengths). That closes the last C++ hop in a dense training step — forward and backward both build their graphs through the python API. aten hands logsumexp back as `(B, H, S)` fp32, which is the dense backward's expected layout modulo the trailing 1 the descriptor declares, so the conversion is an `unsqueeze`. The dense parity test previously asserted `bwd_cpp` incremented; it now asserts the opposite — `bwd` increments and `bwd_cpp` does NOT — so a silent regression back to the C++ path fails the suite instead of passing it. The existing numeric parity checks (against a math reference, with the stock flash backend as the rounding yardstick) are unchanged and still pass, which is what says the op's gradients are as good as the worker's. Verified: provider 17 passed, ops 28 passed. End to end on the POC bridge in pytorch_cudnn_python_api: `bridge calls: fwd=10 bwd=10` (was fwd=10 bwd=0), 10/10 cases passing. Co-Authored-By: Claude Fable 5 --- python/cudnn/torch/sdpa_provider.py | 34 +++++++++++++++++-------- test/python/sdpa/test_torch_provider.py | 11 +++++--- 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/python/cudnn/torch/sdpa_provider.py b/python/cudnn/torch/sdpa_provider.py index 7d0e33161..0a21f41b6 100644 --- a/python/cudnn/torch/sdpa_provider.py +++ b/python/cudnn/torch/sdpa_provider.py @@ -25,8 +25,8 @@ 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 +kernel): attn_bias, dropout_p > 0, and the padded dense backward (per-batch +lengths). Dense and varlen backward both run on the python API. The forward runs on the python API either way, so training still exercises the python fwd path. """ @@ -99,15 +99,27 @@ def _bwd( *, 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, + # Anything cudnn::sdpa_bwd does not serve goes to the C++ worker op + # (bit-exact with the shadowed native kernel): attention bias, dropout, + # and the padded dense path (per-batch lengths). Everything else runs on + # the python API, closing the last C++ hop in a dense training step. + if attn_bias is not None or dropout_p > 0.0 or cum_seq_q is not None: + 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 + + calls["bwd"] += 1 + attn_scale = scale if scale is not None else query.shape[-1] ** -0.5 + # aten hands us logsumexp as (B, H, S) fp32 — exactly the layout the dense + # backward wants, modulo the trailing 1 the descriptor declares. + lse = logsumexp if logsumexp.dim() == 4 else logsumexp.unsqueeze(-1) + return torch.ops.cudnn.sdpa_bwd( + grad_out, query, key, value, out, lse, attn_scale, + is_causal=is_causal, + is_deterministic=torch.are_deterministic_algorithms_enabled(), ) # fmt: skip diff --git a/test/python/sdpa/test_torch_provider.py b/test/python/sdpa/test_torch_provider.py index 765355e9c..360fa15e8 100644 --- a/test/python/sdpa/test_torch_provider.py +++ b/test/python/sdpa/test_torch_provider.py @@ -82,15 +82,18 @@ def test_sdpa_dense_parity(B, Hq, Hkv, Sq, Skv, D, dtype, is_causal, scale, enab 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"] + # Dense: BOTH directions run on the python API now that cudnn::sdpa_bwd + # serves dense. The C++ worker keeps only what the op declines (bias, + # dropout, padded), so bwd_cpp must NOT move here. + fwd_before, bwd_before = provider.calls["fwd"], provider.calls["bwd"] + cpp_before = 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 provider.calls["bwd"] == bwd_before + 1, "provider bwd did not intercept" + assert provider.calls["bwd_cpp"] == cpp_before, "dense backward fell back to the C++ worker" 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) From 7a8e5839079dc71e5973628b6d8519f744267be1 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 1 Sep 2026 00:44:35 -0700 Subject: [PATCH 7/7] docs(agents): record the concrete test traps hit while wiring the torch ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six things that each cost real time, written where the next person will look rather than left in a PR thread: - `test/python/sdpa/` mixes harness modules with tests, so the `test_` prefix is load-bearing — a test file dropped there without it is treated as a helper, never collected, and the suite stays green while asserting nothing. Includes the collect-only check to confirm pickup by the DEFAULT sweep. - Read pytest's own summary: `| grep -c passed` counts collection errors as passes, and `cmd | tail` reports tail's exit status — a failed build behind a pipe reads as success. Both produced wrong 'all green' reports here. - Scale tolerances to the tensor: GQA dK/dV sum over h_q/h_kv query heads, so a fixed absolute bound that passes at h_kv == h_q fails at a group size of 4 while the relative error is unchanged (~0.5%). - A regression test must be seen RED against the unfixed code first. - Seed before allocating, or two runs meant to be compared differ by data. - When removing a fallback, invert its counter assertion rather than deleting it, so a silent regression to the old path fails instead of passing. Companion to the Rule 3 verification recipe in python/cudnn/AGENTS.md. Co-Authored-By: Claude Fable 5 --- test/AGENTS.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/test/AGENTS.md b/test/AGENTS.md index 6bd0fafc1..ff216bfca 100644 --- a/test/AGENTS.md +++ b/test/AGENTS.md @@ -20,7 +20,7 @@ pytest test_conv_fprop.py # one file — note the default -m L0 filter still pytest fe_api/gemm/ # OSS kernel tests ``` -Requirements: `pip install -e ".[cutedsl]"` plus `pytest pytest-xdist looseversion`. `fe_api/` additionally requires an SM90/SM100-class GPU; tests skip (or should skip) on unsupported arch/dtype/backend-version combos rather than fail. +**Read pytest's own summary line; do not post-process the output.** `... | grep -c "passed"` counts a collection error as a pass, and `cmd | tail` reports *tail's* exit status, so a failed build or a failed suite behind a pipe looks like success. Both have produced confidently wrong "all green" reports here. Requirements: `pip install -e ".[cutedsl]"` plus `pytest pytest-xdist looseversion`. `fe_api/` additionally requires an SM90/SM100-class GPU; tests skip (or should skip) on unsupported arch/dtype/backend-version combos rather than fail. ### conftest.py landmines — read before editing @@ -33,6 +33,11 @@ Requirements: `pip install -e ".[cutedsl]"` plus `pytest pytest-xdist looseversi ### Layout - `test/python/test_*.py` — core graph-API tests (conv, matmul, norms, SDPA `test_mhas*.py`, rope, kernel cache, OSS engines `test_sm{90,100}_prefill_oss_engine.py`, ...). Shared SDPA references in `test/python/sdpa/`. +- **`test/python/sdpa/` is a mixed directory and the `test_` prefix is load-bearing.** `fp16.py`, `helpers.py`, `random_config.py` are harness modules the tests import; `sdpa/test_*.py` (and `sdpa/frost/test_*.py`) are collected tests. `pytest.ini` sets no `python_files` override, so a test file dropped there **without** the prefix is silently treated as a helper — it is never collected, and the suite stays green while asserting nothing. After moving or adding a test, confirm it is picked up by the *default* sweep, not just when named directly: + + ```bash + pytest --collect-only -q | grep -c sdpa/test_torch_ops.py + ``` - `test/python/fe_api//` — one subdir per OSS kernel family (`gemm/`, `grouped_gemm/`, `bsa/`, `dsa/`, `nsa/`, `norm/`, `sdpa/`), each with `test_.py` + utils/reference modules. ### Conventions for new tests @@ -40,3 +45,7 @@ Requirements: `pip install -e ".[cutedsl]"` plus `pytest pytest-xdist looseversi - Mark with a level (`@pytest.mark.L0` ... `L4`): L0 must stay fast (default CI smoke); big parameter sweeps go to higher levels. - Gate on capability, don't assume it: skip via `check_support()` failures, `cudnn.backend_version()`, and `torch.cuda.get_device_capability()`. - Compare against a reference implementation (see existing `*_ref.py` / `*_reference.py` patterns) with dtype-appropriate tolerances. +- **Scale the tolerance to the tensor, not to the dtype alone.** A fixed absolute bound quietly becomes wrong when magnitudes grow: GQA dK/dV sum over `h_q/h_kv` query heads, so at a group size of 4 the *relative* error stays ~0.5% while `|dv|` peaks near 9.6 and blows a bound that passed at `h_kv == h_q`. Compare against `TOL * max(|ref|.max(), 1.0)`, or the next GQA ratio someone adds will look like a correctness regression. +- **A regression test must be seen RED.** Before trusting one, run it against the unfixed code — restore the old line, confirm it fails, restore the fix. `test_dsl_sm100_thd_interleaved_kv_views` and `test_varlen_backward_does_not_sync` were both checked this way, and both were genuinely red beforehand; a test written for a bug and never seen to fail is asserting an unknown. +- **Seed before you allocate.** `torch.manual_seed()` after constructing the inputs seeds nothing that matters. Two runs meant to be compared then differ by data, and the assertion fails (or worse, passes) for a reason unrelated to what is under test — if two runs must be comparable, build the inputs once and reuse them. +- **When you remove a fallback, invert its counter assertion — do not delete it.** Tests that asserted `calls["bwd_cpp"]` incremented had to become "`calls["bwd"]` increments **and** `bwd_cpp` does not", so a silent regression to the old path fails the suite instead of passing it.