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: 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/sdpa/fwd/torch_op.py b/python/cudnn/sdpa/fwd/torch_op.py index 805b3fab0..c73761f77 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 @@ -159,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 @@ -343,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] @@ -374,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) @@ -545,8 +557,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 +587,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 +618,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 +664,144 @@ 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) + + # 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: + 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,19 +825,36 @@ 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") - 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] @@ -759,6 +944,7 @@ def _sdpa_bwd_impl( o_stride=o_stride, stats_stride=stats_stride, is_deterministic=is_deterministic, + is_thd=True, ), ) @@ -816,9 +1002,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) @@ -871,31 +1064,65 @@ 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 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, + is_deterministic=torch.are_deterministic_algorithms_enabled(), + ) + 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 # 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/__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..0a21f41b6 --- /dev/null +++ b/python/cudnn/torch/sdpa_provider.py @@ -0,0 +1,310 @@ +# 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 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. +""" + +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, +): + # 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 + + +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. + # (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]), + 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/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. diff --git a/test/python/sdpa/test_torch_ops.py b/test/python/sdpa/test_torch_ops.py index ab0c19e63..f44b8af44 100644 --- a/test/python/sdpa/test_torch_ops.py +++ b/test/python/sdpa/test_torch_ops.py @@ -199,6 +199,159 @@ 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 + @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 + 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.""" diff --git a/test/python/sdpa/test_torch_provider.py b/test/python/sdpa/test_torch_provider.py new file mode 100644 index 000000000..360fa15e8 --- /dev/null +++ b/test/python/sdpa/test_torch_provider.py @@ -0,0 +1,265 @@ +# 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: 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"] == 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) + 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_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, + 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