diff --git a/python/cudnn/sdpa/__init__.py b/python/cudnn/sdpa/__init__.py index 923ee8c54..363d3a895 100644 --- a/python/cudnn/sdpa/__init__.py +++ b/python/cudnn/sdpa/__init__.py @@ -13,7 +13,7 @@ from cudnn.sdpa.fwd import SdpaFwdDslSm100, SdpaFwdDslSm120, SdpaFwdDslSm80 from cudnn.sdpa.fwd import sdpa_fwd_wrapper_dsl_sm100, sdpa_fwd_wrapper_dsl_sm120, sdpa_fwd_wrapper_sm80 - from cudnn.sdpa.bwd import SdpaBwdDslSm120, SdpabwdSm80 + from cudnn.sdpa.bwd import SdpaBwdDslSm120, SdpaBwdDslSm80 from cudnn.sdpa.bwd import sdpa_bwd_wrapper_dsl_sm120, sdpa_bwd_wrapper_sm80 Submodule imports stay lazy there (PEP 562): eager imports used to drag the diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index ae4fefe3f..3df1b4791 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -17,8 +17,8 @@ "SdpaBwdDsl": (".api_dsl", "SdpaBwdDsl"), "SdpaBwdDslSm120": (".api_dsl", "SdpaBwdDslSm120"), "sdpa_bwd_wrapper_dsl_sm120": (".api_dsl", "sdpa_bwd_wrapper_dsl_sm120"), - "SdpabwdSm80": (".api", "SdpabwdSm80"), - "sdpa_bwd_wrapper_sm80": (".api", "sdpa_bwd_wrapper_sm80"), + "SdpaBwdDslSm80": (".api_dsl", "SdpaBwdDslSm80"), + "sdpa_bwd_wrapper_sm80": (".api_dsl", "sdpa_bwd_wrapper_sm80"), } __all__ = list(_LAZY_EXPORTS) diff --git a/python/cudnn/sdpa/bwd/api.py b/python/cudnn/sdpa/bwd/api.py deleted file mode 100644 index 207b6e425..000000000 --- a/python/cudnn/sdpa/bwd/api.py +++ /dev/null @@ -1,746 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from typing import Optional, Tuple -import inspect -import logging -import math - -from cuda.bindings import driver as cuda -import torch - - -from cudnn.api_base import APIBase, TupleDict - -_logger = logging.getLogger(__name__) - - -_KERNEL_MOD = {} - - -def _stream_ctx(current_stream): - """Context manager dispatching onto ``current_stream`` (a ``cuda.CUstream`` - or raw stream int); the kernels launch on torch's current stream, so an - ExternalStream context routes them. ``None`` keeps the current stream, and - a raw handle equal to torch's current/default stream reuses that torch - stream object rather than wrapping it: ``ExternalStream(0)`` breaks - re-execution on some torch builds (NGC), where every launch after the - compile run silently no-ops (all-zero outputs; caught by test_mhas_v2's - determinism re-run). Mirrors gemm/cutedsl/grouped/backend_utils.py.""" - import contextlib - - if current_stream is None: - return contextlib.nullcontext() - handle = int(current_stream) - torch_current = torch.cuda.current_stream() - if handle in (0, 1, 2) or handle == torch_current.cuda_stream: - return contextlib.nullcontext() - torch_default = torch.cuda.default_stream() - if handle == torch_default.cuda_stream: - return torch.cuda.stream(torch_default) - return torch.cuda.stream(torch.cuda.ExternalStream(handle)) - - -# The generic kernel now supports d_qk != d_v (split sub-groups) and d up to -# 256, so the envelope spans gptoss(64,64) / llama(128,128) / dsv3(192,128) / -# qwen(256,256). Kernel constraint: d_qk >= d_v (the per-sub-group split). -# The flavor only sets the d-pad target; the kernel derives qo_stages / drop-sDQ -# from the (padded) d_qk for the A100 SMEM budget. -from ..fwd import config_sm80 as _fwd_config_sm80 - -_FLAVOR_DIMS = { - name: (cfg.D_QK, cfg.D_V) - for name, cfg in ( - ("gptoss", _fwd_config_sm80.GPTOSS_CFG), - ("llama", _fwd_config_sm80.LLAMA_CFG), - ("dsv3", _fwd_config_sm80.DSV3_CFG), - ("qwen", _fwd_config_sm80.QWEN_CFG), - ) -} -_SUPPORTED_FLAVORS = ("gptoss", "llama", "dsv3", "qwen") - - -def _load_kernel_module(key: str = "f16"): - """Lazily import + cache an SM80 BPROP kernel module. - - ``"f16"`` is the GENERIC kernel (``bprop_f16_sm80``): fully parameterized - on d_qk/d_v with the full feature set (masks / bias / dBias / - sink / rope / THD / deterministic). ``"d64"`` is the - dedicated plain-dense d=64 MHA perf variant (~2x faster on A100); it - supports NO features — its ``backward(**_ignored)`` silently swallows - every feature kwarg, so callers must never rely on the signature filter - and only select it through :func:`_d64_fast_path_eligible`. - """ - if key not in _KERNEL_MOD: - if key == "d64": - from .kernels import bprop_d64_f16_sm80 as _mod - else: - from .kernels import bprop_f16_sm80 as _mod - - _KERNEL_MOD[key] = _mod - return _KERNEL_MOD[key] - - -def _d64_fast_path_eligible(*, d_qk, d_v, h_q, h_kv, s_q, s_kv, mask_token, right_bound, causal_bottom_right, bw_kwargs) -> bool: - """Whether the dedicated d=64 kernel can serve this call EXACTLY. - - The perf variant computes a plain dense MHA backward and nothing else; - every condition here guards a feature it would silently ignore. - """ - d64 = _load_kernel_module("d64") - if (d_qk, d_v) != (64, 64) or h_q != h_kv: - return False - if s_q % d64.M_BLOCK != 0 or s_kv % d64.N_BLOCK != 0: - return False - if mask_token != "none" or right_bound != 0 or causal_bottom_right: - return False - for feature in ("seq_kv_lens", "seq_len_q", "bias", "sinks", "rope_freqs"): - if bw_kwargs.get(feature) is not None: - return False - if bw_kwargs.get("deterministic"): - return False - return True - - -def _pick_flavor(d_qk: int, d_v: int) -> str: - """Smallest BPROP flavor whose ``(D_QK, D_V)`` envelope covers - ``(d_qk, d_v)`` (fdqk >= d_qk and fdv >= d_v); the user's heads are padded - up to the flavor dim. The kernel supports d_qk != d_v but requires the - (padded) d_qk >= d_v — the flavor list guarantees this (every flavor has - fdqk >= fdv, and a d_qk < d_v case lands on an equal-d flavor after pad).""" - for flavor in _SUPPORTED_FLAVORS: - fdqk, fdv = _FLAVOR_DIMS[flavor] - if d_qk == fdqk and d_v == fdv: - return flavor - for flavor in _SUPPORTED_FLAVORS: - fdqk, fdv = _FLAVOR_DIMS[flavor] - if d_qk <= fdqk and d_v <= fdv: - return flavor - raise ValueError(f"SM80 BPROP: no flavor envelope covers (D_QK={d_qk}, D_V={d_v}); " f"supported: {_FLAVOR_DIMS}.") - - -def _pad_last_dim(t: torch.Tensor, new_last: int) -> torch.Tensor: - """Zero-pad the trailing dim of an fp16/bf16 tensor up to ``new_last``.""" - old_last = t.shape[-1] - if old_last == new_last: - return t - if old_last > new_last: - raise ValueError(f"_pad_last_dim: tensor's last dim {old_last} exceeds target {new_last}") - pad = torch.zeros((*t.shape[:-1], new_last - old_last), dtype=t.dtype, device=t.device) - return torch.cat([t, pad], dim=-1).contiguous() - - -def _bshd(t: torch.Tensor) -> torch.Tensor: - """BHSD → BSHD (stride-only transpose; contiguous-ify only if needed).""" - x = t.transpose(1, 2) - return x if x.is_contiguous() else x.contiguous() - - -# --------------------------------------------------------------------------- -# APIBase subclass. -# --------------------------------------------------------------------------- -class SdpabwdSm80(APIBase): - """SM80 (A100) SDPA backward. - - Mirrors the SM80 forward adapter. Inputs are the forward activations (Q/K/V/O), the - loss gradient dO, and the forward stats LSE. Outputs dQ/dK/dV (+ dBias when - an additive bias is present). - """ - - def __init__( - self, - sample_q: torch.Tensor, - sample_k: torch.Tensor, - sample_v: torch.Tensor, - sample_o: torch.Tensor, - sample_do: torch.Tensor, - sample_lse: torch.Tensor, - is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - scale_softmax: Optional[float] = None, - causal_bottom_right: bool = False, - has_seq_kv_lens: bool = False, - has_bias: bool = False, - ): - super().__init__() - self._warn_experimental_api() - self._logger.debug("Entering __init__ (bwd)") - - self.q_desc = self._make_tensor_desc(sample_q, name="q") - self.k_desc = self._make_tensor_desc(sample_k, name="k") - self.v_desc = self._make_tensor_desc(sample_v, name="v") - self.o_desc = self._make_tensor_desc(sample_o, name="o") - self.do_desc = self._make_tensor_desc(sample_do, name="dO") - self.lse_desc = self._make_tensor_desc(sample_lse, name="lse") - - self.is_causal = is_causal - self.window_size_left, self.window_size_right = window_size - self.scale_softmax = scale_softmax - self.causal_bottom_right = bool(causal_bottom_right) - self.has_seq_kv_lens = bool(has_seq_kv_lens) - self.has_bias = bool(has_bias) - - # Filled by check_support(). - self.flavor: Optional[str] = None - self.flavor_d_qk: Optional[int] = None - self.flavor_d_v: Optional[int] = None - self.mask_token: Optional[str] = None - self.swa_window_runtime: int = 0 - self.right_bound: int = 0 - self.head_dim_qk: Optional[int] = None - self.head_dim_v: Optional[int] = None - self._logger.debug("__init__ (bwd) completed") - - # ------------------------------------------------------------------ - def check_support(self) -> bool: - self._logger.debug("Entering check_support (bwd)") - - _REQ = (3, 1, 2, 0) - for desc_name in ["q_desc", "k_desc", "v_desc", "o_desc", "do_desc"]: - d = getattr(self, desc_name) - self._value_error_if(d.ndim != 4, f"{d.name} must be rank-4 (B, H, S, D); got {d.ndim}") - _shape = d.shape - _act = tuple(ax for ax in d.stride_order if _shape[ax] != 1) - _exp = tuple(ax for ax in _REQ if _shape[ax] != 1) - self._value_error_if( - _act != _exp, f"{d.name} must have d,h,s,b stride order (3,1,2,0) " f"(size-1 dims wildcarded); got {d.stride_order} shape {_shape}" - ) - - b, h_qo, s_qo, d_qk = self.q_desc.shape - _, h_kv, s_kv, _ = self.k_desc.shape - _, _, _, d_v = self.v_desc.shape - - self._check_tensor_shape(self.q_desc, (b, h_qo, s_qo, d_qk), name="Q") - self._check_tensor_shape(self.k_desc, (b, h_kv, s_kv, d_qk), name="K") - self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_v), name="V") - self._check_tensor_shape(self.o_desc, (b, h_qo, s_qo, d_v), name="O") - self._check_tensor_shape(self.do_desc, (b, h_qo, s_qo, d_v), name="dO") - - for label, val in (("B", b), ("H_q", h_qo), ("H_kv", h_kv), ("S_q", s_qo), ("S_kv", s_kv), ("D_QK", d_qk), ("D_V", d_v)): - self._value_error_if(int(val) <= 0, f"{label} must be > 0; got {val}") - - self._value_error_if(h_qo % h_kv != 0, f"H_q ({h_qo}) must be divisible by H_kv ({h_kv}) for GQA / MQA") - - # Kernel supports d_qk != d_v (split sub-groups) but requires d_qk >= d_v - # (a d_qk < d_v case is padded up to an equal-d flavor by _pick_flavor). - self._value_error_if(d_qk < d_v, f"SM80 BPROP requires D_QK >= D_V; got D_QK={d_qk}, D_V={d_v}") - max_dqk = max(fdqk for fdqk, _ in _FLAVOR_DIMS.values()) - max_dv = max(fdv for _, fdv in _FLAVOR_DIMS.values()) - self._value_error_if( - d_qk > max_dqk or d_v > max_dv, - f"SM80 BPROP: head dim (D_QK={d_qk}, D_V={d_v}) exceeds supported " f"envelope (D_QK<={max_dqk}, D_V<={max_dv}); larger heads not yet ported.", - ) - - self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") - for desc in [self.k_desc, self.v_desc, self.o_desc, self.do_desc]: - self._check_dtype(desc, self.dtype, name=desc.name, extra_error_msg=f"{desc.name} must match Q dtype (FP16/BF16)") - self._check_dtype(self.lse_desc, torch.float32, name="LSE") - self._check_tensor_shape(self.lse_desc, (b, h_qo, s_qo), name="LSE") - self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM80") - - self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM80 BPROP") - device = self.q_desc.device - major, minor = torch.cuda.get_device_capability(device) - self._value_error_if((major, minor) != (8, 0), f"SdpabwdSm80 requires SM80 (A100); found SM{major}{minor} on {device}") - - self.flavor = _pick_flavor(d_qk, d_v) - self.flavor_d_qk, self.flavor_d_v = _FLAVOR_DIMS[self.flavor] - self.head_dim_qk = int(d_qk) - self.head_dim_v = int(d_v) - - # ---- mask token (same resolution as the forward adapter) ------ - swa_left = self.window_size_left - swa_right = self.window_size_right - self.right_bound = 0 - if self.is_causal: - self.mask_token = "causal" if swa_left < 0 else "causal_swa" - self.swa_window_runtime = max(0, swa_left) if swa_left >= 0 else 0 - self.right_bound = max(0, swa_right) - elif swa_left >= 0: - # A left window alone selects SWA; window_size_right is only - # meaningful with is_causal=True. - self._not_implemented_error_if(swa_right > 0, "SM80 BPROP: non-causal SWA with window_size_right > 0 unsupported") - self.mask_token = "swa" - self.swa_window_runtime = swa_left - else: - # window_size=(-1, r) without is_causal: a bare right bound has no - # diagonal to anchor to — reject rather than silently pick a mask - # (mirrors the forward adapter and the THD path). - self._not_implemented_error_if( - swa_right >= 0, - "SM80 BPROP: window_size_right without a left window or is_causal=True has no effect; pass is_causal=True or a left window", - ) - self.mask_token = "none" - self.swa_window_runtime = 0 - - self._value_error_if( - self.causal_bottom_right and not (self.is_causal or self.window_size_left >= 0), - "SM80 BPROP: causal_bottom_right requires is_causal and/or a left window", - ) - - if self.scale_softmax is None or self.scale_softmax == 0.0: - self.scale_softmax = 1.0 / math.sqrt(d_qk) - - self._is_supported = True - self._logger.debug("check_support (bwd) completed") - return True - - # ------------------------------------------------------------------ - def _needs_bshd_stage(self, desc) -> bool: - """Whether ``desc``'s BSHD transpose is non-contiguous — i.e. execute's - kernel-facing view would need a gather into staging.""" - b, h, sq, d = desc.shape - expect = (sq * h * d, d, h * d, 1) # BHSD-logical view of a compact BSHD buffer - return tuple(desc.stride) != expect - - def scratch_workspace_bytes( - self, - *, - has_bias: Optional[bool] = None, - bias_batch: int = 1, - has_sink: bool = False, - deterministic: bool = False, - need_do_dot: bool = True, - ) -> int: - """Per-execute scratch requirement (issue #514): head-dim pad / BSHD - gather staging for Q/K/V/O/dO plus the kernel's internal scratch - (``bprop_f16_sm80.scratch_bytes``; the generic kernel's buffer set - covers the d64 fast path's). The feature flags must match what - execute() will be called with — the engine lowering passes its graph - facts; the default reads the constructor's ``has_bias``.""" - self._ensure_support_checked() - from ..fwd.api_dsl import ws_align - from .kernels import bprop_f16_sm80 as _kmod - - elem = 2 # fp16/bf16 — check_support admits no other input dtype - b, hq, sq, _ = self.q_desc.shape - _, hkv, skv, _ = self.k_desc.shape - fdqk, fdv = self.flavor_d_qk, self.flavor_d_v - pad_qk = self.head_dim_qk < fdqk - pad_v = self.head_dim_v < fdv - if has_bias is None: - has_bias = self.has_bias - total = 0 - # Pad / gather staging, in execute()'s carve order (Q, K, V, O, dO). - for desc, s_len, hh, pad, fd in ( - (self.q_desc, sq, hq, pad_qk, fdqk), - (self.k_desc, skv, hkv, pad_qk, fdqk), - (self.v_desc, skv, hkv, pad_v, fdv), - (self.o_desc, sq, hq, pad_v, fdv), - (self.do_desc, sq, hq, pad_v, fdv), - ): - if pad: - total += ws_align(b * s_len * hh * fd * elem) - elif self._needs_bshd_stage(desc): - total += ws_align(math.prod(desc.shape) * elem) - # Kernel-internal scratch at the PADDED (flavor) head dims. - total += _kmod.scratch_bytes( - B=b, - SQ=sq, - SKV=skv, - H=hq, - Hk=hkv, - d_qk=fdqk, - d_v=fdv, - io_bytes=elem, - deterministic=deterministic, - has_bias=bool(has_bias), - bias_batch=bias_batch, - has_sink=has_sink, - need_do_dot=need_do_dot, - ) - return total - - # ------------------------------------------------------------------ - def compile(self) -> None: - """No-op — the kernel module owns its own per-shape ``lru_cache``; - first ``execute()`` JITs and reuses thereafter.""" - self._logger.debug("Entering compile (bwd, no-op — kernel self-caches)") - self._ensure_support_checked() - self._compiled_kernel = True - self._logger.debug("compile (bwd) completed") - - # ------------------------------------------------------------------ - def execute( - self, - q_tensor: torch.Tensor, - k_tensor: torch.Tensor, - v_tensor: torch.Tensor, - o_tensor: torch.Tensor, - do_tensor: torch.Tensor, - lse_tensor: torch.Tensor, - dq_tensor: torch.Tensor, - dk_tensor: torch.Tensor, - dv_tensor: torch.Tensor, - dbias_tensor: Optional[torch.Tensor] = None, - dsink_tensor: Optional[torch.Tensor] = None, - scale_softmax: Optional[float] = None, - current_stream: Optional[cuda.CUstream] = None, - seq_kv_lens: Optional[torch.Tensor] = None, - seq_len_q: Optional[torch.Tensor] = None, - bias_tensor: Optional[torch.Tensor] = None, - sinks: Optional[torch.Tensor] = None, - rope_freqs: Optional[torch.Tensor] = None, - deterministic: bool = False, - workspace: Optional[torch.Tensor] = None, - ) -> None: - self._logger.debug("Entering execute (bwd)") - if self._compiled_kernel is None: - raise RuntimeError("SdpabwdSm80 is not compiled") - scale_val = self.scale_softmax if (scale_softmax is None or scale_softmax == 0.0) else float(scale_softmax) - - kernel = _load_kernel_module() - - # Per-execute scratch: carved from the caller's workspace when one is - # provided (the engine executor always passes one sized by - # scratch_workspace_bytes(); issue #514), otherwise allocated (the - # standalone wrapper paths). - carver = None - if workspace is not None: - from ..fwd.api_dsl import WorkspaceCarver - - carver = WorkspaceCarver( - workspace, - self.scratch_workspace_bytes( - has_bias=bias_tensor is not None, - bias_batch=(bias_tensor.shape[0] if bias_tensor is not None else 1), - has_sink=sinks is not None, - deterministic=bool(deterministic), - ), - "SdpabwdSm80", - ) - - pad_v = self.head_dim_v < self.flavor_d_v - pad_qk = self.head_dim_qk < self.flavor_d_qk - - def _stage(t: torch.Tensor, pad: bool, fd: int) -> torch.Tensor: - """Kernel-facing BSHD view of BHSD-logical ``t``: zero-copy when the - transpose is contiguous, otherwise gathered (and head-dim padded) - into carved staging — or allocated when no workspace was given.""" - view = t.transpose(1, 2) - d = view.shape[-1] - if pad: - if carver is not None: - bb, ss, hh, _ = view.shape - dst = carver.take(bb * ss * hh * fd, t.dtype).view(bb, ss, hh, fd) - dst[..., :d].copy_(view) - dst[..., d:].zero_() - return dst - return _pad_last_dim(view.contiguous() if not view.is_contiguous() else view, fd) - if view.is_contiguous(): - return view - if carver is not None: - dst = carver.take(t.numel(), t.dtype).view(view.shape) - dst.copy_(view) - return dst - return view.contiguous() - - # BHSD → BSHD for the kernel, in scratch_workspace_bytes()'s sizing - # order (Q, K, V, O, dO). - Q = _stage(q_tensor, pad_qk, self.flavor_d_qk) - K = _stage(k_tensor, pad_qk, self.flavor_d_qk) - V = _stage(v_tensor, pad_v, self.flavor_d_v) - O = _stage(o_tensor, pad_v, self.flavor_d_v) - dO = _stage(do_tensor, pad_v, self.flavor_d_v) - - # Build the feature-kwarg superset; drop any the kernel doesn't accept. - bw_kwargs = dict( - scale=scale_val, - mask=self.mask_token, - swa_window=int(self.swa_window_runtime), - right_bound=int(self.right_bound), - causal_bottom_right=self.causal_bottom_right, - seq_kv_lens=seq_kv_lens, - seq_len_q=seq_len_q, - bias=bias_tensor, - sinks=sinks, - rope_freqs=rope_freqs, - deterministic=bool(deterministic), - # Kernel-internal scratch: the unconsumed workspace tail (issue #514). - workspace=carver.remaining() if carver is not None else None, - ) - # Route plain dense MHA d=64 calls to the dedicated perf kernel - # (~2x faster on A100). The gate must stay exhaustive: the d64 - # kernel's ``backward(**_ignored)`` silently swallows any feature - # kwarg it does not implement, so an under-gated call would produce - # wrong gradients rather than an error. - if _d64_fast_path_eligible( - d_qk=self.head_dim_qk, - d_v=self.head_dim_v, - h_q=q_tensor.shape[1], - h_kv=k_tensor.shape[1], - s_q=q_tensor.shape[2], - s_kv=k_tensor.shape[2], - mask_token=self.mask_token, - right_bound=int(self.right_bound), - causal_bottom_right=self.causal_bottom_right, - bw_kwargs=bw_kwargs, - ): - kernel = _load_kernel_module("d64") - self._logger.debug("execute (bwd): routing to the dedicated d64 kernel") - accepted = inspect.signature(kernel.backward).parameters - bw_kwargs = {kk: vv for kk, vv in bw_kwargs.items() if kk in accepted} - - with _stream_ctx(current_stream): - res = kernel.backward(Q, K, V, dO, O, lse_tensor, **bw_kwargs) - dQ_k, dK_k, dV_k = res[0], res[1], res[2] - # backward() appends optional grads in a FIXED order: dBias (if bias), - # then dSink (if sinks). Reconstruct positions from what we passed. - _idx = 3 - dBias_k = None - dSink_k = None - if bias_tensor is not None: - dBias_k = res[_idx] - _idx += 1 - if sinks is not None: - dSink_k = res[_idx] - _idx += 1 - - # Slice off any d-padding, transpose BSHD → BHSD, copy into user tensors. - if pad_qk: - dQ_k = dQ_k[..., : self.head_dim_qk] - dK_k = dK_k[..., : self.head_dim_qk] - if pad_v: - dV_k = dV_k[..., : self.head_dim_v] - dq_tensor.copy_(dQ_k.transpose(1, 2)) - dk_tensor.copy_(dK_k.transpose(1, 2)) - dv_tensor.copy_(dV_k.transpose(1, 2)) - if dbias_tensor is not None and dBias_k is not None: - # dBias is head-major [., H, SQ, SKV] (like bias) — no transpose. - # copy_ casts in place; a .to() would allocate a staging tensor. - dbias_tensor.copy_(dBias_k) - if dsink_tensor is not None and dSink_k is not None: - dsink_tensor.copy_(dSink_k) - self._logger.debug("execute (bwd) completed") - - -# --------------------------------------------------------------------------- -# THD / varlen backward (mirrors fwd/api.py::_thd_forward). -# --------------------------------------------------------------------------- -def _thd_backward(q, k, v, o, do, lse, *, cu_q, cu_k, scale_softmax, is_causal, window_size, causal_bottom_right, sinks=None, deterministic=False): - """THD / varlen backward: q/k/v/o/do are PACKED ``[1, T, H, D]`` (BSHD, - B==1 — no transpose), ``lse`` is packed ``[1, H, T_q]`` (head-major, - matching the kernel's THD LSE layout), and cu_q/cu_k are ``[n_seq+1]`` - cumulative seqlens. Routes straight to the kernel's THD backward - (over-provisioned grid; MHA-only), reusing the flavor-pick + d-pad. - Returns packed ``[1, T, H, D]`` dQ/dK/dV (BHSD-equivalent for B==1).""" - d_qk = q.shape[-1] - d_v = v.shape[-1] - h_q = q.shape[2] - flavor = _pick_flavor(d_qk, d_v) - fdqk, fdv = _FLAVOR_DIMS[flavor] - # Resolve the default scale from the USER's head dim before padding: the - # kernel would otherwise derive 1/sqrt(D) from the padded flavor width - # (e.g. 1/sqrt(128) for a d=96 llama-flavor call) — silently wrong - # gradients. Mirrors the forward THD path. - if scale_softmax is None or scale_softmax == 0.0: - scale_softmax = 1.0 / math.sqrt(d_qk) - pad_qk = d_qk < fdqk - pad_v = d_v < fdv - if pad_qk: - q = _pad_last_dim(q, fdqk) - k = _pad_last_dim(k, fdqk) - if pad_v: - v = _pad_last_dim(v, fdv) - o = _pad_last_dim(o, fdv) - do = _pad_last_dim(do, fdv) - # mask token from cuDNN's (is_causal, window_size=(left,right)). - wl, wr = window_size - if is_causal and wl >= 0: - mask_token, swa = "causal_swa", wl - elif is_causal: - mask_token, swa = "causal", 0 - elif wl >= 0: - mask_token, swa = "swa", wl - else: - mask_token, swa = "none", 0 - right_bound = wr if (is_causal and wr is not None and wr > 0) else 0 - kernel = _load_kernel_module() - sinks_t = sinks.to(dtype=torch.float32, device=q.device).reshape(h_q).contiguous() if sinks is not None else None - bw_kwargs = dict( - scale=scale_softmax, - mask=mask_token, - swa_window=int(swa), - right_bound=int(right_bound), - causal_bottom_right=bool(causal_bottom_right), - cu_seqlens_q=cu_q, - cu_seqlens_k=cu_k, - sinks=sinks_t, - deterministic=bool(deterministic), - ) - acc = inspect.signature(kernel.backward).parameters - bw_kwargs = {kk: vv for kk, vv in bw_kwargs.items() if kk in acc} - res = kernel.backward(q, k, v, do, o, lse, **bw_kwargs) - dQ_k, dK_k, dV_k = res[0], res[1], res[2] - _idx = 3 - dSink_k = None - if sinks_t is not None: - dSink_k = res[_idx] - _idx += 1 - if pad_qk: - dQ_k = dQ_k[..., :d_qk].contiguous() - dK_k = dK_k[..., :d_qk].contiguous() - if pad_v: - dV_k = dV_k[..., :d_v].contiguous() - out = TupleDict(dq_tensor=dQ_k, dk_tensor=dK_k, dv_tensor=dV_k) - if dSink_k is not None: - out["dsink_tensor"] = dSink_k - return out - - -# --------------------------------------------------------------------------- -# Functional wrapper (mirrors the forward surface). -# --------------------------------------------------------------------------- -_cache_of_objects: dict = {} - - -def sdpa_bwd_wrapper_sm80( - q_tensor: torch.Tensor, - k_tensor: torch.Tensor, - v_tensor: torch.Tensor, - o_tensor: torch.Tensor, - do_tensor: torch.Tensor, - lse_tensor: torch.Tensor, - is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), - scale_softmax: Optional[float] = None, - causal_bottom_right: bool = False, - current_stream: Optional[cuda.CUstream] = None, - seq_kv_lens: Optional[torch.Tensor] = None, - seq_len_q: Optional[torch.Tensor] = None, - bias_tensor: Optional[torch.Tensor] = None, - sinks: Optional[torch.Tensor] = None, - rope_freqs: Optional[torch.Tensor] = None, - cum_seqlen_q_tensor: Optional[torch.Tensor] = None, - cum_seqlen_k_tensor: Optional[torch.Tensor] = None, - deterministic: bool = False, -) -> TupleDict: - """SM80 (A100) SDPA backward. - - Returns ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=... - [, dbias_tensor=...][, dsink_tensor=...])`` — BHSD grads; dBias - head-major [., H, SQ, SKV] when ``bias_tensor`` is given; dSink (H,) - fp32 when ``sinks`` is given (stable order: dq, dk, dv, dbias, dsink). - ALiBi and block_mask are not supported (use the graph API, which routes - them to the cuDNN backend); bias/dBias remain fully served. - """ - # THD / varlen: q/k/v/o/dO are PACKED [1, T, H, D] (BSHD) + cu_seqlens; - # lse is packed [1, H, T_q]. Dedicated path that skips the dense BHSD - # transpose + dense grad alloc (mirrors fwd/api.py's THD branch). - if cum_seqlen_q_tensor is not None: - # Reject dense-only features up front: _thd_backward accepts only - # sinks/deterministic, and silently computing gradients without a - # requested feature is worse than an error. - for label, present in ( - ("bias_tensor", bias_tensor is not None), - ("rope_freqs", rope_freqs is not None), - ("seq_kv_lens", seq_kv_lens is not None), - ("seq_len_q", seq_len_q is not None), - ): - if present: - raise NotImplementedError(f"SM80 SDPA THD (cum_seqlen_*) backward does not support {label}; the dense path serves it") - with _stream_ctx(current_stream): - return _thd_backward( - q_tensor, - k_tensor, - v_tensor, - o_tensor, - do_tensor, - lse_tensor, - cu_q=cum_seqlen_q_tensor, - cu_k=cum_seqlen_k_tensor, - scale_softmax=scale_softmax, - is_causal=is_causal, - window_size=window_size, - causal_bottom_right=causal_bottom_right, - sinks=sinks, - deterministic=deterministic, - ) - for nm, t in (("Q", q_tensor), ("V", v_tensor), ("O", o_tensor), ("dO", do_tensor)): - if t.ndim != 4: - raise ValueError(f"{nm} must be rank-4 BHSD; got {t.ndim}D") - - # Allocate grad outputs in cuDNN-FE BHSD-physical stride order (3,1,2,0): - # contiguous (B, S, H, D) then transpose to a (B, H, S, D) view. - b, h_q, s_q, d_qk = q_tensor.shape - d_v = v_tensor.shape[-1] - dq = torch.empty((b, s_q, h_q, d_qk), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) - # dK/dV take K/V leading shape (GQA: h_kv heads). - h_kv, s_kv = k_tensor.shape[1], k_tensor.shape[2] - dk = torch.empty((b, s_kv, h_kv, d_qk), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) - dv = torch.empty((b, s_kv, h_kv, d_v), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) - # dBias: fp32, same shape as bias ([., H, SQ, SKV]). - dbias = torch.zeros_like(bias_tensor, dtype=torch.float32) if bias_tensor is not None else None - # dSink: fp32 [H] (sink-logit gradient). - dsink = torch.zeros(h_q, dtype=torch.float32, device=q_tensor.device) if sinks is not None else None - - cache_key = ( - q_tensor.shape, - k_tensor.shape, - v_tensor.shape, - q_tensor.stride(), - k_tensor.stride(), - v_tensor.stride(), - q_tensor.dtype, - is_causal, - window_size, - scale_softmax, - causal_bottom_right, - seq_kv_lens is not None, - bias_tensor is not None, - (bias_tensor.dtype if bias_tensor is not None else None), - sinks is not None, - rope_freqs is not None, - q_tensor.device, - ) - sdpa_bwd = _cache_of_objects.get(cache_key) - if sdpa_bwd is None: - _logger.debug("sdpa_bwd_wrapper_sm80: building new SdpabwdSm80") - sdpa_bwd = SdpabwdSm80( - sample_q=q_tensor, - sample_k=k_tensor, - sample_v=v_tensor, - sample_o=o_tensor, - sample_do=do_tensor, - sample_lse=lse_tensor, - is_causal=is_causal, - window_size=window_size, - scale_softmax=scale_softmax, - causal_bottom_right=causal_bottom_right, - has_seq_kv_lens=seq_kv_lens is not None, - has_bias=bias_tensor is not None, - ) - assert sdpa_bwd.check_support(), "Unsupported configuration" - sdpa_bwd.compile() - _cache_of_objects[cache_key] = sdpa_bwd - - sdpa_bwd.execute( - q_tensor=q_tensor, - k_tensor=k_tensor, - v_tensor=v_tensor, - o_tensor=o_tensor, - do_tensor=do_tensor, - lse_tensor=lse_tensor, - dq_tensor=dq, - dk_tensor=dk, - dv_tensor=dv, - dbias_tensor=dbias, - dsink_tensor=dsink, - scale_softmax=scale_softmax, - current_stream=current_stream, - seq_kv_lens=seq_kv_lens, - seq_len_q=seq_len_q, - bias_tensor=bias_tensor, - sinks=sinks, - rope_freqs=rope_freqs, - deterministic=deterministic, - ) - - out = TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) - if dbias is not None: - out["dbias_tensor"] = dbias - if dsink is not None: - out["dsink_tensor"] = dsink - return out diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 686ab20be..3271075e2 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -5,6 +5,7 @@ from __future__ import annotations +import inspect import logging import math import os @@ -27,6 +28,7 @@ padded_head_dims as _sm120_padded_head_dims, ) from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, _torch_stream_context, ws_align +from cudnn.sdpa.fwd import config_sm80 as _fwd_config_sm80 _SM120_KERNEL_FILE = "bprop_f16_sm120.py" _SM120_DTYPE_QKV_CODE = { @@ -880,3 +882,1024 @@ def sdpa_bwd_wrapper_dsl_sm120( if dsink_tensor is not None: out["dsink_tensor"] = dsink_tensor return out + + +# ============================================================================= +# SM80 (A100) backward — SdpaBwdDslSm80 + the sdpa_bwd_wrapper_sm80 entry +# point. Ports the pre-TemplateParams ``SdpabwdSm80`` (bwd/api.py, deleted) +# onto the shared SdpaBwdDsl adapter contract, mirroring the forward port +# (#682): one lowering function (``lower_dsl_bwd``) now drives both backward +# cells. The kernels stay self-caching until the TemplateParams conversion +# (the #689 analogue; issue #604's sym_int THD extents land there). +# ============================================================================= + +_SM80_BWD_KERNEL_MOD = {} + +_SM80_BWD_FLAVOR_DIMS = { + name: (cfg.D_QK, cfg.D_V) + for name, cfg in ( + ("gptoss", _fwd_config_sm80.GPTOSS_CFG), + ("llama", _fwd_config_sm80.LLAMA_CFG), + ("dsv3", _fwd_config_sm80.DSV3_CFG), + ("qwen", _fwd_config_sm80.QWEN_CFG), + ) +} +_SM80_BWD_SUPPORTED_FLAVORS = ("gptoss", "llama", "dsv3", "qwen") + + +def _sm80_bwd_kernel_mod(key: str = "d64"): + """Lazily import + cache the dedicated d=64 SM80 BPROP kernel module. + + ``"d64"`` (the only key) is the plain-dense d=64 MHA perf variant (~2x + faster on A100); it supports NO features — its ``backward(**_ignored)`` + silently swallows every feature kwarg, so callers must never rely on the + signature filter and only select it through + :func:`_sm80_d64_fast_path_eligible`. The GENERIC kernel + (``bprop_f16_sm80``) is a TemplateParams module loaded per-specialization + via :func:`_load_sm80_bwd_module` instead. + """ + assert key == "d64", f"generic SM80 bwd kernels load via _load_sm80_bwd_module; got {key!r}" + if key not in _SM80_BWD_KERNEL_MOD: + from .kernels import bprop_d64_f16_sm80 as _mod + + _SM80_BWD_KERNEL_MOD[key] = _mod + return _SM80_BWD_KERNEL_MOD[key] + + +def _sm80_d64_fast_path_eligible(*, d_qk, d_v, h_q, h_kv, s_q, s_kv, mask_token, right_bound, causal_bottom_right, bw_kwargs) -> bool: + """Whether the dedicated d=64 kernel can serve this call EXACTLY. + + The perf variant computes a plain dense MHA backward and nothing else; + every condition here guards a feature it would silently ignore. + """ + d64 = _sm80_bwd_kernel_mod("d64") + if (d_qk, d_v) != (64, 64) or h_q != h_kv: + return False + if s_q % d64.M_BLOCK != 0 or s_kv % d64.N_BLOCK != 0: + return False + if mask_token != "none" or right_bound != 0 or causal_bottom_right: + return False + for feature in ("seq_kv_lens", "seq_len_q", "bias", "sinks", "rope_freqs"): + if bw_kwargs.get(feature) is not None: + return False + if bw_kwargs.get("deterministic"): + return False + return True + + +def _sm80_bwd_pick_flavor(d_qk: int, d_v: int) -> str: + """Smallest BPROP flavor whose ``(D_QK, D_V)`` envelope covers + ``(d_qk, d_v)`` (fdqk >= d_qk and fdv >= d_v); the user's heads are padded + up to the flavor dim. The kernel supports d_qk != d_v but requires the + (padded) d_qk >= d_v — the flavor list guarantees this (every flavor has + fdqk >= fdv, and a d_qk < d_v case lands on an equal-d flavor after pad).""" + for flavor in _SM80_BWD_SUPPORTED_FLAVORS: + fdqk, fdv = _SM80_BWD_FLAVOR_DIMS[flavor] + if d_qk == fdqk and d_v == fdv: + return flavor + for flavor in _SM80_BWD_SUPPORTED_FLAVORS: + fdqk, fdv = _SM80_BWD_FLAVOR_DIMS[flavor] + if d_qk <= fdqk and d_v <= fdv: + return flavor + raise ValueError(f"SM80 BPROP: no flavor envelope covers (D_QK={d_qk}, D_V={d_v}); " f"supported: {_SM80_BWD_FLAVOR_DIMS}.") + + +def _sm80_bwd_pad_last_dim(t: torch.Tensor, new_last: int) -> torch.Tensor: + """Zero-pad the trailing dim of an fp16/bf16 tensor up to ``new_last``.""" + old_last = t.shape[-1] + if old_last == new_last: + return t + if old_last > new_last: + raise ValueError(f"_sm80_bwd_pad_last_dim: tensor's last dim {old_last} exceeds target {new_last}") + pad = torch.zeros((*t.shape[:-1], new_last - old_last), dtype=t.dtype, device=t.device) + return torch.cat([t, pad], dim=-1).contiguous() + + +def _sm80_thd_backward(q, k, v, o, do, lse, *, cu_q, cu_k, scale_softmax, is_causal, window_size, causal_bottom_right, sinks=None, deterministic=False): + """THD / varlen backward: q/k/v/o/do are PACKED ``[1, T, H, D]`` (BSHD, + B==1 — no transpose), ``lse`` is packed ``[1, H, T_q]`` (head-major, + matching the kernel's THD LSE layout), and cu_q/cu_k are ``[n_seq+1]`` + cumulative seqlens. Loads the THD template specialization (packed token + totals are ``cute.sym_int`` dynamics — one artifact per (params, n_seq), + issue #604) and drives the kernel chain directly (over-provisioned grid; + GQA reduces over the query-head group). Returns packed ``[1, T, H, D]`` + dQ/dK/dV (BHSD-equivalent for B==1). + + Wrapper-only Rule-3 residual: the grid extents (max seqlen per side, the + logical sequence count) are RUNTIME ints read from cu_seqlens on the HOST + — a D2H sync this functional wrapper accepts (the graph/engine path never + routes THD here without materialized host seqlens). + """ + if sinks is not None: + raise NotImplementedError("SM80 THD bprop: attention sinks are dense-only") + if deterministic: + raise NotImplementedError("SM80 THD bprop: deterministic dQ is dense-only (no plan-time semaphore size under sym_int sq)") + d_qk = q.shape[-1] + d_v = v.shape[-1] + h_q = q.shape[2] + h_kv = k.shape[2] + flavor = _sm80_bwd_pick_flavor(d_qk, d_v) + fdqk, fdv = _SM80_BWD_FLAVOR_DIMS[flavor] + # Resolve the default scale from the USER's head dim before padding: the + # kernel would otherwise derive 1/sqrt(D) from the padded flavor width + # (e.g. 1/sqrt(128) for a d=96 llama-flavor call) — silently wrong + # gradients. Mirrors the forward THD path. + if scale_softmax is None or scale_softmax == 0.0: + scale_softmax = 1.0 / math.sqrt(d_qk) + pad_qk = d_qk < fdqk + pad_v = d_v < fdv + if pad_qk: + q = _sm80_bwd_pad_last_dim(q, fdqk) + k = _sm80_bwd_pad_last_dim(k, fdqk) + if pad_v: + v = _sm80_bwd_pad_last_dim(v, fdv) + o = _sm80_bwd_pad_last_dim(o, fdv) + do = _sm80_bwd_pad_last_dim(do, fdv) + # cuDNN's (is_causal, window_size=(left,right)) → mask params. + wl, wr = window_size + has_swa = wl is not None and wl >= 0 + swa = int(wl) if has_swa else 0 + right_bound = int(wr) if (is_causal and wr is not None and wr > 0) else 0 + from cudnn.sdpa.bwd.config_sm80 import bwd_params_for_flavor + + # NOTE: llama-swept tiles always (matching the dense adapter — the gptoss + # wide-Q-tile row stays unwired pending a perf gate); the flavor picks + # only the ENVELOPE dims, which must reach the compiled kernel (a flavor + # name alone would leave the template at its 128/128 defaults while the + # buffers pad to the envelope — OOB at d=64, wrong grads at 192/256). + params = bwd_params_for_flavor( + "llama", + io_bf16=(q.dtype == torch.bfloat16), + d_qk=fdqk, + d_v=fdv, + is_causal=bool(is_causal), + has_swa=has_swa, + causal_bottom_right=bool(causal_bottom_right) and (bool(is_causal) or has_swa), + thd_varlen=True, + sched_policy=_BWD_SCHED_NATURAL, # LPT+THD is a future tweak + ) + mod = _load_sm80_bwd_module(params) + # Host-side grid math (the wrapper-only D2H documented above). + cu_q_host = cu_q.to(dtype=torch.int32, device="cpu") + cu_k_host = cu_k.to(dtype=torch.int32, device="cpu") + n_seq = cu_q_host.numel() - 1 + assert cu_k_host.numel() == n_seq + 1, "cu_seqlens_q / cu_seqlens_k length mismatch" + max_sq = int((cu_q_host[1:] - cu_q_host[:-1]).max()) + max_skv = int((cu_k_host[1:] - cu_k_host[:-1]).max()) + c = mod.compile(1, h_q, h_kv, 0, 0, swa_window=swa, n_batch_logical=n_seq) + t_q = q.shape[1] + t_kv = k.shape[1] + dev = q.device + stream = cuda.CUstream(torch.cuda.current_stream(dev).cuda_stream) + q, k, v, o, do = (t.contiguous() for t in (q, k, v, o, do)) + lse_t = lse.to(dtype=torch.float32, device=dev).contiguous() + cu_q_t = cu_q.to(dtype=torch.int32, device=dev).contiguous() + cu_k_t = cu_k.to(dtype=torch.int32, device=dev).contiguous() + dq_acc = torch.zeros(1, t_q, h_q, fdqk, dtype=torch.float32, device=dev) + dQ_k = torch.empty(1, t_q, h_q, fdqk, dtype=q.dtype, device=dev) + # dK/dV write buffers carry the h_q query heads (one slice per query head); + # MHA: they ARE the outputs. GQA: reduced over the group below. + dk_ws = torch.empty(1, t_kv, h_q, fdqk, dtype=q.dtype, device=dev) + dv_ws = torch.empty(1, t_kv, h_q, fdv, dtype=q.dtype, device=dev) + dot = torch.empty(1, h_q, t_q, dtype=torch.float32, device=dev) + dummy_i32 = torch.zeros(1, dtype=torch.int32, device=dev) + dummy_f32 = torch.zeros(1, dtype=torch.float32, device=dev) + c.do_dot(_fd_tvm(o), _fd_tvm(do), _fd_tvm(dot), _int32(h_q * t_q), stream) + _sm80_bwd_call( + c.main, + q=q, + k=k, + v=v, + do=do, + dq_acc=dq_acc, + dk_ws=dk_ws, + dv_ws=dv_ws, + lse=lse_t, + do_dot=dot, + seq_kv=dummy_i32, + bias=dummy_f32, + dbias=dummy_f32, + rope_cs=dummy_f32, + cu_q=cu_q_t, + cu_k=cu_k_t, + seq_q=dummy_i32, + dq_sem=dummy_i32, + n_q_tiles=(t_q + params.tile_q - 1) // params.tile_q, + scale_log2=float(scale_softmax) * _BWD_LOG2E, + attn_scale=float(scale_softmax), + right_bound=right_bound, + inv_scale=1.0 / float(scale_softmax), + bias_bstride=0, + sem_q_stride=0, + grid_kv_tiles=(max_skv + params.tile_kv - 1) // params.tile_kv, + grid_batch=n_seq, + stream=stream, + ) + c.cast(_fd_tvm(dq_acc), _fd_tvm(dQ_k), _int32((t_q * h_q * fdqk) // 2), stream) + if h_q != h_kv: + dK_k = torch.empty(1, t_kv, h_kv, fdqk, dtype=q.dtype, device=dev) + dV_k = torch.empty(1, t_kv, h_kv, fdv, dtype=q.dtype, device=dev) + c.reduce_k(_fd_tvm(dk_ws), _fd_tvm(dK_k), _int32(t_kv * h_kv * fdqk), stream) + c.reduce_v(_fd_tvm(dv_ws), _fd_tvm(dV_k), _int32(t_kv * h_kv * fdv), stream) + else: + dK_k, dV_k = dk_ws, dv_ws + if pad_qk: + dQ_k = dQ_k[..., :d_qk].contiguous() + dK_k = dK_k[..., :d_qk].contiguous() + if pad_v: + dV_k = dV_k[..., :d_v].contiguous() + return TupleDict(dq_tensor=dQ_k, dk_tensor=dK_k, dv_tensor=dV_k) + + +# --------------------------------------------------------------------------- +# Functional wrapper (mirrors the forward surface). +# --------------------------------------------------------------------------- +_cache_of_objects: dict = {} + + +_SM80_BWD_KERNEL_FILE = "bprop_f16_sm80.py" +# The shared tile_dsl scheduler vocabulary maps identity onto the bwd grid +# decode (NATURAL == plain 3-D == 0, LPT == kv-major == 1). +from cudnn.frost.tile_dsl.constants import SCHED_LPT as _BWD_SCHED_LPT # noqa: E402 +from cudnn.frost.tile_dsl.constants import SCHED_NATURAL as _BWD_SCHED_NATURAL # noqa: E402 + + +def _load_sm80_bwd_module(params): + """Load one uniquely named backward kernel module per parameter set.""" + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels", _SM80_BWD_KERNEL_FILE) + return load_template(path, params, tag="sdpa_bwd_sm80") + + +def _sm80_bwd_call( + compiled, + *, + q, + k, + v, + do, + dq_acc, + dk_ws, + dv_ws, + lse, + do_dot, + seq_kv, + bias, + dbias, + rope_cs, + cu_q, + cu_k, + seq_q, + dq_sem, + n_q_tiles, + scale_log2, + attn_scale, + right_bound, + inv_scale, + bias_bstride, + sem_q_stride, + grid_kv_tiles, + grid_batch, + stream, +): + """Invoke one compiled main-bprop artifact (the traced ``_bprop_host`` + ABI: 17 tensors, then 9 runtime scalars and the launch stream).""" + import cutlass + from cutlass.cute.runtime import from_dlpack as _fd + + def fd(t): + # The kernels compile with --enable-tvm-ffi, so host-side conversions + # must produce TVM-FFI tensors regardless of the env latch. + return _fd(t, enable_tvm_ffi=True) + + compiled( + fd(q), + fd(k), + fd(v), + fd(do), + fd(dq_acc), + fd(dk_ws), + fd(dv_ws), + fd(lse), + fd(do_dot), + fd(seq_kv), + fd(bias), + fd(dbias), + fd(rope_cs), + fd(cu_q), + fd(cu_k), + fd(seq_q), + fd(dq_sem), + cutlass.Int32(n_q_tiles), + cutlass.Float32(scale_log2), + cutlass.Float32(attn_scale), + cutlass.Int32(right_bound), + cutlass.Float32(inv_scale), + cutlass.Int32(bias_bstride), + cutlass.Int32(sem_q_stride), + cutlass.Int32(grid_kv_tiles), + cutlass.Int32(grid_batch), + stream, + ) + + +_BWD_LOG2E = math.log2(math.e) + + +def _int32(v): + import cutlass + + return cutlass.Int32(int(v)) + + +def _fd_tvm(t): + from cutlass.cute.runtime import from_dlpack as _fd + + return _fd(t, enable_tvm_ffi=True) + + +class SdpaBwdDslSm80(SdpaBwdDsl): + """SM80 (A100) SDPA backward via the pre-TemplateParams CuTe-DSL kernels. + + Follows the SM120 adapter lifecycle (check_support → compile → execute) on + top of the self-caching SM80 kernel modules; the TemplateParams conversion + (the #689 analogue, with issue #604's sym_int THD extents) swaps the + kernel seam without touching this contract. SM80-only operands (bias → + dBias, RoPE) arrive as extra optional keywords, as the ``SdpaBwdDsl`` + contract permits. + + Layouts: any dense layout with the head dim innermost-contiguous is + served — non-BSHD-compact operands (dense_flex) gather into carved + staging, a strided stats input gathers to the packed LSE the kernels + read (``Capabilities.strided_stats``), and head dims inside a flavor + envelope pad host-side into the same carved buffers (issue #514: with a + workspace provided, execute allocates nothing). + """ + + def __init__( + self, *args, has_bias: bool = False, bias_is_fp32: bool = True, bias_batch: int = 1, has_rope: bool = False, rope_max_s: int = 0, **kwargs + ) -> None: + # SM80-only plan-time facts (scratch sizing + template identity); the + # base contract carries everything else. + self._has_bias = bool(has_bias) + self._bias_is_fp32 = bool(bias_is_fp32) + self._bias_batch = int(bias_batch) + self._has_rope = bool(has_rope) + self._rope_max_s = int(rope_max_s) + super().__init__(*args, **kwargs) + + def _initialize_implementation(self) -> None: + self.flavor: Optional[str] = None + self.flavor_d_qk: Optional[int] = None + self.flavor_d_v: Optional[int] = None + self.mask_token: Optional[str] = None + self.swa_window_runtime: int = 0 + self.right_bound_runtime: int = 0 + self._dummy_cache: dict = {} + + def _dummy(self, key: str, device: torch.device, factory) -> torch.Tensor: + """A cached device-local dummy for a dead ABI slot (AGENTS.md Rule 1: + no per-execute allocation; the matching has_* Constexpr is False so + the kernel never reads it).""" + cache_key = (key, device) + tensor = self._dummy_cache.get(cache_key) + if tensor is None: + tensor = factory() + self._dummy_cache[cache_key] = tensor + return tensor + + # ------------------------------------------------------------------ + def check_support(self) -> bool: + self._logger.debug("Entering check_support") + + for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc, self.do_desc): + self._value_error_if(desc.ndim != 4, f"{desc.name} must be rank-4 (B, H, S, D); got {desc.ndim}") + + b, h_qo, s_qo, d_qk = self.q_desc.shape + _, h_kv, s_kv, _ = self.k_desc.shape + _, _, _, d_v = self.v_desc.shape + + self._check_tensor_shape(self.q_desc, (b, h_qo, s_qo, d_qk), name="Q") + self._check_tensor_shape(self.k_desc, (b, h_kv, s_kv, d_qk), name="K") + self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_v), name="V") + self._check_tensor_shape(self.o_desc, (b, h_qo, s_qo, d_v), name="O") + self._check_tensor_shape(self.do_desc, (b, h_qo, s_qo, d_v), name="dO") + self._check_tensor_shape(self.dq_desc, (b, h_qo, s_qo, d_qk), name="dQ") + self._check_tensor_shape(self.dk_desc, (b, h_kv, s_kv, d_qk), name="dK") + self._check_tensor_shape(self.dv_desc, (b, h_kv, s_kv, d_v), name="dV") + + for label, val in (("B", b), ("H_q", h_qo), ("H_kv", h_kv), ("S_q", s_qo), ("S_kv", s_kv), ("D_QK", d_qk), ("D_V", d_v)): + self._value_error_if(int(val) <= 0, f"{label} must be > 0; got {val}") + self._value_error_if(h_qo % h_kv != 0, f"H_q ({h_qo}) must be divisible by H_kv ({h_kv}) for GQA / MQA") + + # The kernel supports d_qk != d_v (split sub-groups) but requires + # d_qk >= d_v; head dims inside a flavor envelope pad host-side. + self._value_error_if(d_qk < d_v, f"SM80 BPROP requires D_QK >= D_V; got D_QK={d_qk}, D_V={d_v}") + max_dqk = max(fdqk for fdqk, _ in _SM80_BWD_FLAVOR_DIMS.values()) + max_dv = max(fdv for _, fdv in _SM80_BWD_FLAVOR_DIMS.values()) + self._value_error_if( + d_qk > max_dqk or d_v > max_dv, + f"SM80 BPROP: head dim (D_QK={d_qk}, D_V={d_v}) exceeds supported " f"envelope (D_QK<={max_dqk}, D_V<={max_dv}); larger heads not yet ported.", + ) + + self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") + for desc in (self.k_desc, self.v_desc, self.o_desc, self.do_desc, self.dq_desc, self.dk_desc, self.dv_desc): + self._check_dtype(desc, self.dtype, name=desc.name, extra_error_msg=f"{desc.name} must match Q dtype (FP16/BF16)") + self._check_dtype(self.stats_desc, torch.float32, name="stats") + stats_shape = tuple(self.stats_desc.shape) + self._value_error_if( + stats_shape not in ((b, h_qo, s_qo), (b, h_qo, s_qo, 1)), + f"stats must be (B, H_q, S_q[, 1]) = ({b}, {h_qo}, {s_qo}[, 1]); got {stats_shape}", + ) + # Any stats layout is served: a non-contiguous input gathers into + # carved staging (the kernels read a packed natural-log LSE). + self._stats_needs_stage = not self.stats_desc.is_contiguous() + + self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM80 BPROP") + device = self.q_desc.device + major, minor = torch.cuda.get_device_capability(device) + self._value_error_if((major, minor) != (8, 0), f"SdpaBwdDslSm80 requires SM80 (A100); found SM{major}{minor} on {device}") + + self._value_error_if( + self.tile_m is not None or self.tile_n is not None, + "SM80 BPROP wires no tile knobs; tile_m/tile_n must be unset", + ) + + self.flavor = _sm80_bwd_pick_flavor(d_qk, d_v) + self.flavor_d_qk, self.flavor_d_v = _SM80_BWD_FLAVOR_DIMS[self.flavor] + + # ---- mask token (same resolution as the forward adapter) ---------- + swa_left = -1 if self.window_size_left is None else int(self.window_size_left) + swa_right = 0 if self.window_size_right is None else int(self.window_size_right) + self.right_bound_runtime = 0 + if self.is_causal: + self.mask_token = "causal" if swa_left < 0 else "causal_swa" + self.swa_window_runtime = max(0, swa_left) if swa_left >= 0 else 0 + self.right_bound_runtime = max(0, swa_right) + elif swa_left >= 0: + self._not_implemented_error_if(swa_right > 0, "SM80 BPROP: non-causal SWA with window_size_right > 0 unsupported") + self.mask_token = "swa" + self.swa_window_runtime = swa_left + else: + self._not_implemented_error_if( + swa_right > 0, + "SM80 BPROP: window_size_right without a left window or is_causal=True has no effect; pass is_causal=True or a left window", + ) + self.mask_token = "none" + self.swa_window_runtime = 0 + self._value_error_if( + self.causal_bottom_right and not (self.is_causal or swa_left >= 0), + "SM80 BPROP: causal_bottom_right requires is_causal and/or a left window", + ) + + if self.scale_softmax is None or self.scale_softmax == 0.0: + self.scale_softmax = 1.0 / math.sqrt(d_qk) + + self.batch_size = int(b) + self.s_q_max = int(s_qo) + self.s_k_max = int(s_kv) + self.h_q = int(h_qo) + self.h_kv = int(h_kv) + self.head_dim_qk = int(d_qk) + self.head_dim_v = int(d_v) + + self._is_supported = True + self._logger.debug("check_support completed") + return True + + # ------------------------------------------------------------------ + def compile(self) -> None: + """Plan-time JIT: build the TemplateParams from the plan facts, load + the specialized module via ``frost.template_loader`` (same seam as the + SM120 adapter), and compile the full kernel chain for this shape. + The dedicated plain-dense d=64 fast path keeps its self-caching module + (dense-only — issue #604 concerns the THD extents, which never route + there); its JIT happens on the first execute.""" + self._logger.debug("Entering compile") + self._ensure_support_checked() + from cudnn.sdpa.bwd.config_sm80 import bwd_params_for_flavor + + # Scheduler resolution (the old backward()'s "auto"): kv-major LPT for + # causal load-balance, the plain 3-D grid otherwise; the deterministic + # relay REQUIRES the plain decode (kv_tile == blockIdx.x). + sched = _BWD_SCHED_LPT if (self.is_causal and not self.deterministic) else _BWD_SCHED_NATURAL + # NOTE: the generic pipeline always ran the llama-swept tile point + # regardless of flavor (the old backward() defaults); the gptoss + # wide-Q-tile row stays unwired pending an adapter-level perf gate. + self._params = bwd_params_for_flavor( + "llama", + io_bf16=self.dtype == torch.bfloat16, + d_qk=self.flavor_d_qk, + d_v=self.flavor_d_v, + is_causal=self.is_causal, + has_swa=self.swa_window_runtime > 0 or (self.window_size_left is not None and self.window_size_left >= 0), + causal_bottom_right=self.causal_bottom_right, + has_seq_kv_lens=self.seq_kv_lens_present, + has_seq_q_lens=self.seq_q_lens_present, + has_bias=self._has_bias, + bias_is_fp32=self._bias_is_fp32, + bias_broadcast=self._bias_batch == 1, + has_sink=self.sink_desc is not None, + has_rope=self._has_rope, + deterministic=self.deterministic, + thd_varlen=False, + sched_policy=sched, + ) + # RoPE preconditions the old backward() asserted (the params validator + # covers d_qk <= 128; these two involve the shape, known only here). + if self._has_rope: + self._value_error_if( + self._rope_max_s < max(self.s_q_max, self.s_k_max), + f"rope_freqs rows ({self._rope_max_s}) must cover max(S_q={self.s_q_max}, S_kv={self.s_k_max})", + ) + self._not_implemented_error_if( + bool(self.s_q_max % self._params.tile_q or self.s_k_max % self._params.tile_kv), + "SM80 bprop: RoPE requires S_q/S_kv tile-aligned", + ) + # d64 fast-path gate: every input is plan-time state now. + self._use_d64 = _sm80_d64_fast_path_eligible( + d_qk=self.head_dim_qk, + d_v=self.head_dim_v, + h_q=self.h_q, + h_kv=self.h_kv, + s_q=self.s_q_max, + s_kv=self.s_k_max, + mask_token=self.mask_token, + right_bound=int(self.right_bound_runtime), + causal_bottom_right=self.causal_bottom_right, + bw_kwargs=dict( + seq_kv_lens=object() if self.seq_kv_lens_present else None, + seq_len_q=object() if self.seq_q_lens_present else None, + bias=object() if self._has_bias else None, + sinks=object() if self.sink_desc is not None else None, + rope_freqs=object() if self._has_rope else None, + deterministic=self.deterministic, + ), + ) + if self._use_d64: + self._kmod = None + self._compiled_kernel = True # d64 self-caches on first execute + else: + self._kmod = _load_sm80_bwd_module(self._params) + self._compiled_kernel = self._kmod.compile( + b=self.batch_size, + h=self.h_q, + h_kv=self.h_kv, + sq=self.s_q_max, + skv=self.s_k_max, + swa_window=int(self.swa_window_runtime), + rope_max_s=self._rope_max_s, + n_batch_logical=0, + ) + self._logger.debug("compile completed") + + # ------------------------------------------------------------------ + def _bshd_gather_bytes(self, desc) -> int: + """Bytes to gather ``desc`` into a compact BSHD buffer, or 0 when its + BSHD transpose is already contiguous.""" + b, h, s, d = desc.shape + if tuple(desc.stride) == (s * h * d, d, h * d, 1): + return 0 + return ws_align(b * h * s * d * 2) # fp16/bf16 only on this row + + def scratch_workspace_bytes(self) -> int: + """Per-execute scratch (issue #514): dense_flex gathers / head-dim pad + staging for the five input operands, strided-stats staging, and the + kernel-internal buffers (``bprop_f16_sm80.scratch_bytes``; the generic + kernel's set covers the d64 fast path's). All plan-time state — no + arguments.""" + self._ensure_support_checked() + from .kernels import bprop_f16_sm80 as _kmod + + elem = 2 # fp16/bf16 — check_support admits no other input dtype + b, hq, hkv = self.batch_size, self.h_q, self.h_kv + sq, skv = self.s_q_max, self.s_k_max + fdqk, fdv = self.flavor_d_qk, self.flavor_d_v + pad_qk = self.head_dim_qk < fdqk + pad_v = self.head_dim_v < fdv + total = 0 + # Pad / gather staging, in execute()'s carve order (Q, K, V, O, dO). + for desc, s_len, hh, pad, fd in ( + (self.q_desc, sq, hq, pad_qk, fdqk), + (self.k_desc, skv, hkv, pad_qk, fdqk), + (self.v_desc, skv, hkv, pad_v, fdv), + (self.o_desc, sq, hq, pad_v, fdv), + (self.do_desc, sq, hq, pad_v, fdv), + ): + if pad: + total += ws_align(int(desc.shape[0]) * s_len * hh * fd * elem) + else: + total += self._bshd_gather_bytes(desc) + if self._stats_needs_stage: + total += ws_align(b * hq * sq * 4) + total += _kmod.scratch_bytes( + B=b, + SQ=sq, + SKV=skv, + H=hq, + Hk=hkv, + d_qk=fdqk, + d_v=fdv, + io_bytes=elem, + deterministic=self.deterministic, + has_bias=self._has_bias, + bias_batch=self._bias_batch, + has_sink=self.sink_desc is not None, + need_do_dot=True, + ) + return total + + # ------------------------------------------------------------------ + def execute( + self, + q_tensor: torch.Tensor, + k_tensor: torch.Tensor, + v_tensor: torch.Tensor, + o_tensor: torch.Tensor, + do_tensor: torch.Tensor, + stats_tensor: torch.Tensor, + dq_tensor: torch.Tensor, + dk_tensor: torch.Tensor, + dv_tensor: torch.Tensor, + scale_softmax: Optional[float] = None, + workspace: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + seq_q_lens: Optional[torch.Tensor] = None, + seq_kv_lens: Optional[torch.Tensor] = None, + sink_tensor: Optional[torch.Tensor] = None, + dsink_tensor: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + dbias_tensor: Optional[torch.Tensor] = None, + rope_freqs: Optional[torch.Tensor] = None, + ) -> None: + self._logger.debug("Entering execute") + if self._compiled_kernel is None: + raise RuntimeError("SdpaBwdDslSm80 is not compiled") + + # Init-time flags are compile-time facts; execute must match them + # exactly, in both directions (Hard Rule 1). + self._value_error_if(self._has_bias != (bias_tensor is not None), "bias presence must match the plan (has_bias)") + self._value_error_if(self._has_rope != (rope_freqs is not None), "rope_freqs presence must match the plan (has_rope)") + self._value_error_if((self.sink_desc is not None) != (sink_tensor is not None), "sink presence must match the plan") + self._value_error_if(self.seq_kv_lens_present != (seq_kv_lens is not None), "seq_kv_lens presence must match the plan") + self._value_error_if(self.seq_q_lens_present != (seq_q_lens is not None), "seq_q_lens presence must match the plan") + + scale_val = self.scale_softmax if (scale_softmax is None or scale_softmax == 0.0) else float(scale_softmax) + device = q_tensor.device + launch_stream = self._get_default_stream(current_stream) + + # Per-execute scratch: carved from the caller's workspace when one is + # provided (the engine lowering passes one sized by + # scratch_workspace_bytes(); issue #514), otherwise allocated (the + # standalone wrapper path). Carve order mirrors the sizing order: + # operand staging first, then the kernel-internal buffers. + carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaBwdDslSm80") if workspace is not None else None + + def _take(numel, dtype, zero=False, shape=None): + if carver is not None: + t = carver.take(numel, dtype) + if zero: + t.zero_() + else: + t = (torch.zeros if zero else torch.empty)(numel, dtype=dtype, device=device) + return t.view(shape) if shape is not None else t + + with _torch_stream_context(current_stream, device): + pad_qk = self.head_dim_qk < self.flavor_d_qk + pad_v = self.head_dim_v < self.flavor_d_v + + def _stage(t: torch.Tensor, pad: bool, fd: int) -> torch.Tensor: + """Kernel-facing BSHD view of BHSD-logical ``t``: zero-copy when + the transpose is contiguous, else gathered (and head-dim + padded) into carved staging — or allocated on the wrapper + path.""" + view = t.transpose(1, 2) + d = view.shape[-1] + if pad: + bb, ss, hh, _ = view.shape + dst = _take(bb * ss * hh * fd, t.dtype, shape=(bb, ss, hh, fd)) + dst[..., :d].copy_(view) + dst[..., d:].zero_() + return dst + if view.is_contiguous(): + return view + dst = _take(t.numel(), t.dtype, shape=view.shape) + dst.copy_(view) + return dst + + # BHSD → BSHD, in scratch_workspace_bytes()'s sizing order. + Q = _stage(q_tensor, pad_qk, self.flavor_d_qk) + K = _stage(k_tensor, pad_qk, self.flavor_d_qk) + V = _stage(v_tensor, pad_v, self.flavor_d_v) + O = _stage(o_tensor, pad_v, self.flavor_d_v) + dO = _stage(do_tensor, pad_v, self.flavor_d_v) + + # Stats → packed (B, H, S) LSE; squeeze(-1) is a valid view for + # any (B, H, S, 1) strides; strided inputs gather into staging + # (the kernels READ a packed LSE — raw-pointer addressing). + lse = stats_tensor.squeeze(-1) if stats_tensor.ndim == 4 else stats_tensor + if not lse.is_contiguous(): + lse_stage = _take(lse.numel(), torch.float32, shape=lse.shape) + lse_stage.copy_(lse) + lse = lse_stage + + # --- d64 fast path: the dedicated plain-dense MHA kernel keeps + # its legacy self-caching module (dense-only; #604 is THD-only). + if self._use_d64: + d64 = _sm80_bwd_kernel_mod("d64") + res = d64.backward(Q, K, V, dO, O, lse, scale=scale_val, workspace=carver.remaining() if carver is not None else None) + dQ_k, dK_k, dV_k = res[0], res[1], res[2] + dq_tensor.copy_(dQ_k.transpose(1, 2)) + dk_tensor.copy_(dK_k.transpose(1, 2)) + dv_tensor.copy_(dV_k.transpose(1, 2)) + self._logger.debug("execute completed (d64 fast path)") + return + + p_ = self._params + c = self._compiled_kernel + b, hq, hkv = self.batch_size, self.h_q, self.h_kv + sq, skv = self.s_q_max, self.s_k_max + fdqk, fdv = self.flavor_d_qk, self.flavor_d_v + gqa = hq != hkv + + # Kernel-internal scratch, in the module scratch_bytes() order. + dsink_acc = _take(hq, torch.float32, zero=True) if sink_tensor is not None else None + dq_acc = _take(b * sq * hq * fdqk, torch.float32, zero=True, shape=(b, sq, hq, fdqk)) + dq_ws = _take(b * sq * hq * fdqk, self.dtype, shape=(b, sq, hq, fdqk)) + if p_.deterministic: + sem_units = max(b * hq * c.sem_q_stride, 1) + dq_sem = _take(sem_units, torch.int32, zero=True) + else: + dq_sem = self._dummy("zero_i32", device, lambda: torch.zeros(1, dtype=torch.int32, device=device)) + # dK/dV: MHA at native dims binds the caller's compact-BSHD views + # directly (no copy-back, mirroring the SM120 adapter); anything + # else stages (GQA per-q-head partials, head-dim pads, dense_flex). + dk_view = dk_tensor.transpose(1, 2) + dv_view = dv_tensor.transpose(1, 2) + dk_direct = not gqa and not pad_qk and dk_view.is_contiguous() + dv_direct = not gqa and not pad_v and dv_view.is_contiguous() + dk_ws = dk_view if dk_direct else _take(b * skv * hq * fdqk, self.dtype, shape=(b, skv, hq, fdqk)) + dv_ws = dv_view if dv_direct else _take(b * skv * hq * fdv, self.dtype, shape=(b, skv, hq, fdv)) + dk_out = dv_out = None + if gqa: + dk_out = _take(b * skv * hkv * fdqk, self.dtype, shape=(b, skv, hkv, fdqk)) + dv_out = _take(b * skv * hkv * fdv, self.dtype, shape=(b, skv, hkv, fdv)) + if bias_tensor is not None: + self._value_error_if(not bias_tensor.is_contiguous(), "bias must be contiguous") + self._value_error_if( + tuple(bias_tensor.shape) != (self._bias_batch, hq, sq, skv), + f"bias must be ({self._bias_batch}, {hq}, {sq}, {skv}); got {tuple(bias_tensor.shape)}", + ) + bias_b = bias_tensor + dbias_acc = _take(self._bias_batch * hq * sq * skv, torch.float32, zero=True, shape=(self._bias_batch, hq, sq, skv)) + else: + bias_dt = torch.float32 if p_.bias_is_fp32 else self.dtype + bias_b = self._dummy(f"one_{bias_dt}", device, lambda: torch.ones(1, dtype=bias_dt, device=device)) + dbias_acc = self._dummy("zero_f32", device, lambda: torch.zeros(1, dtype=torch.float32, device=device)) + dot = _take(b * hq * sq, torch.float32, shape=(b, hq, sq)) + if rope_freqs is not None: + # (cos, sin) table build — wrapper-only fusion (the engine row + # never admits RoPE); per-execute by contract, like the caller + # passing fresh angle tables. + d2 = fdqk // 2 + rf = rope_freqs.to(dtype=torch.float32, device=device).reshape(rope_freqs.shape[0], -1) + self._value_error_if(rf.shape[1] < d2, f"rope_freqs last dim ({rf.shape[1]}) must be >= d_qk//2 ({d2})") + self._value_error_if( + rf.shape[0] != self._rope_max_s, f"rope_freqs rows ({rf.shape[0]}) must equal the compiled rope_max_s ({self._rope_max_s})" + ) + angles = rf[:, :d2] + rope_b = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() + else: + rope_b = self._dummy("zero_f32", device, lambda: torch.zeros(1, dtype=torch.float32, device=device)) + if seq_kv_lens is not None: + seq_kv_b = seq_kv_lens.reshape(-1) + self._value_error_if(seq_kv_b.dtype != torch.int32 or not seq_kv_b.is_contiguous(), "seq_kv_lens must be contiguous int32") + else: + seq_kv_b = self._dummy("zero_i32", device, lambda: torch.zeros(1, dtype=torch.int32, device=device)) + if seq_q_lens is not None: + seq_q_b = seq_q_lens.reshape(-1) + self._value_error_if(seq_q_b.dtype != torch.int32 or not seq_q_b.is_contiguous(), "seq_q_lens must be contiguous int32") + else: + seq_q_b = self._dummy("zero_i32", device, lambda: torch.zeros(1, dtype=torch.int32, device=device)) + if sink_tensor is not None: + sinks_b = sink_tensor.reshape(-1) + self._value_error_if(sinks_b.dtype != torch.float32 or not sinks_b.is_contiguous(), "sinks must be contiguous fp32") + cu_dummy = self._dummy("zero_i32", device, lambda: torch.zeros(1, dtype=torch.int32, device=device)) + + # --- launch chain: do_dot → (dSink) → main → dQ cast → (GQA reduce) + c.do_dot(_fd_tvm(O), _fd_tvm(dO), _fd_tvm(dot), _int32(b * hq * sq), launch_stream) + if sink_tensor is not None: + c.dsink(_fd_tvm(lse), _fd_tvm(dot), _fd_tvm(sinks_b), _fd_tvm(dsink_acc), _int32(b * hq), launch_stream) + _sm80_bwd_call( + c.main, + q=Q, + k=K, + v=V, + do=dO, + dq_acc=dq_acc, + dk_ws=dk_ws, + dv_ws=dv_ws, + lse=lse, + do_dot=dot, + seq_kv=seq_kv_b, + bias=bias_b, + dbias=dbias_acc, + rope_cs=rope_b, + cu_q=cu_dummy, + cu_k=cu_dummy, + seq_q=seq_q_b, + dq_sem=dq_sem, + n_q_tiles=(sq + p_.tile_q - 1) // p_.tile_q, + scale_log2=scale_val * _BWD_LOG2E, + attn_scale=scale_val, + right_bound=int(self.right_bound_runtime), + inv_scale=1.0 / float(scale_val), + bias_bstride=0 if self._bias_batch == 1 else hq * sq * skv, + sem_q_stride=c.sem_q_stride, + grid_kv_tiles=0, + grid_batch=0, + stream=launch_stream, + ) + c.cast(_fd_tvm(dq_acc), _fd_tvm(dq_ws), _int32((b * sq * hq * fdqk) // 2), launch_stream) + if gqa: + c.reduce_k(_fd_tvm(dk_ws), _fd_tvm(dk_out), _int32(b * skv * hkv * fdqk), launch_stream) + c.reduce_v(_fd_tvm(dv_ws), _fd_tvm(dv_out), _int32(b * skv * hkv * fdv), launch_stream) + + # Copy-backs: slice any d-padding, transpose BSHD → BHSD into the + # caller's buffers (copy_ casts in place; no .to()). Direct-bound + # dK/dV already landed in place. + dq_src = dq_ws[..., : self.head_dim_qk] if pad_qk else dq_ws + dq_tensor.copy_(dq_src.transpose(1, 2)) + dk_src = dk_out if gqa else dk_ws + dv_src = dv_out if gqa else dv_ws + if not dk_direct: + dk_tensor.copy_((dk_src[..., : self.head_dim_qk] if pad_qk else dk_src).transpose(1, 2)) + if not dv_direct: + dv_tensor.copy_((dv_src[..., : self.head_dim_v] if pad_v else dv_src).transpose(1, 2)) + if dbias_tensor is not None and bias_tensor is not None: + dbias_tensor.copy_(dbias_acc) + if dsink_tensor is not None and dsink_acc is not None: + dsink_tensor.view(-1).copy_(dsink_acc) + self._logger.debug("execute completed") + + +_sm80_bwd_cache: dict = {} + + +def sdpa_bwd_wrapper_sm80( + q_tensor: torch.Tensor, + k_tensor: torch.Tensor, + v_tensor: torch.Tensor, + o_tensor: torch.Tensor, + do_tensor: torch.Tensor, + lse_tensor: torch.Tensor, + is_causal: bool = False, + window_size: "tuple[int, int]" = (-1, -1), + scale_softmax: Optional[float] = None, + causal_bottom_right: bool = False, + current_stream: Optional[cuda.CUstream] = None, + seq_kv_lens: Optional[torch.Tensor] = None, + seq_len_q: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + rope_freqs: Optional[torch.Tensor] = None, + cum_seqlen_q_tensor: Optional[torch.Tensor] = None, + cum_seqlen_k_tensor: Optional[torch.Tensor] = None, + deterministic: bool = False, +) -> TupleDict: + """SM80 (A100) SDPA backward. + + Returns ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=... + [, dbias_tensor=...][, dsink_tensor=...])`` — BHSD grads; dBias + head-major [., H, SQ, SKV] when ``bias_tensor`` is given; dSink (H,) + fp32 when ``sinks`` is given (stable order: dq, dk, dv, dbias, dsink). + ALiBi and block_mask are not supported (use the graph API, which routes + them to the cuDNN backend); bias/dBias remain fully served. + """ + # THD / varlen: q/k/v/o/dO are PACKED [1, T, H, D] (BSHD) + cu_seqlens; + # lse is packed [1, H, T_q]. Dedicated path that skips the dense BHSD + # transpose + dense grad alloc (mirrors the forward THD branch). + if cum_seqlen_q_tensor is not None: + for label, present in ( + ("bias_tensor", bias_tensor is not None), + ("rope_freqs", rope_freqs is not None), + ("seq_kv_lens", seq_kv_lens is not None), + ("seq_len_q", seq_len_q is not None), + ): + if present: + raise NotImplementedError(f"SM80 SDPA THD (cum_seqlen_*) backward does not support {label}; the dense path serves it") + with _torch_stream_context(current_stream, q_tensor.device): + return _sm80_thd_backward( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + do_tensor, + lse_tensor, + cu_q=cum_seqlen_q_tensor, + cu_k=cum_seqlen_k_tensor, + scale_softmax=scale_softmax, + is_causal=is_causal, + window_size=window_size, + causal_bottom_right=causal_bottom_right, + sinks=sinks, + deterministic=deterministic, + ) + for nm, t in (("Q", q_tensor), ("V", v_tensor), ("O", o_tensor), ("dO", do_tensor)): + if t.ndim != 4: + raise ValueError(f"{nm} must be rank-4 BHSD; got {t.ndim}D") + + # Allocate grad outputs in cuDNN-FE BHSD-physical stride order (3,1,2,0): + # contiguous (B, S, H, D) then transpose to a (B, H, S, D) view. + b, h_q, s_q, d_qk = q_tensor.shape + d_v = v_tensor.shape[-1] + dq = torch.empty((b, s_q, h_q, d_qk), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) + h_kv, s_kv = k_tensor.shape[1], k_tensor.shape[2] + dk = torch.empty((b, s_kv, h_kv, d_qk), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) + dv = torch.empty((b, s_kv, h_kv, d_v), dtype=q_tensor.dtype, device=q_tensor.device).transpose(1, 2) + dbias = torch.zeros_like(bias_tensor, dtype=torch.float32) if bias_tensor is not None else None + dsink = torch.zeros(h_q, dtype=torch.float32, device=q_tensor.device) if sinks is not None else None + + cache_key = ( + q_tensor.shape, + k_tensor.shape, + v_tensor.shape, + q_tensor.stride(), + k_tensor.stride(), + v_tensor.stride(), + q_tensor.dtype, + is_causal, + window_size, + scale_softmax, + causal_bottom_right, + deterministic, + seq_kv_lens is not None, + seq_len_q is not None, + bias_tensor is not None, + (bias_tensor.dtype if bias_tensor is not None else None), + (bias_tensor.shape[0] if bias_tensor is not None else None), + sinks is not None, + rope_freqs is not None, + (int(rope_freqs.shape[0]) if rope_freqs is not None else 0), + q_tensor.device, + ) + sdpa_bwd = _sm80_bwd_cache.get(cache_key) + if sdpa_bwd is None: + _logger.debug("sdpa_bwd_wrapper_sm80: building new SdpaBwdDslSm80") + wl, wr = window_size + sdpa_bwd = SdpaBwdDslSm80( + sample_q=q_tensor, + sample_k=k_tensor, + sample_v=v_tensor, + sample_o=o_tensor, + sample_do=do_tensor, + sample_stats=lse_tensor, + sample_dq=dq, + sample_dk=dk, + sample_dv=dv, + sample_sink=sinks, + sample_dsink=dsink, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_size_left=(None if wl is None or wl < 0 else int(wl)), + window_size_right=(None if wr is None or wr < 0 else int(wr)), + deterministic=deterministic, + scale_softmax=scale_softmax, + seq_kv_lens_present=seq_kv_lens is not None, + seq_q_lens_present=seq_len_q is not None, + has_bias=bias_tensor is not None, + bias_is_fp32=(bias_tensor.dtype == torch.float32 if bias_tensor is not None else True), + bias_batch=(int(bias_tensor.shape[0]) if bias_tensor is not None else 1), + has_rope=rope_freqs is not None, + rope_max_s=(int(rope_freqs.shape[0]) if rope_freqs is not None else 0), + ) + assert sdpa_bwd.check_support(), "Unsupported configuration" + sdpa_bwd.compile() + _sm80_bwd_cache[cache_key] = sdpa_bwd + + sdpa_bwd.execute( + q_tensor=q_tensor, + k_tensor=k_tensor, + v_tensor=v_tensor, + o_tensor=o_tensor, + do_tensor=do_tensor, + stats_tensor=lse_tensor, + dq_tensor=dq, + dk_tensor=dk, + dv_tensor=dv, + dbias_tensor=dbias, + dsink_tensor=dsink, + scale_softmax=scale_softmax, + current_stream=current_stream, + seq_kv_lens=seq_kv_lens, + seq_q_lens=seq_len_q, + sink_tensor=sinks, + bias_tensor=bias_tensor, + rope_freqs=rope_freqs, + ) + + out = TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) + if dbias is not None: + out["dbias_tensor"] = dbias + if dsink is not None: + out["dsink_tensor"] = dsink + return out diff --git a/python/cudnn/sdpa/bwd/config_sm80.py b/python/cudnn/sdpa/bwd/config_sm80.py index fc7bd51fb..f0e2de642 100644 --- a/python/cudnn/sdpa/bwd/config_sm80.py +++ b/python/cudnn/sdpa/bwd/config_sm80.py @@ -48,3 +48,111 @@ class Cfg: LLAMA_CFG = Cfg(D_QK=128, D_V=128, TILE_KV=64, TILE_Q=64, WARPS_PER_SG=4) GPTOSS_CFG = Cfg(D_QK=64, D_V=64, TILE_KV=64, TILE_Q=128, WARPS_PER_SG=4) + + +# --------------------------------------------------------------------------- +# TemplateParams — the compile-time identity of one SM80 backward kernel +# specialization. Mirrors fwd/config_sm80.py (#689) and config_sm120: a +# frozen, hashable record injected into the kernel template as the +# ``FROST_TEMPLATE_PARAMS`` module global by ``frost.template_loader``. +# Everything here is PLAN-TIME data — never a runtime tensor value (Hard +# Rule 4). Shape axes (b/h/hk/sq/skv, the actual head dims under the flavor +# envelope, the SWA width) stay arguments of the template module's +# ``compile()`` and its per-shape lru cache; THD packed token totals compile +# DYNAMIC (``cute.sym_int``) there and are never part of any key (issue #604). +# --------------------------------------------------------------------------- +from cudnn.frost.tile_dsl.constants import SCHED_LPT as _SCHED_LPT # noqa: E402 +from cudnn.frost.tile_dsl.constants import SCHED_NATURAL as _SCHED_NATURAL # noqa: E402 + + +@dataclass(frozen=True) +class TemplateParams: + """One SM80 backward kernel specialization (the module-identity axes).""" + + # dtype: fp16 or bf16 I/O (one mma pipeline serves both). + io_bf16: bool = False + # Flavor envelope head dims (the compile-time D box; the actual dims are + # compile() arguments and may be smaller — the host pads to the box). + d_qk: int = 128 + d_v: int = 128 + # Backward pipeline geometry (see the Cfg table above; tile_kv must equal + # warps_per_sg*16 and tile_q a multiple of it). + tile_kv: int = 64 + tile_q: int = 64 + warps_per_sg: int = 4 + # Mask family. right_bound stays a RUNTIME argument (it widens the causal + # band without changing the traced structure). + is_causal: bool = False + has_swa: bool = False + causal_bottom_right: bool = False + # Optional operands / outputs (compile-time ABI presence; a + # missing-but-required or provided-but-uncompiled operand raises at + # execute — Hard Rule 1). + has_seq_kv_lens: bool = False + has_seq_q_lens: bool = False + has_bias: bool = False # bias input => dBias accumulator output + bias_is_fp32: bool = True + bias_broadcast: bool = True # bias batch dim 1 (broadcast) vs B + has_sink: bool = False # sinks input => dSink output (standalone reduction) + has_rope: bool = False + # Deterministic dQ: the kv-ordered gmem-semaphore relay (forces + # SCHED_DEFAULT; the semaphore is carved caller scratch). + deterministic: bool = False + # Packed varlen (wrapper-only today; the engine row declares thd=False). + thd_varlen: bool = False + # Tile-scheduler policy in the SHARED frost vocabulary + # (tile_dsl.constants.SCHED_*): the bwd grid interprets NATURAL as its + # plain kv-major grid and LPT as the kv-major LPT remap (LPT_L2 is a + # forward-only policy today). + sched_policy: int = _SCHED_NATURAL + + +def validate_bwd_params(p: TemplateParams) -> None: + """Raising validator — a failure here means the capability row or the + adapter lied about what this template can serve.""" + if p.tile_kv != p.warps_per_sg * 16: + raise ValueError(f"sm80 bwd: tile_kv ({p.tile_kv}) must equal warps_per_sg*16 ({p.warps_per_sg * 16})") + if p.tile_q % (p.warps_per_sg * 16) != 0: + raise ValueError(f"sm80 bwd: tile_q ({p.tile_q}) must be a multiple of warps_per_sg*16 ({p.warps_per_sg * 16})") + if p.d_qk % 16 != 0 or p.d_v % 16 != 0 or p.d_qk <= 0 or p.d_v <= 0: + raise ValueError(f"sm80 bwd: template head dims must be positive multiples of 16; got ({p.d_qk}, {p.d_v})") + if p.d_qk < p.d_v: + raise ValueError(f"sm80 bwd: d_qk ({p.d_qk}) must be >= d_v ({p.d_v}) (the per-sub-group split)") + if (p.d_qk // 2) % 16 != 0: + raise ValueError(f"sm80 bwd: d_qk//2 ({p.d_qk // 2}) must be a multiple of 16 (ldmatrix.x4 N//8 even)") + if p.d_v % 32 != 0: + raise ValueError(f"sm80 bwd: d_v ({p.d_v}) must be a multiple of 32 (do_dot warp reduce)") + if p.has_rope and p.d_qk > 128: + raise ValueError("sm80 bwd: RoPE requires d_qk <= 128 (the sDQ SMEM staging exceeds the A100 budget beyond that)") + if p.sched_policy not in (_SCHED_NATURAL, _SCHED_LPT): + raise ValueError(f"sm80 bwd: sched_policy must be SCHED_NATURAL or SCHED_LPT; got {p.sched_policy}") + if p.deterministic and p.sched_policy != _SCHED_NATURAL: + raise ValueError("sm80 bwd: deterministic dQ requires SCHED_NATURAL (the kv-ordered semaphore relay)") + if p.causal_bottom_right and not (p.is_causal or p.has_swa): + raise ValueError("sm80 bwd: causal_bottom_right requires is_causal and/or has_swa (nothing to align otherwise)") + if p.thd_varlen and (p.has_bias or p.has_rope or p.has_sink or p.has_seq_kv_lens or p.has_seq_q_lens): + raise ValueError("sm80 bwd: THD carries lengths via cu_seqlens; bias / rope / sink / dense seq-lens are dense-only") + if p.thd_varlen and p.deterministic: + # The dQ-relay semaphore is sized at compile time from the dense sq; + # THD compiles sq as a dynamic sym_int so there is nothing to size it + # from (the FE support surface already rejects deterministic + ragged). + raise ValueError("sm80 bwd: deterministic dQ is dense-only (THD compiles sq dynamic; the relay semaphore has no plan-time size)") + + +def bwd_params_for_flavor(flavor: str, **overrides) -> TemplateParams: + """A backward TemplateParams seeded from one flavor's bprop Cfg row. + + Only llama (the d<=128 shared-pipeline point) and gptoss (the d=64 + wide-Q-tile point) have swept bprop rows; dsv3/qwen ride the llama + pipeline shape with their own envelope dims (the kernel derives + qo_stages / drop-sDQ from d_qk). + """ + cfg = {"llama": LLAMA_CFG, "gptoss": GPTOSS_CFG}.get(flavor, LLAMA_CFG) + # Seed the envelope dims too, so a flavor name alone cannot yield tiles + # from one flavor and head dims from another (callers with a different + # envelope — dsv3/qwen on the llama pipeline — override d_qk/d_v). + base = dict(d_qk=cfg.D_QK, d_v=cfg.D_V, tile_kv=cfg.TILE_KV, tile_q=cfg.TILE_Q, warps_per_sg=cfg.WARPS_PER_SG) + base.update(overrides) + p = TemplateParams(**base) + validate_bwd_params(p) + return p diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index d3427903f..e24f58e4a 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -20,19 +20,23 @@ import logging from dataclasses import dataclass +from functools import partial from typing import Any, Callable, Optional import cudnn from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old from cudnn.sdpa import graph_analyzer as ga - # Lowering dependencies, resolved at build time — see the note in fwd/engines.py: # importing them here would drag the CuTe DSL into every support check. -def _adapter_sm120(): - from cudnn.sdpa.bwd.api_dsl import SdpaBwdDslSm120 +_SM120 = "SdpaBwdDslSm120" +_SM80 = "SdpaBwdDslSm80" + - return SdpaBwdDslSm120 +def _adapter(name: str): + from cudnn.sdpa.bwd import api_dsl + + return getattr(api_dsl, name) def _cuda_driver(): @@ -340,13 +344,17 @@ def build(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None): return spec.lower(spec, facts, knobs) -def lower_dsl_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = None): +def lower_dsl_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = None, api_type: str = _SM120): """Lower the selected SDPA backward engine through its DSL adapter. Descriptor conversion, adapter lifecycle, variant-pack binding, and launch - construction live here; the adapter owns compilation and the three-kernel - execute chain. + construction live here; the adapter owns compilation and the kernel + execute chain. ``EngineSpec.lower`` binds the implementation through + ``api_type`` (mirrors ``lower_dsl_prefill``); SM80-only operands (bias → + dBias, plan-time bias facts) flow only to adapters declaring the matching + constructor / execute keywords. """ + import inspect # Per-port geometry from facts.port_layouts, NOT the live IR tensors: # build_operation_graph rewrites the backward node's K/V ports to @@ -383,7 +391,16 @@ def _desc(geom, dtype, name: str) -> "Any": bias_geom = (tuple(facts.bias_t.get_dim()), tuple(facts.bias_t.get_stride())) if facts.has_bias else None dbias_geom = (tuple(facts.dbias_t.get_dim()), tuple(facts.dbias_t.get_stride())) if facts.has_dbias else None - api = _adapter_sm120()( + adapter_cls = _adapter(api_type) + # SM80-only plan-time facts, forwarded only to adapters declaring them. + _extra_ctor = { + "has_bias": facts.has_bias, + "bias_is_fp32": (facts.bias_t.get_data_type() == cudnn.data_type.FLOAT) if facts.bias_t is not None else True, + "bias_batch": int(facts.bias_t.get_dim()[0]) if facts.bias_t is not None else 1, + "has_rope": False, # RoPE-fused multi-node graphs never reach the analyzer + } + _extra_ctor = {k: v for k, v in _extra_ctor.items() if k in inspect.signature(adapter_cls.__init__).parameters} + api = adapter_cls( sample_q=_desc(q_geom, facts.dtype, "q"), sample_k=_desc(k_geom, facts.dtype, "k"), sample_v=_desc(v_geom, facts.dtype, "v"), @@ -407,6 +424,7 @@ def _desc(geom, dtype, name: str) -> "Any": tile_n=requested.tile_n if requested is not None else None, seq_kv_lens_present=seq_kv_t is not None, seq_q_lens_present=seq_q_t is not None, + **_extra_ctor, ) api.check_support() # raises ValueError / NotImplementedError if unsupported api.compile() @@ -497,8 +515,8 @@ def _execute(variant_pack, workspace=None, stream=None): def _sm80_spec() -> EngineSpec: - """SM80 (A100) backward row: lowers onto the ``cudnn.sdpa`` SM80 APIBase - adapter (``bwd/api.py``), which owns kernel-flavor selection + """SM80 (A100) backward row: lowers through the shared ``lower_dsl_bwd`` + onto ``SdpaBwdDslSm80`` (``bwd/api_dsl.py``), which owns kernel-flavor selection (head-dim envelopes up to (256, 256), incl. rectangular 192/128), host-side head-dim zero-padding, BHSD<->BSHD normalization, per-shape kernel caching, and the dedicated plain-dense d=64 fast path. The @@ -534,176 +552,8 @@ def _sm80_spec() -> EngineSpec: # packed LSE; sm120 reads declared strides natively instead. strided_stats=True, ), - lower=lower_sm80_bwd, - ) - - -def lower_sm80_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = None): - """Lower the SM80 backward row through the ``cudnn.sdpa`` SM80 adapter. - - Built at plan time (issue #514): the adapter is constructed here from the - NORMALIZED buffer descriptors (compact BSHD-physical), its scratch - requirement plus this executor's own dense_flex gather staging is recorded - as ``workspace_bytes``, and execute carves everything from the caller's - workspace — no per-execute allocation on this path. - """ - import dataclasses - - from cudnn.api_base import TensorDesc - from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, ws_align - - from .api import SdpabwdSm80 - - binding = ga.SdpaBinding( - q=facts.q_t, - k=facts.k_t, - v=facts.v_t, - o=facts.o_t, - stats=facts.stats_t, - bias=facts.bias_t, - sink_token=facts.sink_t, - seq_len_kv=facts.seq_kv_t, - seq_len_q=facts.seq_q_t, - do=facts.do_t, - dq=facts.dq_t, - dk=facts.dk_t, - dv=facts.dv_t, - dbias=facts.dbias_t, - dsink=facts.dsink_t, - ) - mask_args = ga.adapter_mask_args(facts) - elem = 2 # fp16/bf16 — mismatch() admits no other input dtype - b, h_q = facts.b, facts.h_q - - def _compact_desc(t, name): - desc = ga.tensor_desc_from_ir(t, name=name) - bb, hh, ss, dd = desc.shape - return dataclasses.replace(desc, stride=(ss * hh * dd, dd, hh * dd, 1), stride_order=(3, 1, 2, 0)) - - def _is_compact_bshd(t) -> bool: - _, h, s, d = tuple(t.get_dim()) - return tuple(t.get_stride()) == (s * h * d, d, h * d, 1) - - # dense_flex gather staging, sized from the PORT layouts (static): a port - # already stored as a compact BSHD-physical allocation is handed through - # zero-copy. The adapter's own scratch then covers head-dim pads and the - # kernel-internal buffers. - ports = ((facts.q_t, "q"), (facts.k_t, "k"), (facts.v_t, "v"), (facts.o_t, "o"), (facts.do_t, "dO")) - - def _port_numel(t) -> int: - n = 1 - for extent in t.get_dim(): - n *= int(extent) - return n - - stage_bytes = {name: (0 if _is_compact_bshd(t) else ws_align(_port_numel(t) * elem)) for t, name in ports} - # Strided stats (Capabilities.strided_stats): the kernels read a PACKED - # (B, H_q, S_q) fp32 LSE, so a stats input with any other declared strides - # is gathered into a carved contiguous chunk at execute. - _stats_contig = facts.stats_t is not None and tuple(facts.stats_t.get_stride()) == (h_q * facts.s_q, facts.s_q, 1, 1) - stats_stage = 0 if (facts.stats_t is None or _stats_contig) else ws_align(b * h_q * facts.s_q * 4) - - q_desc = _compact_desc(facts.q_t, "q") - sample_lse = TensorDesc( - dtype=ga.to_torch_dtype(cudnn.data_type.FLOAT), - shape=(b, h_q, facts.s_q), - stride=(h_q * facts.s_q, facts.s_q, 1), - stride_order=(2, 1, 0), - device=q_desc.device, - name="lse", - ) - api = SdpabwdSm80( - sample_q=q_desc, - sample_k=_compact_desc(facts.k_t, "k"), - sample_v=_compact_desc(facts.v_t, "v"), - sample_o=_compact_desc(facts.o_t, "o"), - sample_do=_compact_desc(facts.do_t, "dO"), - sample_lse=sample_lse, - scale_softmax=facts.scale, - has_seq_kv_lens=facts.seq_kv_t is not None, - has_bias=facts.has_bias, - **mask_args, - ) - if not api.check_support(): - raise ValueError("SdpabwdSm80 declined the normalized graph geometry") - api.compile() - bias_batch = int(facts.bias_t.get_dim()[0]) if facts.bias_t is not None else 1 - api_scratch = api.scratch_workspace_bytes( - has_bias=facts.has_bias, - bias_batch=bias_batch, - has_sink=facts.has_sink, - deterministic=facts.deterministic, + lower=partial(lower_dsl_bwd, api_type=_SM80), ) - total_workspace_bytes = sum(stage_bytes.values()) + stats_stage + api_scratch - - def _normalize(carver, buf, staged: int): - if not staged: - return buf - bb, hh, ss, dd = buf.shape - dst = carver.take(bb * ss * hh * dd, buf.dtype).view(bb, ss, hh, dd) - dst.copy_(buf.permute(0, 2, 1, 3)) - return dst.permute(0, 2, 1, 3) - - def _ir_view(buf, ir_t): - """Reinterpret a variant-pack buffer through the IR tensor's dim/stride. - - cuDNN's execute contract treats variant-pack entries as raw storage - laid out per the IR tensor descriptor — the caller's torch tensor may - be flat or otherwise logically reshaped. The staging/squeeze/copy_ - paths below consume torch views, so rebuild the IR-shaped view instead - of trusting the caller's metadata (mirrors the forward lowering's - ``_ir_view``). INPUT ports only: output-port IR strides are - PROVISIONAL row-major unless the user assigned them (the layout - invariant in docs/python_graph_and_execution_backends.md), so the - gradient outputs below keep the caller tensor's own view — re-striding - them to the provisional layout would scatter the copy-back. - """ - dim, stride = tuple(ir_t.get_dim()), tuple(ir_t.get_stride()) - if tuple(buf.shape) == dim and tuple(buf.stride()) == stride: - return buf - return buf.as_strided(dim, stride) - - def _execute(variant_pack, workspace=None, stream=None): - resolved = ga.resolve_variant_pack(variant_pack, binding) - carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None - # squeeze(-1) is a valid view for ANY (B, H_q, S_q, 1) strides; the - # kernels read a packed LSE, so a strided stats input (strided_stats) - # is gathered into carved contiguous staging first. - lse = _ir_view(resolved[id(facts.stats_t)], facts.stats_t).squeeze(-1) - if stats_stage: - lse_stage = carver.take(b * h_q * facts.s_q, lse.dtype).view(b, h_q, facts.s_q) - lse_stage.copy_(lse) - lse = lse_stage - dbias_buf = resolved.get(id(facts.dbias_t)) if facts.has_dbias and facts.dbias_t is not None else None - dsink_buf = resolved.get(id(facts.dsink_t)) if facts.has_dsink and facts.dsink_t is not None else None - - api.execute( - q_tensor=_normalize(carver, _ir_view(resolved[id(facts.q_t)], facts.q_t), stage_bytes["q"]), - k_tensor=_normalize(carver, _ir_view(resolved[id(facts.k_t)], facts.k_t), stage_bytes["k"]), - v_tensor=_normalize(carver, _ir_view(resolved[id(facts.v_t)], facts.v_t), stage_bytes["v"]), - o_tensor=_normalize(carver, _ir_view(resolved[id(facts.o_t)], facts.o_t), stage_bytes["o"]), - do_tensor=_normalize(carver, _ir_view(resolved[id(facts.do_t)], facts.do_t), stage_bytes["dO"]), - lse_tensor=lse, - dq_tensor=resolved[id(facts.dq_t)], - dk_tensor=resolved[id(facts.dk_t)], - dv_tensor=resolved[id(facts.dv_t)], - dbias_tensor=dbias_buf, - dsink_tensor=dsink_buf.view(-1) if dsink_buf is not None else None, - scale_softmax=facts.scale, - deterministic=facts.deterministic, - # Stream from the caller's handle (ExecutionContext.stream); - # None keeps the current stream. - current_stream=stream, - workspace=carver.remaining() if (carver is not None and api_scratch) else None, - **ga.adapter_feature_buffers(facts, resolved), - ) - return None - - # Executor contract (engine._FrostSdpaBwdPlan): a non-zero workspace_bytes - # means _execute(variant_pack, workspace, stream) with the caller's buffer. - _execute.workspace_bytes = total_workspace_bytes - _execute.binding = binding - return _execute def engine_name(arch: str = "sm120") -> str: diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py index 8d9019915..7f462b0eb 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py @@ -37,7 +37,7 @@ traffic. This is what closes the gap to cuDNN (whose dQ atomic also coalesces): +18 % @ B2H16S4096 (75→89 TFLOPS, ~1.0× cuDNN), +16 % @ B1H16S2048 (1.08×). -**Deterministic dQ** (``backward(deterministic=True)``): the cross-KV-tile dQ +**Deterministic dQ** (``TemplateParams.deterministic``): the cross-KV-tile dQ atomicAdd is order-non-deterministic (fp32 add is non-associative → bitwise varies run-to-run once a sequence spans >1 KV-tile). The deterministic path orders the adds by ``kv_tile`` via a per-(seq,head,q_tile) int32 GMEM semaphore @@ -123,6 +123,14 @@ def from_dlpack(t, **kw): SCHED_DEFAULT = 0 # 3-D grid (kv_tile, head, batch); no reorder (byte-identical) SCHED_LPT = 1 # 1-D kv-major grid for causal load-balance +# TemplateParams injection seam (frost.template_loader): one uniquely named +# module per parameter set, specialized below via cutlass.const_expr folding. +# The shared tile_dsl scheduler vocabulary maps IDENTITY onto the internal +# grid decode (SCHED_NATURAL == SCHED_DEFAULT == 0, SCHED_LPT == 1). +from cudnn.sdpa.bwd.config_sm80 import TemplateParams # noqa: E402 + +PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams()) + def _mask_p(p, kv_abs, q_abs, *, mask_flags: int, swa_window: int, causal_bottom_right: int, causal_diag, eff_skv, right_bound): """Zero a recomputed softmax probability ``p`` if its ``(kv_abs, q_abs)`` @@ -261,8 +269,12 @@ def _bprop_kernel( # LSE/do_dot/bias reads are clamped to = d_v, f"bprop: d_qk ({d_qk}) must be >= d_v ({d_v})" - assert d_qk % 16 == 0 and d_v % 16 == 0, f"bprop: d_qk/d_v must be mult of 16 (got {d_qk}/{d_v})" - # SMEM fit on A100 (164 KiB): drop the dQ-coalescing sDQ staging once d_qk>128, - # and the Q/dO double-buffer (qo_stages 2→1) once d_qk>=256 (qwen). d<=128 - # (llama/gptoss) keeps both (byte-identical fast path). - dq_smem_coalesce = d_qk <= 128 - qo_stages = 1 if d_qk >= 256 else 2 - # ---- THD / varlen: packed [1,T,H,D] Q/K/V/dO/O + cu_seqlens. B==1 packed; - # n_seq logical sequences drive the grid + cu_* sizing. Q.shape[1]/ - # K.shape[1] are the packed totals T_q/T_kv (== the kernel's SQ/SKV). --- - thd = cu_seqlens_q is not None - _carver = None - if workspace is not None: - assert not thd, "workspace carving is dense-only (the engine path; THD comes via the wrappers)" - from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver - - _carver = WorkspaceCarver( - workspace, - scratch_bytes( - B=B, - SQ=SQ, - SKV=SKV, - H=H, - Hk=Hk, - d_qk=d_qk, - d_v=d_v, - io_bytes=Q.element_size(), - deterministic=bool(deterministic), - has_bias=bias is not None, - bias_batch=(bias.shape[0] if bias is not None else 1), - has_sink=sinks is not None, - need_do_dot=do_dot is None, - tile_q=tile_q, - ), - "bprop_f16_sm80", - ) +# =========================================================================== +# Template entry point (the #689 contract, backward flavor): one call per +# shape compiles (or fetches) the FULL kernel chain for this module's PARAMS +# specialization. THD packed token totals compile DYNAMIC (``cute.sym_int``) +# and are never part of the key (issue #604) — callers pass ``sq = skv = 0`` +# there; ``n_batch_logical`` (the logical sequence count) sizes the +# ``cu_seqlens`` ABI and IS plan-time. Launch marshaling lives in the adapter +# (``api_dsl._sm80_bwd_call``); this module holds no host runtime logic. +# The stats input is READ as a packed (B, H, SQ) LSE (raw-pointer packed +# addressing in the device code) — a strided graph stats input is gathered +# into carved staging by the adapter, unlike the forward's #712-style native +# strided WRITES. +# =========================================================================== +from typing import NamedTuple # noqa: E402 - def _scratch(numel, dtype, zero): - if _carver is None: - return (torch.zeros if zero else torch.empty)(numel, dtype=dtype, device=Q.device) - t = _carver.take(numel, dtype) - if zero: - t.zero_() - return t - - if thd: - assert cu_seqlens_k is not None, "THD needs both cu_seqlens_q and cu_seqlens_k" - assert B == 1, f"THD: Q/K/V/dO/O must be packed [1,T,H,D]; got batch dim {B}" - # GQA/MQA work under THD: the per-query-head dK_ws/dV_ws workspace - # ([1,T_kv,H,d]) + the _dkv_reduce over the query-head group are - # layout-agnostic (B=1, SKV=T_kv packed); the kv-tile CTA reads the - # GQA-mapped kv_head and writes its query-head slice exactly as in the - # dense path. (Was conservatively asserted MHA-only.) - cu_q_host = cu_seqlens_q.to(dtype=torch.int32, device="cpu") - cu_k_host = cu_seqlens_k.to(dtype=torch.int32, device="cpu") - n_seq = cu_q_host.numel() - 1 - assert cu_k_host.numel() == n_seq + 1, "cu_seqlens_q / cu_seqlens_k length mismatch" - max_skv = int((cu_k_host[1:] - cu_k_host[:-1]).max()) - max_sq = int((cu_q_host[1:] - cu_q_host[:-1]).max()) - else: - n_seq = B - max_sq = SQ - # ---- mask string → compile-time bitmask ------------------------------- - _MASK_TOK = { - "none": MASK_NONE, - "causal": MASK_CAUSAL, - "swa": MASK_SWA, - "causal_swa": MASK_CAUSAL | MASK_SWA, - } - assert mask in _MASK_TOK, f"backward: mask must be one of {list(_MASK_TOK)}; got {mask!r}" - mask_flags = _MASK_TOK[mask] - # Scheduler: kv-major LPT grid balances the causal load (light high-kv tiles - # land in the last wave). 'auto' uses LPT only when causal (where the q-skip - # makes work uneven); non-causal work is uniform so DEFAULT is fine. - assert sched in ("auto", "default", "natural", "lpt"), f"backward: bad sched {sched!r}" - if sched == "lpt": - sched_policy = SCHED_LPT - elif sched in ("default", "natural"): - sched_policy = SCHED_DEFAULT - else: # auto - sched_policy = SCHED_LPT if (mask_flags & MASK_CAUSAL) else SCHED_DEFAULT - # Deterministic dQ orders the per-(seq,head,q_tile) atomicAdd by kv_tile == the - # 3-D grid's blockIdx.x — so it REQUIRES the SCHED_DEFAULT decode (kv_tile=bx). - # The kv-major LPT 1-D flat grid would remap kv_tile and break both the order - # and the deadlock-freedom (predecessor = lower blockIdx). Force it off. - if deterministic: - sched_policy = SCHED_DEFAULT - has_seq_kv_lens = seq_kv_lens is not None - has_seq_len_q = seq_len_q is not None - if has_seq_kv_lens: - mask_flags |= MASK_PADDED - cbr = 1 if causal_bottom_right else 0 - has_bias = bias is not None - if has_bias: - assert bias.dim() == 4 and bias.shape[1] == H, f"bias must be [1|B, H, SQ, SKV]; got {tuple(bias.shape)}" - bias_is_fp32 = bias.dtype == torch.float32 - bias_batch = bias.shape[0] # 1 (broadcast over B) or B - assert bias_batch in (1, B), f"bias batch dim must be 1 or B={B}; got {bias_batch}" - bias_bstride = 0 if bias_batch == 1 else H * SQ * SKV - else: - bias_is_fp32, bias_batch, bias_bstride = True, 1, 0 - # RoPE: build the [max_s, d_qk//2, 2] (cos,sin) table from the angles - # (mirrors the forward). Backward rotates Q/K on load + un-rotates dQ/dK. - has_rope = rope_freqs is not None - if has_rope: - assert seq_kv_lens is None, "RoPE is dense-only (no THD/padded) on SM80 bprop" - d2 = d_qk // 2 - rf = rope_freqs.to(dtype=torch.float32, device=Q.device).reshape(rope_freqs.shape[0], -1) - rope_max_s = rf.shape[0] - assert rf.shape[1] >= d2, f"rope_freqs last dim ({rf.shape[1]}) must be >= d_qk//2 ({d2})" - assert rope_max_s >= max(SQ, SKV), f"rope_freqs max_s ({rope_max_s}) must cover max(SQ={SQ}, SKV={SKV})" - angles = rf[:, :d2] - rope_cs_t = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() - else: - rope_max_s = 1 - rope_cs_t = _dummy1z(torch.float32, Q.device) - # Attention sink: dQ/dK/dV need NO kernel change (P recomputed from the - # sink-aware LSE the caller passes); only dSink is computed (standalone). - has_sink = sinks is not None - if has_sink: - sinks_t = sinks.to(dtype=torch.float32, device=Q.device).reshape(H).contiguous() - dsink_t = _scratch(H, torch.float32, True) - # THD is dense-feature-only for now (bias/rope/sink/seq_kv_lens are - # dense-only); per-sequence padding is handled by the packed bounds, not - # the PADDED mask. THD uses SCHED_DEFAULT (LPT+THD is a future tweak). - if thd: - assert not ( - has_bias or has_rope or has_sink or has_seq_kv_lens or has_seq_len_q - ), "THD/varlen bprop: bias/rope/sink/seq_kv_lens/seq_len_q not supported yet" - sched_policy = SCHED_DEFAULT - # RoPE staging reuses the sDQ SMEM buffer, which the d_qk > 128 configs - # drop (dq_smem_coalesce) to stay under the 164 KiB dynamic-SMEM cap — - # re-enabling it via has_rope would exceed the budget at launch - # (~48-64 KiB over at d=192/256). - assert not (has_rope and d_qk > 128), "RoPE bprop requires d_qk <= 128 (the sDQ SMEM staging exceeds the A100 budget beyond that)" - # seq_len_q [B] int32 — per-batch live Q length (dense PADDED). Dummy 1-elem - # when unused (kernel never reads it at has_seq_len_q=False). - if has_seq_len_q: - seqq_t = seq_len_q.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seqq_t = _dummy1z(torch.int32, Q.device) - assert d_qk % 2 == 0 - # dQ splits d-cols across the two sub-groups → each reads a DQ_N = d_qk//2 - # column slice of sK. load_b_smem_x4 takes the d-col offset as `col_base` - # (swizzle computed on the TRUE column), so the Swz128B half-column read is - # correct for ANY d — no longer requires d_qk//2 % 64 == 0. The only - # remaining shape constraint is DQ_N % 16 == 0 (load_b_smem_x4 N//8 even): - # d=128→64 ✓, d=64→32 ✓. - assert (d_qk // 2) % 16 == 0, f"d_qk//2 ({d_qk//2}) must be a multiple of 16 (ldmatrix.x4 N//8 even)" - assert d_v % 32 == 0, f"d_v ({d_v}) must be a multiple of 32 (do_dot warp reduce)" - # Partial-tile (arbitrary seqlen): SQ/SKV need NOT be tile multiples. The - # last straddling tile zero-fills OOB rows on load, masks P=0 for kv>=SKV / - # q>=SQ, clamps LSE/do_dot/bias reads, and row-gates the dV/dK/dQ/dBias - # stores (all compile-time gated on SQ%tile_q / SKV%tile_kv → dense path is - # byte-identical). RoPE still requires alignment (see below). - if has_rope and (SQ % tile_q or SKV % tile_kv): - raise NotImplementedError("SM80 bprop: RoPE requires SQ/SKV tile-aligned") - # dQ M-tiling: tile_q rows are covered by warps_per_sg warps × 16 × M_BLOCKS, - # so tile_q must be an exact multiple of warps_per_sg*16 (llama 64→1 block, - # gptoss 128→2). dV/dK/BMM1 keep M=tile_kv=warps_per_sg*16 (1 block). - assert tile_q % (warps_per_sg * 16) == 0, f"tile_q ({tile_q}) must be a multiple of warps_per_sg*16 ({warps_per_sg*16})" - assert tile_kv == warps_per_sg * 16, f"tile_kv ({tile_kv}) must equal warps_per_sg*16 ({warps_per_sg*16})" - if scale is None: - scale = 1.0 / (d_qk**0.5) - scale_log2 = scale * math.log2(math.e) - inv_scale = 1.0 / float(scale) - - dQ_acc = _scratch(B * SQ * H * d_qk, torch.float32, True).view(B, SQ, H, d_qk) - dQ = _scratch(B * SQ * H * d_qk, Q.dtype, False).view(B, SQ, H, d_qk) - # Deterministic-dQ relay counter: one int32 per (seq, head, q_tile), zeroed - # per launch. Stride = ceil(max_SQ/tile_q) so the per-seq q_iter (THD) or the - # dense q_iter both index in-bounds; n_seq sequences (= B dense). 1-elem dummy - # (never touched) on the fast path so it costs nothing. - sem_q_stride = (max_sq + tile_q - 1) // tile_q if deterministic else 0 - sem_units = n_seq * H * sem_q_stride if deterministic else 1 - dq_sem = _scratch(max(sem_units, 1), torch.int32, True) if deterministic else _dummy1z(torch.int32, Q.device) - # dK/dV write buffers have H_q heads (one slice per query head — no atomics). - # MHA (gqa_ratio==1): they ARE the outputs. GQA: a per-query-head workspace - # that a reduction kernel sums over the group → [B,SKV,Hk,d] outputs. - dK_ws = _scratch(B * SKV * H * d_qk, Q.dtype, False).view(B, SKV, H, d_qk) - dV_ws = _scratch(B * SKV * H * d_v, Q.dtype, False).view(B, SKV, H, d_v) - if gqa_ratio == 1: - dK, dV = dK_ws, dV_ws - else: - dK = _scratch(B * SKV * Hk * d_qk, Q.dtype, False).view(B, SKV, Hk, d_qk) - dV = _scratch(B * SKV * Hk * d_v, Q.dtype, False).view(B, SKV, Hk, d_v) - lse_t = lse.to(dtype=torch.float32, device=Q.device).contiguous() - # seq_kv_lens [B] int32 (or 1-elem dummy when not padded — never read). - if has_seq_kv_lens: - seqk_t = seq_kv_lens.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seqk_t = _dummy1z(torch.int32, Q.device) - # cu_seqlens [n_seq+1] int32 (THD) or 1-elem dummy (dense — never read). The - # over-provisioned THD grid covers the longest sequence (ceil(max_skv/tile_kv) - # kv-tiles) × H × n_seq; short sequences early-out per kv-tile. - if thd: - cu_q_t = cu_seqlens_q.to(dtype=torch.int32, device=Q.device).contiguous() - cu_k_t = cu_seqlens_k.to(dtype=torch.int32, device=Q.device).contiguous() - grid_kv_tiles = (max_skv + tile_kv - 1) // tile_kv - grid_batch = n_seq - else: - cu_q_t = _dummy1z(torch.int32, Q.device) - cu_k_t = _dummy1z(torch.int32, Q.device) - grid_kv_tiles = 0 - grid_batch = 0 - # Bias + dBias (fp32 accumulator, same shape as bias; atomicAdd reduces over - # batch when bias is broadcast [1,H,SQ,SKV]). - if has_bias: - bias_t = bias.contiguous() - dbias_t = _scratch(bias_batch * H * SQ * SKV, torch.float32, True).view(bias_batch, H, SQ, SKV) - else: - # Dummy must match the fake tensor _compile_main builds at has_bias=False - # (bias_is_fp32 defaults True → fp32). - bias_t = _dummy1z(torch.float32, Q.device) - dbias_t = _dummy1z(torch.float32, Q.device) - - torch_stream = torch.cuda.current_stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - # ---- do_dot: on-device (default) or caller-supplied ------------------ - if do_dot is None: - dot_t = _scratch(B * H * SQ, torch.float32, False).view(B, H, SQ) - dd_fn = _compile_do_dot(B, H, SQ, d_v, io_is_bf16) - dd_fn(from_dlpack(O), from_dlpack(dO), from_dlpack(dot_t), cutlass.Int32(B * H * SQ), stream) - else: - dot_t = do_dot.to(dtype=torch.float32, device=Q.device).contiguous() - - # dSink (standalone reduction; dQ/dK/dV are already sink-correct via LSE). - if has_sink: - ds_fn = _compile_dsink(B, H, SQ) - ds_fn(from_dlpack(lse_t), from_dlpack(dot_t), from_dlpack(sinks_t), from_dlpack(dsink_t), cutlass.Int32(B * H), stream) - - fn = _compile_main( - B, - H, - SQ, - SKV, - d_qk, - d_v, - tile_kv, - tile_q, - warps_per_sg, - io_is_bf16, - qo_stages=qo_stages, - dq_smem_coalesce=dq_smem_coalesce, - Hk=Hk, - mask_flags=mask_flags, - swa_window=swa_window, - causal_bottom_right=cbr, - has_seq_kv_lens=has_seq_kv_lens, - has_bias=has_bias, - bias_is_fp32=bias_is_fp32, - bias_batch=bias_batch, - has_rope=has_rope, - rope_max_s=rope_max_s, - has_seq_len_q=has_seq_len_q, - thd_varlen=thd, - n_seq=n_seq, - deterministic=deterministic, - sem_units=max(sem_units, 1), - sched_policy=sched_policy, - ) - fn( - from_dlpack(Q), - from_dlpack(K), - from_dlpack(V), - from_dlpack(dO), - from_dlpack(dQ_acc), - from_dlpack(dK_ws), - from_dlpack(dV_ws), - from_dlpack(lse_t), - from_dlpack(dot_t), - from_dlpack(seqk_t), - from_dlpack(bias_t), - from_dlpack(dbias_t), - from_dlpack(rope_cs_t), - from_dlpack(cu_q_t), - from_dlpack(cu_k_t), - from_dlpack(seqq_t), - from_dlpack(dq_sem), - cutlass.Int32((SQ + tile_q - 1) // tile_q), - cutlass.Float32(scale_log2), - cutlass.Float32(scale), - cutlass.Int32(right_bound), - cutlass.Float32(inv_scale), - cutlass.Int32(bias_bstride), - cutlass.Int32(sem_q_stride), - cutlass.Int32(grid_kv_tiles), - cutlass.Int32(grid_batch), - stream, - ) +class CompiledBwd(NamedTuple): + """The compiled artifacts of one backward specialization + shape.""" - cast_fn = _compile_cast(B, H, SQ, d_qk, io_is_bf16) - n_vecs = (B * SQ * H * d_qk) // 2 - cast_fn(from_dlpack(dQ_acc), from_dlpack(dQ), cutlass.Int32(n_vecs), stream) + main: object + do_dot: object + cast: object + reduce_k: object # None unless GQA (h != h_kv) + reduce_v: object # None unless GQA + dsink: object # None unless PARAMS.has_sink + sem_q_stride: int # deterministic-dQ semaphore stride (0 when off) - # GQA: sum dK_ws/dV_ws over each query-head group → [B,SKV,Hk,d] outputs. - if gqa_ratio > 1: - rk_fn = _compile_dkv_reduce(B, SKV, H, Hk, d_qk, io_is_bf16) - rk_fn(from_dlpack(dK_ws), from_dlpack(dK), cutlass.Int32(B * SKV * Hk * d_qk), stream) - rv_fn = _compile_dkv_reduce(B, SKV, H, Hk, d_v, io_is_bf16) - rv_fn(from_dlpack(dV_ws), from_dlpack(dV), cutlass.Int32(B * SKV * Hk * d_v), stream) - # Optional grads appended in a FIXED order (dBias, then dSink); the caller - # reconstructs positions from has_bias / has_sink (which it passed in). - outs = [dQ, dK, dV] - if has_bias: - outs.append(dbias_t) # fp32 [1|B, H, SQ, SKV] - if has_sink: - outs.append(dsink_t) # fp32 [H] - return tuple(outs) +@lru_cache(maxsize=None) +def compile( # noqa: A001 — the template contract's entry point + b: int, + h: int, + h_kv: int, + sq: int, + skv: int, + swa_window: int = 0, + rope_max_s: int = 0, + n_batch_logical: int = 0, +): + """Compile (or fetch) this template specialization for one shape. + + The head dims are PARAMS.d_qk / PARAMS.d_v — the flavor box; the host pads + operands to it, so unlike the forward there is no narrower runtime ``d``. + Dense: ``sq``/``skv`` are the physical extents. THD (PARAMS.thd_varlen): + pass ``b = 1``, ``sq = skv = 0`` — the packed token extents compile as one + ``cute.sym_int`` per ragged group (Q/dO/dQ/LSE/do_dot share t_q; K/V and + the per-query-head dK/dV write buffers share t_kv), so one artifact + re-binds any totals. + """ + p = PARAMS + io_dtype = cutlass.BFloat16 if p.io_bf16 else cutlass.Float16 + mask_flags = (MASK_CAUSAL if p.is_causal else MASK_NONE) | (MASK_SWA if p.has_swa else 0) | (MASK_PADDED if p.has_seq_kv_lens else 0) + # SMEM budget derivations (see the assertions in the device code): drop + # the dQ-coalescing sDQ staging past d_qk 128; single Q/dO buffer at 256. + dq_smem_coalesce = p.d_qk <= 128 + qo_stages = 1 if p.d_qk >= 256 else 2 + gqa = h != h_kv + if p.thd_varlen: + t_q = cute.sym_int(divisibility=1) + t_kv = cute.sym_int(divisibility=1) + _b, _sq, _skv = 1, t_q, t_kv + n_seq = n_batch_logical + else: + _b, _sq, _skv = b, sq, skv + n_seq = b + # Deterministic-dQ relay counter stride: ceil(max_SQ / tile_q). THD + + # deterministic is rejected upstream (the FE support surface), so the + # dense sq is always real here when deterministic is on. + sem_q_stride = ((sq + p.tile_q - 1) // p.tile_q) if p.deterministic else 0 + sem_units = max(n_seq * h * sem_q_stride, 1) + + def _fake(dtype, shape, order, align=16): + return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=order, assumed_align=align) + + r4 = (3, 2, 1, 0) + fq = _fake(io_dtype, (_b, _sq, h, p.d_qk), r4) + fk = _fake(io_dtype, (_b, _skv, h_kv, p.d_qk), r4) + fv = _fake(io_dtype, (_b, _skv, h_kv, p.d_v), r4) + fdo = _fake(io_dtype, (_b, _sq, h, p.d_v), r4) + fdq_acc = _fake(cutlass.Float32, (_b, _sq, h, p.d_qk), r4) + # dK/dV WRITE buffers carry H_q heads (per-query-head; GQA reduces after). + fdk_ws = _fake(io_dtype, (_b, _skv, h, p.d_qk), r4) + fdv_ws = _fake(io_dtype, (_b, _skv, h, p.d_v), r4) + fl = _fake(cutlass.Float32, (_b, h, _sq), (2, 1, 0)) + fdt = _fake(cutlass.Float32, (_b, h, _sq), (2, 1, 0)) + fsk = _fake(cutlass.Int32, (b if p.has_seq_kv_lens else 1,), (0,), align=4) + bias_dtype = cutlass.Float32 if p.bias_is_fp32 else io_dtype + bias_b = 1 if p.bias_broadcast else b + fbias = _fake(bias_dtype, ((bias_b, h, sq, skv) if p.has_bias else (1,)), (r4 if p.has_bias else (0,))) + fdbias = _fake(cutlass.Float32, ((bias_b, h, sq, skv) if p.has_bias else (1,)), (r4 if p.has_bias else (0,))) + frope = _fake(cutlass.Float32, ((rope_max_s, p.d_qk // 2, 2) if p.has_rope else (1,)), ((2, 1, 0) if p.has_rope else (0,))) + _cu_len = (n_batch_logical + 1) if p.thd_varlen else 1 + fcuq = _fake(cutlass.Int32, (_cu_len,), (0,), align=4) + fcuk = _fake(cutlass.Int32, (_cu_len,), (0,), align=4) + fsq = _fake(cutlass.Int32, (b if p.has_seq_q_lens else 1,), (0,), align=4) + fsem = _fake(cutlass.Int32, (sem_units,), (0,), align=4) + fstream = cuda.CUstream(0) + + main = cute.compile( + _bprop_host, + fq, + fk, + fv, + fdo, + fdq_acc, + fdk_ws, + fdv_ws, + fl, + fdt, + fsk, + fbias, + fdbias, + frope, + fcuq, + fcuk, + fsq, + fsem, + p.d_qk, + p.d_v, + p.tile_kv, + p.tile_q, + p.warps_per_sg, + int(qo_stages), + bool(dq_smem_coalesce), + io_dtype, + int(mask_flags), + int(swa_window), + int(1 if p.causal_bottom_right else 0), + bool(p.has_seq_kv_lens), + bool(p.has_bias), + bool(p.bias_is_fp32), + bool(p.has_rope), + bool(p.has_seq_q_lens), + bool(p.thd_varlen), + bool(p.deterministic), + int(p.sched_policy), + cutlass.Int32(0), + cutlass.Float32(0.0), + cutlass.Float32(0.0), + cutlass.Int32(0), + cutlass.Float32(0.0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + fstream, + options="--enable-tvm-ffi", + ) + fo = _fake(io_dtype, (_b, _sq, h, p.d_v), r4) + do_dot = cute.compile(_do_dot_host, fo, fdo, fdt, p.d_v, io_dtype, cutlass.Int32(0), fstream, options="--enable-tvm-ffi") + fdq_out = _fake(io_dtype, (_b, _sq, h, p.d_qk), r4) + cast = cute.compile(_cast_host, fdq_acc, fdq_out, io_dtype, cutlass.Int32(0), fstream, options="--enable-tvm-ffi") + reduce_k = reduce_v = None + if gqa: + fdk_out = _fake(io_dtype, (_b, _skv, h_kv, p.d_qk), r4) + fdv_out = _fake(io_dtype, (_b, _skv, h_kv, p.d_v), r4) + reduce_k = cute.compile(_dkv_reduce_host, fdk_ws, fdk_out, p.d_qk, h, h_kv, io_dtype, cutlass.Int32(0), fstream, options="--enable-tvm-ffi") + reduce_v = cute.compile(_dkv_reduce_host, fdv_ws, fdv_out, p.d_v, h, h_kv, io_dtype, cutlass.Int32(0), fstream, options="--enable-tvm-ffi") + dsink = None + if p.has_sink: + fsinks = _fake(cutlass.Float32, (h,), (0,), align=4) + fdsink = _fake(cutlass.Float32, (h,), (0,), align=4) + dsink = cute.compile(_dsink_host, fl, fdt, fsinks, fdsink, sq, cutlass.Int32(0), fstream, options="--enable-tvm-ffi") + return CompiledBwd(main=main, do_dot=do_dot, cast=cast, reduce_k=reduce_k, reduce_v=reduce_v, dsink=dsink, sem_q_stride=sem_q_stride) diff --git a/test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py b/test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py index 0c645f214..f65f266f4 100644 --- a/test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py +++ b/test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py @@ -182,17 +182,19 @@ def _run(): @pytest.mark.L0 @torch_fork_set_rng(seed=0) -def test_sdpa_bwd_sm80_d64_fast_path(): +def test_sdpa_bwd_sm80_d64_fast_path(monkeypatch): """The dedicated d=64 kernel routes only for plain dense MHA and agrees - with the generic kernel on the same inputs.""" + with the generic kernel on the same inputs (the generic side runs through + the adapter with the d64 gate forced off — the generic module is a + TemplateParams template with no standalone entry point).""" try: - from cudnn.sdpa.bwd import api as api_sm80 - from cudnn.sdpa.bwd.kernels import bprop_d64_f16_sm80 as d64, bprop_f16_sm80 as gen + from cudnn.sdpa.bwd import api_dsl as api_sm80 + from cudnn.sdpa.bwd.kernels import bprop_d64_f16_sm80 as d64 except ImportError as e: pytest.skip(f"SM80 SDPA API not available: {e}") common = dict(d_qk=64, d_v=64, h_q=8, h_kv=8, s_q=512, s_kv=512, mask_token="none", right_bound=0, causal_bottom_right=False, bw_kwargs={}) - assert api_sm80._d64_fast_path_eligible(**common) + assert api_sm80._sm80_d64_fast_path_eligible(**common) # every gated condition individually disqualifies for override in ( dict(d_qk=48, d_v=48), # padded flavor @@ -204,15 +206,155 @@ def test_sdpa_bwd_sm80_d64_fast_path(): dict(bw_kwargs={"bias": object()}), dict(bw_kwargs={"deterministic": True}), ): - assert not api_sm80._d64_fast_path_eligible(**{**common, **override}), override + assert not api_sm80._sm80_d64_fast_path_eligible(**{**common, **override}), override b, h, s, d = 2, 8, 512, 64 q = torch.randn(b, s, h, d, dtype=torch.float16, device="cuda") # BSHD (kernel layout) k, v, do, o = (torch.randn_like(q) for _ in range(4)) lse = torch.randn(b, h, s, dtype=torch.float32, device="cuda").abs() + 5 scale = 1.0 / math.sqrt(d) - dq_g, dk_g, dv_g = gen.backward(q, k, v, do, o, lse, scale=scale, mask="none") + # Generic path: build the adapter directly (BHSD-logical views of the same + # BSHD storage) with the d64 gate forced off, so it compiles + launches + # the generic TemplateParams module. + monkeypatch.setattr(api_sm80, "_sm80_d64_fast_path_eligible", lambda **kw: False) + qb, kb, vb, ob, dob = (t.transpose(1, 2) for t in (q, k, v, o, do)) + dq_g = torch.empty(b, s, h, d, dtype=q.dtype, device="cuda").transpose(1, 2) + dk_g = torch.empty_like(dq_g) + dv_g = torch.empty_like(dq_g) + eng = api_sm80.SdpaBwdDslSm80( + sample_q=qb, + sample_k=kb, + sample_v=vb, + sample_o=ob, + sample_do=dob, + sample_stats=lse, + sample_dq=dq_g, + sample_dk=dk_g, + sample_dv=dv_g, + is_causal=False, + scale_softmax=scale, + ) + assert eng.check_support() + eng.compile() + assert not eng._use_d64 + eng.execute( + q_tensor=qb, k_tensor=kb, v_tensor=vb, o_tensor=ob, do_tensor=dob, stats_tensor=lse, dq_tensor=dq_g, dk_tensor=dk_g, dv_tensor=dv_g, scale_softmax=scale + ) + dq_g, dk_g, dv_g = (t.transpose(1, 2) for t in (dq_g, dk_g, dv_g)) # back to BSHD dq_d, dk_d, dv_d = d64.backward(q, k, v, do, o, lse, scale=scale) torch.testing.assert_close(dq_d.float(), dq_g.float(), rtol=2e-2, atol=2e-2) torch.testing.assert_close(dk_d.float(), dk_g.float(), rtol=2e-2, atol=2e-2) torch.testing.assert_close(dv_d.float(), dv_g.float(), rtol=2e-2, atol=2e-2) + + +def _thd_run_and_check(lens, h, d_qk, d_v, *, check_grads=True): + """Run the THD fwd+bwd wrappers on packed random inputs; when + ``check_grads``, compare each sequence's dQ/dK/dV slice against the dense + fp32 autograd reference (the existing ``_ref_grads`` pattern).""" + import itertools + + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_sm80 + from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80 + + t = int(sum(lens)) + cu = torch.tensor([0] + list(itertools.accumulate(lens)), dtype=torch.int32, device="cuda") + q = torch.randn(1, t, h, d_qk, dtype=torch.float16, device="cuda") + k = torch.randn(1, t, h, d_qk, dtype=torch.float16, device="cuda") + v = torch.randn(1, t, h, d_v, dtype=torch.float16, device="cuda") + do = torch.randn(1, t, h, d_v, dtype=torch.float16, device="cuda") + fwd = sdpa_fwd_wrapper_sm80(q, k, v, is_causal=True, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu, max_s_q=int(max(lens))) + out = sdpa_bwd_wrapper_sm80(q, k, v, fwd["o_tensor"], do, fwd["lse_tensor"], is_causal=True, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu) + if check_grads: + for i in range(len(lens)): + lo, hi = int(cu[i]), int(cu[i + 1]) + + # packed [T, H, D] slice -> BHSD [1, H, S, D] + def _bhsd(x, _lo=lo, _hi=hi): + return x[0, _lo:_hi].permute(1, 0, 2).unsqueeze(0) + + _, dq_ref, dk_ref, dv_ref = _ref_grads(_bhsd(q), _bhsd(k), _bhsd(v), _bhsd(do), is_causal=True, window_left=-1, scale=1.0 / math.sqrt(d_qk)) + torch.testing.assert_close(_bhsd(out["dq_tensor"]).to(torch.float32), dq_ref, rtol=3e-2, atol=3e-2) + torch.testing.assert_close(_bhsd(out["dk_tensor"]).to(torch.float32), dk_ref, rtol=3e-2, atol=3e-2) + torch.testing.assert_close(_bhsd(out["dv_tensor"]).to(torch.float32), dv_ref, rtol=3e-2, atol=3e-2) + return out + + +@pytest.mark.L0 +@pytest.mark.parametrize("d_qk,d_v", [(64, 64), (192, 128)], ids=["gptoss_env", "dsv3_env"]) +@torch_fork_set_rng(seed=0) +def test_sm80_bwd_thd_flavor_envelope_dims(d_qk, d_v): + """Regression: THD must compile the kernel at the FLAVOR ENVELOPE dims, + not the template's 128/128 defaults — a flavor-name-only params build + wrote dQ/dK/dV out of bounds at d=64 and returned wrong gradients at + 192/128 (CodeRabbit critical on the TemplateParams port).""" + try: + _thd_run_and_check([96, 160], h=4, d_qk=d_qk, d_v=d_v) + except ImportError as e: + pytest.skip(f"SM80 SDPA API not available: {e}") + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_sm80_bwd_thd_compile_key_plan_time_only(): + """Issue #604 regression (backward): the packed THD token totals are + RUNTIME values, so two varlen backward calls with different totals must + re-bind ONE compiled artifact (the bprop template's per-shape lru sees a + single miss) — never mint a compile per step, the continuous-batching + pathology no correctness test catches.""" + from cudnn.frost import template_loader + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_sm80 + from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80 + + H, D = 4, 128 + + def varlen(lens): + import itertools + + t = int(sum(lens)) + cu = torch.tensor([0] + list(itertools.accumulate(lens)), dtype=torch.int32, device="cuda") + q = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + k = torch.randn_like(q) + v = torch.randn_like(q) + do = torch.randn_like(q) + fwd = sdpa_fwd_wrapper_sm80(q, k, v, is_causal=True, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu, max_s_q=int(max(lens))) + return sdpa_bwd_wrapper_sm80( + q, + k, + v, + fwd["o_tensor"], + do, + fwd["lse_tensor"], + is_causal=True, + cum_seqlen_q_tensor=cu, + cum_seqlen_k_tensor=cu, + ) + + def cache_totals(): + # Count ONLY the bprop template's per-shape lru (the fwd wrapper runs + # too, and its counters are covered by the forward's twin test); the + # counters are session-global, so assert on DELTAS across our calls. + mods = [m for (path, _params), m in template_loader._MODULES.items() if "bprop" in str(path)] + infos = [m.compile.cache_info() for m in mods if hasattr(m.compile, "cache_info")] + return sum(i.misses for i in infos), sum(i.hits for i in infos) + + varlen([96, 160]) # first call: one compile + n_modules_before = len(template_loader._MODULES) + misses_0, hits_0 = cache_totals() + varlen([128, 64, 320]) # different totals AND batch count + # Different logical batch counts legitimately re-specialize (the cu fake + # length is plan-time); different TOKEN TOTALS at the same batch count + # must not — and the re-bound artifact must still be CORRECT, so this + # call's gradients are validated against the dense per-sequence reference. + # The counter window brackets THIS call alone: a pure cache hit, zero + # misses (netting against the n_seqs=3 call could mask a leak). + misses_pre_rebind, hits_pre_rebind = cache_totals() + _thd_run_and_check([64, 256], h=H, d_qk=D, d_v=D) # same n_seqs as call 1, different totals + misses_post_rebind, hits_post_rebind = cache_totals() + assert misses_post_rebind == misses_pre_rebind, "the same-batch-count re-bind minted a compile (token totals leaked into the key)" + assert hits_post_rebind > hits_pre_rebind, "expected the same-batch-count re-bind to cache-hit" + assert len(template_loader._MODULES) == n_modules_before, "a new template specialization was minted by runtime data" + misses_1, hits_1 = cache_totals() + # Call 2 (n_seqs=3) may legitimately re-specialize once; call 3 shares + # call 1's key (n_seqs=2, different token totals) and MUST cache-hit. + assert misses_1 - misses_0 <= 1, f"THD bprop compile key leaked runtime data: {misses_1 - misses_0} new misses" + assert hits_1 - hits_0 >= 1, "expected a cache hit on the same-batch-count re-call" diff --git a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py index 57f28dbf2..9d11315ea 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py @@ -208,8 +208,10 @@ def test_bwd_engine_end_to_end(): dob = gb.tensor(name="dO", dim=(B, H, S, D), stride=st, data_type=_HALF) statsb = gb.tensor(name="stats", dim=(B, H, S, 1), stride=(H * S, S, 1, 1), data_type=cudnn.data_type.FLOAT) dq, dk, dv = gb.sdpa_backward(q=qb, k=kb, v=vb, o=ob, dO=dob, stats=statsb, attn_scale=_SCALE, use_causal_mask=True) + # Output layout is honored only when DECLARED (the graph invariant: + # IR-inferred output strides are provisional) — bind BSHD-physical. for t in (dq, dk, dv): - t.set_output(True).set_data_type(_HALF) + t.set_output(True).set_data_type(_HALF).set_stride(st) _native_then_pin(gb, _BWD) do_buf = _buf() @@ -358,8 +360,10 @@ def test_engine_execute_does_not_allocate(): dob = gb.tensor(name="dO", dim=(B, H, S, D), stride=st, data_type=_HALF) statsb = gb.tensor(name="stats", dim=(B, H, S, 1), stride=stats_stride, data_type=cudnn.data_type.FLOAT) dq, dk, dv = gb.sdpa_backward(q=qb, k=kb, v=vb, o=ob, dO=dob, stats=statsb, attn_scale=_SCALE, use_causal_mask=True) - for t in (dq, dk, dv): - t.set_output(True).set_data_type(_HALF) + # Declared BSHD-physical output layouts (IR-inferred strides are provisional). + dq.set_output(True).set_data_type(_HALF).set_stride(st) + dk.set_output(True).set_data_type(_HALF).set_stride(st_kv) + dv.set_output(True).set_data_type(_HALF).set_stride(st_kv) _native_then_pin(gb, _BWD) assert gb.get_workspace_size() > 0, "the SM80 bwd executor must report its carved scratch" do_buf = _buf() diff --git a/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py b/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py index fce8f38ce..6a007ab33 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py @@ -146,8 +146,10 @@ def test_sm80_bwd_respects_handle_stream_and_is_capturable(): dob = gb.tensor(dim=dims, stride=_STRIDES, data_type=_HALF, name="dO") statsb = gb.tensor(dim=(_B, _H, _S, 1), stride=(_H * _S, _S, 1, 1), data_type=_F32, name="stats") dq, dk, dv = gb.sdpa_backward(q=qb, k=kb, v=vb, o=ob, dO=dob, stats=statsb, attn_scale=1.0 / math.sqrt(_D), use_causal_mask=True) + # Output layout is honored only when DECLARED (IR-inferred strides are + # provisional) — bind BSHD-physical to match the buffers below. for x in (dq, dk, dv): - x.set_output(True).set_data_type(_HALF) + x.set_output(True).set_data_type(_HALF).set_stride(_STRIDES) _build_and_pin(gb, "sdpa_bwd_sm80") do_gpu = _mk_buf()