From 08284d2bb783099638c3e37467960145ecddf390 Mon Sep 17 00:00:00 2001 From: barretw Date: Mon, 17 Aug 2026 21:54:51 -0700 Subject: [PATCH 1/2] add MLA --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 40 +- python/cudnn/sdpa/bwd/api_dsl.py | 124 ++--- python/cudnn/sdpa/bwd/config_sm120.py | 15 +- python/cudnn/sdpa/bwd/engines.py | 3 + .../cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 433 +++++++++++------- .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 213 ++++++++- 6 files changed, 568 insertions(+), 260 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 6aefb562a..4a9831aba 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -145,16 +145,25 @@ warp-partition triple: | 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. +SMEM per CTA is `Q_STAGES·tile_q·d_qk + tile_q·d_v + tile_kv·d_qk + +max(tile_kv·d_v, 2·tile_q·tile_kv)` 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. + +The Q/K head dim may exceed the V head dim (MLA: DeepSeek-V3 and Kimi-K2.6 +train at 192/128). `d_qk` sizes Q/K/dQ/dK and `d_v` sizes V/O/dO/dV: GEMM1 +contracts over `d_qk`, GEMM2 over `d_v`, and dK/dV share one warp partition +with per-side column slices. Tile defaults come from `d_qk`. Unequal dims +must both be multiples of 64 (one smem page/swizzle); the adapter pads each +side to its own native kernel size and raises the VO side to at least 64 +when the sizes differ. ### dQ scatter and the scrambled workspace @@ -236,7 +245,9 @@ only the unused relay operand remains in the kernel ABI. - 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`) + served by zero-padding D to the next supported size. Rectangular + `d_qk >= d_v` (MLA, e.g. 192/128) is supported: each side pads to its own + native size (the VO side raises to >= 64 when the sizes differ) - Masks: none, causal (top-left or bottom-right), right-band-widened causal (`diagonal_band_right_bound` > 0, the causal diagonal shifted right by a compile-time R), sliding window (left-window offset, with or without @@ -251,6 +262,7 @@ only the unused relay operand remains in the kernel ABI. - No dropout / bias / ALiBi / softcap / THD - Workspace (carved from the caller's buffer): fp32 `delta` and `dq_accum` scratch plus int32 relay-counter storage (reserved in both modes); GQA adds - the io-dtype `dk_ws`/`dv_ws` partials buffers (`B·S_kv·H_q·D_padded` - elements each, where `D_padded` is the adapter's zero-padded head - dimension); use `scratch_workspace_bytes()` for the exact total + the io-dtype `dk_ws`/`dv_ws` partials buffers (`B·S_kv·H_q·d_qk_padded` and + `B·S_kv·H_q·d_v_padded` elements, where `d_*_padded` are the adapter's + zero-padded head dimensions); use `scratch_workspace_bytes()` for the exact + total diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 7702552fc..668e0be4b 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -23,6 +23,7 @@ SUPPORTED_HEAD_DIMS as _SM120_SUPPORTED_HEAD_DIMS, TemplateParams as Sm120TemplateParams, padded_head_dim as _sm120_padded_head_dim, + padded_head_dims as _sm120_padded_head_dims, ) from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, _torch_stream_context, ws_align @@ -111,7 +112,8 @@ def __init__( 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.head_dim_qk: Optional[int] = None + self.head_dim_v: Optional[int] = None self.dtype: Optional[torch.dtype] = None self._initialize_implementation() self._logger.debug("__init__ completed") @@ -157,8 +159,11 @@ 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 + # Kernel-facing head dims per side (QK: Q/K/dQ/dK, V: V/O/dO/dV). + # When one differs from its D, that side's operands stage through + # zero-padded compact copies. + self.head_dim_qk_padded: Optional[int] = None + self.head_dim_v_padded: Optional[int] = None # name -> staging number-of-elements for each non-BSHD-compact port self._staging_numels: dict[str, int] = {} @@ -189,29 +194,41 @@ def check_support(self) -> bool: b, h_q, s_q, d_qk = self.q_desc.shape _, h_kv, s_kv, _ = self.k_desc.shape + d_v = int(self.v_desc.shape[3]) 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.v_desc, (b, h_kv, s_kv, d_v), name="V") + self._check_tensor_shape(self.o_desc, (b, h_q, s_q, d_v), name="O") + self._check_tensor_shape(self.do_desc, (b, h_q, s_q, d_v), 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)): + for label, val in (("B", b), ("H_q", h_q), ("H_kv", h_kv), ("S_q", s_q), ("S_kv", s_kv), ("D_QK", d_qk), ("D_V", d_v)): self._value_error_if(int(val) <= 0, f"{label} must be > 0; got {val}") self._value_error_if( h_q % h_kv != 0, f"SM120 DSL SDPA backward requires H_q to be a multiple of H_kv (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( - self.head_dim_padded is None, - f"D ({d_qk}) must be a multiple of 8 and <= {max(_SM120_SUPPORTED_HEAD_DIMS)}", + d_v > d_qk, + f"SM120 DSL SDPA backward requires D_QK >= D_V (MLA-style rectangular head dims); got D_QK={d_qk}, D_V={d_v}", ) - # 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( + d_qk % 8 != 0 or _sm120_padded_head_dim(int(d_qk)) is None, + f"D_QK ({d_qk}) must be a multiple of 8 and <= {max(_SM120_SUPPORTED_HEAD_DIMS)}", + ) + self._value_error_if( + d_v % 8 != 0 or _sm120_padded_head_dim(int(d_v)) is None, + f"D_V ({d_v}) must be a multiple of 8 and <= {max(_SM120_SUPPORTED_HEAD_DIMS)}", + ) + self.head_dim_qk_padded, self.head_dim_v_padded = _sm120_padded_head_dims(int(d_qk), int(d_v)) + # A side's operands all stage when its D pads; otherwise only the non-BSHD-compact ones. + for desc in (self.q_desc, self.k_desc, self.dq_desc, self.dk_desc): + if self.head_dim_qk_padded != d_qk or not self._bshd_physical_ok(desc): + self._staging_numels[desc.name] = math.prod(desc.shape[:-1]) * self.head_dim_qk_padded + for desc in (self.v_desc, self.o_desc, self.do_desc, self.dv_desc): + if self.head_dim_v_padded != d_v or not self._bshd_physical_ok(desc): + self._staging_numels[desc.name] = math.prod(desc.shape[:-1]) * self.head_dim_v_padded self._value_error_if( self.stats_desc.ndim != 4 or tuple(self.stats_desc.shape) != (b, h_q, s_q, 1), @@ -304,7 +321,8 @@ def check_support(self) -> bool: 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.head_dim_qk = int(d_qk) + self.head_dim_v = int(d_v) self._sq_rounded = _round_up(self.s_q_max, _SM120_ROW_ROUND) self._is_supported = True @@ -340,8 +358,9 @@ def compile(self) -> None: qh=self.h_q, sq=self.s_q_max, skv=self.s_k_max, - d=self.head_dim_padded, + d_qk=self.head_dim_qk_padded, kvh=self.h_kv, + d_v=self.head_dim_v_padded, ) self._logger.debug("compile completed") @@ -350,13 +369,14 @@ def _dq_sem_len(self) -> int: return self.batch_size * self.h_q * _round_up(self.s_q_max, _SM120_MIN_Q_TILE) // _SM120_MIN_Q_TILE - def _dkv_ws_elems(self) -> int: - """io-dtype elements of each GQA partials buffer (dk_ws / dv_ws); - 0 for MHA, where they alias the dk/dv outputs.""" + def _dkv_ws_elems(self) -> tuple[int, int]: + """io-dtype elements of the GQA partials buffers (dk_ws, dv_ws); + (0, 0) for MHA, where they alias the dk/dv outputs.""" if self.h_q == self.h_kv: - return 0 - return self.batch_size * self.s_k_max * self.h_q * self.head_dim_padded + return (0, 0) + rows = self.batch_size * self.s_k_max * self.h_q + return (rows * self.head_dim_qk_padded, rows * self.head_dim_v_padded) def _checked_seq_lens(self, seq_lens: torch.Tensor, name: str) -> torch.Tensor: """Validate per-batch lengths and return a (B,) int32 view (never a copy/cast).""" @@ -379,16 +399,17 @@ def _checked_seq_lens(self, seq_lens: torch.Tensor, name: str) -> torch.Tensor: return seq_lens.reshape(-1) 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_QK]) + dq_sem (int32 flat [B*H*ceil(SQ/32)], deterministic relay counters) - + dk_ws/dv_ws (io [B, SKV, H_q, D] each, per-q-head partials, GQA only) - + one compact staging copy per non-BSHD-compact operand.""" + + dk_ws/dv_ws (io [B, SKV, H_q, D_QK] / [B, SKV, H_q, D_V], per-q-head + partials, GQA only) + 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_padded * 4) + dq_accum_bytes = ws_align(self.batch_size * self._sq_rounded * self.h_q * self.head_dim_qk_padded * 4) dq_sem_bytes = ws_align(self._dq_sem_len() * 4) - dkv_ws_bytes = 2 * ws_align(self._dkv_ws_elems() * self.dtype.itemsize) + dkv_ws_bytes = sum(ws_align(elems * self.dtype.itemsize) for elems in self._dkv_ws_elems()) 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 + dkv_ws_bytes + staging_bytes @@ -446,7 +467,7 @@ 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_padded, torch.float32) + dq_accum = carver.take(self.batch_size * self._sq_rounded * self.h_q * self.head_dim_qk_padded, torch.float32) dq_sem = carver.take(self._dq_sem_len(), torch.int32) if current_stream is None: @@ -457,26 +478,27 @@ def execute( import cutlass - # 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 + # Non-compact operands (a whole side when its D pads) stage through + # workspace-carved compact copies with zero-filled pad columns. The QK + # side (Q/dQ/K/dK) and the VO side (V/O/dO/dV) pad independently. + d_qk, dqk_pad = self.head_dim_qk, self.head_dim_qk_padded + d_v, dv_pad = self.head_dim_v, self.head_dim_v_padded - def _staged_bshd(tensor: torch.Tensor) -> torch.Tensor: + def _staged_bshd(tensor: torch.Tensor, d_orig: int, d_pad: int) -> torch.Tensor: view = tensor.transpose(1, 2) - if not pads and view.is_contiguous(): + if d_pad == d_orig 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) + if d_pad != d_orig: + staged[..., d_orig:].zero_() + staged[..., :d_orig].copy_(view) return staged - def _staged_out_bshd(tensor: torch.Tensor): + def _staged_out_bshd(tensor: torch.Tensor, d_orig: int, d_pad: int): """(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(): + if d_pad == d_orig 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 @@ -485,14 +507,14 @@ def _staged_out_bshd(tensor: torch.Tensor): seq_kv_t = self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else None 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) + q = _staged_bshd(q_tensor, d_qk, dqk_pad) + k = _staged_bshd(k_tensor, d_qk, dqk_pad) + v = _staged_bshd(v_tensor, d_v, dv_pad) + o = _staged_bshd(o_tensor, d_v, dv_pad) + do = _staged_bshd(do_tensor, d_v, dv_pad) + dq, dq_user = _staged_out_bshd(dq_tensor, d_qk, dqk_pad) + dk, dk_user = _staged_out_bshd(dk_tensor, d_qk, dqk_pad) + dv, dv_user = _staged_out_bshd(dv_tensor, d_v, dv_pad) lse = stats_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) kernels = self._compiled_kernel @@ -501,9 +523,9 @@ def _staged_out_bshd(tensor: torch.Tensor): if self.h_q == self.h_kv: dk_ws, dv_ws = dk, dv else: - ws_shape = (self.batch_size, self.s_k_max, self.h_q, self.head_dim_padded) - dk_ws = carver.take(self._dkv_ws_elems(), self.dtype).view(ws_shape) - dv_ws = carver.take(self._dkv_ws_elems(), self.dtype).view(ws_shape) + dkw_elems, dvw_elems = self._dkv_ws_elems() + dk_ws = carver.take(dkw_elems, self.dtype).view(self.batch_size, self.s_k_max, self.h_q, dqk_pad) + dv_ws = carver.take(dvw_elems, self.dtype).view(self.batch_size, self.s_k_max, self.h_q, dv_pad) # Kernel chain (dot -> main -> [reduce] -> cvt) kernels.dot(o, do, delta, dq_accum, dq_sem, current_stream) @@ -530,9 +552,9 @@ def _staged_out_bshd(tensor: torch.Tensor): kernels.cvt(dq_accum, dq, cutlass.Float32(scale_val), current_stream) if kernels.dsink is not None: kernels.dsink(lse, delta, sink_tensor.reshape(self.h_q), dsink_tensor.reshape(self.h_q), seq_q_t, current_stream) - for user_view, staged in ((dq_user, dq), (dk_user, dk), (dv_user, dv)): + for user_view, staged, d_orig in ((dq_user, dq, d_qk), (dk_user, dk, d_qk), (dv_user, dv, d_v)): if user_view is not None: - user_view.copy_(staged[..., : self.head_dim]) + user_view.copy_(staged[..., :d_orig]) def _tensor_signature(tensor: torch.Tensor) -> tuple: diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py index 7db5a1386..0cb006b58 100644 --- a/python/cudnn/sdpa/bwd/config_sm120.py +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -15,11 +15,24 @@ def padded_head_dim(d: int) -> "int | None": - """Smallest native bin >= ``d``, or ``None`` when ``d`` exceeds every bin.""" + """Smallest native kernel head-dim size >= ``d``, or ``None`` when ``d`` exceeds them all.""" return min((b for b in SUPPORTED_HEAD_DIMS if b >= d), default=None) +def padded_head_dims(d_qk: int, d_v: int) -> "tuple[int, int] | None": + """Native kernel head-dim sizes for a head-dim pair.""" + + d_qk_pad = padded_head_dim(d_qk) + d_v_pad = padded_head_dim(d_v) + if d_qk_pad is None or d_v_pad is None: + return None + # Unequal dims must both be multiples of 64 (one smem swizzle). + if d_v_pad != d_qk_pad: + d_v_pad = max(d_v_pad, 64) + return d_qk_pad, d_v_pad + + @dataclass(frozen=True) class TemplateParams: """Per-graph parameters that change the traced SM120 backward kernel. diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index dc1298c77..f26e8581a 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -179,6 +179,8 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", requested: return f"head dims (D_QK={facts.d_qk}, D_V={facts.d_v}) exceed the {max(capabilities.d)} envelope" elif facts.d_qk not in capabilities.d: return f"serves D in {sorted(capabilities.d)}; graph has D={facts.d_qk}" + elif facts.d_v not in capabilities.d: + return f"serves D_V in {sorted(capabilities.d)}; graph has D_V={facts.d_v}" 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: @@ -261,6 +263,7 @@ def _sm120_spec() -> EngineSpec: sm_hi=_BLACKWELL_GEFORCE[1], # Any head size multipled of 8 d=frozenset(range(8, max(_SM120_HEAD_DIMS) + 1, 8)), + dqk_ge_dv=True, dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), gqa=True, causal=True, diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index e26e2728d..8434a0f99 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -66,21 +66,23 @@ def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b -def largest_warp_partition(m_dim: int, n_dim: int) -> int: +def largest_warp_partition(m_dim: int, *n_dims: 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. + pairs). The dK/dV pair shares one partition, so it passes both head + dims as ``n_dims`` and A must validate for each. 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: + if m_dim % (16 * a) == 0 and all((n_dim * a // 8) % 16 == 0 for n_dim in n_dims): return a - raise ValueError(f"no valid warp partition for M{m_dim} N{n_dim}") + raise ValueError(f"no valid warp partition for M{m_dim} N{n_dims}") @cute.jit @@ -323,7 +325,7 @@ def _bwd_dq_scatter( DQ_REPS: cutlass.Constexpr[int], DQ_NF: cutlass.Constexpr[int], M: cutlass.Constexpr[int], - d: cutlass.Constexpr[int], + d_qk: cutlass.Constexpr[int], ): """dQ accumulate into the scrambled dq_accum workspace.""" t_r = math_tidx // 32 @@ -332,12 +334,12 @@ def _bwd_dq_scatter( 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): + if cutlass.const_expr(d_qk >= 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 + addr = dqa_base + (t_r + jm * 8) * (H * d_qk) + t_c * 2 + jn * 64 else: - addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d) + (t_c % 16) * 2 + addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d_qk) + (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]) @@ -357,7 +359,7 @@ def _bwd_gemm5_dk( M: cutlass.Constexpr[int], PDS: cutlass.Constexpr[int], PAGE: cutlass.Constexpr[int], - DKV_PER: cutlass.Constexpr[int], + DK_PER: cutlass.Constexpr[int], io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], ): """GEMM 5: acc_dk += dS^T @ Q (the iteration's last sQ reader).""" @@ -379,13 +381,13 @@ def _bwd_gemm5_dk( sQ_st, b_k_step=kc, M=16 * DKV_REPS, - N=DKV_PER, + N=DK_PER, b_trans=True, b_rows=M, b_page=PAGE, lane=lane, ab_dtype=io_dtype, - col_base=wd_k * DKV_PER, + col_base=wd_k * DK_PER, ) @@ -446,7 +448,7 @@ class SM120FusedMultiHeadAttentionFP16Backward: 192: (32, 64), 256: (32, 64), } - # (d, q_tile, kv_tile) -> (warps_m_sdp, warps_m_dkv, warps_m_dq): for each + # (d_qk, 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 = { @@ -467,7 +469,8 @@ def __init__( window_size_left: int | None = None, window_size_right: int | None = None, deterministic: bool = False, - head_dim: int = 128, + head_dim_qk: int = 128, + head_dim_v: int = 0, # 0 = same as head_dim_qk. use_pdl: bool = True, q_tile: int = 0, kv_tile: int = 0, @@ -486,9 +489,13 @@ def __init__( self.seq_q_lens_present = bool(seq_q_lens_present) # sink LSE is finite on padded rows; trim them explicitly (LSE := +inf, P = 0) self.trim_q_rows = bool(sink_present) and self.seq_q_lens_present - self.d = head_dim + self.d_qk = head_dim_qk + self.d_v = int(head_dim_v) or head_dim_qk + # current MLA requires both to be multiples of 64 so one smem swizzle serves every tile. + if self.d_v != head_dim_qk and (head_dim_qk % 64 or self.d_v % 64): + raise ValueError(f"unequal head dims must both be multiples of 64; got d_qk={head_dim_qk}, d_v={self.d_v}") self.use_pdl = bool(use_pdl) - self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim] + self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim_qk] if q_tile: self.q_tile = int(q_tile) if kv_tile: @@ -499,34 +506,37 @@ def __init__( 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)) + tuned = self.CONFIG.get((head_dim_qk, self.q_tile, self.kv_tile)) if tuned is not None: self.warps_m_sdp, self.warps_m_dkv, self.warps_m_dq = tuned + # The tuned dK partition does not tile the dV head dim. + if (self.d_v * self.warps_m_dkv // 8) % 16: + self.warps_m_dkv = largest_warp_partition(self.kv_tile, head_dim_qk, self.d_v) 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 + self.warps_m_dkv = largest_warp_partition(self.kv_tile, head_dim_qk, self.d_v) + self.warps_m_dq = largest_warp_partition(self.q_tile, head_dim_qk) + M_, N_, d_qk_ = self.q_tile, self.kv_tile, head_dim_qk 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"), + (self.warps_m_dkv, N_, d_qk_, "warps_m_dkv"), + (self.warps_m_dkv, N_, self.d_v, "warps_m_dkv (dV)"), + (self.warps_m_dq, M_, d_qk_, "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 + raise ValueError(f"invalid {tag}={a_} for M{M_} N{N_} d_qk{d_qk_} d_v{self.d_v}") + self.page = 64 if head_dim_qk % 64 == 0 and self.d_v % 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 + M, N, d_qk, d_v = self.q_tile, self.kv_tile, self.d_qk, self.d_v # 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) + smem_elems = q_stages * M * d_qk + M * d_v + N * d_qk + max(N * d_v, 2 * M * N) if smem_elems * in_dtype.bytes <= cap: break else: @@ -534,18 +544,18 @@ def __init__( # 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_sdO = q_stages * M * d_qk + self.off_sK = self.off_sdO + M * d_v + self.off_sV = self.off_sK + N * d_qk self.off_sdS = self.off_sV # aliases sV (V is in regs) self.off_sP = self.off_sV + M * N self.smem_elems = smem_elems @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.""" + def load_tma_tile(self, s_dst, tma_desc, mbar, batch, head, seq, rows: cutlass.Constexpr[int], cols: cutlass.Constexpr[int]): + """Load one paged/swizzled `(rows, cols)` tile with TMA.""" elems_per_page = rows * self.page - for pg in cutlass.range_constexpr(self.tma_copy_iters): + for pg in cutlass.range_constexpr(cols // self.page): if prims.elect_sync(): prims.cp_async_bulk_tensor_shared_cta_global( s_dst.subview(pg * elems_per_page), @@ -559,14 +569,14 @@ def kernel( self, q: cute.Tensor, # [B, SQ, HQ, D] io dtype (BSHD) k: cute.Tensor, # [B, SKV, HKV, D] - v: cute.Tensor, # [B, SKV, HKV, D] - do: cute.Tensor, # [B, SQ, HQ, D] + v: cute.Tensor, # [B, SKV, HKV, DV] + do: cute.Tensor, # [B, SQ, HQ, DV] lse: cute.Tensor, # [B, HQ, SQ] fp32 (natural-log LSE) delta: cute.Tensor, # [B, HQ, SQ_r128] fp32 (dot_do_o output) dq_accum: cute.Tensor, # [B*SQ_r128*HQ*D] fp32 (scrambled, zeroed) dq_sem: cute.Tensor, # [B*HQ*num_q_tiles] int32 relay turn counters, one per (batch, head, q-tile); zeroed by dot (deterministic only) dk_ws: cute.Tensor, # [B, SKV, HQ, D] dK destination: dk itself when MHA; per-q-head partials summed by _dkv_reduce_kernel when GQA - dv_ws: cute.Tensor, # [B, SKV, HQ, D] dV destination (same as dK) + dv_ws: cute.Tensor, # [B, SKV, HQ, DV] dV destination (same as dK) seq_q_lens: Optional[cute.Tensor], # [B] int32 per-batch Q lengths; None unless seq_q_lens_present seq_kv_lens: Optional[cute.Tensor], # [B] int32 per-batch KV lengths; None unless seq_kv_lens_present tma_q_desc: cutlass.GridConstant[cuda.TensorMap], @@ -577,7 +587,8 @@ def kernel( attn_scale: cutlass.Float32, # linear scale (dq/dk output) ) -> None: io_dtype = self.in_dtype - d = self.d + d_qk = self.d_qk + d_v = self.d_v M = self.q_tile N = self.kv_tile PAGE = self.page @@ -594,14 +605,17 @@ def kernel( 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 + DK_PER = d_qk * WM_DKV // 8 + DK_NF = DK_PER // 8 + DV_PER = d_v * WM_DKV // 8 + DV_NF = DV_PER // 8 # dQ: warp (wq, wdq). DQ_REPS = M // (16 * WM_DQ) - DQ_PER = d * WM_DQ // 8 + DQ_PER = d_qk * WM_DQ // 8 DQ_NF = DQ_PER // 8 - D_CHUNKS = d // 16 # SdP k-reduce + DQK_CHUNKS = d_qk // 16 # S = Q @ K^T k-reduce (D_QK) + DV_CHUNKS = d_v // 16 # dP = dO @ V^T k-reduce (D_V) 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 @@ -620,7 +634,8 @@ def kernel( GROUP = HQ // HKV # query heads per KV head (1 = plain MHA) SQ_R = ((SQ + 127) // 128) * 128 kv_base = n_block * N - q_row_stride = HQ * d # row stride of HQ-headed BSHD tensors (Q side; also dk_ws/dv_ws, whose head axis is HQ) + qk_row_stride = HQ * d_qk # row stride of HQ-headed D_QK-wide BSHD tensors (Q; also dk_ws, whose head axis is HQ) + v_row_stride = HQ * d_v # row stride of HQ-headed D_V-wide BSHD tensors (dO; also dv_ws) # Per-batch actual lengths (Padding mask) seqlen_q = SQ @@ -671,10 +686,10 @@ def kernel( n_iters = cutlass.Int32(0) smem = cutlass.Array(io_dtype, self.smem_elems, space=cutlass.AddressSpace.smem, alignment=128) - 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 + sQ = smem # Q_STAGES * M * d_qk + sdO = smem.subview(self.off_sdO) # M * d_qk + sK = smem.subview(self.off_sK) # N * d_qk + sV = smem.subview(self.off_sV) # N * d_qk 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) @@ -706,17 +721,17 @@ def kernel( 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, kv_head, kv_base, rows=N) - self.load_tma_tile(sK, tma_k_desc, k_mbar, batch, kv_head, kv_base, rows=N) + prims.mbarrier_arrive_expect_tx(v_mbar, N * d_v * io_dtype.bytes) + prims.mbarrier_arrive_expect_tx(k_mbar, N * d_qk * io_dtype.bytes) + self.load_tma_tile(sV, tma_v_desc, v_mbar, batch, kv_head, kv_base, rows=N, cols=d_v) + self.load_tma_tile(sK, tma_k_desc, k_mbar, batch, kv_head, kv_base, rows=N, cols=d_qk) 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, q_head, (m_block_max - 1) * M, rows=M) + prims.mbarrier_arrive_expect_tx(q_full, M * d_qk * io_dtype.bytes) + self.load_tma_tile(sQ, tma_q_desc, q_full, batch, q_head, (m_block_max - 1) * M, rows=M, cols=d_qk) if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx(do_full, M * d * io_dtype.bytes) + prims.mbarrier_arrive_expect_tx(do_full, M * d_v * io_dtype.bytes) self.load_tma_tile( sdO, tma_do_desc, @@ -725,6 +740,7 @@ def kernel( q_head, (m_block_max - 1) * M, rows=M, + cols=d_v, ) while not prims.mbarrier_try_wait_parity(v_mbar, cutlass.Int32(0)): pass @@ -751,29 +767,30 @@ def kernel( 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) + prims.mbarrier_arrive_expect_tx(next_q_full, M * d_qk * io_dtype.bytes) self.load_tma_tile( - sQ.subview(next_stage * M * d), + sQ.subview(next_stage * M * d_qk), tma_q_desc, next_q_full, batch, q_head, next_m * M, rows=M, + cols=d_qk, ) # 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, q_head, next_m * M, rows=M) + prims.mbarrier_arrive_expect_tx(do_full, M * d_v * io_dtype.bytes) + self.load_tma_tile(sdO, tma_do_desc, do_full, batch, q_head, next_m * M, rows=M, cols=d_v) 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, q_head, next_m * M, rows=M) + prims.mbarrier_arrive_expect_tx(q_full, M * d_qk * io_dtype.bytes) + self.load_tma_tile(sQ, tma_q_desc, q_full, batch, q_head, next_m * M, rows=M, cols=d_qk) elif warp < self.load_warp_id: prims.setmaxregister(240, prims.SetMaxRegisterAction.INCREASE) @@ -826,8 +843,8 @@ def kernel( pass # V -> registers. - v_persist = cutlass.Array(cutlass.Int32, D_CHUNKS * VREG_PAIRS * 4, alignment=16) - for kc in cutlass.range_constexpr(D_CHUNKS): + v_persist = cutlass.Array(cutlass.Int32, DV_CHUNKS * VREG_PAIRS * 4, alignment=16) + for kc in cutlass.range_constexpr(DV_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 @@ -847,10 +864,11 @@ def kernel( # 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 = cutlass.Array(cutlass.Float32, DKV_REPS * DK_NF * 4, alignment=16) + acc_dv = cutlass.Array(cutlass.Float32, DKV_REPS * DV_NF * 4, alignment=16) + for i in cutlass.range_constexpr(DKV_REPS * DK_NF * 4): acc_dk[i] = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(DKV_REPS * DV_NF * 4): acc_dv[i] = cutlass.Float32(0.0) wq = math_warp % WM_DQ @@ -871,7 +889,7 @@ def kernel( stage = j & cutlass.Int32(1) else: stage = cutlass.Int32(0) - sQ_st = sQ.subview(stage * M * d) + sQ_st = sQ.subview(stage * M * d_qk) q_row0 = m_block * M cute.arch.barrier(barrier_id=3, number_of_threads=288) @@ -885,7 +903,7 @@ def kernel( # 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): + for kc in cutlass.range_constexpr(DQK_CHUNKS): af = [] for rep in cutlass.range_constexpr(SDP_REPS): qf = load_a_frag( @@ -994,7 +1012,7 @@ def kernel( # 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): + for kc in cutlass.range_constexpr(DV_CHUNKS): af = [] for rep in cutlass.range_constexpr(SDP_REPS): dof = load_a_frag( @@ -1065,13 +1083,13 @@ def kernel( sdO, b_k_step=kc, M=16 * DKV_REPS, - N=DKV_PER, + N=DV_PER, b_trans=True, b_rows=M, b_page=PAGE, lane=lane, ab_dtype=io_dtype, - col_base=wd_k * DKV_PER, + col_base=wd_k * DV_PER, ) # GEMM3 is the final dO consumer; this rendezvous lets @@ -1109,7 +1127,7 @@ def kernel( M=M, PDS=PDS, PAGE=PAGE, - DKV_PER=DKV_PER, + DK_PER=DK_PER, io_dtype=io_dtype, ) cute.arch.barrier(barrier_id=5, number_of_threads=288) @@ -1149,10 +1167,10 @@ def kernel( if nq0 + r_loc >= seqlen_q: val = cutlass.Float32(float("inf")) lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) - dqa_base = ((batch * SQ_R + q_row0) * HQ + q_head) * d + dqa_base = ((batch * SQ_R + q_row0) * HQ + q_head) * d_qk 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, HQ, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d=d) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, HQ, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d_qk=d_qk) if cutlass.const_expr(self.deterministic): _bwd_det_release(det_sem, m_block, det_turn, warp) else: @@ -1193,10 +1211,10 @@ def kernel( if nq0 + r_loc >= seqlen_q: val = cutlass.Float32(float("inf")) lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) - dqa_base = ((batch * SQ_R + q_row0) * HQ + q_head) * d + dqa_base = ((batch * SQ_R + q_row0) * HQ + q_head) * d_qk 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, HQ, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d=d) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, HQ, DQ_REPS=DQ_REPS, DQ_NF=DQ_NF, M=M, d_qk=d_qk) if cutlass.const_expr(self.deterministic): _bwd_det_release(det_sem, m_block, det_turn, warp) _bwd_gemm5_dk( @@ -1212,7 +1230,7 @@ def kernel( M=M, PDS=PDS, PAGE=PAGE, - DKV_PER=DKV_PER, + DK_PER=DK_PER, io_dtype=io_dtype, ) @@ -1226,42 +1244,68 @@ def kernel( 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 + for nf in cutlass.range_constexpr(max(DK_NF, DV_NF)): 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, - ) + if cutlass.const_expr(nf < DK_NF): + off = (rep * DK_NF + nf) * 4 + c0 = wd_k * DK_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) + if cutlass.const_expr(nf < DV_NF): + off_v = (rep * DV_NF + nf) * 4 + c0_v = wd_k * DV_PER + nf * 8 + 2 * p_lane + tile_ptr(sdV, r0, c0_v, page=PAGE, rows=N).store( + pack_half2(acc_dv[off_v + 0], acc_dv[off_v + 1], io_dtype), + alignment=4, + ) + tile_ptr(sdV, r8, c0_v, page=PAGE, rows=N).store( + pack_half2(acc_dv[off_v + 2], acc_dv[off_v + 3], io_dtype), + alignment=4, + ) cute.arch.barrier(barrier_id=2, number_of_threads=256) # smem -> gmem. dk_ws/dv_ws rows are HQ-headed: dk/dv themselves # when MHA (HQ == HKV and q_head == kv_head), one slot per q head # under GQA — the same addressing covers both. - chunks_per_row = d // _COPY_ELEMS - total = N * chunks_per_row - # workspace's head-dim base offset - whd_base = (batch * SKV + kv_base) * q_row_stride + q_head * d - 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): - w_off = whd_base + row * q_row_stride + col - copy16_smem_to_gmem(tile_ptr(sdK, row, col, page=PAGE, rows=N), dkws_ptr + w_off) - copy16_smem_to_gmem(tile_ptr(sdV, row, col, page=PAGE, rows=N), dvws_ptr + w_off) + if cutlass.const_expr(d_qk == d_v): + chunks_per_row = d_qk // _COPY_ELEMS + total = N * chunks_per_row + # workspace's head-dim base offset + whd_base = (batch * SKV + kv_base) * qk_row_stride + q_head * d_qk + 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): + w_off = whd_base + row * qk_row_stride + col + copy16_smem_to_gmem(tile_ptr(sdK, row, col, page=PAGE, rows=N), dkws_ptr + w_off) + copy16_smem_to_gmem(tile_ptr(sdV, row, col, page=PAGE, rows=N), dvws_ptr + w_off) + else: + # Unequal head dims: dK rows are D_QK-wide and dV rows + # D_V-wide, so the chunk->(row, col) maps differ per side. + k_chunks_per_row = d_qk // _COPY_ELEMS + whd_base_k = (batch * SKV + kv_base) * qk_row_stride + q_head * d_qk + for i in cutlass.range_constexpr(N * k_chunks_per_row // 256): + chunk = i * 256 + math_tidx + row = chunk // k_chunks_per_row + col = (chunk % k_chunks_per_row) * _COPY_ELEMS + if (not cutlass.const_expr(PARTIAL_KV)) or (kv_base + row < SKV): + w_off = whd_base_k + row * qk_row_stride + col + copy16_smem_to_gmem(tile_ptr(sdK, row, col, page=PAGE, rows=N), dkws_ptr + w_off) + v_chunks_per_row = d_v // _COPY_ELEMS + whd_base_v = (batch * SKV + kv_base) * v_row_stride + q_head * d_v + for i in cutlass.range_constexpr(N * v_chunks_per_row // 256): + chunk = i * 256 + math_tidx + row = chunk // v_chunks_per_row + col = (chunk % v_chunks_per_row) * _COPY_ELEMS + if (not cutlass.const_expr(PARTIAL_KV)) or (kv_base + row < SKV): + w_off = whd_base_v + row * v_row_stride + col + copy16_smem_to_gmem(tile_ptr(sdV, row, col, page=PAGE, rows=N), dvws_ptr + w_off) else: prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) @@ -1327,13 +1371,14 @@ def __call__( @cute.kernel def _dot_do_o_kernel( - o: cute.Tensor, # [B, SQ, H, D] - do: cute.Tensor, # [B, SQ, H, D] + o: cute.Tensor, # [B, SQ, H, DV] + do: cute.Tensor, # [B, SQ, H, DV] 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], + d_qk: cutlass.Constexpr[int], # D_QK: dq_accum's head dim + d_v: cutlass.Constexpr[int], # D_V: O/dO's head dim page: cutlass.Constexpr[int], use_pdl: cutlass.Constexpr[bool], deterministic: cutlass.Constexpr[bool], @@ -1345,7 +1390,7 @@ def _dot_do_o_kernel( SQ = o.shape[1] H = o.shape[2] SQ_R = ((SQ + 127) // 128) * 128 - row_stride = H * d + row_stride = H * d_v M = q_tile o_ptr = o.iterator.raw_ptr() @@ -1353,7 +1398,7 @@ def _dot_do_o_kernel( dd_ptr = delta.iterator.raw_ptr() dqa_ptr = dq_accum.iterator.raw_ptr() - base = ((batch * SQ + m_block * M) * H + head) * d + base = ((batch * SQ + m_block * M) * H + head) * d_v dd_base = (batch * H + head) * SQ_R + m_block * M q_left = SQ - m_block * M @@ -1361,7 +1406,7 @@ def _dot_do_o_kernel( rows_per_pass = 256 // tpr col0 = (tidx % tpr) * _COPY_ELEMS row0 = tidx // tpr - n_pages = d // page + n_pages = d_v // page for rp in cutlass.range_constexpr(M // rows_per_pass): row = row0 + rp * rows_per_pass acc = cutlass.Float32(0.0) @@ -1401,10 +1446,10 @@ def _dot_do_o_kernel( ), cutlass.Float32, ) - dqa_base = ((batch * SQ_R + m_block * M) * H + head) * d + dqa_base = ((batch * SQ_R + m_block * M) * H + head) * d_qk 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 + for jn in cutlass.range_constexpr(d_qk // (ztpr * 4)): + addr = dqa_base + (zr0 + im * zrows) * (H * d_qk) + zc0 + jn * ztpr * 4 (dqa_ptr + addr).store(zero4, alignment=16) if cutlass.const_expr(deterministic): @@ -1424,14 +1469,15 @@ def _dot_do_o_host( dq_accum: cute.Tensor, dq_sem: cute.Tensor, q_tile: cutlass.Constexpr[int], - d: cutlass.Constexpr[int], + d_qk: cutlass.Constexpr[int], + d_v: 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, dq_sem, q_tile, d, page, use_pdl, deterministic).launch( + _dot_do_o_kernel(o, do, delta, dq_accum, dq_sem, q_tile, d_qk, d_v, page, use_pdl, deterministic).launch( grid=(m_blocks, o.shape[2], o.shape[0]), block=(256, 1, 1), stream=stream, @@ -1449,7 +1495,7 @@ 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], + d_qk: cutlass.Constexpr[int], page: cutlass.Constexpr[int], warps_m_dq: cutlass.Constexpr[int], attn_scale: cutlass.Float32, @@ -1471,7 +1517,7 @@ def _convert_dq_kernel( M = q_tile WM_DQ = warps_m_dq DQ_REPS = M // (16 * WM_DQ) - DQ_PER = d * WM_DQ // 8 + DQ_PER = d_qk * WM_DQ // 8 DQ_NF = DQ_PER // 8 wq = warp % WM_DQ wd_q = warp // WM_DQ @@ -1479,22 +1525,22 @@ def _convert_dq_kernel( 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) + sdQ = cutlass.Array(io_dtype, M * d_qk, 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 + dqa_base = ((batch * SQ_R + m_block * M) * H + head) * d_qk 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): + if cutlass.const_expr(d_qk >= 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 + addr = dqa_base + (t_r + jm * 8) * (H * d_qk) + t_c * 2 + jn * 64 else: - addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d) + (t_c % 16) * 2 + addr = dqa_base + (t_r + (t_c // 16) * 8 + i_pair * 16) * (H * d_qk) + (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 @@ -1506,9 +1552,9 @@ def _convert_dq_kernel( 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 + row_stride = H * d_qk + g_base = ((batch * SQ + m_block * M) * H + head) * d_qk + chunks_per_row = d_qk // _COPY_ELEMS for i in cutlass.range_constexpr(M * chunks_per_row // 256): chunk = i * 256 + tidx row = chunk // chunks_per_row @@ -1525,7 +1571,7 @@ def _convert_dq_host( dq_accum: cute.Tensor, dq: cute.Tensor, q_tile: cutlass.Constexpr[int], - d: cutlass.Constexpr[int], + d_qk: cutlass.Constexpr[int], page: cutlass.Constexpr[int], warps_m_dq: cutlass.Constexpr[int], attn_scale: cutlass.Float32, @@ -1534,7 +1580,7 @@ def _convert_dq_host( 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( + _convert_dq_kernel(dq_accum, dq, q_tile, d_qk, 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, @@ -1547,13 +1593,46 @@ def _convert_dq_host( # --------------------------------------------------------------------------- +@cute.jit +def _reduce_group_vec( + ws_ptr, + out_ptr, + idx, + hkv, + hq, + *, + d: cutlass.Constexpr[int], + group: cutlass.Constexpr[int], + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +): + """Sum one 16 B output vector over the group's q-head partials (fp32, + fixed order -> deterministic) and store it in the io dtype.""" + VEC = 8 # 8 elements per vector (16 bytes) + pos = idx * VEC + col = pos % d + rowh = pos // d # (b*SKV + s)*HKV + kv_head + kh = rowh % hkv + bs = rowh // hkv + in0 = (bs * hq + kh * group) * d + col + acc = cutlass.Array(cutlass.Float32, VEC) + for e in cutlass.range_constexpr(VEC): + acc[e] = cutlass.Float32(0.0) + for g in cutlass.range_constexpr(group): + w = (ws_ptr + in0 + g * d).load(count=VEC) + for e in cutlass.range_constexpr(VEC): + acc[e] = acc[e] + w[e].to(cutlass.Float32) + vec = cutlass.Vector.from_elements(tuple(acc[e].to(io_dtype) for e in range(VEC)), io_dtype) + (out_ptr + pos).store(vec, alignment=16) + + @cute.kernel def _dkv_reduce_kernel( dk_ws: cute.Tensor, # [B, SKV, HQ, D] io dtype (one dK partial per q head) - dv_ws: cute.Tensor, # [B, SKV, HQ, D] io dtype (one dV partial per q head) + dv_ws: cute.Tensor, # [B, SKV, HQ, DV] io dtype (one dV partial per q head) dk: cute.Tensor, # [B, SKV, HKV, D] io dtype out - dv: cute.Tensor, # [B, SKV, HKV, D] io dtype out - d: cutlass.Constexpr[int], + dv: cute.Tensor, # [B, SKV, HKV, DV] io dtype out + d_qk: cutlass.Constexpr[int], + d_v: cutlass.Constexpr[int], group: cutlass.Constexpr[int], io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], use_pdl: cutlass.Constexpr[bool], @@ -1569,34 +1648,26 @@ def _dkv_reduce_kernel( HKV = dk.shape[2] HQ = HKV * group VEC = 8 # 8 elements per vector (16 bytes) - OUT_VECS = B * SKV * HKV * d // VEC dkws_ptr = dk_ws.iterator.raw_ptr() dvws_ptr = dv_ws.iterator.raw_ptr() dk_ptr = dk.iterator.raw_ptr() dv_ptr = dv.iterator.raw_ptr() gidx = bidx * 256 + tidx # host launch 256 threads - if gidx < OUT_VECS: - output_flat_position = gidx * VEC - col = output_flat_position % d - rowh = output_flat_position // d # (b*SKV + s)*HKV + kv_head - kh = rowh % HKV - bs = rowh // HKV - in0 = (bs * HQ + kh * group) * d + col - acc_k = cutlass.Array(cutlass.Float32, VEC) - acc_v = cutlass.Array(cutlass.Float32, VEC) - for e in cutlass.range_constexpr(VEC): - acc_k[e] = cutlass.Float32(0.0) - acc_v[e] = cutlass.Float32(0.0) - for g in cutlass.range_constexpr(group): - kw = (dkws_ptr + in0 + g * d).load(count=VEC) - vw = (dvws_ptr + in0 + g * d).load(count=VEC) - for e in cutlass.range_constexpr(VEC): - acc_k[e] = acc_k[e] + kw[e].to(cutlass.Float32) - acc_v[e] = acc_v[e] + vw[e].to(cutlass.Float32) - kvec = cutlass.Vector.from_elements(tuple(acc_k[e].to(io_dtype) for e in range(VEC)), io_dtype) - vvec = cutlass.Vector.from_elements(tuple(acc_v[e].to(io_dtype) for e in range(VEC)), io_dtype) - (dk_ptr + output_flat_position).store(kvec, alignment=16) - (dv_ptr + output_flat_position).store(vvec, alignment=16) + if cutlass.const_expr(d_qk == d_v): + OUT_VECS = B * SKV * HKV * d_qk // VEC + if gidx < OUT_VECS: + _reduce_group_vec(dkws_ptr, dk_ptr, gidx, HKV, HQ, d=d_qk, group=group, io_dtype=io_dtype) + _reduce_group_vec(dvws_ptr, dv_ptr, gidx, HKV, HQ, d=d_qk, group=group, io_dtype=io_dtype) + else: + # Unequal head dims: dK and dV vectors index different row widths, so + # the flat thread range covers dK's vectors first, then dV's. + K_VECS = B * SKV * HKV * d_qk // VEC + V_VECS = B * SKV * HKV * d_v // VEC + if gidx < K_VECS: + _reduce_group_vec(dkws_ptr, dk_ptr, gidx, HKV, HQ, d=d_qk, group=group, io_dtype=io_dtype) + else: + if gidx < K_VECS + V_VECS: + _reduce_group_vec(dvws_ptr, dv_ptr, gidx - K_VECS, HKV, HQ, d=d_v, group=group, io_dtype=io_dtype) @cute.jit @@ -1605,14 +1676,19 @@ def _dkv_reduce_host( dv_ws: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, - d: cutlass.Constexpr[int], + d_qk: cutlass.Constexpr[int], + d_v: cutlass.Constexpr[int], group: cutlass.Constexpr[int], io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], use_pdl: cutlass.Constexpr[bool], stream: cuda_driver.CUstream, ): - out_vecs = cute.ceil_div(dk.shape[0] * dk.shape[1] * dk.shape[2] * d, 8) - _dkv_reduce_kernel(dk_ws, dv_ws, dk, dv, d, group, io_dtype, use_pdl).launch( + if cutlass.const_expr(d_qk == d_v): + out_vecs = cute.ceil_div(dk.shape[0] * dk.shape[1] * dk.shape[2] * d_qk, 8) + else: + # Split index space: one thread per dK vector plus one per dV vector. + out_vecs = cute.ceil_div(dk.shape[0] * dk.shape[1] * dk.shape[2] * (d_qk + d_v), 8) + _dkv_reduce_kernel(dk_ws, dv_ws, dk, dv, d_qk, d_v, group, io_dtype, use_pdl).launch( grid=(cute.ceil_div(out_vecs, 256), 1, 1), block=(256, 1, 1), stream=stream, @@ -1700,16 +1776,16 @@ def compile( # noqa: A001 qh: int = 1, sq: int = 128, skv: int = 128, - d: int = 128, - kvh: int = 0, + d_qk: int = 128, + d_v: int = 0, # the V/O/dO/dV head dim (0 = ``d_qk``; unequal dims serve MLA). + kvh: int = 0, # the KV head count for GQA/MQA (0 = ``qh``, plain MHA). ) -> SimpleNamespace: """Compile and cache the backward chain for one compact BSHD shape (dot, main, cvt, plus the group-reduce kernel when GQA). - - ``kvh`` is the KV head count for GQA/MQA (0 = ``qh``, plain MHA). """ kvh = int(kvh) or int(qh) + d_v = int(d_v) or int(d_qk) if qh % kvh: raise ValueError(f"GQA requires qh to be a multiple of kvh; got qh={qh}, kvh={kvh}") bwd = SM120FusedMultiHeadAttentionFP16Backward( @@ -1719,7 +1795,8 @@ def compile( # noqa: A001 window_size_left=PARAMS.window_size_left, window_size_right=PARAMS.window_size_right, deterministic=PARAMS.deterministic, - head_dim=d, + head_dim_qk=d_qk, + head_dim_v=d_v, use_pdl=PARAMS.use_pdl, q_tile=PARAMS.q_tile, kv_tile=PARAMS.kv_tile, @@ -1737,25 +1814,25 @@ def _fake(dtype, shape): assumed_align=16, ) - fake_q = _fake(STORAGE_DTYPE, (b, sq, qh, d)) - fake_k = _fake(STORAGE_DTYPE, (b, skv, kvh, d)) - fake_v = _fake(STORAGE_DTYPE, (b, skv, kvh, 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, kvh, d)) - fake_dv = _fake(STORAGE_DTYPE, (b, skv, kvh, d)) + fake_q = _fake(STORAGE_DTYPE, (b, sq, qh, d_qk)) + fake_k = _fake(STORAGE_DTYPE, (b, skv, kvh, d_qk)) + fake_v = _fake(STORAGE_DTYPE, (b, skv, kvh, d_v)) + fake_o = _fake(STORAGE_DTYPE, (b, sq, qh, d_v)) + fake_do = _fake(STORAGE_DTYPE, (b, sq, qh, d_v)) + fake_dq = _fake(STORAGE_DTYPE, (b, sq, qh, d_qk)) + fake_dk = _fake(STORAGE_DTYPE, (b, skv, kvh, d_qk)) + fake_dv = _fake(STORAGE_DTYPE, (b, skv, kvh, d_v)) 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_dq_accum = _fake(cutlass.Float32, (b * sq_r * qh * d_qk,)) # 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),)) - # Main-kernel dK/dV destinations, always HQ-headed: alias dk/dv for MHA + # Main-kernel dK/dV destinations, always HQ-headed: alias dk/d_v for MHA # (qh == kvh); per-q-head partials summed by _dkv_reduce_kernel for GQA. has_gqa = kvh != qh - fake_dk_ws = _fake(STORAGE_DTYPE, (b, skv, qh, d)) - fake_dv_ws = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_dk_ws = _fake(STORAGE_DTYPE, (b, skv, qh, d_qk)) + fake_dv_ws = _fake(STORAGE_DTYPE, (b, skv, qh, d_v)) fake_seq_q_lens = _fake(cutlass.Int32, (b,)) if PARAMS.seq_q_lens_present else None fake_seq_kv_lens = _fake(cutlass.Int32, (b,)) if PARAMS.seq_kv_lens_present else None fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -1769,7 +1846,8 @@ def _fake(dtype, shape): fake_dq_accum, fake_dq_sem, bwd.q_tile, - d, + d_qk, + d_v, bwd.page, bwd.use_pdl, bwd.deterministic, @@ -1800,7 +1878,7 @@ def _fake(dtype, shape): fake_dq_accum, fake_dq, bwd.q_tile, - d, + d_qk, bwd.page, bwd.warps_m_dq, cutlass.Float32(1.0), @@ -1817,7 +1895,8 @@ def _fake(dtype, shape): fake_dv_ws, fake_dk, fake_dv, - d, + d_qk, + d_v, qh // kvh, STORAGE_DTYPE, bwd.use_pdl, 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 b07d161bc..39d9bdc99 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -115,24 +115,33 @@ def _expected_workspace_bytes( h_q: int, s_q: int, head_dim: int, - staged: tuple[torch.Tensor, ...] = (), + staged_qk: tuple[torch.Tensor, ...] = (), + staged_vo: tuple[torch.Tensor, ...] = (), h_kv: int | None = None, s_kv: int | None = None, io_itemsize: int = 2, + head_dim_v: int | None = None, ) -> int: - from cudnn.sdpa.bwd.config_sm120 import padded_head_dim + from cudnn.sdpa.bwd.config_sm120 import padded_head_dims from cudnn.sdpa.fwd.api_dsl import ws_align - d_pad = padded_head_dim(head_dim) + head_dim_v = head_dim if head_dim_v is None else head_dim_v + # Per-side native kernel head-dim sizes — same helper the adapter uses. + d_pad, dv_pad = padded_head_dims(head_dim, head_dim_v) sq_r = -(-s_q // 128) * 128 h_kv = h_q if h_kv is None else h_kv s_kv_eff = s_kv if s_kv is not None else s_q dq_sem = batch * h_q * (-(-s_q // 32)) # int32 relay counters (min q-tile 32) # dk_ws/dv_ws GQA partials buffers in the io dtype (none carved for MHA, where the main kernel writes dk/dv directly) - dkv_ws_half = 0 if h_kv == h_q else batch * s_kv_eff * h_q * d_pad * io_itemsize - base = ws_align(batch * h_q * sq_r * 4) + ws_align(batch * sq_r * h_q * d_pad * 4) + ws_align(dq_sem * 4) + 2 * ws_align(dkv_ws_half) - # 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()) + dkv_ws = 0 + if h_kv != h_q: + dkv_ws = ws_align(batch * s_kv_eff * h_q * d_pad * io_itemsize) + ws_align(batch * s_kv_eff * h_q * dv_pad * io_itemsize) + base = ws_align(batch * h_q * sq_r * 4) + ws_align(batch * sq_r * h_q * d_pad * 4) + ws_align(dq_sem * 4) + dkv_ws + # Padded-width staging copy per non-BSHD-compact operand (a whole side when its D pads). + staging = sum(ws_align(t.numel() // head_dim * d_pad * t.element_size()) for t in staged_qk if d_pad != head_dim or not t.transpose(1, 2).is_contiguous()) + staging += sum( + ws_align(t.numel() // head_dim_v * dv_pad * t.element_size()) for t in staged_vo if dv_pad != head_dim_v or not t.transpose(1, 2).is_contiguous() + ) return base + staging @@ -167,9 +176,10 @@ 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 + head_dim_v = v_gpu.shape[3] 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) + dv_gpu = _bhsd(batch, h_kv, v_gpu.shape[2], head_dim_v, dtype, empty=True, layout=grad_layout) graph = cudnn.pygraph( io_data_type=io_dtype, @@ -250,9 +260,19 @@ def _run_bwd_graph( workspace_size = graph.get_workspace_size() if plan_name == ENGINE: - staged = (q_gpu, k_gpu, v_gpu, o_gpu, do_gpu, dq_gpu, dk_gpu, dv_gpu) + staged_qk = (q_gpu, k_gpu, dq_gpu, dk_gpu) + staged_vo = (v_gpu, o_gpu, do_gpu, dv_gpu) assert workspace_size == _expected_workspace_bytes( - batch, h_q, q_gpu.shape[2], head_dim, staged, h_kv=h_kv, s_kv=k_gpu.shape[2], io_itemsize=q_gpu.element_size() + batch, + h_q, + q_gpu.shape[2], + head_dim, + staged_qk, + staged_vo, + h_kv=h_kv, + s_kv=k_gpu.shape[2], + io_itemsize=q_gpu.element_size(), + head_dim_v=head_dim_v, ) workspace = torch.empty(max(workspace_size, 1), dtype=torch.uint8, device="cuda") @@ -288,6 +308,7 @@ def _run_case( s_q: int = 512, s_kv: int = 512, head_dim: int = 64, + head_dim_v: int | None = None, dtype: torch.dtype = torch.float16, is_causal: bool = False, causal_bottom_right: bool = False, @@ -303,11 +324,12 @@ def _run_case( sink: bool = False, ) -> str: h_kv = h_q if h_kv is None else h_kv + head_dim_v = head_dim if head_dim_v is None else head_dim_v scale = 1.0 / math.sqrt(head_dim) q = _bhsd(batch, h_q, s_q, head_dim, dtype, layout=layout) k = _bhsd(batch, h_kv, s_kv, head_dim, dtype, layout=layout) - v = _bhsd(batch, h_kv, s_kv, head_dim, dtype, layout=layout) - do = _bhsd(batch, h_q, s_q, head_dim, dtype, layout=layout) + v = _bhsd(batch, h_kv, s_kv, head_dim_v, dtype, layout=layout) + do = _bhsd(batch, h_q, s_q, head_dim_v, dtype, layout=layout) sink_gpu = torch.randn(1, h_q, 1, 1, dtype=torch.float32, device="cuda") if sink else None o, stats, dq_ref, dk_ref, dv_ref, dsink_ref = _ref_bwd( q, @@ -322,7 +344,7 @@ def _run_case( padding=padding, sink_token=sink_gpu, ) - o = _bhsd(batch, h_q, s_q, head_dim, dtype, empty=True, layout=layout).copy_(o) + o = _bhsd(batch, h_q, s_q, head_dim_v, dtype, empty=True, layout=layout).copy_(o) seq_q_lens = seq_kv_lens = None if padding is not None: seq_q_lens = torch.tensor(padding[0], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) @@ -900,20 +922,21 @@ def test_sdpa_bwd_dsl_sm120_deterministic_bitwise_padding(): _run_bitwise_case(s_q=1024, s_kv=1024, head_dim=64, is_causal=True, padding=([1024, 300], [1000, 128])) -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): +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, head_dim_v: int | None = None): """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 + head_dim_v = head_dim if head_dim_v is None else head_dim_v 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) + v = _bhsd(batch, heads, s_kv, head_dim_v, dtype) + do = _bhsd(batch, heads, s_q, head_dim_v, 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) + o = _bhsd(batch, heads, s_q, head_dim_v, 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) @@ -979,3 +1002,159 @@ def test_sdpa_bwd_dsl_sm120_tile_knobs(head_dim: int, q_tile: int, kv_tile: int) """ _run_case(head_dim=head_dim, is_causal=True, q_tile=q_tile, kv_tile=kv_tile) + + +# --------------------------------------------------------------------------- +# Rectangular head dims (D_QK > D_V): MLA training shapes. +# --------------------------------------------------------------------------- + + +def _run_mla_wrapper_case( + *, + head_dim: int, + head_dim_v: int, + h_q: int = 4, + h_kv: int | None = None, + s_q: int = 512, + s_kv: int = 512, + dtype: torch.dtype = torch.float16, + is_causal: bool = False, + causal_bottom_right: bool = False, +) -> None: + """Rectangular-dims case through the direct wrapper, vs the torch ref.""" + + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch = 2 + h_kv = h_q if h_kv is None else h_kv + scale = 1.0 / math.sqrt(head_dim) + q = _bhsd(batch, h_q, s_q, head_dim, dtype) + k = _bhsd(batch, h_kv, s_kv, head_dim, dtype) + v = _bhsd(batch, h_kv, s_kv, head_dim_v, dtype) + do = _bhsd(batch, h_q, s_q, head_dim_v, 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, h_q, s_q, head_dim_v, 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("mask", ["dense", "causal_tl", "causal_br"]) +@torch_fork_set_rng(seed=50) +def test_sdpa_bwd_dsl_sm120_mla_192_128(mask: str): + """DeepSeek-V3 / Kimi-K2.6 MLA training shape: D_QK=192 (128 nope + 64 + rope), D_V=128. D_QK > 128 has no graph surface, so the graph build must + be rejected and the direct wrapper serves it (same fallback pattern as + the large-D tests).""" + + _require_dsl() + import cudnn + + is_causal = mask != "dense" + causal_bottom_right = mask == "causal_br" + try: + _run_case(head_dim=192, head_dim_v=128, 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}" + + _run_mla_wrapper_case(head_dim=192, head_dim_v=128, is_causal=is_causal, causal_bottom_right=causal_bottom_right) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=51) +def test_sdpa_bwd_dsl_sm120_mla_192_128_gqa_bf16(): + """MLA dims + GQA + bf16 through the direct wrapper: the split-index + group-reduce sums the D_QK-wide dK partials and D_V-wide dV partials.""" + + _require_dsl() + _run_mla_wrapper_case(head_dim=192, head_dim_v=128, h_q=8, h_kv=2, dtype=torch.bfloat16, is_causal=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=52) +def test_sdpa_bwd_dsl_sm120_mla_192_128_tails(): + """MLA dims with non-tile-multiple sequence tails (partial Q/KV tiles).""" + + _require_dsl() + _run_mla_wrapper_case(head_dim=192, head_dim_v=128, s_q=193, s_kv=257, is_causal=True) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("head_dim", "head_dim_v"), + [(128, 64), (96, 40)], + ids=["native_128_64", "padded_96_40"], +) +@pytest.mark.parametrize("is_causal", [False, True], ids=["dense", "causal"]) +@torch_fork_set_rng(seed=53) +def test_sdpa_bwd_dsl_sm120_rect_head_dims_graph(head_dim: int, head_dim_v: int, is_causal: bool): + """Rectangular D_QK > D_V through the graph path: native 128/64, and the + padded envelope (96/40 stages into the native 128/64 sizes; the VO side + raises to 64 so both kernel dims stay multiples of 64).""" + + _run_case(head_dim=head_dim, head_dim_v=head_dim_v, is_causal=is_causal) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=54) +def test_sdpa_bwd_dsl_sm120_rect_head_dims_gqa_graph(): + """GQA + rectangular dims via the graph path: the group-reduce kernel's + split index space (dK vectors then dV vectors) covers both partials.""" + + _run_case(h_q=8, h_kv=2, head_dim=128, head_dim_v=64, is_causal=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=55) +def test_sdpa_bwd_dsl_sm120_rect_head_dims_padding_graph(): + """Rectangular dims compose with the padding mask (per-batch seq lens).""" + + _run_case(head_dim=128, head_dim_v=64, is_causal=True, padding=([512, 300], [512, 128])) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=56) +def test_sdpa_bwd_dsl_sm120_mla_deterministic_bitwise(): + """Repeated deterministic MLA (192/128) runs are bitwise identical.""" + + _require_dsl() + runs, (dq_ref, dk_ref, dv_ref) = _run_wrapper_det_case( + 192, + head_dim_v=128, + s_q=1024, + s_kv=1024, + is_causal=True, + window_size_left=None, + n_runs=3, + ) + 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) + for run_i, out in enumerate(runs[1:], start=1): + for grad in ("dq_tensor", "dk_tensor", "dv_tensor"): + assert torch.equal(out[grad], runs[0][grad]), f"run {run_i}: {grad} is not bitwise reproducible (MLA 192/128)" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=57) +def test_sdpa_bwd_dsl_sm120_rect_rejects_dv_gt_dqk(): + """D_V > D_QK is out of the dqk_ge_dv envelope: the adapter rejects it.""" + + _require_dsl() + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, s, dtype = 2, 4, 256, torch.float16 + q = _bhsd(batch, heads, s, 64, dtype) + k = _bhsd(batch, heads, s, 64, dtype) + v = _bhsd(batch, heads, s, 128, dtype) + do = _bhsd(batch, heads, s, 128, dtype) + o = torch.zeros_like(do) # never consumed: check_support rejects first + stats = torch.zeros(batch, heads, s, 1, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="D_QK >= D_V"): + sdpa_bwd_wrapper_dsl_sm120(q, k, v, o, do, stats) From 139b911349735a776699615a1fc98f28565e0a8e Mon Sep 17 00:00:00 2001 From: barretw Date: Mon, 17 Aug 2026 22:51:21 -0700 Subject: [PATCH 2/2] fix comments --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 2 +- python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 4a9831aba..f94b2e597 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -148,7 +148,7 @@ warp-partition triple: SMEM per CTA is `Q_STAGES·tile_q·d_qk + tile_q·d_v + tile_kv·d_qk + max(tile_kv·d_v, 2·tile_q·tile_kv)` 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 +when it does not 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 diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index c3844cd15..b908cc440 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -687,9 +687,9 @@ def kernel( smem = cutlass.Array(io_dtype, self.smem_elems, space=cutlass.AddressSpace.smem, alignment=128) sQ = smem # Q_STAGES * M * d_qk - sdO = smem.subview(self.off_sdO) # M * d_qk + sdO = smem.subview(self.off_sdO) # M * d_v sK = smem.subview(self.off_sK) # N * d_qk - sV = smem.subview(self.off_sV) # N * d_qk + sV = smem.subview(self.off_sV) # N * d_v (region max(N * d_v, 2 * M * N)) 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)