diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md new file mode 100644 index 000000000..ce0c68862 --- /dev/null +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -0,0 +1,200 @@ +# SDPA Backward (SM120) + +**This is an experimental API and subject to change.** + +## Overview + +**SDPA backward** pass for the NVIDIA Blackwell GeForce line (`SM120` / +`SM121`: RTX 50-series, RTX PRO 6000 Blackwell, DGX Spark), implemented with +CuTe DSL primitives using TMA loads and a warp-specialized producer/consumer +schedule. Consumes the forward activations (`Q/K/V/O`), the loss gradient +`dO`, and the forward `LSE`; produces `dQ/dK/dV`. + +Two integration surfaces are provided: + +* a standalone wrapper (documented below), `cudnn.sdpa_bwd_wrapper_dsl_sm120`, and +* a FROST engine (`sdpa_bwd_sm120`, see `cudnn.sdpa.bwd.engines`) that serves + single-node `sdpa_backward()` graphs built with `cudnn.pygraph` when + selected from the ranked plan list (`graph.plans` / + `graph.select_plan(i)`) with `CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1`. + +The kernel lives at `python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py` +(one fused five-GEMM main kernel, plus a `dot` preprocess and a `cvt` +dQ-finalize kernel per call). + +## Requirements + +The `cutedsl` optional dependency (`nvidia-cutlass-dsl` + `apache-tvm-ffi`) +and an SM120 or SM121 device. + +## API Usage + +```python +from cudnn.sdpa import sdpa_bwd_wrapper_dsl_sm120 + +grads = sdpa_bwd_wrapper_dsl_sm120( + q_tensor=q, k_tensor=k, v_tensor=v, + o_tensor=o, do_tensor=do, stats_tensor=stats, # from the forward pass + is_causal=True, + causal_bottom_right=False, + window_size_left=None, # W: keys with k < q + diag - W are masked + deterministic=False, # ordered dQ KV-tile reduction (bitwise-reproducible) + scale_softmax=None, # None -> 1/sqrt(D) +) +dq, dk, dv = grads["dq_tensor"], grads["dk_tensor"], grads["dv_tensor"] +``` + +Tensors are logical `(B, H, S, D)`; any dense layout with the head dim +innermost-contiguous is accepted (`dense_flex`) — non-BSHD-compact operands +are staged through workspace copies. `stats` is the natural-log forward LSE, +fp32 `(B, H, S_q, 1)` contiguous. + +Through the graph API, per-plan sequence-tile-width knobs can be requested via +`SdpaBwdKnobs`: `tile_m` controls the Q tile (`q_tile`), and `tile_n` controls +the KV tile (`kv_tile`). + +## Determinism + +By default the dQ accumulation across KV tiles uses fp32 atomics, so the result +can be bitwise non-deterministic when a q-tile receives contributions from +multiple KV-tile CTAs. `deterministic=True` serializes the per-`(batch, head, +q_tile)` additions in ascending KV-tile order through a GMEM turn-counter +array (the FlashAttention-style ordered-reduction relay). dK/dV are +deterministic in both modes. This maps to the graph API's +`use_deterministic_algorithm` and costs a shape-dependent slowdown of the main +kernel while keeping the workspace linear in sequence length. + +## Kernel design and optimizations + +### The three-kernel chain + +One backward call is three launches, overlapped with programmatic dependent +launch (PDL) so each kernel's prologue runs under its predecessor's tail: + +``` +dot delta = rowsum(O ∘ dO); zeroes dq_accum (and, when deterministic, the relay counters) +main the fused five-GEMM pass; writes dK/dV, accumulates dQ into dq_accum +cvt dq_accum (fp32, scrambled) -> dQ (io dtype), applying attn_scale +``` + +### Main-kernel pipeline: KV-stationary, five chained GEMMs + +Grid is `(num_kv_tiles, H, B)` — one CTA owns one KV tile, loads K/V **once**, +and walks every q-tile of its (batch, head) in descending order. Per q-tile +iteration: + +``` +GEMM1 S = Q · Kᵀ (K streamed from SMEM) + P = exp2((scale·S − LSE) · log2(e)) replay from natural-log LSE +GEMM2 dP = dO · Vᵀ (V resident in registers after one ldmatrix pass) + dS = P ∘ (dP − delta) in fp32 accumulators; I/O-dtype copy -> SMEM +GEMM3 dV += Pᵀ · dO (P read back transposed via ldmatrix.trans) +GEMM4 dQ = dS · K -> fp32 atomic scatter into dq_accum +GEMM5 dK += dSᵀ · Q (the iteration's last sQ reader) +``` + +dK/dV live in registers across the whole pass (CTA-private KV rows — no +atomics) and are written once in the epilogue through SMEM buffers that alias +the dead sK/sV regions. dQ is the transposed case — every KV tile contributes +to every q-tile row — hence the cross-CTA atomic workspace. + +Key register/SMEM economies: P stays in fp32 accumulator registers for the +dS pointwise (no SMEM round trip for the P→dS chain); V is register-resident +so `sdS` aliases `sV`. The `CONFIG` table selects three 2-D partitions of the +8 math warps per `(D, q_tile, kv_tile)`: GEMM1/2 share the S/dP partition, +GEMM3/5 share the dK/dV partition, and GEMM4 uses the dQ partition. This keeps +MMA fragments `ldmatrix`-legal and balances the accumulators. + +### Warp specialization + +384 threads = 12 warps: **8 math warps** (`setmaxregister` up to 240), **1 TMA +producer warp** (down to 24 registers), and 3 register-donor warps (down to +24; they exist only to hand their registers to the math warps). The producer +prefetches the tensormaps, issues the one-time K/V TMA loads, then streams +Q/dO tiles through an mbarrier `expect_tx` ring — double-buffered +(`Q_STAGES == 2`) where SMEM allows, so the next tile's Q is in flight while +the current one computes. Producer and consumers rendezvous on 288-thread +named barriers (loop-top ready/consumed, post-GEMM3 dO release, and the +single-buffer Q refill); math-only synchronization uses separate 256-thread +barriers that exclude the producer. + +### Head-size support and tile configs + +Each head dim selects a sweep-tuned default `(q_tile, kv_tile)` and +warp-partition triple: + +| D | q_tile × kv_tile | note | +|---|---|---| +| 32 | 128 × 64 | | +| 64 | 64 × 128 | wide KV tile: fewer CTAs, halves Q/dO re-reads | +| 128 | 64 × 64 | | +| 192 | 32 × 64 | double-buffered Q at the default config | +| 256 | 32 × 64 | SMEM-bound: single-buffered Q at the default config | + +SMEM per CTA is `(Q_STAGES + 1)·M·D + N·D + max(N·D, 2·M·N)` elements against +the ~99 KB SM120 cap. The constructor tries `Q_STAGES = 2` and falls back to a +**single Q buffer** when it doesn't fit; among the default configs, this occurs +at D=256. In the single-buffer branch the iteration reorders GEMM5 *before* +GEMM4 (GEMM5 is sQ's last reader), so the Q refill for the next tile hides +behind GEMM4 and the dQ scatter instead of stalling the loop. Head dims that +are multiples of 8 but are not native sizes are zero-padded by the adapter +(pad columns contribute nothing anywhere in the chain). Explicit +`tile_m`/`tile_n` knobs override the Q/KV tile defaults; off-table combinations +derive their warp partitions from a largest-valid rule. + +### dQ scatter and the scrambled workspace + +Naive per-element atomics from the MMA fragment layout produce scattered +addresses. Instead `dq_accum` uses a fragment-order ("scrambled") layout in +which the 32 lanes of a warp each reduce an adjacent fp32 pair, covering 64 +consecutive floats per `red.global.add.v2.f32` invocation. The coalesced layout +reduces dQ atomic traffic, and the `cvt` kernel un-scrambles it while converting +to the I/O dtype. `dot` pre-zeroes the workspace (fused with the delta +reduction; PDL orders it before `main`'s first add). + +### Causal and sliding-window masks + +Masking is applied twice, cheaply: + +* **Loop bounds** do the heavy lifting: causal clamps the first q-tile + (`q_block_min`, bottom-right via `diag_off = S_kv − S_q`), a left window + clamps the last (`q_block_max`) — fully-masked tiles are never visited, so + square causal attention runs roughly half as many tile iterations. +* **In-register score masking** runs only on tiles that straddle a mask edge + (`do_mask_causal` / `do_mask_window` gates); interior tiles skip it. + +The softmax replay guards fully-masked rows (forward `LSE = −inf`) by +substituting `+inf`, reconstructing `P = 0` instead of NaN. Non-tile-multiple +sequence tails are handled with load clamps and store row-gates. + +### Deterministic vs. non-deterministic dQ + +The default path's relaxed atomics make the fp32 add order — hence the +bitwise result — scheduling-dependent. Deterministic mode serializes each +(batch, head, q-tile)'s adds in ascending KV-tile order through an int32 +turn-counter array: one elected lane spins on an acquire load until the +counter equals the CTA's turn, a math-warps-only barrier releases the scatter, +and after a second barrier a single `st.release.gpu` publishes `turn + 1`. +The release store's own fence drains the adds; in the double-buffered branch, +the following GEMM5 overlaps that drain. +Correctness rests on CTAs dispatching in ascending `blockIdx.x` (so the +awaited predecessor is always resident or done) and on the mask loop bounds +making each q-tile's visitors contiguous in KV index — under a sliding window +the turn subtracts the first visitor, +`kv_lo = max((q_block·tile_q + diag − W) // tile_kv, 0)`. +The relay instructions fold out under `const_expr` when determinism is off; +only the unused relay operand remains in the kernel ABI. + +## Support surface and constraints + +- SM120 and SM121, e.g. RTX 5090, RTX PRO 6000 Blackwell, and DGX Spark +- Dtypes: FP16 / BF16 (LSE fp32) +- Head dims: 32/64/128/192/256 natively; any other multiple of 8 up to 256 is + served by zero-padding D to the next supported size (`d_qk == d_v`) +- Masks: none, causal (top-left or bottom-right), sliding window + (left-window offset, with or without causal) +- Equal Q/KV head counts (no GQA/MQA); no dropout / bias / ALiBi / sinks / + softcap / THD +- Workspace (carved from the caller's buffer): fp32 `delta` and `dq_accum` + scratch plus int32 relay-counter storage (reserved in both modes); padded-D + and non-compact layouts add staging copies diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 70df78a4c..3661dc40e 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -29,6 +29,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Forward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-forward-fe-oss-sm100-d256) - [SDPA Backward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-backward-fe-oss-sm100-d256) +- [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) ## Installation and setup diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index b40bf4268..91ec3e4fe 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -22,8 +22,9 @@ SEQ_Q_TILES as _SM120_Q_TILES, SUPPORTED_HEAD_DIMS as _SM120_SUPPORTED_HEAD_DIMS, TemplateParams as Sm120TemplateParams, + padded_head_dim as _sm120_padded_head_dim, ) -from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, ws_align +from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, _torch_stream_context, ws_align _SM120_KERNEL_FILE = "bprop_f16_sm120.py" _SM120_DTYPE_QKV_CODE = { @@ -33,6 +34,9 @@ # 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 +# dq_sem is sized for the smallest legal q-tile so one formula covers every +# tile choice; must match the template's fake_dq_sem sizing. +_SM120_MIN_Q_TILE = 32 _logger = logging.getLogger(__name__) @@ -65,6 +69,7 @@ def __init__( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: Optional[int] = None, + deterministic: bool = False, scale_softmax: Optional[float] = None, tile_m: Optional[int] = None, tile_n: Optional[int] = None, @@ -86,6 +91,7 @@ def __init__( self.is_causal = bool(is_causal) self.causal_bottom_right = bool(causal_bottom_right) self.window_size_left = None if window_size_left is None else int(window_size_left) + self.deterministic = bool(deterministic) 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) @@ -104,24 +110,6 @@ def __init__( 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.""" @@ -155,6 +143,10 @@ def _initialize_implementation(self) -> None: self.compute_capability: Optional[tuple[int, int]] = None self._k_mod = None self._sq_rounded: Optional[int] = None + # Kernel-facing head dim. When it differs from D, operands stage through zero-padded compact copies + self.head_dim_padded: Optional[int] = None + # name -> staging number-of-elements for each non-BSHD-compact port + self._staging_numels: dict[str, int] = {} @staticmethod def _bshd_physical_ok(desc: TensorDesc) -> bool: @@ -166,16 +158,19 @@ def _bshd_physical_ok(desc: TensorDesc) -> bool: def check_support(self) -> bool: self._logger.debug("Entering check_support") + from cudnn.sdpa.graph_analyzer import dense_layout_ok + + self._staging_numels = {} 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}", + not dense_layout_ok(tuple(desc.shape), tuple(desc.stride)), + f"{desc.name} must have the head dim innermost-contiguous (stride 1) and " + f"non-broadcast, non-overlapping strides (any B/H/S order, padded " + f"strides allowed); got stride {desc.stride} shape {desc.shape}", ) b, h_q, s_q, d_qk = self.q_desc.shape @@ -194,10 +189,15 @@ def check_support(self) -> bool: h_q != h_kv, f"SM120 DSL SDPA backward does not implement GQA / MQA; got H_q={h_q}, H_kv={h_kv}", ) + self.head_dim_padded = _sm120_padded_head_dim(int(d_qk)) if d_qk % 8 == 0 else None 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.head_dim_padded is None, + f"D ({d_qk}) must be a multiple of 8 and <= {max(_SM120_SUPPORTED_HEAD_DIMS)}", ) + # All operands stage when D pads; otherwise only non-BSHD-compact ones. + 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): + if self.head_dim_padded != d_qk or not self._bshd_physical_ok(desc): + self._staging_numels[desc.name] = math.prod(desc.shape[:-1]) * self.head_dim_padded self._value_error_if( self.stats_desc.ndim != 4 or tuple(self.stats_desc.shape) != (b, h_q, s_q, 1), @@ -278,6 +278,7 @@ def compile(self) -> None: is_causal=self.is_causal, causal_top_left=self.is_causal and not self.causal_bottom_right, window_size_left=self.window_size_left, + deterministic=self.deterministic, q_tile=self.q_tile, kv_tile=self.kv_tile, ) @@ -288,17 +289,26 @@ def compile(self) -> None: qh=self.h_q, sq=self.s_q_max, skv=self.s_k_max, - d=self.head_dim, + d=self.head_dim_padded, ) self._logger.debug("compile completed") + def _dq_sem_len(self) -> int: + """Element count of the dq_sem relay-counter buffer (int32).""" + + return self.batch_size * self.h_q * _round_up(self.s_q_max, _SM120_MIN_Q_TILE) // _SM120_MIN_Q_TILE + def scratch_workspace_bytes(self) -> int: - """delta (fp32 [B, H, SQ_r128]) + dq_accum (fp32 flat [B*SQ_r128*H*D]).""" + """delta (fp32 [B, H, SQ_r128]) + dq_accum (fp32 flat [B*SQ_r128*H*D]) + + dq_sem (int32 flat [B*H*ceil(SQ/32)], deterministic relay counters) + + one compact staging copy per non-BSHD-compact operand.""" 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 + dq_accum_bytes = ws_align(self.batch_size * self._sq_rounded * self.h_q * self.head_dim_padded * 4) + dq_sem_bytes = ws_align(self._dq_sem_len() * 4) + staging_bytes = sum(ws_align(numel * self.dtype.itemsize) for numel in self._staging_numels.values()) + return delta_bytes + dq_accum_bytes + dq_sem_bytes + staging_bytes def execute( self, @@ -325,7 +335,8 @@ def execute( 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) + dq_accum = carver.take(self.batch_size * self._sq_rounded * self.h_q * self.head_dim_padded, torch.float32) + dq_sem = carver.take(self._dq_sem_len(), torch.int32) if current_stream is None: # Direct call (no dispatch-forwarded stream): fall back to torch's @@ -335,34 +346,63 @@ def execute( 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) + # Non-compact operands (all operands when D pads) stage through + # workspace-carved compact copies with zero-filled pad columns. + d_pad = self.head_dim_padded + pads = d_pad != self.head_dim + + def _staged_bshd(tensor: torch.Tensor) -> torch.Tensor: + view = tensor.transpose(1, 2) + if not pads and view.is_contiguous(): + return view + b, s, h, _ = view.shape + staged = carver.take(b * s * h * d_pad, self.dtype).view(b, s, h, d_pad) + if pads: + staged[..., self.head_dim :].zero_() + staged[..., : self.head_dim].copy_(view) + return staged + + def _staged_out_bshd(tensor: torch.Tensor): + """(kernel-facing compact BSHD buffer, user view to scatter back into or None).""" + view = tensor.transpose(1, 2) + if not pads and view.is_contiguous(): + return view, None + b, s, h, _ = view.shape + return carver.take(b * s * h * d_pad, self.dtype).view(b, s, h, d_pad), view + + with _torch_stream_context(current_stream, q_tensor.device): + q = _staged_bshd(q_tensor) + k = _staged_bshd(k_tensor) + v = _staged_bshd(v_tensor) + o = _staged_bshd(o_tensor) + do = _staged_bshd(do_tensor) + dq, dq_user = _staged_out_bshd(dq_tensor) + dk, dk_user = _staged_out_bshd(dk_tensor) + dv, dv_user = _staged_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, dq_sem, current_stream) + kernels.main( + q, + k, + v, + do, + lse, + delta, + dq_accum, + dq_sem, + dk, + dv, + cutlass.Float32(scale_log2), + cutlass.Float32(scale_val), + current_stream, + ) + kernels.cvt(dq_accum, dq, cutlass.Float32(scale_val), current_stream) + for user_view, staged in ((dq_user, dq), (dk_user, dk), (dv_user, dv)): + if user_view is not None: + user_view.copy_(staged[..., : self.head_dim]) def _tensor_signature(tensor: torch.Tensor) -> tuple: @@ -383,6 +423,7 @@ def sdpa_bwd_wrapper_dsl_sm120( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: Optional[int] = None, + deterministic: bool = False, scale_softmax: Optional[float] = None, ) -> TupleDict: """Run SM120 SDPA backward and return ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=...)``.""" @@ -404,6 +445,7 @@ def sdpa_bwd_wrapper_dsl_sm120( bool(is_causal), bool(causal_bottom_right), window_size_left, + bool(deterministic), scale_softmax, ) api = _wrapper_api_cache.get(cache_key) @@ -421,6 +463,7 @@ def sdpa_bwd_wrapper_dsl_sm120( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + deterministic=deterministic, scale_softmax=scale_softmax, ) api.check_support() diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py index 4c06b9148..93285c363 100644 --- a/python/cudnn/sdpa/bwd/config_sm120.py +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -9,9 +9,15 @@ from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 -SEQ_Q_TILES = (64, 128) +SEQ_Q_TILES = (32, 64, 128) SEQ_KV_TILES = (64, 128) -SUPPORTED_HEAD_DIMS = (32, 64, 128) +SUPPORTED_HEAD_DIMS = (32, 64, 128, 192, 256) + + +def padded_head_dim(d: int) -> "int | None": + """Smallest native bin >= ``d``, or ``None`` when ``d`` exceeds every bin.""" + + return min((b for b in SUPPORTED_HEAD_DIMS if b >= d), default=None) @dataclass(frozen=True) @@ -28,6 +34,7 @@ class TemplateParams: is_causal: bool = False causal_top_left: bool = False window_size_left: int | None = None + deterministic: bool = False use_pdl: bool = True q_tile: int = 0 kv_tile: int = 0 diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index cd03ef2ea..79c318866 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -113,6 +113,17 @@ class Capabilities: thd: bool = False cu_seq_len: bool = False # cu_seq_len_q / cu_seq_len_kv prefix sums (no row serves these yet) + # Dense layout envelope this engine accepts: + # "bshd" — Q/K/V/O must be BSHD-physical (stride order 3,1,2,0). + # "dense_flex" — any B/H/S stride permutation, padded (oversized) + # strides included, as long as the head dim is + # innermost-contiguous (stride 1) and the strides are + # non-broadcast / non-overlapping (facts.dense_layout; + # see graph_analyzer.dense_layout_ok). The DSL executor + # normalizes such tensors to the kernel's canonical + # BSHD-compact buffers (zero-copy when already BSHD). + layouts: frozenset[str] = frozenset({"bshd"}) + # Tuning-knob domains this engine's lowering honors (see SdpaBwdKnobs). tile_ms: frozenset[int] = frozenset() tile_ns: frozenset[int] = frozenset() @@ -159,7 +170,13 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", requested: 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: + if "dense_flex" in capabilities.layouts: + if not facts.dense_layout: + return ( + "Q/K/V/O/dO/dQ/dK/dV must have the head dim innermost-contiguous (stride 1) and " + "non-broadcast, non-overlapping strides (any B/H/S order, padded strides allowed)" + ) + elif 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 ( @@ -223,11 +240,14 @@ def _sm120_spec() -> EngineSpec: capabilities=Capabilities( sm_lo=_BLACKWELL_GEFORCE[0], sm_hi=_BLACKWELL_GEFORCE[1], - d=frozenset(_SM120_HEAD_DIMS), + # Any head size multipled of 8 + d=frozenset(range(8, max(_SM120_HEAD_DIMS) + 1, 8)), dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), causal=True, bottom_right=True, swa=True, + layouts=frozenset({"bshd", "dense_flex"}), + deterministic=True, tile_ms=frozenset(_SM120_Q_TILES), tile_ns=frozenset(_SM120_KV_TILES), ), @@ -268,17 +288,13 @@ def lower_dsl_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = 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) + # Per-port geometry from facts.port_layouts, NOT the live IR tensors: + # build_operation_graph rewrites the backward node's K/V ports to + # transposed (B, H, D, S) views; the analyzer captured the geometry with + # that rewrite undone. + ports = {name: (tuple(dim), tuple(stride)) for name, dim, stride in facts.port_layouts} + q_geom, k_geom, v_geom, o_geom = ports["q"], ports["k"], ports["v"], ports["o"] + do_geom, dq_geom, dk_geom, dv_geom = ports["dO"], ports["dQ"], ports["dK"], ports["dV"] stats_geom = ((facts.b, facts.h_q, facts.s_q, 1), (facts.h_q * facts.s_q, facts.s_q, 1, 1)) import torch @@ -301,17 +317,18 @@ def _desc(geom, dtype, name: str) -> "Any": api = _adapter_sm120()( 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_k=_desc(k_geom, facts.dtype, "k"), + sample_v=_desc(v_geom, facts.dtype, "v"), + sample_o=_desc(o_geom, facts.dtype, "o"), + sample_do=_desc(do_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"), + sample_dq=_desc(dq_geom, facts.dtype, "dQ"), + sample_dk=_desc(dk_geom, facts.dtype, "dK"), + sample_dv=_desc(dv_geom, facts.dtype, "dV"), is_causal=facts.causal, causal_bottom_right=facts.bottom_right, window_size_left=facts.window_left, + deterministic=facts.deterministic, 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, @@ -338,12 +355,12 @@ def _desc(geom, dtype, name: str) -> "Any": ) def _canonical_view(buf, geom): - """Reinterpret a variant-pack buffer through the canonical geometry. + """Reinterpret a variant-pack buffer through the port's 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. + executor consumes torch views, so rebuild the port-shaped view here. No-op when the caller already passed a matching view. """ dim, stride = geom @@ -355,14 +372,14 @@ 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), + k_tensor=_canonical_view(resolved[id(binding.k)], k_geom), + v_tensor=_canonical_view(resolved[id(binding.v)], v_geom), + o_tensor=_canonical_view(resolved[id(binding.o)], o_geom), + do_tensor=_canonical_view(resolved[id(binding.do)], do_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), + dq_tensor=_canonical_view(resolved[id(binding.dq)], dq_geom), + dk_tensor=_canonical_view(resolved[id(binding.dk)], dk_geom), + dv_tensor=_canonical_view(resolved[id(binding.dv)], dv_geom), scale_softmax=facts.scale, # Scratch comes from the CALLER's workspace (never allocated # here): the dispatch sized/validated it against workspace_bytes; diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index 8a9a96456..c9ad5aa3b 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -18,12 +18,11 @@ Constraints: * Supported input dtypes: Float16 and BFloat16 (output dtype matches) -* Head dimension must be one of 32, 64, or 128 +* Head dimension must be one of 32, 64, 128, 192, or 256; the adapter + serves any other multiple of 8 up to 256 by zero-padding D. * Equal Q/KV head counts (no GQA); no dropout/alibi/softcap * Optional causal (top-left or bottom-right) and sliding-window masks -* 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 = @@ -266,6 +265,169 @@ def mma_abregs( acc[s + 3] = c3 +@cute.jit +def _bwd_gemm4_dq( + acc_dq, + sdS, + sK, + wq, + wd_q, + lane, + *, + DQ_REPS: cutlass.Constexpr[int], + DQ_NF: cutlass.Constexpr[int], + KV_CHUNKS: cutlass.Constexpr[int], + WM_DQ: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + PDS: cutlass.Constexpr[int], + PAGE: cutlass.Constexpr[int], + DQ_PER: cutlass.Constexpr[int], + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """GEMM 4: acc_dq = dS @ K^T (reads only sdS/sK, never sQ).""" + 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, + ) + + +@cute.jit +def _bwd_dq_scatter( + acc_dq, + dqa_ptr, + dqa_base, + math_tidx, + H, + *, + DQ_REPS: cutlass.Constexpr[int], + DQ_NF: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + d: cutlass.Constexpr[int], +): + """dQ accumulate into the scrambled dq_accum workspace.""" + t_r = math_tidx // 32 + t_c = math_tidx % 32 + 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]) + + +@cute.jit +def _bwd_gemm5_dk( + acc_dk, + sdS, + sQ_st, + wn_k, + wd_k, + lane, + *, + Q_CHUNKS: cutlass.Constexpr[int], + DKV_REPS: cutlass.Constexpr[int], + WM_DKV: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + PDS: cutlass.Constexpr[int], + PAGE: cutlass.Constexpr[int], + DKV_PER: cutlass.Constexpr[int], + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """GEMM 5: acc_dk += dS^T @ Q (the iteration's last sQ reader).""" + 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, + ) + + +@cute.jit +def _bwd_det_wait(det_sem, m_block, det_turn, warp): + """Deterministic-relay entry: block the 8 compute warps until it is this + CTA's turn for q-tile ``m_block`` (FA3 / cuDNN-SM90 STAGES=4 scheme). + + One elected lane of warp 0 spins on an acquire load of the turn counter; + barrier 6 (compute warps only — never the producer's 288-thread + barriers) releases the other warps into the dQ scatter.""" + if warp == 0: + if prims.elect_sync(): + while ( + prims.load_ext( + det_sem + m_block, + order=prims.MemOrder.ACQUIRE, + scope=prims.MemScope.GPU, + ) + != det_turn + ): + pass + cute.arch.barrier(barrier_id=6, number_of_threads=256) + + +@cute.jit +def _bwd_det_release(det_sem, m_block, det_turn, warp): + """Deterministic-relay exit: pass the turn for ``m_block`` to the next + kv tile once every compute warp has issued its dQ reds. + + The release store alone orders the relaxed reds before the handoff: the + barrier sequences the other warps' reds against this thread (CTA scope) + and st.release makes them cumulatively visible at GPU scope (its own + membar; an extra fence here doubled the per-handoff drain cost).""" + cute.arch.barrier(barrier_id=6, number_of_threads=256) + if warp == 0: + if prims.elect_sync(): + prims.store_ext( + (det_turn + 1).ir_value(), + det_sem + m_block, + order=prims.MemOrder.RELEASE, + scope=prims.MemScope.GPU, + ) + + # --------------------------------------------------------------------------- # Main kernel. # --------------------------------------------------------------------------- @@ -278,6 +440,8 @@ class SM120FusedMultiHeadAttentionFP16Backward: 32: (128, 64), 64: (64, 128), 128: (64, 64), + 192: (32, 64), + 256: (32, 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 @@ -288,6 +452,8 @@ class SM120FusedMultiHeadAttentionFP16Backward: (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 + (192, 32, 64): (2, 4, 2), # default for d192 + (256, 32, 64): (2, 4, 2), # default for d256 } def __init__( @@ -296,6 +462,7 @@ def __init__( is_causal: bool = False, causal_top_left: bool = False, window_size_left: int | None = None, + deterministic: bool = False, head_dim: int = 128, use_pdl: bool = True, q_tile: int = 0, @@ -305,6 +472,7 @@ def __init__( self.is_causal = is_causal self.causal_top_left = bool(causal_top_left) self.window_size_left = window_size_left + self.deterministic = bool(deterministic) self.d = head_dim self.use_pdl = bool(use_pdl) self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim] @@ -341,18 +509,24 @@ def __init__( 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 + # Double-buffer Q when SMEM allows (prefetch hides the TMA latency); + # single-buffered Q pays an end-of-iteration rendezvous (d256's only fit). + cap = cutlass.utils.get_smem_capacity_in_bytes("sm_120") + for q_stages in (2, 1): + smem_elems = (q_stages + 1) * M * d + N * d + max(N * d, 2 * M * N) + if smem_elems * in_dtype.bytes <= cap: + break + else: + raise ValueError(f"smem {smem_elems * in_dtype.bytes} bytes exceeds the sm_120 cap of {cap} bytes even single-buffered") + # smem element offsets. + self.q_stages = q_stages + self.off_sQ = 0 # `q_stages` buffers + self.off_sdO = q_stages * M * d + self.off_sK = self.off_sdO + 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}") + self.smem_elems = smem_elems @cute.jit def load_tma_tile(self, s_dst, tma_desc, mbar, batch, head, seq, rows: cutlass.Constexpr[int]): @@ -377,6 +551,7 @@ def kernel( 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) + dq_sem: cute.Tensor, # [B*H*num_q_tiles] int32 relay turn counters, one per (batch, head, q-tile); zeroed by dot (deterministic only) dk: cute.Tensor, # [B, SKV, H, D] output dv: cute.Tensor, # [B, SKV, H, D] output tma_q_desc: cutlass.GridConstant[cuda.TensorMap], @@ -391,6 +566,7 @@ def kernel( M = self.q_tile N = self.kv_tile PAGE = self.page + Q_STAGES = self.q_stages 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 @@ -434,6 +610,8 @@ def kernel( dqa_ptr = dq_accum.iterator.raw_ptr() dk_ptr = dk.iterator.raw_ptr() dv_ptr = dv.iterator.raw_ptr() + if cutlass.const_expr(self.deterministic): + dqsem_ptr = dq_sem.iterator.raw_ptr() PARTIAL_Q = (SQ % M) != 0 PARTIAL_KV = (SKV % N) != 0 @@ -456,7 +634,7 @@ def kernel( n_iters = cute.math.max(m_block_max - m_block_min, cutlass.Int32(0)) smem = cutlass.Array(io_dtype, self.smem_elems, space=cutlass.AddressSpace.smem, alignment=128) - sQ = smem # 2 * M * d (double buffer) + sQ = smem # Q_STAGES * M * d sdO = smem.subview(self.off_sdO) # M * d sK = smem.subview(self.off_sK) # N * d sV = smem.subview(self.off_sV) # N * d @@ -487,6 +665,8 @@ def kernel( khd_base = (batch * SKV + kv_base) * row_stride + head * d lse_base = (batch * H + head) * SQ dd_base = (batch * H + head) * SQ_R + if cutlass.const_expr(self.deterministic): + det_sem = dqsem_ptr + (batch * H + head) * ((SQ + M - 1) // M) if warp == self.load_warp_id: prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) @@ -516,8 +696,12 @@ def kernel( 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) + if cutlass.const_expr(Q_STAGES == 2): + load_stage = load_j & cutlass.Int32(1) + q_phase_p = (load_j // 2) & cutlass.Int32(1) + else: + load_stage = cutlass.Int32(0) + q_phase_p = load_j & 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) @@ -526,26 +710,35 @@ def kernel( # 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. + if cutlass.const_expr(Q_STAGES == 2): + # when double-buffer for q, prefetch for the next iteration + 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 (dV += P^T*dO): 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) + if cutlass.const_expr(Q_STAGES == 1): + # Post-GEMM5 (dK += dS^T*Q): every consumer is done with sQ. + cute.arch.barrier(barrier_id=5, number_of_threads=288) + if load_j + 1 < n_iters: + 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, next_m * M, rows=M) elif warp < self.load_warp_id: prims.setmaxregister(240, prims.SetMaxRegisterAction.INCREASE) @@ -635,7 +828,10 @@ def kernel( j = cutlass.Int32(0) while j < n_iters: m_block = m_block_max - 1 - j - stage = j & cutlass.Int32(1) + if cutlass.const_expr(Q_STAGES == 2): + stage = j & cutlass.Int32(1) + else: + stage = cutlass.Int32(0) sQ_st = sQ.subview(stage * M * d) q_row0 = m_block * M @@ -831,86 +1027,134 @@ def kernel( # 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( + # Deterministic relay turn for this q-tile: the dQ adds of a + # (batch, head, q-tile) happen in ascending kv-tile order. + if cutlass.const_expr(self.deterministic): + if cutlass.const_expr(self.window_size_left is not None): + # SWA clamps m_block_max, so a q-tile's visitors start + # at kv tile n_lo = max((m_block*M + diag_off - W) // N, 0) + # (inverts the clamp); count turns from there. + det_turn = n_block - cute.math.max( + (m_block * M + diag_off - self.window_size_left) // N, + cutlass.Int32(0), + ) + else: + det_turn = n_block + + if cutlass.const_expr(Q_STAGES == 1): + # Single Q buffer: GEMM5 (sQ's last reader) first, so the + # Q refill hides behind GEMM4 + the dQ scatter. (2-stage + # keeps GEMM4-first: atomics drain during GEMM5 instead.) + _bwd_gemm5_dk( + acc_dk, + sdS, + sQ_st, + wn_k, + wd_k, + lane, + Q_CHUNKS=Q_CHUNKS, + DKV_REPS=DKV_REPS, + WM_DKV=WM_DKV, + M=M, + PDS=PDS, + PAGE=PAGE, + DKV_PER=DKV_PER, + io_dtype=io_dtype, + ) + cute.arch.barrier(barrier_id=5, number_of_threads=288) + _bwd_gemm4_dq( acc_dq, - af, + sdS, 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, + wq, + wd_q, + lane, + DQ_REPS=DQ_REPS, + DQ_NF=DQ_NF, + KV_CHUNKS=KV_CHUNKS, + WM_DQ=WM_DQ, + M=M, + N=N, + PDS=PDS, + PAGE=PAGE, + DQ_PER=DQ_PER, + io_dtype=io_dtype, ) - - # 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 - val = (lse_ptr + lse_base + nq0 + r_loc).load() - if cutlass.const_expr((self.is_causal and not self.causal_top_left) or self.window_size_left is not None): - # A fully masked row has LSE = -inf, which can produce - # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. - if val == cutlass.Float32(float("-inf")): - val = cutlass.Float32(float("inf")) - lse_r[rep * 2 + hf] = val * 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( + # Reload LSE for the next (lower) m-block: overlaps the + # in-flight Q refill issued at barrier 5. + 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 + val = (lse_ptr + lse_base + nq0 + r_loc).load() + if cutlass.const_expr((self.is_causal and not self.causal_top_left) or self.window_size_left is not None): + # A fully masked row has LSE = -inf, which can produce + # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. + if val == cutlass.Float32(float("-inf")): + val = cutlass.Float32(float("inf")) + lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) + dqa_base = ((batch * SQ_R + q_row0) * H + head) * d + if cutlass.const_expr(self.deterministic): + _bwd_det_wait(det_sem, m_block, det_turn, warp) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, H, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d=d) + if cutlass.const_expr(self.deterministic): + _bwd_det_release(det_sem, m_block, det_turn, warp) + else: + _bwd_gemm4_dq( + acc_dq, + sdS, + sK, + wq, + wd_q, + lane, + DQ_REPS=DQ_REPS, + DQ_NF=DQ_NF, + KV_CHUNKS=KV_CHUNKS, + WM_DQ=WM_DQ, + M=M, + N=N, + PDS=PDS, + PAGE=PAGE, + DQ_PER=DQ_PER, + io_dtype=io_dtype, + ) + # Reload LSE for the next (lower) m-block (develop-exact + # position: between GEMM4 and the scatter, hiding the + # global-load latency behind the atomic drain + GEMM5). + 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 + val = (lse_ptr + lse_base + nq0 + r_loc).load() + if cutlass.const_expr((self.is_causal and not self.causal_top_left) or self.window_size_left is not None): + # A fully masked row has LSE = -inf, which can produce + # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. + if val == cutlass.Float32(float("-inf")): + val = cutlass.Float32(float("inf")) + lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) + dqa_base = ((batch * SQ_R + q_row0) * H + head) * d + if cutlass.const_expr(self.deterministic): + _bwd_det_wait(det_sem, m_block, det_turn, warp) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, H, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d=d) + if cutlass.const_expr(self.deterministic): + _bwd_det_release(det_sem, m_block, det_turn, warp) + _bwd_gemm5_dk( acc_dk, - af, + sdS, 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, + wn_k, + wd_k, + lane, + Q_CHUNKS=Q_CHUNKS, + DKV_REPS=DKV_REPS, + WM_DKV=WM_DKV, + M=M, + PDS=PDS, + PAGE=PAGE, + DKV_PER=DKV_PER, + io_dtype=io_dtype, ) j += 1 @@ -969,6 +1213,7 @@ def __call__( lse: cute.Tensor, delta: cute.Tensor, dq_accum: cute.Tensor, + dq_sem: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, softmax_scale_log2: cutlass.Float32, @@ -990,6 +1235,7 @@ def __call__( lse, delta, dq_accum, + dq_sem, dk, dv, tma_q_desc, @@ -1008,7 +1254,7 @@ def __call__( # --------------------------------------------------------------------------- -# Preprocess kernel: delta = rowsum(dO * O) + dq_accum zeroing +# Preprocess kernel: delta = rowsum(dO * O) + dq_accum / dq_sem zeroing # --------------------------------------------------------------------------- @@ -1018,10 +1264,12 @@ def _dot_do_o_kernel( 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) + dq_sem: cute.Tensor, # [B*H*num_q_tiles] int32 relay turn counters (zeroed here when deterministic) q_tile: cutlass.Constexpr[int], d: cutlass.Constexpr[int], page: cutlass.Constexpr[int], use_pdl: cutlass.Constexpr[bool], + deterministic: cutlass.Constexpr[bool], ): if cutlass.const_expr(use_pdl): cute.arch.griddepcontrol_launch_dependents() @@ -1092,6 +1340,14 @@ def _dot_do_o_kernel( addr = dqa_base + (zr0 + im * zrows) * (H * d) + zc0 + jn * ztpr * 4 (dqa_ptr + addr).store(zero4, alignment=16) + if cutlass.const_expr(deterministic): + # Reset this q-tile's relay turn counter (PDL-ordered before the main + # kernel's first acquire, like the dq_accum zeroing above). + if tidx == 0: + num_q_tiles = (SQ + M - 1) // M + sem_ptr = dq_sem.iterator.raw_ptr() + (sem_ptr + (batch * H + head) * num_q_tiles + m_block).store(cutlass.Int32(0)) + @cute.jit def _dot_do_o_host( @@ -1099,14 +1355,16 @@ def _dot_do_o_host( do: cute.Tensor, delta: cute.Tensor, dq_accum: cute.Tensor, + dq_sem: cute.Tensor, q_tile: cutlass.Constexpr[int], d: cutlass.Constexpr[int], page: cutlass.Constexpr[int], use_pdl: cutlass.Constexpr[bool], + deterministic: 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( + _dot_do_o_kernel(o, do, delta, dq_accum, dq_sem, q_tile, d, page, use_pdl, deterministic).launch( grid=(m_blocks, o.shape[2], o.shape[0]), block=(256, 1, 1), stream=stream, @@ -1233,6 +1491,7 @@ def compile( # noqa: A001 is_causal=PARAMS.is_causal, causal_top_left=PARAMS.causal_top_left, window_size_left=PARAMS.window_size_left, + deterministic=PARAMS.deterministic, head_dim=d, use_pdl=PARAMS.use_pdl, q_tile=PARAMS.q_tile, @@ -1259,6 +1518,9 @@ def _fake(dtype, shape): 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,)) + # Sized for the smallest legal q-tile (32) so one formula covers every + # tile choice; must match the adapter's carve (scratch_workspace_bytes). + fake_dq_sem = _fake(cutlass.Int32, (b * qh * ceil_div(sq, 32),)) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) options = "--enable-tvm-ffi" @@ -1268,10 +1530,12 @@ def _fake(dtype, shape): fake_do, fake_delta, fake_dq_accum, + fake_dq_sem, bwd.q_tile, d, bwd.page, bwd.use_pdl, + bwd.deterministic, fake_stream, options=options, ) @@ -1284,6 +1548,7 @@ def _fake(dtype, shape): fake_lse, fake_delta, fake_dq_accum, + fake_dq_sem, fake_dk, fake_dv, cutlass.Float32(1.0), diff --git a/python/cudnn/sdpa/graph_analyzer.py b/python/cudnn/sdpa/graph_analyzer.py index a98d6462f..e5ae06022 100644 --- a/python/cudnn/sdpa/graph_analyzer.py +++ b/python/cudnn/sdpa/graph_analyzer.py @@ -198,6 +198,9 @@ class SdpaGraphFacts: # strides — any B/H/S order, padded strides allowed (see dense_layout_ok). # The actual per-tensor stride tuples stay available via q_t/k_t/v_t/o_t. dense_layout: bool = True + # Backward only: (name, dim, stride) per rank-4 layout port, with the + # K/V transposed-port rewrite undone; () on forward graphs. + port_layouts: tuple = () is_mxfp8: bool = False # block-scale MXFP8 (FP8 Q/K/V + per-32-block E8M0 SF) is_fp8: bool = False # per-tensor FP8 (FP8 Q/K/V + scalar descales) dtype_o: Optional[Any] = None # O dtype as cudnn.data_type @@ -538,6 +541,7 @@ def _square_transposed(dim: tuple, stride: tuple) -> bool: uniform_dtype=uniform, bshd_layout=bshd, dense_layout=dense_layout, + port_layouts=(tuple((name, dims[name], strides[name]) for name, _ in rank4_ports) if is_backward else ()), is_mxfp8=is_mxfp8, is_fp8=is_fp8, dtype_o=(o_dtype if _fp8_family else q_dtype), diff --git a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py index 582dc9dd1..f198257d0 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -49,10 +49,12 @@ def _require_dsl() -> None: from frost_test_utils import select_engine as _select_engine # noqa: F401 -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.""" +def _bhsd(batch: int, heads: int, sequence: int, head_dim: int, dtype: torch.dtype, empty: bool = False, layout: str = "bshd") -> torch.Tensor: + """Logical BHSD over compact BSHD storage, or BHSD-contiguous for layout="bhsd".""" factory = torch.empty if empty else torch.randn + if layout == "bhsd": + return factory(batch, heads, sequence, head_dim, dtype=dtype, device="cuda") return factory(batch, sequence, heads, head_dim, dtype=dtype, device="cuda").transpose(1, 2) @@ -83,11 +85,17 @@ def _ref_bwd( return o_ref.to(q.dtype), stats_ref.contiguous(), dq.to(q.dtype), dk.to(q.dtype), dv.to(q.dtype) -def _expected_workspace_bytes(batch: int, heads: int, s_q: int, head_dim: int) -> int: +def _expected_workspace_bytes(batch: int, heads: int, s_q: int, head_dim: int, staged: tuple[torch.Tensor, ...] = ()) -> int: + from cudnn.sdpa.bwd.config_sm120 import padded_head_dim from cudnn.sdpa.fwd.api_dsl import ws_align + d_pad = padded_head_dim(head_dim) sq_r = -(-s_q // 128) * 128 - return ws_align(batch * heads * sq_r * 4) + ws_align(batch * sq_r * heads * head_dim * 4) + dq_sem = batch * heads * (-(-s_q // 32)) # int32 relay counters (min q-tile 32) + base = ws_align(batch * heads * sq_r * 4) + ws_align(batch * sq_r * heads * d_pad * 4) + ws_align(dq_sem * 4) + # Padded-width staging copy per non-BSHD-compact operand (all when D pads). + staging = sum(ws_align(t.numel() // head_dim * d_pad * t.element_size()) for t in staged if d_pad != head_dim or not t.transpose(1, 2).is_contiguous()) + return base + staging def _run_bwd_graph( @@ -102,9 +110,11 @@ def _run_bwd_graph( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + deterministic: bool = False, select: bool = True, q_tile: int | None = None, kv_tile: int | None = None, + grad_layout: str = "bshd", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]: """Build and execute the SM120 FROST backward graph; returns (dq, dk, dv, plan_name).""" @@ -115,9 +125,9 @@ def _run_bwd_graph( 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) + dq_gpu = _bhsd(batch, h_q, q_gpu.shape[2], head_dim, dtype, empty=True, layout=grad_layout) + dk_gpu = _bhsd(batch, h_kv, k_gpu.shape[2], head_dim, dtype, empty=True, layout=grad_layout) + dv_gpu = _bhsd(batch, h_kv, v_gpu.shape[2], head_dim, dtype, empty=True, layout=grad_layout) graph = cudnn.pygraph( io_data_type=io_dtype, @@ -147,6 +157,8 @@ def _run_bwd_graph( bwd_kwargs["use_causal_mask"] = True if window_size_left is not None: bwd_kwargs["sliding_window_length"] = window_size_left + 1 + if deterministic: + bwd_kwargs["use_deterministic_algorithm"] = True dq, dk, dv = graph.sdpa_backward(**bwd_kwargs) dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride()) @@ -178,7 +190,8 @@ def _run_bwd_graph( 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) + staged = (q_gpu, k_gpu, v_gpu, o_gpu, do_gpu, dq_gpu, dk_gpu, dv_gpu) + assert workspace_size == _expected_workspace_bytes(batch, h_q, q_gpu.shape[2], head_dim, staged) workspace = torch.empty(max(workspace_size, 1), dtype=torch.uint8, device="cuda") variant_pack = { @@ -212,19 +225,22 @@ def _run_case( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + deterministic: bool = False, select: bool = True, q_tile: int | None = None, kv_tile: int | None = None, + layout: str = "bshd", + grad_layout: str = "bshd", ) -> 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) + q = _bhsd(batch, heads, s_q, head_dim, dtype, layout=layout) + k = _bhsd(batch, heads, s_kv, head_dim, dtype, layout=layout) + v = _bhsd(batch, heads, s_kv, head_dim, dtype, layout=layout) + do = _bhsd(batch, heads, s_q, head_dim, dtype, layout=layout) 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, window_size_left=window_size_left ) - o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True, layout=layout).copy_(o) dq, dk, dv, plan_name = _run_bwd_graph( q, k, @@ -236,9 +252,11 @@ def _run_case( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + deterministic=deterministic, select=select, q_tile=q_tile, kv_tile=kv_tile, + grad_layout=grad_layout, ) tol = _tolerances(dtype) torch.testing.assert_close(dq.float(), dq_ref.float(), **tol) @@ -368,6 +386,86 @@ def test_sdpa_bwd_dsl_sm120_sliding_window_tails(): _run_case(s_q=193, s_kv=257, head_dim=128, is_causal=True, window_size_left=16) +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [192, 256]) +@pytest.mark.parametrize("mask", ["dense", "causal_tl", "causal_br"]) +@torch_fork_set_rng(seed=12) +def test_sdpa_bwd_dsl_sm120_large_d_wrapper(mask: str, head_dim: int): + """D>128: the graph API's hidden-dim surface stops at 128, so the graph + build must be rejected and the direct wrapper serves it (same fallback + pattern as the sq_gt_skv bottom-right case above).""" + + _require_dsl() + import cudnn + + is_causal = mask != "dense" + causal_bottom_right = mask == "causal_br" + try: + _run_case(head_dim=head_dim, is_causal=is_causal, causal_bottom_right=causal_bottom_right) + return + except cudnn.cudnnGraphNotSupportedError as exc: + assert "hidden_dim" in str(exc), f"unexpected graph rejection: {exc}" + + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, s_q, s_kv, dtype = 2, 4, 512, 512, torch.float16 + 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) + out = sdpa_bwd_wrapper_dsl_sm120(q, k, v, o, do, stats, is_causal=is_causal, causal_bottom_right=causal_bottom_right, scale_softmax=scale) + tol = _tolerances(dtype) + torch.testing.assert_close(out["dq_tensor"].float(), dq_ref.float(), **tol) + torch.testing.assert_close(out["dk_tensor"].float(), dk_ref.float(), **tol) + torch.testing.assert_close(out["dv_tensor"].float(), dv_ref.float(), **tol) + + +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [40, 72, 120]) +@pytest.mark.parametrize("is_causal", [False, True], ids=["dense", "causal"]) +@torch_fork_set_rng(seed=13) +def test_sdpa_bwd_dsl_sm120_padded_head_dim(head_dim: int, is_causal: bool): + """Non-bin head dims run zero-padded in the next bin via the graph path.""" + + _run_case(head_dim=head_dim, is_causal=is_causal) + + +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [136, 200]) +@torch_fork_set_rng(seed=14) +def test_sdpa_bwd_dsl_sm120_padded_head_dim_wrapper(head_dim: int): + """Non-bin head dims above the graph API's 128 cap: graph build is + rejected, the direct wrapper serves them zero-padded.""" + + _require_dsl() + import cudnn + + try: + _run_case(head_dim=head_dim, is_causal=True) + return + except cudnn.cudnnGraphNotSupportedError as exc: + assert "hidden_dim" in str(exc), f"unexpected graph rejection: {exc}" + + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, s_q, s_kv, dtype = 2, 4, 512, 512, torch.float16 + 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=True) + o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + out = sdpa_bwd_wrapper_dsl_sm120(q, k, v, o, do, stats, is_causal=True, scale_softmax=scale) + tol = _tolerances(dtype) + torch.testing.assert_close(out["dq_tensor"].float(), dq_ref.float(), **tol) + torch.testing.assert_close(out["dk_tensor"].float(), dk_ref.float(), **tol) + torch.testing.assert_close(out["dv_tensor"].float(), dv_ref.float(), **tol) + + @pytest.mark.L0 @torch_fork_set_rng(seed=5) def test_sdpa_bwd_dsl_sm120_auto_routing(): @@ -377,6 +475,146 @@ def test_sdpa_bwd_dsl_sm120_auto_routing(): assert plan_name == ENGINE +@pytest.mark.L0 +@torch_fork_set_rng(seed=8) +def test_sdpa_bwd_dsl_sm120_dense_flex_bhsd_contiguous(): + """BHSD-contiguous operands (dense_flex): served via workspace-carved compact + staging copies — a gather for the inputs, a scatter-back for dQ/dK/dV.""" + + _run_case(head_dim=64, is_causal=True, layout="bhsd", grad_layout="bhsd") + + +@pytest.mark.L0 +@pytest.mark.parametrize("mask", ["dense", "causal_tl", "causal_br", "swa"]) +@torch_fork_set_rng(seed=12) +def test_sdpa_bwd_dsl_sm120_deterministic_numeric(mask: str): + """use_deterministic_algorithm=True routes to the engine and keeps parity + (default tolerances) for every mask family the kernel serves.""" + + _run_case( + s_q=384 if mask == "causal_br" else 512, + s_kv=1024 if mask == "causal_br" else 512, + head_dim=64, + is_causal=mask != "dense", + causal_bottom_right=mask == "causal_br", + window_size_left=127 if mask == "swa" else None, + deterministic=True, + ) + + +def _run_bitwise_case(n_runs: int = 3, **case_kwargs) -> None: + """Same inputs, ``n_runs`` independent graph runs: outputs must be bitwise equal.""" + + batch, heads, dtype = 2, 4, torch.float16 + s_q = case_kwargs.pop("s_q", 1024) + s_kv = case_kwargs.pop("s_kv", 1024) + head_dim = case_kwargs.pop("head_dim", 64) + 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, _, _, _ = _ref_bwd( + q, + k, + v, + do, + scale=scale, + is_causal=case_kwargs.get("is_causal", False), + causal_bottom_right=case_kwargs.get("causal_bottom_right", False), + window_size_left=case_kwargs.get("window_size_left"), + ) + o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + runs = [_run_bwd_graph(q, k, v, o, do, stats, scale=scale, deterministic=True, **case_kwargs) for _ in range(n_runs)] + dq0, dk0, dv0, _ = runs[0] + for run_i, (dq, dk, dv, _) in enumerate(runs[1:], start=1): + assert torch.equal(dq, dq0), f"run {run_i}: dQ is not bitwise reproducible" + assert torch.equal(dk, dk0), f"run {run_i}: dK is not bitwise reproducible" + assert torch.equal(dv, dv0), f"run {run_i}: dV is not bitwise reproducible" + + +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("mask", ["dense", "causal", "swa"]) +@torch_fork_set_rng(seed=13) +def test_sdpa_bwd_dsl_sm120_deterministic_bitwise(head_dim: int, mask: str): + """Repeated deterministic runs are bitwise identical (dQ relay ordering).""" + + _run_bitwise_case( + head_dim=head_dim, + is_causal=mask != "dense", + window_size_left=127 if mask == "swa" else None, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=14) +def test_sdpa_bwd_dsl_sm120_deterministic_bitwise_tails_knobs(): + """Bitwise reproducibility with partial tails and a non-default tile knob.""" + + _run_bitwise_case(s_q=1000, s_kv=999, head_dim=64, is_causal=True, q_tile=128, kv_tile=64) + + +def _run_wrapper_det_case(head_dim: int, *, s_q: int, s_kv: int, is_causal: bool, window_size_left: int | None, n_runs: int = 1): + """Deterministic run(s) through the direct wrapper (D>128 has no graph + surface); returns (outputs per run, references).""" + + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, dtype = 2, 4, torch.float16 + 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, window_size_left=window_size_left) + o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + runs = [ + sdpa_bwd_wrapper_dsl_sm120(q, k, v, o, do, stats, is_causal=is_causal, window_size_left=window_size_left, deterministic=True, scale_softmax=scale) + for _ in range(n_runs) + ] + return runs, (dq_ref, dk_ref, dv_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=15) +def test_sdpa_bwd_dsl_sm120_deterministic_large_d_numeric(): + """Deterministic relay on the single-Q-buffer branch (D=256 -> q32, + Q_STAGES == 1): causal + sliding window + non-tile-multiple tails, vs ref.""" + + _require_dsl() + runs, (dq_ref, dk_ref, dv_ref) = _run_wrapper_det_case(256, s_q=1000, s_kv=1000, is_causal=True, window_size_left=127) + tol = _tolerances(torch.float16) + out = runs[0] + torch.testing.assert_close(out["dq_tensor"].float(), dq_ref.float(), **tol) + torch.testing.assert_close(out["dk_tensor"].float(), dk_ref.float(), **tol) + torch.testing.assert_close(out["dv_tensor"].float(), dv_ref.float(), **tol) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("head_dim", "mask"), + [(192, "causal"), (256, "causal"), (256, "swa")], +) +@torch_fork_set_rng(seed=16) +def test_sdpa_bwd_dsl_sm120_deterministic_large_d_bitwise(head_dim: int, mask: str): + """Repeated deterministic runs are bitwise identical on the q32 large-D path.""" + + _require_dsl() + runs, _ = _run_wrapper_det_case( + head_dim, + s_q=1024, + s_kv=1024, + is_causal=mask != "dense", + window_size_left=127 if mask == "swa" else None, + n_runs=3, + ) + first = runs[0] + for run_i, out in enumerate(runs[1:], start=1): + for grad in ("dq_tensor", "dk_tensor", "dv_tensor"): + assert torch.equal(out[grad], first[grad]), f"run {run_i}: {grad} is not bitwise reproducible (D={head_dim}, {mask})" + + @pytest.mark.L0 @pytest.mark.parametrize( ("head_dim", "q_tile", "kv_tile"), diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 66aff1bb7..ce2bdceef 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -959,6 +959,10 @@ def test_bwd_facts_kv_transposed_view_canonicalized(): facts = _facts(_mk_bwd_graph(kv_transposed_view=True)) assert (facts.s_kv, facts.d_qk) == (S, _BWD_D) assert facts.bshd_layout + # port_layouts (what bwd lowering consumes) has the rewrite undone too. + ports = {name: (dim, stride) for name, dim, stride in facts.port_layouts} + assert ports["k"] == ((B, H, S, _BWD_D), _bshd_strides(H, S, _BWD_D)) + assert ports["v"] == ((B, H, S, _BWD_D), _bshd_strides(H, S, _BWD_D)) def test_bwd_probe_accepts(monkeypatch): @@ -987,9 +991,10 @@ def test_bwd_probe_rejects_gqa(monkeypatch): def test_bwd_probe_rejects_unsupported_head_dim(monkeypatch): + # Envelope: any multiple of 8 up to 256 (adapter pads); reject the rest. 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)) + assert not _bwd_eligible(_mk_bwd_graph(d=100)) + assert not _bwd_eligible(_mk_bwd_graph(d=264)) def test_bwd_probe_causal_notches(monkeypatch): @@ -1002,9 +1007,10 @@ def test_bwd_probe_causal_notches(monkeypatch): assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(s_q=S // 2, use_causal_mask_bottom_right=True, sliding_window_length=64)) -def test_bwd_probe_rejects_deterministic(monkeypatch): +def test_bwd_probe_accepts_deterministic(monkeypatch): + # use_deterministic_algorithm is served by the ordered-relay dQ path. monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) - assert not _bwd_eligible(_mk_bwd_graph(use_deterministic_algorithm=True)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(use_deterministic_algorithm=True)) def test_bwd_probe_rejects_bias(monkeypatch): @@ -1017,12 +1023,15 @@ def test_bwd_probe_rejects_dbias(monkeypatch): assert not _bwd_eligible(_mk_bwd_graph(dbias=True)) -def test_bwd_probe_rejects_non_bshd_layout(monkeypatch): +def test_bwd_probe_accepts_dense_flex_layouts(monkeypatch): + # Same dense_flex envelope as the forward rows: any B/H/S order with the + # head dim innermost; the adapter stages to compact BSHD. 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)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(grad_strides=bhsd_contig)) + # Head dim NOT innermost (S innermost instead) is outside dense_flex. + s_innermost = (H * S * _BWD_D, S * _BWD_D, 1, S) + assert not _bwd_eligible(_mk_bwd_graph(grad_strides=s_innermost)) def test_bwd_probe_rejects_strided_stats(monkeypatch): @@ -1050,12 +1059,10 @@ def test_bwd_mismatch_reason_strings(monkeypatch): 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(d=100))) + assert reason is not None and "100" in reason reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(use_causal_mask=True, use_alibi_mask=True))) assert reason is not None and "ALiBi" 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))