diff --git a/python/cudnn/engines/engine_ids.py b/python/cudnn/engines/engine_ids.py index 832f4a840..e75748250 100644 --- a/python/cudnn/engines/engine_ids.py +++ b/python/cudnn/engines/engine_ids.py @@ -33,7 +33,7 @@ LINEAR_ATTENTION_ID_BASE = PYTHON_ENGINE_ID_BASE + 100 # 20_100..20_199 FROST_GEMM_ID_BASE = PYTHON_ENGINE_ID_BASE + 200 # 20_200..20_299 FROST_SDPA_FWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 300 # 20_300..20_399 -FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 400 # reserved +FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 400 # 20_400..20_499 OUT_OF_TREE_ID_BASE = PYTHON_ENGINE_ID_BASE + 10_000 # 30_000+, register_backend() # The delegating entry: the backend picks among candidates it holds but does not diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index 91ebb22e2..6bd92f5eb 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -43,7 +43,7 @@ from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple -from .engine_ids import FROST_GEMM_ID_BASE, FROST_SDPA_FWD_ID_BASE, LINEAR_ATTENTION_ID_BASE +from .engine_ids import FROST_GEMM_ID_BASE, FROST_SDPA_BWD_ID_BASE, FROST_SDPA_FWD_ID_BASE, LINEAR_ATTENTION_ID_BASE _LOG = logging.getLogger("cudnn.engines.manifest") @@ -101,6 +101,7 @@ def matches(self, node_types: frozenset, sm: Optional[int]) -> bool: _GEMM_ANCHOR = frozenset({"MATMUL", "MATMUL_FP8", "MOE_GROUPED_MATMUL"}) _GEMM_CLOSURE = _GEMM_ANCHOR | frozenset({"POINTWISE", "REDUCTION", "RESHAPE", "BLOCK_SCALE_QUANTIZE", "BLOCK_SCALE_DEQUANTIZE"}) _SDPA_FWD = frozenset({"SDPA", "SDPA_FP8", "SDPA_MXFP8"}) +_SDPA_BWD = frozenset({"SDPA_BWD"}) _GDN = frozenset({"GDN", "GDN_BWD"}) _GDN2 = frozenset({"GDN2", "GDN2_BWD"}) _KDA = frozenset({"KDA", "KDA_BWD"}) @@ -176,6 +177,18 @@ def matches(self, node_types: frozenset, sm: Optional[int]) -> bool: sm_lo=100, opt_in=True, ), + EngineRow( + FROST_SDPA_BWD_ID_BASE + 0, + "frost_sdpa_bwd", + "cudnn.sdpa.bwd.engine", + "FrostSdpaBwdEngines", + _SDPA_BWD, + id_hi=FROST_SDPA_BWD_ID_BASE + 100, + # TODO: widen when an SM100/SM80 spec lands + sm_lo=120, + sm_hi=121, + opt_in=True, + ), ) diff --git a/python/cudnn/frost/tile_dsl/mma.py b/python/cudnn/frost/tile_dsl/mma.py index 94b0c337c..66340c6f2 100644 --- a/python/cudnn/frost/tile_dsl/mma.py +++ b/python/cudnn/frost/tile_dsl/mma.py @@ -12,6 +12,36 @@ from .swizzle import swizzle_xor_128b, swizzle_lin_128b +@cute.jit +def ptx_mma_m16n8k16_f32( + a0: cutlass.Int32, + a1: cutlass.Int32, + a2: cutlass.Int32, + a3: cutlass.Int32, + b0: cutlass.Int32, + b1: cutlass.Int32, + c0: cutlass.Float32, + c1: cutlass.Float32, + c2: cutlass.Float32, + c3: cutlass.Float32, + ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> tuple[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32]: + """``mma.sync.aligned.m16n8k16.row.col.f32.{f16|bf16}.{f16|bf16}.f32``.""" + if cutlass.const_expr(ab_dtype != cutlass.Float16 and ab_dtype != cutlass.BFloat16): + raise TypeError(f"Invalid A/B dtype: {ab_dtype}") + ab_tag = "f16" if cutlass.const_expr(ab_dtype == cutlass.Float16) else "bf16" + return cute.arch.inline_ptx( + f"mma.sync.aligned.m16n8k16.row.col.f32.{ab_tag}.{ab_tag}.f32 {{$0,$1,$2,$3}}, {{$4,$5,$6,$7}}, {{$8,$9}}, {{$10,$11,$12,$13}};", + write_only_types=[ + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ], + read_only_args=[a0, a1, a2, a3, b0, b1, c0, c1, c2, c3], + ) + + @cute.jit def mma_ss(desc, desc_a_base, desc_b_base, tmem_c, tmem_sf_a=None, tmem_sf_b=None, accumulate: bool = False, k_start: int = 0, k_count=None): if cutlass.const_expr(desc.cta_group == 1): diff --git a/python/cudnn/frost/tile_dsl/swizzle.py b/python/cudnn/frost/tile_dsl/swizzle.py index 93c4f7148..e8b16adea 100644 --- a/python/cudnn/frost/tile_dsl/swizzle.py +++ b/python/cudnn/frost/tile_dsl/swizzle.py @@ -15,6 +15,49 @@ def swizzle_xor_128b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2): return swz_chunk * chunk_elems + in_chunk +@cute.jit +def swizzle_xor_64b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2): + chunk_elems = 16 // elem_bytes + chunk_idx = col_elem // chunk_elems + in_chunk = col_elem % chunk_elems + swz_chunk = chunk_idx ^ ((row >> 1) & 3) + return swz_chunk * chunk_elems + in_chunk + + +@cute.jit +def swizzle_xor_32b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2): + chunk_elems = 16 // elem_bytes + chunk_idx = col_elem // chunk_elems + in_chunk = col_elem % chunk_elems + swz_chunk = chunk_idx ^ ((row >> 2) & 1) + return swz_chunk * chunk_elems + in_chunk + + +@cute.jit +def swizzle_xor( + row: cutlass.Int32, + col: cutlass.Int32, + row_stride: cutlass.Constexpr[int], + elem_bytes: cutlass.Constexpr[int], +) -> cutlass.Int32: + """Return the physical SMEM column for an XOR-swizzled row-major tile. + + The XOR is applied at the 16-byte boundary for all element widths. + ``elem_bytes`` selects the element-domain shift and swizzle chunk size. + """ + row_stride_bytes = cutlass.const_expr(row_stride * elem_bytes) + if cutlass.const_expr(row_stride_bytes % 128 == 0): + chunk_elems = 128 // elem_bytes + swizzled = swizzle_xor_128b(row, col % chunk_elems, elem_bytes=elem_bytes) + elif cutlass.const_expr(row_stride_bytes % 64 == 0): + chunk_elems = 64 // elem_bytes + swizzled = swizzle_xor_64b(row, col % chunk_elems, elem_bytes=elem_bytes) + else: + chunk_elems = 32 // elem_bytes + swizzled = swizzle_xor_32b(row, col % chunk_elems, elem_bytes=elem_bytes) + return (col // chunk_elems) * chunk_elems + swizzled + + @cute.jit def swizzle_lin_128b(lin, *, row_stride_log2: cutlass.Constexpr[int], elem_bytes: cutlass.Constexpr[int] = 2): chunk_log2 = cutlass.const_expr((16 // elem_bytes).bit_length() - 1) diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index 21e4a53fc..05f918123 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -2,8 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 from .api import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 +from .api_dsl import SdpaBwdDsl, SdpaBwdDslSm120, sdpa_bwd_wrapper_dsl_sm120 __all__ = [ "SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256", + "SdpaBwdDsl", + "SdpaBwdDslSm120", + "sdpa_bwd_wrapper_dsl_sm120", ] diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py new file mode 100644 index 000000000..6e0360658 --- /dev/null +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -0,0 +1,447 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""cuDNN-frontend adapter over the FROST DSL SDPA backward kernels.""" + +from __future__ import annotations + +import logging +import math +import os +from abc import abstractmethod +from typing import Optional + +import torch +from cuda.bindings import driver as cuda + +from cudnn.api_base import APIBase, TensorDesc, TupleDict +from cudnn.frost.template_loader import load_template +from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 +from cudnn.sdpa.bwd.config_sm120 import ( + SEQ_KV_TILES as _SM120_KV_TILES, + SEQ_Q_TILES as _SM120_Q_TILES, + SUPPORTED_HEAD_DIMS as _SM120_SUPPORTED_HEAD_DIMS, + TemplateParams as Sm120TemplateParams, +) +from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, ws_align + +_SM120_KERNEL_FILE = "bprop_f16_sm120.py" +_SM120_DTYPE_QKV_CODE = { + torch.bfloat16: DTYPE_BF16, + torch.float16: DTYPE_FP16, +} +# delta / dq_accum rows are padded to multiples of 128 (the kernel's +# dq_accum layout contract: tile_q must divide 128). +_SM120_ROW_ROUND = 128 + +_logger = logging.getLogger(__name__) + + +def _load_sm120_kernel_module(params: Sm120TemplateParams): + """Load one uniquely named backward kernel module per parameter set.""" + + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels", _SM120_KERNEL_FILE) + return load_template(path, params, tag="sdpa_bwd_sm120") + + +def _round_up(x: int, a: int) -> int: + return -(-int(x) // a) * a + + +class SdpaBwdDsl(APIBase): + """Implementation-agnostic interface for FROST DSL SDPA-backward kernels.""" + + def __init__( + self, + sample_q: torch.Tensor | TensorDesc, + sample_k: torch.Tensor | TensorDesc, + sample_v: torch.Tensor | TensorDesc, + sample_o: torch.Tensor | TensorDesc, + sample_do: torch.Tensor | TensorDesc, + sample_stats: torch.Tensor | TensorDesc, + sample_dq: torch.Tensor | TensorDesc, + sample_dk: torch.Tensor | TensorDesc, + sample_dv: torch.Tensor | TensorDesc, + is_causal: bool = False, + causal_bottom_right: bool = False, + scale_softmax: Optional[float] = None, + tile_m: Optional[int] = None, + tile_n: Optional[int] = None, + ) -> None: + super().__init__() + self._warn_experimental_api() + self._logger.debug("Entering __init__") + + 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.stats_desc = self._make_tensor_desc(sample_stats, name="stats") + self.dq_desc = self._make_tensor_desc(sample_dq, name="dQ") + self.dk_desc = self._make_tensor_desc(sample_dk, name="dK") + self.dv_desc = self._make_tensor_desc(sample_dv, name="dV") + + self.is_causal = bool(is_causal) + self.causal_bottom_right = bool(causal_bottom_right) + self.scale_softmax = scale_softmax + self.tile_m = None if tile_m is None else int(tile_m) + self.tile_n = None if tile_n is None else int(tile_n) + + self.batch_size: Optional[int] = None + self.s_q_max: Optional[int] = None + self.s_k_max: Optional[int] = None + self.h_q: Optional[int] = None + self.h_kv: Optional[int] = None + self.head_dim: Optional[int] = None + self.dtype: Optional[torch.dtype] = None + self._initialize_implementation() + self._logger.debug("__init__ completed") + + @abstractmethod + def _initialize_implementation(self) -> None: + """Initialize state private to specific implementations.""" + + @staticmethod + def _to_bshd(tensor: torch.Tensor) -> torch.Tensor: + """Return the compact kernel-facing BSHD tensor for a logical-BHSD INPUT.""" + + view = tensor.transpose(1, 2) + return view if view.is_contiguous() else view.contiguous() + + @staticmethod + def _out_bshd(tensor: torch.Tensor) -> torch.Tensor: + """The compact BSHD view of a logical-BHSD OUTPUT, or raise.""" + + view = tensor.transpose(1, 2) + if not view.is_contiguous(): + raise ValueError( + "output tensor must be logical (B, H, S, D) over compact BSHD storage; " f"got stride {tuple(tensor.stride())} shape {tuple(tensor.shape)}" + ) + return view + + @abstractmethod + def scratch_workspace_bytes(self) -> int: + """Return the per-execution scratch requirement for this implementation.""" + + @abstractmethod + 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, + ) -> None: + """Execute the compiled kernel chain using the common operand set.""" + + +class SdpaBwdDslSm120(SdpaBwdDsl): + """Compile and execute fixed-length SM120/SM121 SDPA backward.""" + + def _initialize_implementation(self) -> None: + # 0 = the kernel's per-head-dim CONFIG default. + self.q_tile = 0 if self.tile_m is None else int(self.tile_m) + self.kv_tile = 0 if self.tile_n is None else int(self.tile_n) + self.compute_capability: Optional[tuple[int, int]] = None + self._k_mod = None + self._sq_rounded: Optional[int] = None + + @staticmethod + def _bshd_physical_ok(desc: TensorDesc) -> bool: + """True when the logical-BHSD desc sits on compact BSHD storage.""" + + b, h, s, d = desc.shape + return tuple(desc.stride) == (s * h * d, d, h * d, 1) + + 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.dq_desc, self.dk_desc, self.dv_desc): + self._value_error_if( + desc.ndim != 4, + f"{desc.name} must be rank-4 (B, H, S, D); got {desc.ndim}", + ) + self._value_error_if( + not self._bshd_physical_ok(desc), + f"{desc.name} must be logical (B, H, S, D) over compact BSHD storage " + f"(the SM120 backward kernels hard-code the H*D row stride); got " + f"stride {desc.stride} shape {desc.shape}", + ) + + b, h_q, s_q, d_qk = self.q_desc.shape + _, h_kv, s_kv, _ = self.k_desc.shape + 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_qk), name="V") + self._check_tensor_shape(self.o_desc, (b, h_q, s_q, d_qk), name="O") + self._check_tensor_shape(self.do_desc, (b, h_q, s_q, d_qk), name="dO") + self._check_tensor_shape(self.dq_desc, tuple(self.q_desc.shape), name="dQ") + self._check_tensor_shape(self.dk_desc, tuple(self.k_desc.shape), name="dK") + self._check_tensor_shape(self.dv_desc, tuple(self.v_desc.shape), name="dV") + + for label, val in (("B", b), ("H_q", h_q), ("H_kv", h_kv), ("S_q", s_q), ("S_kv", s_kv), ("D", d_qk)): + self._value_error_if(int(val) <= 0, f"{label} must be > 0; got {val}") + self._not_implemented_error_if( + h_q != h_kv, + f"SM120 DSL SDPA backward does not implement GQA / MQA; got H_q={h_q}, H_kv={h_kv}", + ) + self._value_error_if( + d_qk not in _SM120_SUPPORTED_HEAD_DIMS, + f"D ({d_qk}) must be one of {_SM120_SUPPORTED_HEAD_DIMS}", + ) + + self._value_error_if( + self.stats_desc.ndim != 4 or tuple(self.stats_desc.shape) != (b, h_q, s_q, 1), + f"stats must be (B, H_q, S_q, 1); got {tuple(self.stats_desc.shape)}", + ) + self._value_error_if( + not self.stats_desc.is_contiguous(), + f"stats must be contiguous; got stride {self.stats_desc.stride}", + ) + self._check_dtype(self.stats_desc, torch.float32, name="stats") + + 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", + ) + self._value_error_if( + desc.device != self.q_desc.device, + f"{desc.name} must be on device {self.q_desc.device}, got {desc.device}", + ) + self._value_error_if( + self.q_desc.device.type != "cuda", + f"Q must be a CUDA tensor, got device {self.q_desc.device}", + ) + + self._value_error_if( + self.q_tile not in (0,) + _SM120_Q_TILES, + f"q_tile must be one of {(0,) + _SM120_Q_TILES} (0 = per-head-dim default); got {self.q_tile}", + ) + self._value_error_if( + self.kv_tile not in (0,) + _SM120_KV_TILES, + f"kv_tile must be one of {(0,) + _SM120_KV_TILES} (0 = per-head-dim default); got {self.kv_tile}", + ) + self._value_error_if( + self.causal_bottom_right and not self.is_causal, + "causal_bottom_right requires is_causal=True", + ) + # The kernel's diagonal is bottom-right (FA2). Top-left causal is + # identical when S_q == S_kv, which is what the engine gates on. + self._value_error_if( + self.is_causal and not self.causal_bottom_right and s_q != s_kv, + "top-left causal with S_q != S_kv is not supported (the kernel diagonal is bottom-right)", + ) + # Bottom-right causal with S_q > S_kv produces fully-masked query + # rows whose forward stats are -inf; the backward exp2 replay is not + # specified for those rows. Conservatively rejected (engine gates on + # this too). + self._value_error_if( + self.is_causal and s_q > s_kv, + "causal with S_q > S_kv is not supported (fully-masked query rows)", + ) + + self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") + self.compute_capability = torch.cuda.get_device_capability(self.q_desc.device) + self._runtime_error_if( + self.compute_capability not in {(12, 0), (12, 1)}, + f"SdpaBwdDslSm120 requires SM120 or SM121, found SM{self.compute_capability[0]}{self.compute_capability[1]}", + ) + + 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_q) + self.s_k_max = int(s_kv) + self.h_q = int(h_q) + self.h_kv = int(h_kv) + self.head_dim = int(d_qk) + self._sq_rounded = _round_up(self.s_q_max, _SM120_ROW_ROUND) + self._is_supported = True + + self._logger.debug("check_support completed successfully") + return True + + def compile(self) -> None: + """Compile the shape-specialized SM120 FROST backward template.""" + + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + params = Sm120TemplateParams( + dtype_qkv=_SM120_DTYPE_QKV_CODE[self.dtype], + is_causal=self.is_causal, + q_tile=self.q_tile, + kv_tile=self.kv_tile, + ) + self._k_mod = _load_sm120_kernel_module(params) + self._compiled_kernel = self._k_mod.compile( + compute_capability=self.compute_capability, + b=self.batch_size, + qh=self.h_q, + sq=self.s_q_max, + skv=self.s_k_max, + d=self.head_dim, + ) + self._logger.debug("compile completed") + + def scratch_workspace_bytes(self) -> int: + """delta (fp32 [B, H, SQ_r128]) + dq_accum (fp32 flat [B*SQ_r128*H*D]).""" + + self._ensure_support_checked() + delta_bytes = ws_align(self.batch_size * self.h_q * self._sq_rounded * 4) + dq_accum_bytes = ws_align(self.batch_size * self._sq_rounded * self.h_q * self.head_dim * 4) + return delta_bytes + dq_accum_bytes + + 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, + ) -> None: + """Execute tensors matching the compiled specialization.""" + + if self._compiled_kernel is None: + raise RuntimeError("SdpaBwdDslSm120 kernel is not compiled") + + scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) + scale_log2 = scale_val * math.log2(math.e) + + carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "sdpa_bwd_sm120") + delta = carver.take(self.batch_size * self.h_q * self._sq_rounded, torch.float32).reshape(self.batch_size, self.h_q, self._sq_rounded) + dq_accum = carver.take(self.batch_size * self._sq_rounded * self.h_q * self.head_dim, torch.float32) + + if current_stream is None: + # Direct call (no dispatch-forwarded stream): fall back to torch's + # current stream. A stream forwarded from the execute-time handle + # is respected rather than clobbered. + current_stream = cuda.CUstream(torch.cuda.current_stream(q_tensor.device).cuda_stream) + + import cutlass + + q = self._to_bshd(q_tensor) + k = self._to_bshd(k_tensor) + v = self._to_bshd(v_tensor) + o = self._to_bshd(o_tensor) + do = self._to_bshd(do_tensor) + dq = self._out_bshd(dq_tensor) + dk = self._out_bshd(dk_tensor) + dv = self._out_bshd(dv_tensor) + lse = stats_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) + + kernels = self._compiled_kernel + # Three-kernel chain + kernels.dot(o, do, delta, dq_accum, current_stream) + kernels.main( + q, + k, + v, + do, + lse, + delta, + dq_accum, + dk, + dv, + cutlass.Float32(scale_log2), + cutlass.Float32(scale_val), + current_stream, + ) + kernels.cvt(dq_accum, dq, cutlass.Float32(scale_val), current_stream) + + +def _tensor_signature(tensor: torch.Tensor) -> tuple: + """(shape, stride, dtype, device) — everything the specialization keys on.""" + return (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype, tensor.device) + + +_wrapper_api_cache: dict[tuple, SdpaBwdDslSm120] = {} + + +def sdpa_bwd_wrapper_dsl_sm120( + q_tensor: torch.Tensor, + k_tensor: torch.Tensor, + v_tensor: torch.Tensor, + o_tensor: torch.Tensor, + do_tensor: torch.Tensor, + stats_tensor: torch.Tensor, + is_causal: bool = False, + causal_bottom_right: bool = False, + scale_softmax: Optional[float] = None, +) -> TupleDict: + """Run SM120 SDPA backward and return ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=...)``.""" + + dq_tensor = torch.empty_strided(q_tensor.shape, q_tensor.stride(), dtype=q_tensor.dtype, device=q_tensor.device) + dk_tensor = torch.empty_strided(k_tensor.shape, k_tensor.stride(), dtype=k_tensor.dtype, device=k_tensor.device) + dv_tensor = torch.empty_strided(v_tensor.shape, v_tensor.stride(), dtype=v_tensor.dtype, device=v_tensor.device) + + # check_support()/compile() run only on a miss, so the key must carry the + # full signature of every operand the specialization depends on (dq/dk/dv + # are derived from q/k/v above). + cache_key = ( + _tensor_signature(q_tensor), + _tensor_signature(k_tensor), + _tensor_signature(v_tensor), + _tensor_signature(o_tensor), + _tensor_signature(do_tensor), + _tensor_signature(stats_tensor), + bool(is_causal), + bool(causal_bottom_right), + scale_softmax, + ) + api = _wrapper_api_cache.get(cache_key) + if api is None: + api = SdpaBwdDslSm120( + sample_q=q_tensor, + sample_k=k_tensor, + sample_v=v_tensor, + sample_o=o_tensor, + sample_do=do_tensor, + sample_stats=stats_tensor, + sample_dq=dq_tensor, + sample_dk=dk_tensor, + sample_dv=dv_tensor, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + scale_softmax=scale_softmax, + ) + api.check_support() + api.compile() + _wrapper_api_cache[cache_key] = api + + workspace = torch.empty(api.scratch_workspace_bytes(), dtype=torch.uint8, device=q_tensor.device) + api.execute( + q_tensor=q_tensor, + k_tensor=k_tensor, + v_tensor=v_tensor, + o_tensor=o_tensor, + do_tensor=do_tensor, + stats_tensor=stats_tensor, + dq_tensor=dq_tensor, + dk_tensor=dk_tensor, + dv_tensor=dv_tensor, + scale_softmax=scale_softmax, + workspace=workspace, + ) + return TupleDict(dq_tensor=dq_tensor, dk_tensor=dk_tensor, dv_tensor=dv_tensor) diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py new file mode 100644 index 000000000..194e7401f --- /dev/null +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Compile-time configuration for the FROST SM120 SDPA backward template.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 + +SEQ_Q_TILES = (64, 128) +SEQ_KV_TILES = (64, 128) +SUPPORTED_HEAD_DIMS = (32, 64, 128) + + +@dataclass(frozen=True) +class TemplateParams: + """Per-graph parameters that change the traced SM120 backward kernel. + + Tensor geometry deliberately stays out of this record. Scalar dimensions + are inputs to the template module's per-shape ``compile()`` cache, while + strides follow the fixed compact-BSHD kernel contract. This frozen record + identifies the import-time specialization shared by all compatible shapes. + """ + + dtype_qkv: int = DTYPE_FP16 + is_causal: bool = False + use_pdl: bool = True + q_tile: int = 0 + kv_tile: int = 0 + + +def validate_params(params: TemplateParams) -> None: + """Validate the SM120 backward template specialization. + + Reachable failures should already have been rejected by the engine + capabilities or adapter support checks; this validation is a backstop for + direct template use. + """ + + if params.dtype_qkv not in (DTYPE_BF16, DTYPE_FP16): + raise ValueError(f"SM120 SDPA bwd: dtype_qkv must be DTYPE_BF16 ({DTYPE_BF16}) or DTYPE_FP16 ({DTYPE_FP16}); got {params.dtype_qkv}") + if params.q_tile not in (0,) + SEQ_Q_TILES: + raise ValueError(f"SM120 SDPA bwd: q_tile must be one of {(0,) + SEQ_Q_TILES} (0 = per-head-dim default); got {params.q_tile}") + if params.kv_tile not in (0,) + SEQ_KV_TILES: + raise ValueError(f"SM120 SDPA bwd: kv_tile must be one of {(0,) + SEQ_KV_TILES} (0 = per-head-dim default); got {params.kv_tile}") diff --git a/python/cudnn/sdpa/bwd/engine.py b/python/cudnn/sdpa/bwd/engine.py new file mode 100644 index 000000000..7f42269c7 --- /dev/null +++ b/python/cudnn/sdpa/bwd/engine.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The FROST SDPA-backward engines: one BaseEngine per capability cell. + +Listed in ``cudnn/engines/manifest.py`` as ONE row owning the +``FROST_SDPA_BWD_ID_BASE`` block, so ``FrostSdpaBwdEngines()`` returns the whole +family and a graph containing an sdpa_backward() node reaches them through the +ordinary lifecycle — no registration call. The row is opt-in +(``CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1``) until these engines have the arch +coverage to serve graphs unasked. + +The capability table, the probe and the lowering stay in ``engines.py`` +(``ENGINE_SPECS`` / ``analyze_for`` / ``build``); this file is only the engine +contract around them. +""" + +from typing import TYPE_CHECKING, Any, List, Optional + +from cudnn.engines.base import BaseEngine, CompiledPlan, ExecutionContext, PlanConfig +from cudnn.engines.engine_ids import FROST_SDPA_BWD_ID_BASE + +if TYPE_CHECKING: + from cudnn._pygraph import pygraph + + from .engines import EngineSpec + + +def _check_workspace(workspace, required: int, name: str) -> None: + """A FROST executor carves its scratch out of the CALLER's workspace: no + hidden per-execute allocation, stable pointers, CUDA-graph friendly.""" + if workspace is None: + raise ValueError(f"{name} needs a {required}-byte workspace; execute() got none — allocate graph.get_workspace_size() bytes and pass it") + available = workspace.numel() * workspace.element_size() if hasattr(workspace, "numel") else len(workspace) + if available < required: + raise ValueError(f"{name} needs a {required}-byte workspace; the buffer provides {available}") + + +class _FrostSdpaBwdPlan(CompiledPlan): + """A compiled SDPA-backward executor plus the graph binding it was compiled for.""" + + def __init__(self, name: str, compiled: Any): + self._name = name + self._compiled = compiled + # The kernel is bound to specific graph tensors; the variant pack the + # graph API hands us covers every IO tensor of the graph, so key the + # kernel's own operands out of it by uid (uids are eager and unique). + self._tensors = list(compiled.binding.bound_tensors()) + + def get_workspace_size(self) -> int: + return int(getattr(self._compiled, "workspace_bytes", 0) or 0) + + def execute(self, graph: "pygraph", uid_to_data, ctx: ExecutionContext) -> None: + # Keyed by IR tensor object: that is the binding's own identity, and the + # only key resolve_variant_pack() accepts for an auto-assigned uid. + pack = {} + missing = [] + for t in self._tensors: + buf = uid_to_data.get(t.get_uid()) + if buf is None: + missing.append(t.get_name() or t.get_uid()) + else: + pack[t] = buf + if missing: + raise ValueError(f"{self._name}: the variant pack is missing buffers for {missing}") + required = self.get_workspace_size() + if required: + _check_workspace(ctx.workspace, required, self._name) + self._compiled(pack, ctx.workspace, stream=ctx.stream) + else: + self._compiled(pack, stream=ctx.stream) + + +class FrostSdpaBwdEngine(BaseEngine): + """One SDPA-backward capability cell (arch x geometry). + + Wraps a single :class:`~cudnn.sdpa.bwd.engines.EngineSpec`: ``name`` is the + spec's shipped name and ``engine_id`` is its fixed offset in the family's id + block (see :func:`FrostSdpaBwdEngines`). + """ + + def __init__(self, spec: "EngineSpec", offset: int): + super().__init__() + self._spec = spec + self.name = spec.name + self.engine_id = FROST_SDPA_BWD_ID_BASE + offset + + def _decline_reason(self, graph: "pygraph", knobs) -> Optional[str]: + from .engines import analyze_for + + try: + _, reason = analyze_for(self._spec, graph, knobs) + except ValueError as exc: + # ValueError is the analyzer's internal "cannot express this graph"; + # at the engine boundary that is a decline, not a user error. + return str(exc) + return reason + + def check_support(self, graph: "pygraph") -> None: + reason = self._decline_reason(graph, None) + if reason is not None: + raise NotImplementedError(f"{self.name}: {reason}") + + def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: + # One plan, no knobs: nothing proposes a tuning request today, so the + # engine runs at its capability row's default tiles. A knob search + # (SdpaBwdKnobs over Capabilities.tile_ms/tile_ns) becomes several + # entries here; each one's knobs reach build_plan verbatim. + self.check_support(graph) + return [PlanConfig(self.engine_id, self.default_knobs)] + + def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: ExecutionContext = None) -> CompiledPlan: + from .engines import build + + knobs = plan.knobs if plan is not None else None + try: + return _FrostSdpaBwdPlan(self.name, build(self._spec, graph, knobs)) + except (NotImplementedError, ValueError) as exc: + raise NotImplementedError(f"{self.name}: {exc}") from exc + + +# engine_id = FROST_SDPA_BWD_ID_BASE + offset. An offset is FIXED FOREVER: an +# autotune result is (engine_id, knobs) and must replay across versions. +# Appending a spec takes the next free offset; offsets are never reordered or +# reused. Keyed by the spec's shipped name rather than by its position in +# ENGINE_SPECS, because that position is the PREFERENCE order and may change. +_ID_OFFSETS = { + "sdpa_bwd_sm120": 0, +} + + +def FrostSdpaBwdEngines() -> List[FrostSdpaBwdEngine]: + """The SDPA-backward engine family, in ENGINE_SPECS (= preference) order.""" + from .engines import ENGINE_SPECS + + engines = [] + for spec in ENGINE_SPECS: + if spec.name not in _ID_OFFSETS: + raise KeyError(f"engine spec {spec.name!r} has no engine-id offset; allocate the next free one in engine._ID_OFFSETS (never reuse)") + engines.append(FrostSdpaBwdEngine(spec, _ID_OFFSETS[spec.name])) + return engines diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py new file mode 100644 index 000000000..dbb020f75 --- /dev/null +++ b/python/cudnn/sdpa/bwd/engines.py @@ -0,0 +1,374 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""FROST SDPA-backward engine registry: capability declarations + spec table. + +One registered engine per architecture, named ``sdpa_bwd_sm`` (dtype is +NOT part of the identity: a cell's engine serves every dtype its kernel +handles — fp16 and bf16 today — via ``Capabilities.dtypes``). + +The backward opset keeps its own :class:`Capabilities` record rather than +reusing the forward one: the feature model differs (gradient side outputs, +determinism, no phase axis) and the shared analyzer facts already carry +everything both need. The shared analyzer +(``cudnn.sdpa.graph_analyzer.analyze``) parses the graph once into +:class:`SdpaGraphFacts`; each engine's probe is a cheap field-by-field +candidate match against its row below. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Callable, Optional + +import cudnn +import torch +from cuda.bindings import driver as _cuda_driver + +from cudnn.sdpa import graph_analyzer as ga +from cudnn.sdpa.bwd.api_dsl import SdpaBwdDslSm120 +from cudnn.sdpa.bwd.config_sm120 import ( + SEQ_KV_TILES as _SM120_KV_TILES, + SEQ_Q_TILES as _SM120_Q_TILES, + SUPPORTED_HEAD_DIMS as _SM120_HEAD_DIMS, +) + +_LOG = logging.getLogger(__name__) + +_BLACKWELL_GEFORCE_ARCHES = frozenset[tuple[int, int]]({(12, 0), (12, 1)}) + + +@dataclass(frozen=True) +class SdpaBwdKnobs: + """Per-plan tuning request for the SDPA-backward engines. + + This is the operation's knob *vocabulary* — typed fields, no global enum. + ``None`` means "no preference". Travels as ``PlanConfig.knobs``; each + engine's :class:`Capabilities` row advertises the domain it honors, and the + probe rejects the engine for any request outside that domain (a knob is + honored or the engine is ineligible — never silently degraded). + """ + + tile_m: Optional[int] = None # Q sequence tile width (q_tile) + tile_n: Optional[int] = None # KV sequence tile width (kv_tile) + + +@dataclass(frozen=True) +class Capabilities: + """What one backward ENGINE can serve — the envelope of graphs its + lowering can honor. Compared field-by-field against SdpaGraphFacts in the + probe.""" + + arches: frozenset[tuple[int, int]] + d: frozenset[int] # supported head dims (d_qk == d_v required) + dtypes: frozenset[torch.dtype] = frozenset({torch.float16, torch.bfloat16}) + + # optional features a backward graph may request + gqa: bool = False # h_q != h_kv + causal: bool = False + bottom_right: bool = False + deterministic: bool = False # use_deterministic_algorithm=True + dbias: bool = False # dBias output + dsink: bool = False # dSink_token output + bias: bool = False + dropout: bool = False + score_mod: bool = False + paged_kv: bool = False + alibi: bool = False + block_mask: bool = False + rng_dump: bool = False + score_max: bool = False + score_sum_exp: bool = False + dynamic_scale: bool = False + unfuse_fma: bool = False + seq_q_trim: bool = False + right_band_widening: bool = False + swa: bool = False + padded: bool = False + sink: bool = False + thd: bool = False + + # Tuning-knob domains this engine's lowering honors (see SdpaBwdKnobs). + tile_ms: frozenset[int] = frozenset() + tile_ns: frozenset[int] = frozenset() + + +def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", requested: Any = None) -> Optional[str]: + """First reason this engine is not a candidate for these facts and tuning + knobs, or ``None`` when lowering should perform the final feasibility check. + """ + if facts.invalid: + return facts.invalid + if not facts.is_backward: + return "this engine serves sdpa_backward() graphs only" + if requested is not None: + if not isinstance(requested, SdpaBwdKnobs): + return f"knob request is a {type(requested).__name__}, not SdpaBwdKnobs — wrong operation's vocabulary" + for value, domain, label in ( + (requested.tile_m, capabilities.tile_ms, "tile_m"), + (requested.tile_n, capabilities.tile_ns, "tile_n"), + ): + if value is not None and value not in domain: + return f"requested {label}={value} is outside this engine's domain {sorted(domain)}" + if facts.device_cc not in capabilities.arches: + required = " or ".join(f"SM{major}{minor}" for major, minor in sorted(capabilities.arches)) + return f"requires {required}; current device is {facts.device_cc}" + if facts.is_mxfp8 or facts.is_fp8: + return "this engine serves only half (fp16/bf16) sdpa_backward graphs" + if facts.d_qk != facts.d_v: + return f"D_QK must equal D_V; graph has D_QK={facts.d_qk}/D_V={facts.d_v}" + if facts.d_qk not in capabilities.d: + return f"serves D in {sorted(capabilities.d)}; graph has D={facts.d_qk}" + if facts.dtype not in capabilities.dtypes: + return f"dtype {facts.dtype} not in {sorted(str(d) for d in capabilities.dtypes)}" + if not facts.uniform_dtype: + return "K/V/O/dO/dQ/dK/dV dtypes must match Q" + if facts.h_q != facts.h_kv and not capabilities.gqa: + return f"GQA / MQA is not supported (H_q={facts.h_q}, H_kv={facts.h_kv})" + if not facts.bshd_layout: + return "Q/K/V/O/dO/dQ/dK/dV must be BSHD-physical (stride order 3,1,2,0)" + + for fact, cap, label in ( + (facts.deterministic, capabilities.deterministic, "use_deterministic_algorithm (dQ accumulates through fp32 atomics)"), + (facts.has_dbias, capabilities.dbias, "dBias output"), + (facts.has_dsink, capabilities.dsink, "dSink_token output"), + (facts.has_bias, capabilities.bias, "bias"), + (facts.has_dropout, capabilities.dropout, "dropout"), + (facts.has_score_mod, capabilities.score_mod, "score_mod"), + (facts.has_paged_kv, capabilities.paged_kv, "paged attention"), + (facts.has_alibi, capabilities.alibi, "ALiBi"), + (facts.has_block_mask, capabilities.block_mask, "block_mask"), + (facts.has_rng_dump, capabilities.rng_dump, "rng_dump"), + (facts.has_score_max, capabilities.score_max, "score_max"), + (facts.has_score_sum_exp, capabilities.score_sum_exp, "score_sum_exp"), + (facts.dynamic_scale, capabilities.dynamic_scale, "tensor attn_scale"), + (facts.has_unfuse_fma, capabilities.unfuse_fma, "unfuse_fma"), + (facts.seq_q_trim, capabilities.seq_q_trim, "seq_len_q without padding mask"), + (facts.right_band_widening, capabilities.right_band_widening, "causal right-band widening"), + (facts.window_left is not None, capabilities.swa, "sliding window"), + (facts.padded, capabilities.padded, "padding mask"), + (facts.has_sink, capabilities.sink, "sink token"), + (facts.thd, capabilities.thd, "THD / ragged"), + (facts.causal, capabilities.causal, "causal mask"), + ): + if fact and not cap: + return f"graph uses {label}, which this engine does not support" + + if facts.bottom_right and not facts.causal: + return "bottom-right alignment requires a causal upper bound" + if facts.causal: + # The kernel's diagonal is bottom-right-aligned (FA2 convention: + # masked iff k > q + S_kv - S_q). Top-left causal is identical when + # S_q == S_kv. + if facts.bottom_right and not capabilities.bottom_right: + return "graph uses bottom-right causal, which this engine does not support" + if not facts.bottom_right and facts.s_q != facts.s_kv: + return "top-left causal with S_q != S_kv is not supported (the kernel diagonal is bottom-right)" + # Bottom-right causal with S_q > S_kv produces fully-masked query + # rows whose forward stats are -inf; the backward exp2 replay is not + # specified for those rows. Conservatively rejected. + if facts.s_q > facts.s_kv: + return "causal with S_q > S_kv is not supported (fully-masked query rows)" + + # The kernel consumes the forward stats as a contiguous natural-log LSE + # (fp32 (B, H_q, S_q, 1)); a strided stats view has no zero-copy reshape. + if facts.stats_t is not None: + if facts.stats_t.get_data_type() != cudnn.data_type.FLOAT: + return f"stats must be fp32; got {facts.stats_t.get_data_type()}" + s_dim = tuple(facts.stats_t.get_dim()) + s_stride = tuple(facts.stats_t.get_stride()) + expect_dim = (facts.b, facts.h_q, facts.s_q, 1) + expect_stride = (facts.h_q * facts.s_q, facts.s_q, 1, 1) + if s_dim != expect_dim: + return f"stats must be (B, H_q, S_q, 1) = {expect_dim}; got {s_dim}" + if s_stride != expect_stride: + return f"stats must be contiguous {expect_stride}; got stride {s_stride}" + return None + + +@dataclass(frozen=True) +class EngineSpec: + name: str + capabilities: Capabilities + lower: "Callable[[EngineSpec, ga.SdpaGraphFacts, Any], Any]" + + +def _sm120_spec() -> EngineSpec: + return EngineSpec( + name="sdpa_bwd_sm120", + capabilities=Capabilities( + arches=_BLACKWELL_GEFORCE_ARCHES, + d=frozenset(_SM120_HEAD_DIMS), + dtypes=frozenset({torch.float16, torch.bfloat16}), + causal=True, + bottom_right=True, + tile_ms=frozenset(_SM120_Q_TILES), + tile_ns=frozenset(_SM120_KV_TILES), + ), + lower=lower_dsl_bwd, + ) + + +def analyze_for(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None): + """``(facts, reason)``: the parsed graph and the first reason ``spec`` + cannot serve it under ``knobs`` (``None`` when it can). + + The single eligibility entry point, shared by :func:`probe`, :func:`build` + and ``engine.FrostSdpaBwdEngine.check_support``. ``knobs`` is the plan's + tuning request (``PlanConfig.knobs``), ``None`` for no preference. + """ + facts = ga.analyze(graph) + if facts is None: + return None, "graph is not a single sdpa_backward() node" + return facts, mismatch(spec.capabilities, facts, knobs) + + +def probe(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None) -> bool: + _, reason = analyze_for(spec, graph, knobs) + if reason is not None: + _LOG.debug("cudnn.sdpa: %s ineligible: %s", spec.name, reason) + return False + return True + + +def build(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None): + """Lower ``spec`` for ``graph``, or raise the bare ineligibility reason (the + caller — the engine — names itself in the message).""" + facts, reason = analyze_for(spec, graph, knobs) + if reason is not None: + raise ValueError(reason) + return spec.lower(spec, facts, knobs) + + +def lower_dsl_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = None): + """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. + """ + + # Canonical BSHD-physical geometry, fixed at build time from the facts. + # Deliberately NOT read back from the IR tensors at execute: + # ``build_operation_graph`` rewrites the backward node's K/V ports to + # transposed (B, H, D, S) views, so the live ``get_dim()`` after a native + # build would describe the transposed view while the underlying buffer + # keeps the user's canonical layout (which the bshd gate already proved). + def _bshd_geometry(b: int, h: int, s: int, d: int) -> tuple[tuple[int, ...], tuple[int, ...]]: + return (b, h, s, d), (s * h * d, d, h * d, 1) + + q_geom = _bshd_geometry(facts.b, facts.h_q, facts.s_q, facts.d_qk) + kv_geom = _bshd_geometry(facts.b, facts.h_kv, facts.s_kv, facts.d_qk) + stats_geom = ((facts.b, facts.h_q, facts.s_q, 1), (facts.h_q * facts.s_q, facts.s_q, 1, 1)) + + def _desc(geom, dtype: torch.dtype, name: str) -> "Any": + from cudnn.api_base import TensorDesc + + dim, stride = geom + return TensorDesc( + dtype=dtype, + shape=dim, + stride=stride, + stride_order=TensorDesc._compute_stride_order(dim, stride), + device=torch.device("cuda", torch.cuda.current_device()), + name=name, + ) + + api = SdpaBwdDslSm120( + sample_q=_desc(q_geom, facts.dtype, "q"), + sample_k=_desc(kv_geom, facts.dtype, "k"), + sample_v=_desc(kv_geom, facts.dtype, "v"), + sample_o=_desc(q_geom, facts.dtype, "o"), + sample_do=_desc(q_geom, facts.dtype, "dO"), + sample_stats=_desc(stats_geom, torch.float32, "stats"), + sample_dq=_desc(q_geom, facts.dtype, "dQ"), + sample_dk=_desc(kv_geom, facts.dtype, "dK"), + sample_dv=_desc(kv_geom, facts.dtype, "dV"), + is_causal=facts.causal, + causal_bottom_right=facts.bottom_right, + scale_softmax=facts.scale, + tile_m=requested.tile_m if requested is not None else None, + tile_n=requested.tile_n if requested is not None else None, + ) + api.check_support() # raises ValueError / NotImplementedError if unsupported + api.compile() + + # Workspace requirement for the compiled geometry (the FROST executor + # contract, see engine._check_workspace): the delta and dq_accum fp32 + # scratch is carved from the CALLER's workspace at execute, so its size is + # fixed here and recorded on the executor as ``workspace_bytes``. + total_workspace_bytes = api.scratch_workspace_bytes() + + binding = ga.SdpaBinding( + q=facts.q_t, + k=facts.k_t, + v=facts.v_t, + o=facts.o_t, + stats=facts.stats_t, + do=facts.do_t, + dq=facts.dq_t, + dk=facts.dk_t, + dv=facts.dv_t, + ) + + def _canonical_view(buf: torch.Tensor, geom) -> torch.Tensor: + """Reinterpret a variant-pack buffer through the canonical geometry. + + cuDNN's execute contract treats variant-pack entries as raw storage + laid out per the IR tensor descriptor — callers may hand in a torch + tensor whose logical shape is anything with the right bytes. The DSL + executor consumes torch views, so rebuild the canonical view here. + No-op when the caller already passed a matching view. + """ + dim, stride = geom + 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) + api.execute( + q_tensor=_canonical_view(resolved[id(binding.q)], q_geom), + k_tensor=_canonical_view(resolved[id(binding.k)], kv_geom), + v_tensor=_canonical_view(resolved[id(binding.v)], kv_geom), + o_tensor=_canonical_view(resolved[id(binding.o)], q_geom), + do_tensor=_canonical_view(resolved[id(binding.do)], q_geom), + stats_tensor=_canonical_view(resolved[id(binding.stats)], stats_geom), + dq_tensor=_canonical_view(resolved[id(binding.dq)], q_geom), + dk_tensor=_canonical_view(resolved[id(binding.dk)], kv_geom), + dv_tensor=_canonical_view(resolved[id(binding.dv)], kv_geom), + scale_softmax=facts.scale, + # Scratch comes from the CALLER's workspace (never allocated + # here): the dispatch sized/validated it against workspace_bytes; + # the adapter's carver re-validates so a direct call cannot + # silently corrupt memory. + workspace=workspace, + # Stream from the execute-time context (raw CUstream int, + # engine plan passes ctx.stream); None keeps the adapter's + # torch-current-stream fallback. + current_stream=_cuda_driver.CUstream(stream) if stream is not None else None, + ) + return None + + # Executor contract (engine._FrostSdpaBwdPlan): a non-zero workspace_bytes + # means the plan calls _execute(variant_pack, workspace) with the caller's + # buffer; 0 means _execute(variant_pack) and the buffer is never touched. + # ``binding`` lets the plan key this executor's operands out of the graph's + # variant pack (the pack covers every IO tensor of the graph). + _execute.workspace_bytes = total_workspace_bytes + _execute.binding = binding + return _execute + + +def engine_name(arch: str = "sm120") -> str: + """The shipped engine name for a coverage cell (test/user convenience).""" + + return f"sdpa_bwd_{arch}" + + +# Preference order: the ranked plan list offers these in this order (see +# cudnn/sdpa/bwd/engine.py, which wraps each spec as a BaseEngine). +ENGINE_SPECS = (_sm120_spec(),) + +__all__ = ["ENGINE_SPECS", "Capabilities", "EngineSpec", "SdpaBwdKnobs", "analyze_for", "engine_name", "mismatch"] diff --git a/python/cudnn/sdpa/bwd/kernels/__init__.py b/python/cudnn/sdpa/bwd/kernels/__init__.py new file mode 100644 index 000000000..ce2b556ee --- /dev/null +++ b/python/cudnn/sdpa/bwd/kernels/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""DSL SDPA-backward kernel templates. + +Filenames encode the coverage matrix: ``bprop__sm.py``. + +Every template specializes on its architecture's frozen ``TemplateParams`` at +import time (module global ``FROST_TEMPLATE_PARAMS``, injected by +``cudnn.frost.template_loader.load_template``). Tensor geometry remains an +input to each module's cached ``compile()`` function. Import a template +directly only for its all-defaults standalone path. +""" diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py new file mode 100644 index 000000000..c3850b210 --- /dev/null +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -0,0 +1,1263 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""FROST SM120 SDPA backward kernel template (fp16 / bf16). + +A fused multi-head attention (FMHA) backward for the NVIDIA Blackwell +GeForce SM120 family (SM120 and SM121) using TMA loads and a +warp-specialized producer/consumer schedule. + +The backward follows the FlashAttention-2 algorithm (Dao, 2023; +https://github.com/Dao-AILab/flash-attention, BSD-3-Clause): a +KV-stationary, seq-KV-parallel single pass in which each CTA owns one KV +tile and walks the query tiles in descending order, computing the five +chained GEMMs (S = Q*K^T, dP = dO*V^T, dV += P^T*dO, dQ = dS*K, +dK += dS^T*Q) with the softmax VJP fused in registers. dK/dV accumulate +in registers across the whole pass (no atomics); dQ is reduced through an +fp32 workspace and finalized by a small convert kernel. + +Constraints: +* Supported input dtypes: Float16 and BFloat16 (output dtype matches) +* Head dimension must be one of 32, 64, or 128 +* Equal Q/KV head counts (no GQA); no dropout/alibi/local/softcap +* Q/K/V/O/dO/dQ/dK/dV use compact BSHD storage +* LSE input is the natural-log forward stats, fp32 (B, H, SQ) contiguous +* dQ accumulation uses fp32 atomics (not bitwise deterministic) + +One backward call is three kernel launches through the per-shape +``compile()`` cache at the bottom of this module: ``dot`` (delta = +rowsum(dO*O), also zeroes the dq_accum workspace), ``main`` (the fused +five-GEMM pass writing dK/dV), and ``cvt`` (dq_accum fp32 -> dQ io dtype). +""" + +from functools import lru_cache +from types import SimpleNamespace +from typing import Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.utils +import cutlass.experimental.cuda as cuda +import cutlass.cute as cute +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream +from cutlass.experimental import primitives as prims +from cutlass._mlir.dialects import arith + +from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 +from cudnn.frost.tile_dsl.mma import ptx_mma_m16n8k16_f32 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor +from cudnn.sdpa.bwd.config_sm120 import TemplateParams, validate_params + +# The FROST loader injects one immutable specialization before executing this +# module. A direct import uses the dense FP16 defaults. +PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams()) +validate_params(PARAMS) + +STORAGE_DTYPE = {DTYPE_FP16: cutlass.Float16, DTYPE_BF16: cutlass.BFloat16}[PARAMS.dtype_qkv] + +_LOG2E = 1.4426950408889634 +_COPY_ELEMS = 8 # 16-byte gmem<->smem chunk (8 fp16/bf16) + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def largest_warp_partition(m_dim: int, n_dim: int) -> int: + """Largest valid 2-D warp-partition factor A for one GEMM's (M, N) pair. + + The warp grid is (A, 8 // A): the A warps split m_dim into 16-row MMA + (m16n8k16) blocks, so m_dim must be a multiple of 16 * A, and the + per-warp N slice (n_dim * A / 8) must be a multiple of 16 (ldmatrix.x4 + pairs). Used when a macro-tile override deviates from the per-head-dim + CONFIG default, whose hand-tuned partitions only validate for the + default tiles; "largest valid" is a heuristic, not a sweep winner. + """ + + for a in (8, 4, 2, 1): + if m_dim % (16 * a) == 0 and (n_dim * a // 8) % 16 == 0: + return a + raise ValueError(f"no valid warp partition for M{m_dim} N{n_dim}") + + +@cute.jit +def tile_ptr( + sbuf, + row: cutlass.Int32, + col: cutlass.Int32, + *, + page: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], +): + """Element pointer into a paged+swizzled smem tile.""" + pg = col // page + in_col = col % page + off = pg * (rows * page) + row * page + swizzle_xor(row, in_col, page, 2) + return sbuf.subview(off).data_ptr() + + +@cute.jit +def pack_half2(lo, hi, dtype: cutlass.Constexpr[Type[cutlass.Numeric]]): + """Pack two fp32 into a 2-element io-dtype vector (one 4 B store).""" + return cutlass.Vector.from_elements((lo.to(dtype), hi.to(dtype)), dtype) + + +@cute.jit +def _red_add_f32x2(ptr, v0: cutlass.Float32, v1: cutlass.Float32) -> None: + """One red.global.add.v2.f32 covering a thread's adjacent (c0,c1) pair.""" + prims.inline_ptx( + "red.global.add.v2.f32 [$0], {$1, $2};", + read_only_args=[ptr, v0, v1], + ) + + +@cute.jit +def load_a_frag( + sbuf, + kc: cutlass.Constexpr[int], + row0, + lane, + *, + rows: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], +): + """ldmatrix.x4 one (16 x 16) row-major A fragment.""" + row = row0 + lane % 16 + col = kc * 16 + (lane // 16) * 8 + return prims.ldmatrix(tile_ptr(sbuf, row, col, page=page, rows=rows), 4, prims.MMALayout.ROW) + + +@cute.jit +def load_a_frag_transposed( + sbuf, + kc: cutlass.Constexpr[int], + col0, + lane, + *, + rows: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], +): + """ldmatrix.trans.x4: A[M, K] from a tile stored physically [K, M].""" + row = kc * 16 + lane % 16 + col = col0 + (lane // 16) * 8 + return prims.ldmatrix(tile_ptr(sbuf, row, col, page=page, rows=rows), 4, prims.MMALayout.COL) + + +@cute.jit +def copy16_smem_to_gmem(sptr, gptr): + """One 16-byte smem->gmem chunk.""" + v = sptr.load(count=8) + gptr.store(v, alignment=16) + + +@cute.jit +def mma_bstream( + acc, + a_frag, + sB, + *, + b_k_step: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + b_trans: cutlass.Constexpr[bool], + b_rows: cutlass.Constexpr[int], + b_page: cutlass.Constexpr[int], + lane, + ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + col_base=0, + row_base=0, +): + """One k=16 step of an (M x N) MMA, B streamed from smem via + ldmatrix.x4 (2 adjacent 8-column n-frags per fetch). + + acc: ``(M//16) * (N//8) * 4`` fp32, m-rep major then n-frag. + a_frag: ``(M//16) * 4`` Int32 (one 16x16 A fragment per m-rep). + """ + M_BLOCKS = M // 16 + N_FRAGS = N // 8 + PAIRS = N_FRAGS // 2 + a_stride = len(a_frag) // M_BLOCKS + + if cutlass.const_expr(b_trans): + b_row = lane % 16 + n_offset = lane // 16 + layout_flag = prims.MMALayout.COL + else: + b_row = lane % 8 + b_col_subchunk = (lane // 8) % 2 + n_offset = lane // 16 + layout_flag = prims.MMALayout.ROW + + for pair in cutlass.range_constexpr(PAIRS): + n_frag = pair * 2 + if cutlass.const_expr(b_trans): + row = b_k_step * 16 + b_row + col = (n_frag + n_offset) * 8 + col_base + else: + row = row_base + (n_frag + n_offset) * 8 + b_row + col = b_k_step * 16 + b_col_subchunk * 8 + col_base + b_ptr = tile_ptr(sB, row, col, page=b_page, rows=b_rows) + b_v = prims.ldmatrix(b_ptr, 4, layout_flag) + for m_block in cutlass.range_constexpr(M_BLOCKS): + a_off = m_block * a_stride + for half in cutlass.range_constexpr(2): + s = (m_block * N_FRAGS + n_frag + half) * 4 + c0, c1, c2, c3 = ptx_mma_m16n8k16_f32( + a_frag[a_off + 0], + a_frag[a_off + 1], + a_frag[a_off + 2], + a_frag[a_off + 3], + b_v[half * 2 + 0], + b_v[half * 2 + 1], + acc[s + 0], + acc[s + 1], + acc[s + 2], + acc[s + 3], + ab_dtype, + ) + acc[s + 0] = c0 + acc[s + 1] = c1 + acc[s + 2] = c2 + acc[s + 3] = c3 + + +@cute.jit +def mma_abregs( + acc, + a_frag, + b_frag, + *, + b_k_step: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """One k=16 MMA step with both operands resident in registers.""" + M_BLOCKS = M // 16 + N_FRAGS = N // 8 + PAIRS = N_FRAGS // 2 + a_stride = len(a_frag) // M_BLOCKS + b_k_stride = PAIRS * 4 + + for pair in cutlass.range_constexpr(PAIRS): + n_frag = pair * 2 + b_off = b_k_step * b_k_stride + pair * 4 + for m_block in cutlass.range_constexpr(M_BLOCKS): + a_off = m_block * a_stride + for half in cutlass.range_constexpr(2): + s = (m_block * N_FRAGS + n_frag + half) * 4 + c0, c1, c2, c3 = ptx_mma_m16n8k16_f32( + a_frag[a_off + 0], + a_frag[a_off + 1], + a_frag[a_off + 2], + a_frag[a_off + 3], + b_frag[b_off + half * 2 + 0], + b_frag[b_off + half * 2 + 1], + acc[s + 0], + acc[s + 1], + acc[s + 2], + acc[s + 3], + ab_dtype, + ) + acc[s + 0] = c0 + acc[s + 1] = c1 + acc[s + 2] = c2 + acc[s + 3] = c3 + + +# --------------------------------------------------------------------------- +# Main kernel. +# --------------------------------------------------------------------------- + + +class SM120FusedMultiHeadAttentionFP16Backward: + """Configure and launch the SM120 FMHA backward kernel chain.""" + + DEFAULT_TILES = { + 32: (128, 64), + 64: (64, 128), + 128: (64, 64), + } + # (d, q_tile, kv_tile) -> (warps_m_sdp, warps_m_dkv, warps_m_dq): for each + # GEMM the 8 compute warps form an (A, 8 // A) grid; the value is A, the + # warp count along that GEMM's own M (row) axis. + CONFIG = { + (32, 128, 64): (4, 4, 8), # default for d32 + (32, 128, 128): (4, 8, 4), # for very long S + (64, 64, 128): (4, 8, 4), # default for d64 + (64, 128, 64): (8, 2, 4), # For underfilled grids, kv64 can double CTA counts + (128, 64, 64): (2, 1, 4), # default for d128 + } + + def __init__( + self, + in_dtype: Type[cutlass.Numeric] = cutlass.Float16, + is_causal: bool = False, + head_dim: int = 128, + use_pdl: bool = True, + q_tile: int = 0, + kv_tile: int = 0, + ): + self.in_dtype = in_dtype + self.is_causal = is_causal + self.d = head_dim + self.use_pdl = bool(use_pdl) + self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim] + if q_tile: + self.q_tile = int(q_tile) + if kv_tile: + self.kv_tile = int(kv_tile) + if 128 % self.q_tile: + # dq_accum is scrambled in 128-row blocks (SQ_R is rounded to 128 + # and the convert kernel unscrambles per block). + raise ValueError(f"q_tile must divide 128; got {self.q_tile}") + # Warp layouts: the sweep-tuned triple for this exact tile choice + # when we have one, else the largest-valid derivation + tuned = self.CONFIG.get((head_dim, self.q_tile, self.kv_tile)) + if tuned is not None: + self.warps_m_sdp, self.warps_m_dkv, self.warps_m_dq = tuned + else: + self.warps_m_sdp = largest_warp_partition(self.q_tile, self.kv_tile) + self.warps_m_dkv = largest_warp_partition(self.kv_tile, head_dim) + self.warps_m_dq = largest_warp_partition(self.q_tile, head_dim) + M_, N_, d_ = self.q_tile, self.kv_tile, head_dim + for a_, m_dim, n_dim, tag in ( + (self.warps_m_sdp, M_, N_, "warps_m_sdp"), + (self.warps_m_dkv, N_, d_, "warps_m_dkv"), + (self.warps_m_dq, M_, d_, "warps_m_dq"), + ): + if 8 % a_ or m_dim % (16 * a_) or (n_dim * a_ // 8) % 16: + raise ValueError(f"invalid {tag}={a_} for M{M_} N{N_} d{d_}") + self.page = 64 if head_dim % 64 == 0 else 32 + self.threads = 384 + self.num_consumer_warps = 8 + self.load_warp_id = self.num_consumer_warps + self.tma_copy_iters = head_dim // self.page + self.tma_swizzle = cuda.TensorMapSwizzle.s128b if self.page == 64 else cuda.TensorMapSwizzle.s64b + + M, N, d = self.q_tile, self.kv_tile, self.d + # smem element offsets + self.off_sQ = 0 # 2 buffers + self.off_sdO = 2 * M * d + self.off_sK = 3 * M * d + self.off_sV = self.off_sK + N * d + self.off_sdS = self.off_sV # aliases sV (V is in regs) + self.off_sP = self.off_sV + M * N + self.smem_elems = self.off_sV + max(N * d, 2 * M * N) + smem_bytes = self.smem_elems * in_dtype.bytes + cap = cutlass.utils.get_smem_capacity_in_bytes("sm_120") + if smem_bytes > cap: + raise ValueError(f"smem {smem_bytes} B exceeds sm_120 cap {cap}") + + @cute.jit + def load_tma_tile(self, s_dst, tma_desc, mbar, batch, head, seq, rows: cutlass.Constexpr[int]): + """Load one paged/swizzled `(rows, d)` tile with TMA.""" + elems_per_page = rows * self.page + for pg in cutlass.range_constexpr(self.tma_copy_iters): + if prims.elect_sync(): + prims.cp_async_bulk_tensor_shared_cta_global( + s_dst.subview(pg * elems_per_page), + tma_desc.get_ptr(), + (pg * self.page, head, seq, batch), + mbar, + ) + + @cute.kernel + def kernel( + self, + q: cute.Tensor, # [B, SQ, H, D] io dtype (BSHD) + k: cute.Tensor, # [B, SKV, H, D] + v: cute.Tensor, # [B, SKV, H, D] + do: cute.Tensor, # [B, SQ, H, D] + lse: cute.Tensor, # [B, H, SQ] fp32 (natural-log LSE) + delta: cute.Tensor, # [B, H, SQ_r128] fp32 (dot_do_o output) + dq_accum: cute.Tensor, # [B*SQ_r128*H*D] fp32 (scrambled, zeroed) + dk: cute.Tensor, # [B, SKV, H, D] output + dv: cute.Tensor, # [B, SKV, H, D] output + tma_q_desc: cutlass.GridConstant[cuda.TensorMap], + tma_k_desc: cutlass.GridConstant[cuda.TensorMap], + tma_v_desc: cutlass.GridConstant[cuda.TensorMap], + tma_do_desc: cutlass.GridConstant[cuda.TensorMap], + softmax_scale_log2: cutlass.Float32, # scale * log2(e) + attn_scale: cutlass.Float32, # linear scale (dq/dk output) + ) -> None: + io_dtype = self.in_dtype + d = self.d + M = self.q_tile + N = self.kv_tile + PAGE = self.page + PDS = 64 if self.kv_tile >= 64 else self.kv_tile + WM_SDP = self.warps_m_sdp # S/dP warp grid (WM_SDP, 8//WM_SDP), WM_SDP along q rows + WM_DKV = self.warps_m_dkv # dK/dV warp grid (WM_DKV, 8//WM_DKV), WM_DKV along kv rows + WM_DQ = self.warps_m_dq # dQ warp grid (WM_DQ, 8//WM_DQ), WM_DQ along q rows + + # SdP: warp (wm, wn); wm steps over SDP_REPS 16-row MMA blocks + # interleaved by 16*WM_SDP; wn covers SDP_NPER contiguous kv columns. + SDP_REPS = M // (16 * WM_SDP) + SDP_NPER = N * WM_SDP // 8 + SDP_NF = SDP_NPER // 8 + # dKV: warp (wn2, wd); DKV_REPS 16-row MMA blocks interleaved by 16*WM_DKV. + DKV_REPS = N // (16 * WM_DKV) + DKV_PER = d * WM_DKV // 8 + DKV_NF = DKV_PER // 8 + # dQ: warp (wq, wdq). + DQ_REPS = M // (16 * WM_DQ) + DQ_PER = d * WM_DQ // 8 + DQ_NF = DQ_PER // 8 + + D_CHUNKS = d // 16 # SdP k-reduce + Q_CHUNKS = M // 16 # dK/dV k-reduce + KV_CHUNKS = N // 16 # dQ k-reduce + VREG_PAIRS = SDP_NPER // 16 # V-in-regs frag pairs / chunk + + tidx, _, _ = cute.arch.thread_idx() + n_block, head, batch = cute.arch.block_idx() + lane = tidx % 32 + warp = cute.arch.warp_idx() + g_lane = lane // 4 + p_lane = lane % 4 + + SQ = q.shape[1] + SKV = k.shape[1] + H = q.shape[2] + SQ_R = ((SQ + 127) // 128) * 128 + kv_base = n_block * N + row_stride = H * d # BSHD gmem row stride + + lse_ptr = lse.iterator.raw_ptr() + dd_ptr = delta.iterator.raw_ptr() + dqa_ptr = dq_accum.iterator.raw_ptr() + dk_ptr = dk.iterator.raw_ptr() + dv_ptr = dv.iterator.raw_ptr() + + PARTIAL_Q = (SQ % M) != 0 + PARTIAL_KV = (SKV % N) != 0 + + m_block_max = (SQ + M - 1) // M + if cutlass.const_expr(self.is_causal): + m_block_min = cute.math.max(kv_base + SQ - SKV, cutlass.Int32(0)) // M + else: + m_block_min = cutlass.Int32(0) + n_iters = m_block_max - m_block_min + + smem = cutlass.Array(io_dtype, self.smem_elems, space=cutlass.AddressSpace.smem, alignment=128) + sQ = smem # 2 * M * d (double buffer) + sdO = smem.subview(self.off_sdO) # M * d + sK = smem.subview(self.off_sK) # N * d + sV = smem.subview(self.off_sV) # N * d + sdS = smem.subview(self.off_sdS) # M * N (aliases sV) + sP = smem.subview(self.off_sP) # M * N + tma_mbar = cutlass.Array(cutlass.Int64, 5, space=cutlass.AddressSpace.smem, alignment=8) + k_mbar = tma_mbar + v_mbar = tma_mbar.subview(1) + q_full = tma_mbar.subview(2) + do_full = tma_mbar.subview(4) + + if warp == self.load_warp_id: + if prims.elect_sync(): + prims.prefetch_tensormap(tma_k_desc.get_ptr()) + prims.prefetch_tensormap(tma_v_desc.get_ptr()) + prims.prefetch_tensormap(tma_q_desc.get_ptr()) + prims.prefetch_tensormap(tma_do_desc.get_ptr()) + prims.mbarrier_init(k_mbar, 1) + prims.mbarrier_init(v_mbar, 1) + prims.mbarrier_init(q_full, 1) + prims.mbarrier_init(q_full.subview(1), 1) + prims.mbarrier_init(do_full, 1) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + # gmem tile bases for this (batch, head, n_block / m_block). + qhd_base = (batch * SQ) * row_stride + head * d + khd_base = (batch * SKV + kv_base) * row_stride + head * d + lse_base = (batch * H + head) * SQ + dd_base = (batch * H + head) * SQ_R + + if warp == self.load_warp_id: + prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(v_mbar, N * d * io_dtype.bytes) + prims.mbarrier_arrive_expect_tx(k_mbar, N * d * io_dtype.bytes) + self.load_tma_tile(sV, tma_v_desc, v_mbar, batch, head, kv_base, rows=N) + self.load_tma_tile(sK, tma_k_desc, k_mbar, batch, head, kv_base, rows=N) + + if n_iters > 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(q_full, M * d * io_dtype.bytes) + self.load_tma_tile(sQ, tma_q_desc, q_full, batch, head, (m_block_max - 1) * M, rows=M) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(do_full, M * d * io_dtype.bytes) + self.load_tma_tile( + sdO, + tma_do_desc, + do_full, + batch, + head, + (m_block_max - 1) * M, + rows=M, + ) + while not prims.mbarrier_try_wait_parity(v_mbar, cutlass.Int32(0)): + pass + while not prims.mbarrier_try_wait_parity(k_mbar, cutlass.Int32(0)): + pass + for load_j in cutlass.range(n_iters, unroll=1): + load_stage = load_j & cutlass.Int32(1) + q_phase_p = (load_j // 2) & cutlass.Int32(1) + while not prims.mbarrier_try_wait_parity(q_full.subview(load_stage), q_phase_p): + pass + do_phase_p = load_j & cutlass.Int32(1) + while not prims.mbarrier_try_wait_parity(do_full, do_phase_p): + pass + # Loop-top: stage load_j ready AND load_j-1 consumed. + cute.arch.barrier(barrier_id=3, number_of_threads=288) + next_m = m_block_max - 2 - load_j + if load_j + 1 < n_iters: + next_stage = (load_j + 1) & cutlass.Int32(1) + next_q_full = q_full.subview(next_stage) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(next_q_full, M * d * io_dtype.bytes) + self.load_tma_tile( + sQ.subview(next_stage * M * d), + tma_q_desc, + next_q_full, + batch, + head, + next_m * M, + rows=M, + ) + # Post-GEMM3: every consumer is done with sdO. + cute.arch.barrier(barrier_id=4, number_of_threads=288) + if load_j + 1 < n_iters: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(do_full, M * d * io_dtype.bytes) + self.load_tma_tile(sdO, tma_do_desc, do_full, batch, head, next_m * M, rows=M) + + elif warp < self.load_warp_id: + prims.setmaxregister(240, prims.SetMaxRegisterAction.INCREASE) + # LSE for the first (highest) m-block: per-thread direct loads at + # this thread's C-fragment rows + math_warp = warp + math_tidx = tidx + m_block = m_block_max - 1 + wm_s = math_warp % WM_SDP + wn_s = math_warp // WM_SDP + lse_r = cutlass.Array(cutlass.Float32, 2 * SDP_REPS) + dd_r = cutlass.Array(cutlass.Float32, 2 * SDP_REPS) + for rep in cutlass.range_constexpr(SDP_REPS): + for hf in cutlass.range_constexpr(2): + r_loc = wm_s * 16 + rep * 16 * WM_SDP + g_lane + hf * 8 + r_abs = m_block * M + r_loc + if cutlass.const_expr(PARTIAL_Q): + r_cl = cute.math.min(r_abs, SQ - 1) + val = (lse_ptr + lse_base + r_cl).load() + inf = cutlass.Float32(float("inf")) + # branchless (r_abs < SQ) ? 1 : 0 via arith.select + ok32 = cutlass.Int32( + arith.select( + (r_abs < SQ).ir_value(), + cutlass.Int32(1).ir_value(), + cutlass.Int32(0).ir_value(), + ) + ) + if ok32 == 0: + val = inf + lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) + else: + val = (lse_ptr + lse_base + r_abs).load() + lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) + + while not prims.mbarrier_try_wait_parity(v_mbar, cutlass.Int32(0)): + pass + while not prims.mbarrier_try_wait_parity(k_mbar, cutlass.Int32(0)): + pass + + # V -> registers. + v_persist = cutlass.Array(cutlass.Int32, D_CHUNKS * VREG_PAIRS * 4, alignment=16) + for kc in cutlass.range_constexpr(D_CHUNKS): + for pair in cutlass.range_constexpr(VREG_PAIRS): + n_frag = pair * 2 + row = wn_s * SDP_NPER + (n_frag + lane // 16) * 8 + lane % 8 + col = kc * 16 + ((lane // 8) % 2) * 8 + vf = prims.ldmatrix( + tile_ptr(sV, row, col, page=PAGE, rows=N), + 4, + prims.MMALayout.ROW, + ) + v_off = (kc * VREG_PAIRS + pair) * 4 + v_persist[v_off + 0] = vf[0] + v_persist[v_off + 1] = vf[1] + v_persist[v_off + 2] = vf[2] + v_persist[v_off + 3] = vf[3] + cute.arch.barrier(barrier_id=1, number_of_threads=256) + + # dK/dV accumulators. + wn_k = math_warp % WM_DKV + wd_k = math_warp // WM_DKV + acc_dk = cutlass.Array(cutlass.Float32, DKV_REPS * DKV_NF * 4, alignment=16) + acc_dv = cutlass.Array(cutlass.Float32, DKV_REPS * DKV_NF * 4, alignment=16) + for i in cutlass.range_constexpr(DKV_REPS * DKV_NF * 4): + acc_dk[i] = cutlass.Float32(0.0) + acc_dv[i] = cutlass.Float32(0.0) + + wq = math_warp % WM_DQ + wd_q = math_warp // WM_DQ + + acc_s = cutlass.Array(cutlass.Float32, SDP_REPS * SDP_NF * 4, alignment=16) + acc_dp = cutlass.Array(cutlass.Float32, SDP_REPS * SDP_NF * 4, alignment=16) + acc_dq = cutlass.Array(cutlass.Float32, DQ_REPS * DQ_NF * 4, alignment=16) + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + # ---- main loop: m_block descending -------------------------------- + j = cutlass.Int32(0) + while j < n_iters: + m_block = m_block_max - 1 - j + stage = j & cutlass.Int32(1) + sQ_st = sQ.subview(stage * M * d) + q_row0 = m_block * M + + cute.arch.barrier(barrier_id=3, number_of_threads=288) + + # dP_sum per-thread loads (delta buffer is 128-rounded). + for rep in cutlass.range_constexpr(SDP_REPS): + for hf in cutlass.range_constexpr(2): + r_loc = wm_s * 16 + rep * 16 * WM_SDP + g_lane + hf * 8 + dd_r[rep * 2 + hf] = (dd_ptr + dd_base + q_row0 + r_loc).load() + + # GEMM 1: acc_s = Q @ K^T. + for i in cutlass.range_constexpr(SDP_REPS * SDP_NF * 4): + acc_s[i] = cutlass.Float32(0.0) + for kc in cutlass.range_constexpr(D_CHUNKS): + af = [] + for rep in cutlass.range_constexpr(SDP_REPS): + qf = load_a_frag( + sQ_st, + kc, + wm_s * 16 + rep * 16 * WM_SDP, + lane, + rows=M, + page=PAGE, + ) + af = af + [qf[0], qf[1], qf[2], qf[3]] + mma_bstream( + acc_s, + af, + sK, + b_k_step=kc, + M=16 * SDP_REPS, + N=SDP_NPER, + b_trans=False, + b_rows=N, + b_page=PAGE, + lane=lane, + ab_dtype=io_dtype, + row_base=wn_s * SDP_NPER, + ) + + # Mask + softmax (scores -> P, unscaled by attn_scale) and the + # P store to smem. + do_mask_causal = (m_block * M) < (kv_base + N + SQ - SKV) + neg_inf = cutlass.Float32(float("-inf")) + for rep in cutlass.range_constexpr(SDP_REPS): + for nf in cutlass.range_constexpr(SDP_NF): + off = (rep * SDP_NF + nf) * 4 + kv_c0 = wn_s * SDP_NPER + nf * 8 + 2 * p_lane + kv_a0 = kv_base + kv_c0 + kv_a1 = kv_a0 + 1 + r0 = q_row0 + wm_s * 16 + rep * 16 * WM_SDP + g_lane + r8 = r0 + 8 + s0 = acc_s[off + 0] + s1 = acc_s[off + 1] + s2 = acc_s[off + 2] + s3 = acc_s[off + 3] + if cutlass.const_expr(self.is_causal): + if do_mask_causal: + if kv_a0 > r0 + SKV - SQ: + s0 = neg_inf + if kv_a1 > r0 + SKV - SQ: + s1 = neg_inf + if kv_a0 > r8 + SKV - SQ: + s2 = neg_inf + if kv_a1 > r8 + SKV - SQ: + s3 = neg_inf + if cutlass.const_expr(PARTIAL_KV): + if kv_a0 >= SKV: + s0 = neg_inf + s2 = neg_inf + if kv_a1 >= SKV: + s1 = neg_inf + s3 = neg_inf + lse0 = lse_r[rep * 2 + 0] + lse8 = lse_r[rep * 2 + 1] + p0 = cute.math.exp2(s0 * softmax_scale_log2 - lse0, fastmath=True) + p1 = cute.math.exp2(s1 * softmax_scale_log2 - lse0, fastmath=True) + p2 = cute.math.exp2(s2 * softmax_scale_log2 - lse8, fastmath=True) + p3 = cute.math.exp2(s3 * softmax_scale_log2 - lse8, fastmath=True) + acc_s[off + 0] = p0 + acc_s[off + 1] = p1 + acc_s[off + 2] = p2 + acc_s[off + 3] = p3 + # sP store: each (2p, 2p+1) pair packed into one 4 B + # store to the swizzled tile + pr0 = wm_s * 16 + rep * 16 * WM_SDP + g_lane + pr8 = pr0 + 8 + sw0 = tile_ptr(sP, pr0, kv_c0, page=PDS, rows=M) + sw8 = tile_ptr(sP, pr8, kv_c0, page=PDS, rows=M) + sw0.store(pack_half2(p0, p1, io_dtype), alignment=4) + sw8.store(pack_half2(p2, p3, io_dtype), alignment=4) + + # GEMM 2: acc_dp = dO @ V^T (V in registers). + for i in cutlass.range_constexpr(SDP_REPS * SDP_NF * 4): + acc_dp[i] = cutlass.Float32(0.0) + for kc in cutlass.range_constexpr(D_CHUNKS): + af = [] + for rep in cutlass.range_constexpr(SDP_REPS): + dof = load_a_frag( + sdO, + kc, + wm_s * 16 + rep * 16 * WM_SDP, + lane, + rows=M, + page=PAGE, + ) + af = af + [dof[0], dof[1], dof[2], dof[3]] + mma_abregs( + acc_dp, + af, + v_persist, + b_k_step=kc, + M=16 * SDP_REPS, + N=SDP_NPER, + ab_dtype=io_dtype, + ) + + # dS = P * (dP - dP_sum) + for rep in cutlass.range_constexpr(SDP_REPS): + for nf in cutlass.range_constexpr(SDP_NF): + off = (rep * SDP_NF + nf) * 4 + dd0 = dd_r[rep * 2 + 0] + dd8 = dd_r[rep * 2 + 1] + acc_dp[off + 0] = acc_s[off + 0] * (acc_dp[off + 0] - dd0) + acc_dp[off + 1] = acc_s[off + 1] * (acc_dp[off + 1] - dd0) + acc_dp[off + 2] = acc_s[off + 2] * (acc_dp[off + 2] - dd8) + acc_dp[off + 3] = acc_s[off + 3] * (acc_dp[off + 3] - dd8) + + # dS -> fp16 -> sdS. + for rep in cutlass.range_constexpr(SDP_REPS): + for nf in cutlass.range_constexpr(SDP_NF): + off = (rep * SDP_NF + nf) * 4 + kv_c0 = wn_s * SDP_NPER + nf * 8 + 2 * p_lane + pr0 = wm_s * 16 + rep * 16 * WM_SDP + g_lane + pr8 = pr0 + 8 + sw0 = tile_ptr(sdS, pr0, kv_c0, page=PDS, rows=M) + sw8 = tile_ptr(sdS, pr8, kv_c0, page=PDS, rows=M) + sw0.store( + pack_half2(acc_dp[off + 0], acc_dp[off + 1], io_dtype), + alignment=4, + ) + sw8.store( + pack_half2(acc_dp[off + 2], acc_dp[off + 3], io_dtype), + alignment=4, + ) + cute.arch.barrier(barrier_id=1, number_of_threads=256) + + # GEMM 3: acc_dv += P^T @ dO. + for kc in cutlass.range_constexpr(Q_CHUNKS): + af = [] + for rep in cutlass.range_constexpr(DKV_REPS): + pf = load_a_frag_transposed( + sP, + kc, + wn_k * 16 + rep * 16 * WM_DKV, + lane, + rows=M, + page=PDS, + ) + af = af + [pf[0], pf[2], pf[1], pf[3]] + mma_bstream( + acc_dv, + af, + sdO, + b_k_step=kc, + M=16 * DKV_REPS, + N=DKV_PER, + b_trans=True, + b_rows=M, + b_page=PAGE, + lane=lane, + ab_dtype=io_dtype, + col_base=wd_k * DKV_PER, + ) + + # GEMM3 is the final dO consumer; this rendezvous lets + # the producer refill the single dO buffer. + cute.arch.barrier(barrier_id=4, number_of_threads=288) + + # GEMM 4: acc_dq = dS @ K^T. + for i in cutlass.range_constexpr(DQ_REPS * DQ_NF * 4): + acc_dq[i] = cutlass.Float32(0.0) + for kc in cutlass.range_constexpr(KV_CHUNKS): + af = [] + for rep in cutlass.range_constexpr(DQ_REPS): + sf = load_a_frag(sdS, kc, wq * 16 + rep * 16 * WM_DQ, lane, rows=M, page=PDS) + af = af + [sf[0], sf[1], sf[2], sf[3]] + mma_bstream( + acc_dq, + af, + sK, + b_k_step=kc, + M=16 * DQ_REPS, + N=DQ_PER, + b_trans=True, + b_rows=N, + b_page=PAGE, + lane=lane, + ab_dtype=io_dtype, + col_base=wd_q * DQ_PER, + ) + + # Reload LSE for the next (lower) m-block. + if j + 1 < n_iters: + nq0 = (m_block - 1) * M + for rep in cutlass.range_constexpr(SDP_REPS): + for hf in cutlass.range_constexpr(2): + r_loc = wm_s * 16 + rep * 16 * WM_SDP + g_lane + hf * 8 + lse_r[rep * 2 + hf] = (lse_ptr + lse_base + nq0 + r_loc).load() * cutlass.Float32(_LOG2E) + + # dQ accumulate into the scrambled dq_accum. + t_r = math_tidx // 32 + t_c = math_tidx % 32 + dqa_base = ((batch * SQ_R + q_row0) * H + head) * d + for rep in cutlass.range_constexpr(DQ_REPS): + for nf in cutlass.range_constexpr(DQ_NF): + for hv in cutlass.range_constexpr(2): + i_pair = hv + rep * 2 + nf * 2 * DQ_REPS + if cutlass.const_expr(d >= 64): + jm = i_pair % (M // 8) + jn = i_pair // (M // 8) + addr = dqa_base + (t_r + jm * 8) * (H * d) + t_c * 2 + jn * 64 + else: + addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d) + (t_c % 16) * 2 + poff = (rep * DQ_NF + nf) * 4 + hv * 2 + _red_add_f32x2(dqa_ptr + addr, acc_dq[poff + 0], acc_dq[poff + 1]) + + # GEMM 5: acc_dk += dS^T @ Q. + for kc in cutlass.range_constexpr(Q_CHUNKS): + af = [] + for rep in cutlass.range_constexpr(DKV_REPS): + sf = load_a_frag_transposed( + sdS, + kc, + wn_k * 16 + rep * 16 * WM_DKV, + lane, + rows=M, + page=PDS, + ) + af = af + [sf[0], sf[2], sf[1], sf[3]] + mma_bstream( + acc_dk, + af, + sQ_st, + b_k_step=kc, + M=16 * DKV_REPS, + N=DKV_PER, + b_trans=True, + b_rows=M, + b_page=PAGE, + lane=lane, + ab_dtype=io_dtype, + col_base=wd_k * DKV_PER, + ) + + j += 1 + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + # ---- epilogue: dK/dV through smem (sdK aliases sK, sdV aliases sV) -------- + cute.arch.barrier(barrier_id=2, number_of_threads=256) + sdK = sK + sdV = sV + for rep in cutlass.range_constexpr(DKV_REPS): + for nf in cutlass.range_constexpr(DKV_NF): + off = (rep * DKV_NF + nf) * 4 + r0 = wn_k * 16 + rep * 16 * WM_DKV + g_lane + r8 = r0 + 8 + c0 = wd_k * DKV_PER + nf * 8 + 2 * p_lane + dk0 = acc_dk[off + 0] * attn_scale + dk1 = acc_dk[off + 1] * attn_scale + dk2 = acc_dk[off + 2] * attn_scale + dk3 = acc_dk[off + 3] * attn_scale + tile_ptr(sdK, r0, c0, page=PAGE, rows=N).store(pack_half2(dk0, dk1, io_dtype), alignment=4) + tile_ptr(sdK, r8, c0, page=PAGE, rows=N).store(pack_half2(dk2, dk3, io_dtype), alignment=4) + tile_ptr(sdV, r0, c0, page=PAGE, rows=N).store( + pack_half2(acc_dv[off + 0], acc_dv[off + 1], io_dtype), + alignment=4, + ) + tile_ptr(sdV, r8, c0, page=PAGE, rows=N).store( + pack_half2(acc_dv[off + 2], acc_dv[off + 3], io_dtype), + alignment=4, + ) + cute.arch.barrier(barrier_id=2, number_of_threads=256) + + # smem -> gmem + chunks_per_row = d // _COPY_ELEMS + total = N * chunks_per_row + for i in cutlass.range_constexpr(total // 256): + chunk = i * 256 + math_tidx + row = chunk // chunks_per_row + col = (chunk % chunks_per_row) * _COPY_ELEMS + if (not cutlass.const_expr(PARTIAL_KV)) or (kv_base + row < SKV): + g_off = khd_base + row * row_stride + col + copy16_smem_to_gmem(tile_ptr(sdK, row, col, page=PAGE, rows=N), dk_ptr + g_off) + copy16_smem_to_gmem(tile_ptr(sdV, row, col, page=PAGE, rows=N), dv_ptr + g_off) + + else: + prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + do: cute.Tensor, + lse: cute.Tensor, + delta: cute.Tensor, + dq_accum: cute.Tensor, + dk: cute.Tensor, + dv: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + attn_scale: cutlass.Float32, + stream: cuda_driver.CUstream, + ) -> None: + box_kv = (1, self.kv_tile, 1, self.page) + box_q = (1, self.q_tile, 1, self.page) + tma_q_desc = cuda.create_tensor_map_tiled_from_view(q, box_dims=box_q, stride_order=(3, 2, 1, 0), swizzle=self.tma_swizzle) + tma_k_desc = cuda.create_tensor_map_tiled_from_view(k, box_dims=box_kv, stride_order=(3, 2, 1, 0), swizzle=self.tma_swizzle) + tma_v_desc = cuda.create_tensor_map_tiled_from_view(v, box_dims=box_kv, stride_order=(3, 2, 1, 0), swizzle=self.tma_swizzle) + tma_do_desc = cuda.create_tensor_map_tiled_from_view(do, box_dims=box_q, stride_order=(3, 2, 1, 0), swizzle=self.tma_swizzle) + n_blocks = cute.ceil_div(k.shape[1], self.kv_tile) + self.kernel( + q, + k, + v, + do, + lse, + delta, + dq_accum, + dk, + dv, + tma_q_desc, + tma_k_desc, + tma_v_desc, + tma_do_desc, + softmax_scale_log2, + attn_scale, + ).launch( + grid=(n_blocks, q.shape[2], q.shape[0]), + block=(self.threads, 1, 1), + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.use_pdl, + ) + + +# --------------------------------------------------------------------------- +# Preprocess kernel: delta = rowsum(dO * O) + dq_accum zeroing +# --------------------------------------------------------------------------- + + +@cute.kernel +def _dot_do_o_kernel( + o: cute.Tensor, # [B, SQ, H, D] + do: cute.Tensor, # [B, SQ, H, D] + delta: cute.Tensor, # [B, H, SQ_r128] fp32 out + dq_accum: cute.Tensor, # [B*SQ_r128*H*D] fp32 (zeroed here) + q_tile: cutlass.Constexpr[int], + d: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], +): + if cutlass.const_expr(use_pdl): + cute.arch.griddepcontrol_launch_dependents() + m_block, head, batch = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + SQ = o.shape[1] + H = o.shape[2] + SQ_R = ((SQ + 127) // 128) * 128 + row_stride = H * d + M = q_tile + + o_ptr = o.iterator.raw_ptr() + do_ptr = do.iterator.raw_ptr() + dd_ptr = delta.iterator.raw_ptr() + dqa_ptr = dq_accum.iterator.raw_ptr() + + base = ((batch * SQ + m_block * M) * H + head) * d + dd_base = (batch * H + head) * SQ_R + m_block * M + q_left = SQ - m_block * M + + tpr = page // _COPY_ELEMS + rows_per_pass = 256 // tpr + col0 = (tidx % tpr) * _COPY_ELEMS + row0 = tidx // tpr + n_pages = d // page + for rp in cutlass.range_constexpr(M // rows_per_pass): + row = row0 + rp * rows_per_pass + acc = cutlass.Float32(0.0) + if row < q_left: + g_off = base + row * row_stride + col0 + for pg in cutlass.range_constexpr(n_pages): + ov = (o_ptr + g_off + pg * page).load(count=_COPY_ELEMS) + dov = (do_ptr + g_off + pg * page).load(count=_COPY_ELEMS) + for kk in cutlass.range_constexpr(_COPY_ELEMS): + acc = acc + ov[kk].to(cutlass.Float32) * dov[kk].to(cutlass.Float32) + # Allreduce over the tpr threads sharing the row (lane-contiguous). + n_sh = 3 if cutlass.const_expr(tpr == 8) else 2 + for sh in cutlass.range_constexpr(n_sh): + acc = acc + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=acc, + offset=1 << (n_sh - 1 - sh), + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + if tidx % tpr == 0: + (dd_ptr + dd_base + row).store(acc) + + if cutlass.const_expr(use_pdl): + cute.arch.griddepcontrol_wait() + + zrows = 32 if cutlass.const_expr(page == 32) else 16 + ztpr = 256 // zrows + zr0 = tidx // ztpr + zc0 = (tidx % ztpr) * 4 + zero4 = cutlass.Vector.from_elements( + ( + cutlass.Float32(0.0), + cutlass.Float32(0.0), + cutlass.Float32(0.0), + cutlass.Float32(0.0), + ), + cutlass.Float32, + ) + dqa_base = ((batch * SQ_R + m_block * M) * H + head) * d + for im in cutlass.range_constexpr(M // zrows): + for jn in cutlass.range_constexpr(d // (ztpr * 4)): + addr = dqa_base + (zr0 + im * zrows) * (H * d) + zc0 + jn * ztpr * 4 + (dqa_ptr + addr).store(zero4, alignment=16) + + +@cute.jit +def _dot_do_o_host( + o: cute.Tensor, + do: cute.Tensor, + delta: cute.Tensor, + dq_accum: cute.Tensor, + q_tile: cutlass.Constexpr[int], + d: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +): + m_blocks = cute.ceil_div(o.shape[1], q_tile) + _dot_do_o_kernel(o, do, delta, dq_accum, q_tile, d, page, use_pdl).launch( + grid=(m_blocks, o.shape[2], o.shape[0]), + block=(256, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + +# --------------------------------------------------------------------------- +# Convert kernel: scrambled dq_accum (fp32) -> dQ (io dtype) +# --------------------------------------------------------------------------- + + +@cute.kernel +def _convert_dq_kernel( + dq_accum: cute.Tensor, # [B*SQ_r128*H*D] fp32 + dq: cute.Tensor, # [B, SQ, H, D] io dtype out + q_tile: cutlass.Constexpr[int], + d: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], + warps_m_dq: cutlass.Constexpr[int], + attn_scale: cutlass.Float32, + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + use_pdl: cutlass.Constexpr[bool], +): + if cutlass.const_expr(use_pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + m_block, head, batch = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % 32 + warp = cute.arch.warp_idx() + g_lane = lane // 4 + p_lane = lane % 4 + SQ = dq.shape[1] + H = dq.shape[2] + SQ_R = ((SQ + 127) // 128) * 128 + M = q_tile + WM_DQ = warps_m_dq + DQ_REPS = M // (16 * WM_DQ) + DQ_PER = d * WM_DQ // 8 + DQ_NF = DQ_PER // 8 + wq = warp % WM_DQ + wd_q = warp // WM_DQ + + dqa_ptr = dq_accum.iterator.raw_ptr() + dq_ptr = dq.iterator.raw_ptr() + + sdQ = cutlass.Array(io_dtype, M * d, space=cutlass.AddressSpace.smem, alignment=128) + + t_r = tidx // 32 + t_c = tidx % 32 + dqa_base = ((batch * SQ_R + m_block * M) * H + head) * d + for rep in cutlass.range_constexpr(DQ_REPS): + for nf in cutlass.range_constexpr(DQ_NF): + frag = cutlass.Array(cutlass.Float32, 4) + for hv in cutlass.range_constexpr(2): + i_pair = hv + rep * 2 + nf * 2 * DQ_REPS + if cutlass.const_expr(d >= 64): + jm = i_pair % (M // 8) + jn = i_pair // (M // 8) + addr = dqa_base + (t_r + jm * 8) * (H * d) + t_c * 2 + jn * 64 + else: + addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d) + (t_c % 16) * 2 + pv = (dqa_ptr + addr).load(count=2) + frag[hv * 2 + 0] = pv[0] * attn_scale + frag[hv * 2 + 1] = pv[1] * attn_scale + r0 = wq * 16 + rep * 16 * WM_DQ + g_lane + r8 = r0 + 8 + c0 = wd_q * DQ_PER + nf * 8 + 2 * p_lane + tile_ptr(sdQ, r0, c0, page=page, rows=M).store(pack_half2(frag[0], frag[1], io_dtype), alignment=4) + tile_ptr(sdQ, r8, c0, page=page, rows=M).store(pack_half2(frag[2], frag[3], io_dtype), alignment=4) + prims.barrier_cta_sync(0) + + q_left = SQ - m_block * M + row_stride = H * d + g_base = ((batch * SQ + m_block * M) * H + head) * d + chunks_per_row = d // _COPY_ELEMS + for i in cutlass.range_constexpr(M * chunks_per_row // 256): + chunk = i * 256 + tidx + row = chunk // chunks_per_row + col = (chunk % chunks_per_row) * _COPY_ELEMS + if row < q_left: + copy16_smem_to_gmem( + tile_ptr(sdQ, row, col, page=page, rows=M), + dq_ptr + g_base + row * row_stride + col, + ) + + +@cute.jit +def _convert_dq_host( + dq_accum: cute.Tensor, + dq: cute.Tensor, + q_tile: cutlass.Constexpr[int], + d: cutlass.Constexpr[int], + page: cutlass.Constexpr[int], + warps_m_dq: cutlass.Constexpr[int], + attn_scale: cutlass.Float32, + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +): + m_blocks = cute.ceil_div(dq.shape[1], q_tile) + _convert_dq_kernel(dq_accum, dq, q_tile, d, page, warps_m_dq, attn_scale, io_dtype, use_pdl).launch( + grid=(m_blocks, dq.shape[2], dq.shape[0]), + block=(256, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + +@lru_cache(maxsize=None) +def compile( # noqa: A001 + compute_capability: tuple[int, int], + b: int = 1, + qh: int = 1, + sq: int = 128, + skv: int = 128, + d: int = 128, +) -> SimpleNamespace: + """Compile and cache the three-kernel backward chain for one compact BSHD shape.""" + + bwd = SM120FusedMultiHeadAttentionFP16Backward( + in_dtype=STORAGE_DTYPE, + is_causal=PARAMS.is_causal, + head_dim=d, + use_pdl=PARAMS.use_pdl, + q_tile=PARAMS.q_tile, + kv_tile=PARAMS.kv_tile, + ) + sq_r = ceil_div(sq, 128) * 128 + + def _fake(dtype, shape): + return make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(range(len(shape) - 1, -1, -1)), + assumed_align=16, + ) + + fake_q = _fake(STORAGE_DTYPE, (b, sq, qh, d)) + fake_k = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_v = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_o = _fake(STORAGE_DTYPE, (b, sq, qh, d)) + fake_do = _fake(STORAGE_DTYPE, (b, sq, qh, d)) + fake_dq = _fake(STORAGE_DTYPE, (b, sq, qh, d)) + fake_dk = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_dv = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_lse = _fake(cutlass.Float32, (b, qh, sq)) + fake_delta = _fake(cutlass.Float32, (b, qh, sq_r)) + fake_dq_accum = _fake(cutlass.Float32, (b * sq_r * qh * d,)) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + options = "--enable-tvm-ffi" + + compiled_dot = cute.compile( + _dot_do_o_host, + fake_o, + fake_do, + fake_delta, + fake_dq_accum, + bwd.q_tile, + d, + bwd.page, + bwd.use_pdl, + fake_stream, + options=options, + ) + compiled_main = cute.compile( + bwd, + fake_q, + fake_k, + fake_v, + fake_do, + fake_lse, + fake_delta, + fake_dq_accum, + fake_dk, + fake_dv, + cutlass.Float32(1.0), + cutlass.Float32(1.0), + fake_stream, + options=options, + ) + compiled_cvt = cute.compile( + _convert_dq_host, + fake_dq_accum, + fake_dq, + bwd.q_tile, + d, + bwd.page, + bwd.warps_m_dq, + cutlass.Float32(1.0), + STORAGE_DTYPE, + bwd.use_pdl, + fake_stream, + options=options, + ) + return SimpleNamespace(dot=compiled_dot, main=compiled_main, cvt=compiled_cvt) diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index c0cdbcf35..6673b8ef4 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -179,6 +179,8 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti """ if facts.invalid: return facts.invalid + if facts.is_backward: + return "this engine serves sdpa() forward graphs only" if knobs is not None: if not isinstance(knobs, SdpaFwdKnobs): return f"knob request is a {type(knobs).__name__}, not SdpaFwdKnobs — wrong operation's vocabulary" diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 3c1664065..31cdd259a 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -49,6 +49,8 @@ from cutlass.experimental import primitives as prims from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 +from cudnn.frost.tile_dsl.mma import ptx_mma_m16n8k16_f32 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor from cudnn.sdpa.fwd.config_sm120 import ( SEQ_KV_TILES as _SEQ_KV_TILES, SEQ_Q_TILES as _SEQ_Q_TILES, @@ -115,69 +117,6 @@ def nvvm_threadquad_reduction_sum(val: cutlass.Float32) -> cutlass.Float32: return val -@cute.jit -def ptx_mma_m16n8k16_f32( - a0: cutlass.Int32, - a1: cutlass.Int32, - a2: cutlass.Int32, - a3: cutlass.Int32, - b0: cutlass.Int32, - b1: cutlass.Int32, - c0: cutlass.Float32, - c1: cutlass.Float32, - c2: cutlass.Float32, - c3: cutlass.Float32, - ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], -) -> tuple[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32]: - """``mma.sync.aligned.m16n8k16.row.col.f32.{f16|bf16}.{f16|bf16}.f32``.""" - if cutlass.const_expr(ab_dtype != cutlass.Float16 and ab_dtype != cutlass.BFloat16): - raise TypeError(f"Invalid A/B dtype: {ab_dtype}") - ab_tag = "f16" if cutlass.const_expr(ab_dtype == cutlass.Float16) else "bf16" - return cute.arch.inline_ptx( - f"mma.sync.aligned.m16n8k16.row.col.f32.{ab_tag}.{ab_tag}.f32 {{$0,$1,$2,$3}}, {{$4,$5,$6,$7}}, {{$8,$9}}, {{$10,$11,$12,$13}};", - write_only_types=[ - cutlass.Float32, - cutlass.Float32, - cutlass.Float32, - cutlass.Float32, - ], - read_only_args=[a0, a1, a2, a3, b0, b1, c0, c1, c2, c3], - ) - - -@cute.jit -def get_swizzled_col( - row: cutlass.Int32, - col: cutlass.Int32, - row_stride: cutlass.Constexpr[int], - elem_bytes: cutlass.Constexpr[int], -) -> cutlass.Int32: - """Return the physical SMEM column for an XOR-swizzled row-major tile. - - The XOR is applied at the 16-byte boundary for all element widths. - ``elem_bytes`` selects the element-domain shift and swizzle chunk size. - """ - row_stride_bytes = row_stride * elem_bytes - chunk_bytes = 32 - sw_bits = 1 - row_shift = 2 - if row_stride_bytes % 128 == 0: - chunk_bytes = 128 - sw_bits = 3 - row_shift = 0 - elif row_stride_bytes % 64 == 0: - chunk_bytes = 64 - sw_bits = 2 - row_shift = 1 - chunk_size = chunk_bytes // elem_bytes - elems_per_16b = 16 // elem_bytes - sw_base = elems_per_16b.bit_length() - 1 - chunk = col // chunk_size - col_in_chunk = col % chunk_size - bit_msk = (1 << sw_bits) - 1 - return chunk * chunk_size + (col_in_chunk ^ (((row >> row_shift) & bit_msk) << sw_base)) - - @cute.jit def pack_to_i32( src: tuple, @@ -457,7 +396,7 @@ def load_k_frag_pair(k_frag_pair: cutlass.Constexpr[int], d_frag: cutlass.Conste k_smem_ptr = ( mma_params.sK.data_ptr() + k_physical_row * self.tma_swizzle_chunk_elems - + get_swizzled_col( + + swizzle_xor( k_physical_row, k_col_in_chunk, self.tma_swizzle_chunk_elems, @@ -671,7 +610,7 @@ def load_v_frag_pair(v_frag: cutlass.Constexpr[int], d_frag_pair: cutlass.Conste sV_ptr = ( mma_params.sV.data_ptr() + v_physical_row * self.tma_swizzle_chunk_elems - + get_swizzled_col( + + swizzle_xor( v_physical_row, v_col_in_chunk, self.tma_swizzle_chunk_elems, diff --git a/python/cudnn/sdpa/graph_analyzer.py b/python/cudnn/sdpa/graph_analyzer.py index d81547011..da7ae5a8f 100644 --- a/python/cudnn/sdpa/graph_analyzer.py +++ b/python/cudnn/sdpa/graph_analyzer.py @@ -156,6 +156,11 @@ class SdpaGraphFacts: has_unfuse_fma: bool = False has_block_mask: bool = False has_rng_dump: bool = False + is_backward: bool = False # sdpa_backward() node (NodeType.SDPA_BWD) + right_bound: Optional[int] = None # raw resolved right band (0 == causal) + deterministic: bool = False # sdpa_backward(use_deterministic_algorithm=True) + has_dbias: bool = False # dBias output requested (backward) + has_dsink: bool = False # dSink_token output requested (backward) has_score_max: bool = False # per-row/tile score-max side output requested has_score_sum_exp: bool = False # per-row/tile sum-of-exp side output requested dynamic_scale: bool = False # attn_scale passed as a tensor @@ -179,6 +184,13 @@ class SdpaGraphFacts: sink_t: Any = None seq_kv_t: Any = None seq_q_t: Any = None + # backward-only refs + do_t: Any = None + dq_t: Any = None + dk_t: Any = None + dv_t: Any = None + dbias_t: Any = None + dsink_t: Any = None # MXFP8 block-scale (descale) tensors + Amax_O output. sf_q_t: Any = None sf_k_t: Any = None @@ -193,7 +205,8 @@ class SdpaGraphFacts: def _single_sdpa_node(graph: "cudnn.pygraph") -> Optional[Any]: - """The graph's sole SDPA-forward node, or None if the graph is anything else.""" + """The graph's sole SDPA node (forward, backward, or an FP8/MXFP8 flavor), + or None if the graph is anything else.""" try: nodes = graph.nodes except Exception: # noqa: BLE001 — non-IR graph objects @@ -201,7 +214,7 @@ def _single_sdpa_node(graph: "cudnn.pygraph") -> Optional[Any]: if len(nodes) != 1: return None node = nodes[0] - if node.node_type not in (cudnn.NodeType.SDPA, cudnn.NodeType.SDPA_MXFP8, cudnn.NodeType.SDPA_FP8): + if node.node_type not in (cudnn.NodeType.SDPA, cudnn.NodeType.SDPA_BWD, cudnn.NodeType.SDPA_MXFP8, cudnn.NodeType.SDPA_FP8): return None return node @@ -216,8 +229,18 @@ def _record_from_node(node: Any) -> dict: rec: dict = dict(node.params) for port, t in node.inputs.items(): rec.setdefault(port, t) - rec["o"] = node.outputs.get("O") - rec["stats"] = node.outputs.get("Stats") + # Forward: O / Stats are node outputs. Backward: o / stats are INPUT + # ports (already folded above) — don't clobber them with the absent + # forward output ports. + if node.outputs.get("O") is not None: + rec["o"] = node.outputs.get("O") + if node.outputs.get("Stats") is not None: + rec["stats"] = node.outputs.get("Stats") + rec["_is_backward"] = node.node_type == cudnn.NodeType.SDPA_BWD + if rec["_is_backward"]: + for port in ("dQ", "dK", "dV", "dBias", "dSink_token"): + if rec.get(port) is None: + rec[port] = node.outputs.get(port) # Output-style kwargs (passed as sdpa() arguments but recorded in # node.outputs): fold each one in so engines see every requested output. # Missing one here lets an engine that never writes it pass the probe and @@ -250,18 +273,55 @@ def _first_not_none(*vals): def _extract_facts(rec: dict) -> SdpaGraphFacts: + is_backward = bool(rec.get("_is_backward")) q, k, v, o = rec.get("q"), rec.get("k"), rec.get("v"), rec.get("o") if q is None or k is None or v is None or o is None: - return _invalid("missing q/k/v/o on the sdpa() node") - - q_dim, q_stride = tuple(q.get_dim()), tuple(q.get_stride()) - k_dim, k_stride = tuple(k.get_dim()), tuple(k.get_stride()) - v_dim, v_stride = tuple(v.get_dim()), tuple(v.get_stride()) - o_dim, o_stride = tuple(o.get_dim()), tuple(o.get_stride()) - if len({len(q_dim), len(k_dim), len(v_dim), len(o_dim), 4}) != 1: - return _invalid("Q/K/V/O must all be rank-4 (B, H, S, D)") + return _invalid("missing q/k/v/o on the sdpa node") + + rank4_ports = [("q", q), ("k", k), ("v", v), ("o", o)] + if is_backward: + for name in ("dO", "dQ", "dK", "dV"): + t = rec.get(name) + if t is None: + return _invalid(f"missing {name} on the sdpa_backward node") + rank4_ports.append((name, t)) + + dims = {} + strides = {} + for name, t in rank4_ports: + d = tuple(t.get_dim()) + if len(d) != 4: + return _invalid(f"{name} must be rank-4 (B, H, S, D); got rank {len(d)}") + dims[name] = d + strides[name] = tuple(t.get_stride()) + q_dim, q_stride = dims["q"], strides["q"] + k_dim, k_stride = dims["k"], strides["k"] + v_dim, v_stride = dims["v"], strides["v"] + o_dim, o_stride = dims["o"], strides["o"] b, h_q, s_q, d_qk = q_dim + + # ``build_operation_graph`` rewrites the BACKWARD node's K / V ports to + # transposed (B, H, D, S) views (the K^T / V^T the lowering consumes); + # the underlying buffer keeps the user's (B, H, S, D) shape. Canonicalize + # so probing works both before and after the native build. + if is_backward: + + def _square_transposed(dim: tuple, stride: tuple) -> bool: + """Square (S == D) rewritten views are extent-ambiguous; the stride + order disambiguates: the transposed view keeps the buffer's unit + stride, which lands on axis 2 instead of axis 3.""" + return dim[2] == dim[3] != 1 and stride[2] == 1 and stride[3] != 1 + + if (k_dim[3] != d_qk and k_dim[2] == d_qk) or _square_transposed(k_dim, k_stride): + k_dim = (k_dim[0], k_dim[1], k_dim[3], k_dim[2]) + k_stride = (k_stride[0], k_stride[1], k_stride[3], k_stride[2]) + _, h_kv, s_kv, _ = k_dim + if (v_dim[2] != s_kv and v_dim[3] == s_kv) or _square_transposed(v_dim, v_stride): + v_dim = (v_dim[0], v_dim[1], v_dim[3], v_dim[2]) + v_stride = (v_stride[0], v_stride[1], v_stride[3], v_stride[2]) + dims["k"], dims["v"] = k_dim, v_dim + strides["k"], strides["v"] = k_stride, v_stride _, h_kv, s_kv, _ = k_dim d_v = v_dim[-1] if k_dim != (b, h_kv, s_kv, d_qk): @@ -270,6 +330,11 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: return _invalid(f"V shape mismatch (k_dim={k_dim}, v_dim={v_dim})") if o_dim != (b, h_q, s_q, d_v): return _invalid(f"O shape mismatch (q_dim={q_dim}, o_dim={o_dim})") + if is_backward: + if dims["dO"] != (b, h_q, s_q, d_v): + return _invalid("dO shape mismatch") + if dims["dQ"] != dims["q"] or dims["dK"] != dims["k"] or dims["dV"] != dims["v"]: + return _invalid("dQ/dK/dV must match Q/K/V shapes") if any(x <= 0 for x in (b, h_q, h_kv, s_q, s_kv, d_qk, d_v)): return _invalid("B/H/S/D must all be > 0") if h_q % h_kv != 0: @@ -284,9 +349,13 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: # FP8 in: O dtype is independent of the input; only K/V must match Q. uniform = all(_DTYPE_FROM_CUDNN.get(t.get_data_type()) == q_dtype for t in (k, v)) else: - uniform = all(_DTYPE_FROM_CUDNN.get(t.get_data_type()) == q_dtype for t in (k, v, o)) - bshd = all(bshd_layout_ok(d, s) for d, s in ((q_dim, q_stride), (k_dim, k_stride), (v_dim, v_stride), (o_dim, o_stride))) - dense_layout = all(dense_layout_ok(d, s) for d, s in ((q_dim, q_stride), (k_dim, k_stride), (v_dim, v_stride), (o_dim, o_stride))) + _uniform_ports = [k, v, o] + ([rec["dO"], rec["dQ"], rec["dK"], rec["dV"]] if is_backward else []) + uniform = all(_DTYPE_FROM_CUDNN.get(t.get_data_type()) == q_dtype for t in _uniform_ports) + _layout_ports = [(q_dim, q_stride), (k_dim, k_stride), (v_dim, v_stride), (o_dim, o_stride)] + if is_backward: + _layout_ports += [(dims[name], strides[name]) for name in ("dO", "dQ", "dK", "dV")] + bshd = all(bshd_layout_ok(d, s) for d, s in _layout_ports) + dense_layout = all(dense_layout_ok(d, s) for d, s in _layout_ports) # descale_q/k/v are the block-scale SF tensors for MXFP8, or scalar per-tensor # descales for FP8. Both arrive on the same-named node.inputs ports. @@ -356,17 +425,22 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: if sink_dim != (1, h_q, 1, 1): return _invalid(f"sink_token must be (1, H_q, 1, 1); got {sink_dim}") - generate_stats = rec.get("generate_stats") - is_inference = rec.get("is_inference") - if generate_stats is not None: - wants_stats = bool(generate_stats) - elif is_inference is not None: - wants_stats = not bool(is_inference) - else: - wants_stats = False stats = rec.get("stats") - if wants_stats and stats is None: - return _invalid("generate_stats=True but no Stats tensor was returned") + if is_backward: + wants_stats = False + if stats is None: + return _invalid("sdpa_backward requires the forward stats tensor") + else: + generate_stats = rec.get("generate_stats") + is_inference = rec.get("is_inference") + if generate_stats is not None: + wants_stats = bool(generate_stats) + elif is_inference is not None: + wants_stats = not bool(is_inference) + else: + wants_stats = False + if wants_stats and stats is None: + return _invalid("generate_stats=True but no Stats tensor was returned") attn_scale = rec.get("attn_scale") dynamic_scale = attn_scale is not None and not isinstance(attn_scale, (int, float)) @@ -391,6 +465,11 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: bottom_right=bool(align_is_br), window_left=window_left, right_band_widening=right_widening, + is_backward=is_backward, + right_bound=resolved_right, + deterministic=bool(rec.get("use_deterministic_algorithm", False)), + has_dbias=rec.get("dBias") is not None, + has_dsink=rec.get("dSink_token") is not None, has_bias=rec.get("bias") is not None, has_dropout=rec.get("dropout") is not None, has_score_mod=rec.get("fn") is not None, @@ -414,7 +493,13 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: k_t=k, v_t=v, o_t=o, - stats_t=(stats if wants_stats else None), + stats_t=(stats if (wants_stats or is_backward) else None), + do_t=rec.get("dO"), + dq_t=rec.get("dQ"), + dk_t=rec.get("dK"), + dv_t=rec.get("dV"), + dbias_t=rec.get("dBias"), + dsink_t=rec.get("dSink_token"), sink_t=sink_token, seq_kv_t=seq_len_kv, seq_q_t=seq_len_q, @@ -483,6 +568,17 @@ class SdpaBinding: descale_v: Any = None scale_o: Any = None amax_s: Any = None + # SM80 feature operands + backward ports. + bias: Any = None + block_mask: Any = None + score_max: Any = None + score_sum_exp: Any = None + do: Any = None + dq: Any = None + dk: Any = None + dv: Any = None + dbias: Any = None + dsink: Any = None def bound_tensors(self) -> list: return [ @@ -505,6 +601,16 @@ def bound_tensors(self) -> list: self.descale_v, self.scale_o, self.amax_s, + self.bias, + self.block_mask, + self.score_max, + self.score_sum_exp, + self.do, + self.dq, + self.dk, + self.dv, + self.dbias, + self.dsink, ) if t is not None ] diff --git a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py new file mode 100644 index 000000000..1dec74d80 --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""End-to-end tests for the FROST SM120 DSL SDPA-backward engine against a torch reference.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from test_utils import torch_fork_set_rng + +ENGINE = "sdpa_bwd_sm120" + + +def _is_sm120() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) + return (major, minor) in {(12, 0), (12, 1)} + + +def _dsl_deps_available() -> bool: + try: + import cutlass # noqa: F401 + except ImportError: + return False + return True + + +pytestmark = pytest.mark.skipif( + not _is_sm120(), + reason="SM120 DSL SDPA backward engine requires an SM120 or SM121 device.", +) + + +@pytest.fixture(autouse=True) +def _enable_frost(monkeypatch): + """FROST engines resolve only under the env opt-in (read live per call).""" + + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + + +def _require_dsl() -> None: + try: + import cudnn # noqa: F401 + import cudnn.sdpa # noqa: F401 + except ImportError as exc: + pytest.skip(f"SM120 DSL engine not available: {exc}") + if not _dsl_deps_available(): + pytest.skip("cutlass/dsl not installed") + + +def _select_engine(graph, name): + """Pin the ranked entry named ``name`` (graph.plans holds the backend's + plans and the python engines' in one list). A pin is strict: check_support / + build_plans raise if that engine declines the graph.""" + names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] + assert name in names, f"engine {name!r} did not claim this graph; plans={names}" + graph.select_plan(names.index(name)) + return graph + + +def _bhsd(batch: int, heads: int, sequence: int, head_dim: int, dtype: torch.dtype, empty: bool = False) -> torch.Tensor: + """Return logical BHSD backed by compact BSHD physical storage.""" + + factory = torch.empty if empty else torch.randn + return factory(batch, sequence, heads, head_dim, dtype=dtype, device="cuda").transpose(1, 2) + + +def _ref_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + *, + scale: float, + is_causal: bool = False, + causal_bottom_right: bool = False, +): + """FP32 autograd reference: returns (o, stats, dq, dk, dv). + + ``stats`` is the natural-log LSE reshaped to the graph's (B, H, S_q, 1) + layout — exactly what the forward pass would have produced. + """ + + s_q, s_kv = q.shape[2], k.shape[2] + q32 = q.detach().float().requires_grad_() + k32 = k.detach().float().requires_grad_() + v32 = v.detach().float().requires_grad_() + scores = torch.matmul(q32, k32.transpose(-1, -2)) * scale + if is_causal: + diagonal = s_kv - s_q if causal_bottom_right else 0 + mask = torch.ones(s_q, s_kv, device="cuda", dtype=torch.bool).tril(diagonal=diagonal) + scores = scores.masked_fill(~mask, float("-inf")) + stats = torch.logsumexp(scores, dim=-1, keepdim=True).contiguous() # (B, H, S_q, 1) + p32 = torch.softmax(scores, dim=-1) + o32 = torch.matmul(p32, v32) + o32.backward(do.float()) + o = o32.detach().to(q.dtype) + return o, stats, q32.grad.to(q.dtype), k32.grad.to(q.dtype), v32.grad.to(q.dtype) + + +def _expected_workspace_bytes(batch: int, heads: int, s_q: int, head_dim: int) -> int: + from cudnn.sdpa.fwd.api_dsl import ws_align + + sq_r = -(-s_q // 128) * 128 + return ws_align(batch * heads * sq_r * 4) + ws_align(batch * sq_r * heads * head_dim * 4) + + +def _run_bwd_graph( + q_gpu: torch.Tensor, + k_gpu: torch.Tensor, + v_gpu: torch.Tensor, + o_gpu: torch.Tensor, + do_gpu: torch.Tensor, + stats_gpu: torch.Tensor, + *, + scale: float, + is_causal: bool = False, + causal_bottom_right: bool = False, + select: bool = True, + q_tile: int | None = None, + kv_tile: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]: + """Build and execute the SM120 FROST backward graph; returns (dq, dk, dv, plan_name).""" + + _require_dsl() + import cudnn + + dtype = q_gpu.dtype + io_dtype = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16 + batch, h_q, _, head_dim = q_gpu.shape + _, h_kv, _, _ = k_gpu.shape + dq_gpu = _bhsd(batch, h_q, q_gpu.shape[2], head_dim, dtype, empty=True) + dk_gpu = _bhsd(batch, h_kv, k_gpu.shape[2], head_dim, dtype, empty=True) + dv_gpu = _bhsd(batch, h_kv, v_gpu.shape[2], head_dim, dtype, empty=True) + + graph = cudnn.pygraph( + io_data_type=io_dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + q = graph.tensor_like(q_gpu, name="q") + k = graph.tensor_like(k_gpu, name="k") + v = graph.tensor_like(v_gpu, name="v") + o = graph.tensor_like(o_gpu, name="o") + do = graph.tensor_like(do_gpu, name="dO") + stats = graph.tensor_like(stats_gpu, name="stats") + + bwd_kwargs = { + "name": "sdpa_backward", + "q": q, + "k": k, + "v": v, + "o": o, + "dO": do, + "stats": stats, + "attn_scale": scale, + } + if causal_bottom_right: + bwd_kwargs["use_causal_mask_bottom_right"] = True + elif is_causal: + bwd_kwargs["use_causal_mask"] = True + + dq, dk, dv = graph.sdpa_backward(**bwd_kwargs) + dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride()) + dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride()) + dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride()) + + graph.validate() + graph.build_operation_graph() + if q_tile is not None or kv_tile is not None: + # A knob request rides on a plan entry (PlanConfig.knobs): append + # exactly one (engine_id, knobs) plan — the deterministic-replay path. + from cudnn.engines.engine_ids import FROST_SDPA_BWD_ID_BASE + from cudnn.sdpa.bwd.engines import SdpaBwdKnobs + + graph.create_execution_plan(FROST_SDPA_BWD_ID_BASE + 0, SdpaBwdKnobs(tile_m=q_tile, tile_n=kv_tile)) + graph.select_plan(0) + else: + graph.create_execution_plans([cudnn.heur_mode.A]) + if select: + _select_engine(graph, ENGINE) + graph.check_support() + graph.build_plans() + # What actually runs, not what merely ranked first: build_plans settles the + # plan index on the entry that built. + engine = graph.selected_engine + plan_name = engine.name if engine is not None else "backend" + if select or q_tile is not None or kv_tile is not None: + assert plan_name == ENGINE, f"pinned {ENGINE} but {plan_name} would run" + + workspace_size = graph.get_workspace_size() + if plan_name == ENGINE: + assert workspace_size == _expected_workspace_bytes(batch, h_q, q_gpu.shape[2], head_dim) + workspace = torch.empty(max(workspace_size, 1), dtype=torch.uint8, device="cuda") + + variant_pack = { + q: q_gpu, + k: k_gpu, + v: v_gpu, + o: o_gpu, + do: do_gpu, + stats: stats_gpu, + dq: dq_gpu, + dk: dk_gpu, + dv: dv_gpu, + } + graph.execute(variant_pack, workspace) + torch.cuda.synchronize() + return dq_gpu, dk_gpu, dv_gpu, plan_name + + +def _tolerances(dtype: torch.dtype) -> dict: + return {"atol": 2e-2 if dtype == torch.float16 else 5e-2, "rtol": 5e-2} + + +def _run_case( + *, + batch: int = 2, + heads: int = 4, + s_q: int = 512, + s_kv: int = 512, + head_dim: int = 64, + dtype: torch.dtype = torch.float16, + is_causal: bool = False, + causal_bottom_right: bool = False, + select: bool = True, + q_tile: int | None = None, + kv_tile: int | None = None, +) -> str: + scale = 1.0 / math.sqrt(head_dim) + q = _bhsd(batch, heads, s_q, head_dim, dtype) + k = _bhsd(batch, heads, s_kv, head_dim, dtype) + v = _bhsd(batch, heads, s_kv, head_dim, dtype) + do = _bhsd(batch, heads, s_q, head_dim, dtype) + o, stats, dq_ref, dk_ref, dv_ref = _ref_bwd(q, k, v, do, scale=scale, is_causal=is_causal, causal_bottom_right=causal_bottom_right) + o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + dq, dk, dv, plan_name = _run_bwd_graph( + q, + k, + v, + o, + do, + stats, + scale=scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + select=select, + q_tile=q_tile, + kv_tile=kv_tile, + ) + tol = _tolerances(dtype) + torch.testing.assert_close(dq.float(), dq_ref.float(), **tol) + torch.testing.assert_close(dk.float(), dk_ref.float(), **tol) + torch.testing.assert_close(dv.float(), dv_ref.float(), **tol) + return plan_name + + +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [32, 64, 128]) +@pytest.mark.parametrize("is_causal", [False, True], ids=["dense", "causal"]) +@torch_fork_set_rng(seed=0) +def test_sdpa_bwd_dsl_sm120_graph_api(head_dim: int, is_causal: bool): + """FP16 numeric parity per head dim, dense and top-left causal (S_q == S_kv).""" + + _run_case(head_dim=head_dim, is_causal=is_causal) + + +@pytest.mark.L0 +@pytest.mark.parametrize("is_causal", [False, True], ids=["dense", "causal"]) +@torch_fork_set_rng(seed=1) +def test_sdpa_bwd_dsl_sm120_bf16(is_causal: bool): + """BF16 numeric parity at d=64.""" + + _run_case(dtype=torch.bfloat16, head_dim=64, is_causal=is_causal) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=2) +def test_sdpa_bwd_dsl_sm120_cross_seqlen_causal_br(): + """Bottom-right causal with S_q < S_kv (the decode-style tail).""" + + _run_case(s_q=384, s_kv=1024, head_dim=64, is_causal=True, causal_bottom_right=True) + + +@pytest.mark.L0 +@pytest.mark.parametrize("is_causal", [False, True], ids=["dense", "causal_br"]) +@torch_fork_set_rng(seed=4) +def test_sdpa_bwd_dsl_sm120_sequence_tails(is_causal: bool): + """Non-tile-multiple sequence tails exercise the partial-Q/KV predicates.""" + + _run_case( + s_q=193, + s_kv=257, + head_dim=128, + is_causal=is_causal, + causal_bottom_right=is_causal, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=5) +def test_sdpa_bwd_dsl_sm120_auto_routing(): + """Without an explicit select, the eligible graph auto-routes to the engine.""" + + plan_name = _run_case(head_dim=64, is_causal=True, select=False) + assert plan_name == ENGINE + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("head_dim", "q_tile", "kv_tile"), + [ + (64, 128, 64), # sweep-tuned non-default entry (CONFIG hit) + (64, 64, 64), # not in CONFIG (largest_warp_partition fallback) + ], +) +@torch_fork_set_rng(seed=7) +def test_sdpa_bwd_dsl_sm120_tile_knobs(head_dim: int, q_tile: int, kv_tile: int): + """Explicit macro-tile knobs override the per-head-dim CONFIG default. + + One case per warp-layout source: the sweep-tuned CONFIG entry and the + largest_warp_partition fallback. (SMEM-infeasible combinations — e.g. + any d128 non-default — correctly fail the strict-select build in the + kernel constructor instead.) + """ + + _run_case(head_dim=head_dim, is_causal=True, q_tile=q_tile, kv_tile=kv_tile) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index bac480886..812780fd6 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -418,7 +418,7 @@ def _run_thd_case( graph.validate() graph.build_operation_graph() graph.create_execution_plans([cudnn.heur_mode.A]) - graph.select_engines([engine_name(arch="sm120")]) + _select_engine(graph, engine_name(arch="sm120")) graph.check_support() graph.build_plans() workspace = torch.empty(max(1, graph.get_workspace_size()), dtype=torch.uint8, device="cuda") diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index aa58a51b3..472fc4ce6 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -10,6 +10,7 @@ import torch from cudnn.sdpa import graph_analyzer as ga +from cudnn.sdpa.bwd import engines as bwd_engines from cudnn.sdpa.fwd import engines @@ -680,3 +681,201 @@ def test_sm120_knob_domains(monkeypatch): assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(tile_m=64, tile_n=64, cga=1)) assert not _eligible(g, engines.SdpaFwdKnobs(cga=2)) assert not _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) + + +# --------------------------------------------------------------------------- +# SDPA_BWD: facts extraction + sdpa_bwd_sm120 probe gating. Executable +# coverage lives in test_sdpa_bwd_dsl_sm120.py. +# --------------------------------------------------------------------------- + +_BWD_ENGINE = "sdpa_bwd_sm120" +_BWD_D = 64 + + +def _bwd_eligible(graph, knobs=None): + """Names of the FROST SDPA-backward engines whose caps match this graph.""" + return {s.name for s in bwd_engines.ENGINE_SPECS if bwd_engines.probe(s, graph, knobs)} + + +def _bshd_strides(h: int, s: int, d: int) -> tuple[int, int, int, int]: + return (s * h * d, d, h * d, 1) + + +def _mk_bwd_graph( + d: int = _BWD_D, + h_kv: int = H, + s_q: int = S, + s_kv: int = S, + kv_transposed_view: bool = False, + stats_stride: tuple | None = None, + grad_strides: tuple | None = None, + bias: bool = False, + dbias: bool = False, + **bwd_kwargs, +): + g = _mk_graph() + q_dims, q_strides = (B, H, s_q, d), _bshd_strides(H, s_q, d) + kv_dims, kv_strides = (B, h_kv, s_kv, d), _bshd_strides(h_kv, s_kv, d) + if kv_transposed_view: + # Mimic the post-build_operation_graph state: the backward node's K/V + # ports are rewritten to transposed (B, H, D, S) views of the same + # canonical BSHD buffer. + kv_dims = (B, h_kv, d, s_kv) + kv_strides = (s_kv * h_kv * d, d, 1, h_kv * d) + q = g.tensor(dim=q_dims, stride=q_strides, data_type=DTYPE, name="q") + k = g.tensor(dim=kv_dims, stride=kv_strides, data_type=DTYPE, name="k") + v = g.tensor(dim=kv_dims, stride=kv_strides, data_type=DTYPE, name="v") + o = g.tensor(dim=q_dims, stride=_bshd_strides(H, s_q, d), data_type=DTYPE, name="o") + do = g.tensor(dim=q_dims, stride=_bshd_strides(H, s_q, d), data_type=DTYPE, name="dO") + stats = g.tensor( + dim=(B, H, s_q, 1), + stride=stats_stride or (H * s_q, s_q, 1, 1), + data_type=cudnn.data_type.FLOAT, + name="stats", + ) + if bias: + bias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="bias") + bwd_kwargs.update(bias=bias_t) + if dbias: + dbias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="dBias") + bwd_kwargs.update(dBias=dbias_t) + dq, dk, dv = g.sdpa_backward(name="sb", q=q, k=k, v=v, o=o, dO=do, stats=stats, attn_scale=0.125, **bwd_kwargs) + _finish_output(dq, q_dims, grad_strides or _bshd_strides(H, s_q, d)) + _finish_output(dk, (B, h_kv, s_kv, d), grad_strides or _bshd_strides(h_kv, s_kv, d)) + _finish_output(dv, (B, h_kv, s_kv, d), grad_strides or _bshd_strides(h_kv, s_kv, d)) + return g + + +def test_bwd_engines_registered(): + from cudnn.engines import MANIFEST, is_python_engine + + (row,) = [r for r in MANIFEST if r.factory == "FrostSdpaBwdEngines"] + assert is_python_engine(row.engine_id) + assert row.id_end - row.engine_id >= len(bwd_engines.ENGINE_SPECS) + assert bwd_engines.engine_name() == _BWD_ENGINE + + +def test_bwd_facts_extracted(): + g = _mk_bwd_graph(use_causal_mask=True) + facts = _facts(g) + assert facts.is_backward + assert facts.causal and not facts.bottom_right + assert facts.right_bound == 0 + assert not facts.deterministic and not facts.has_dbias and not facts.has_dsink + assert (facts.b, facts.h_q, facts.h_kv, facts.s_q, facts.s_kv, facts.d_qk, facts.d_v) == (B, H, H, S, S, _BWD_D, _BWD_D) + assert facts.dtype == torch.float16 and facts.uniform_dtype + assert facts.bshd_layout + for ref in (facts.do_t, facts.dq_t, facts.dk_t, facts.dv_t, facts.stats_t): + assert ref is not None + assert facts.scale == 0.125 + + +def test_bwd_facts_kv_transposed_view_canonicalized(): + # After build_operation_graph the bwd node's K/V ports describe transposed + # (B, H, D, S) views; the analyzer canonicalizes dims AND strides back so + # geometry and the BSHD layout gate hold before and after the native build. + facts = _facts(_mk_bwd_graph(kv_transposed_view=True)) + assert (facts.s_kv, facts.d_qk) == (S, _BWD_D) + assert facts.bshd_layout + + +def test_bwd_probe_accepts(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph()) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(use_causal_mask=True)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(s_q=S // 2, use_causal_mask_bottom_right=True)) + for d in (32, 128): + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(d=d)) + + +def test_bwd_probe_rejects_forward_graph(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g, d=_BWD_D) + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) + _finish_output(o, dims, strides) + assert not _bwd_eligible(g) + # ... and symmetrically, the forward engines decline a backward graph. + assert not _eligible(_mk_bwd_graph()) + + +def test_bwd_probe_rejects_gqa(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(h_kv=H // 2)) + + +def test_bwd_probe_rejects_unsupported_head_dim(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(d=96)) + assert not _bwd_eligible(_mk_bwd_graph(d=256)) + + +def test_bwd_probe_causal_notches(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + # Top-left causal requires S_q == S_kv (the kernel diagonal is bottom-right). + assert not _bwd_eligible(_mk_bwd_graph(s_q=S // 2, use_causal_mask=True)) + # Causal with S_q > S_kv has fully-masked query rows (stats = -inf). + assert not _bwd_eligible(_mk_bwd_graph(s_q=2 * S, use_causal_mask_bottom_right=True)) + # Sliding window is not supported. + assert not _bwd_eligible(_mk_bwd_graph(use_causal_mask=True, sliding_window_length=64)) + + +def test_bwd_probe_rejects_deterministic(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(use_deterministic_algorithm=True)) + + +def test_bwd_probe_rejects_bias(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(bias=True)) + + +def test_bwd_probe_rejects_dbias(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(dbias=True)) + + +def test_bwd_probe_rejects_non_bshd_layout(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + # BHSD-contiguous gradients are outside the strict-BSHD envelope (no + # normalization copy on the backward path, unlike the forward dense_flex). + bhsd_contig = (H * S * _BWD_D, S * _BWD_D, _BWD_D, 1) + assert not _bwd_eligible(_mk_bwd_graph(grad_strides=bhsd_contig)) + + +def test_bwd_probe_rejects_strided_stats(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + # A padded stats stride has no zero-copy (B, H, S) reshape. + assert not _bwd_eligible(_mk_bwd_graph(stats_stride=(2 * H * S, 2 * S, 2, 1))) + + +def test_bwd_knob_domains(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + g = _mk_bwd_graph() + # In-domain requests are eligible (final per-head-dim feasibility is the + # kernel constructor's, at build). + assert _BWD_ENGINE in _bwd_eligible(g, bwd_engines.SdpaBwdKnobs(tile_m=64, tile_n=128)) + assert _BWD_ENGINE in _bwd_eligible(g, bwd_engines.SdpaBwdKnobs()) # all-None = no preference + # Out-of-domain values are rejected. + assert not _bwd_eligible(g, bwd_engines.SdpaBwdKnobs(tile_m=48)) + assert not _bwd_eligible(g, bwd_engines.SdpaBwdKnobs(tile_n=32)) + # Another operation's vocabulary is rejected wholesale. + assert not _bwd_eligible(g, engines.SdpaFwdKnobs(tile_m=64)) + + +def test_bwd_mismatch_reason_strings(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + caps = bwd_engines.ENGINE_SPECS[0].capabilities + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(h_kv=H // 2))) + assert reason is not None and "GQA" in reason + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(d=96))) + assert reason is not None and "96" in reason + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(s_q=S // 2, use_causal_mask=True))) + assert reason is not None and "top-left" in reason + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(use_deterministic_algorithm=True))) + assert reason is not None and "deterministic" in reason + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph()), engines.SdpaFwdKnobs(tile_m=64)) + assert reason is not None and "knob" in reason + reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph()), bwd_engines.SdpaBwdKnobs(tile_m=48)) + assert reason is not None and "tile_m=48" in reason + assert bwd_engines.mismatch(caps, _facts(_mk_bwd_graph()), bwd_engines.SdpaBwdKnobs(tile_m=64, tile_n=128)) is None