From 113d71b56ffee4239a22e6a662749237e712a494 Mon Sep 17 00:00:00 2001 From: barretw Date: Tue, 11 Aug 2026 06:29:12 -0700 Subject: [PATCH 1/6] add QGA --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 55 +++-- python/cudnn/sdpa/bwd/api_dsl.py | 37 ++- python/cudnn/sdpa/bwd/engines.py | 1 + .../cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 211 +++++++++++++----- .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 106 ++++++++- .../sdpa/frost/test_sdpa_graph_analyzer.py | 10 +- 6 files changed, 331 insertions(+), 89 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index ce0c68862..29fa43f5f 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -20,7 +20,8 @@ Two integration surfaces are provided: 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). +dQ-finalize kernel per call; GQA/MQA adds a fourth `reduce` kernel that sums +the per-q-head dK/dV partials). ## Requirements @@ -47,7 +48,8 @@ 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. +fp32 `(B, H, S_q, 1)` contiguous. GQA/MQA is expressed through the head +counts: `H_kv` may be any divisor of `H_q` (K/V and dK/dV carry `H_kv` heads). 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 @@ -60,28 +62,32 @@ 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 +deterministic in both modes — including under GQA, where the group reduction +runs in a fixed q-head order. 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 +### The 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: +One backward call is three launches (four under GQA), 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 +dot delta = rowsum(O ∘ dO); zeroes dq_accum (and, when deterministic, the relay counters) +main the fused five-GEMM pass; writes dK/dV into dk_ws/dv_ws (aliased to the dk/dv + outputs for MHA, per-q-head partial buffers for GQA); accumulates dQ into dq_accum +reduce GQA only: dK/dV = fixed-order sum of each KV head's group of q-head partials +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: +Grid is `(num_kv_tiles, H_q, B)` — one CTA owns one KV tile of one **query** +head (its KV head is `q_head // group`), 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) @@ -152,6 +158,22 @@ 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). +### GQA/MQA group reduction + +Under GQA the chain rule sums each KV head's gradient over its group of +`group = H_q / H_kv` query heads: `dK[kv_head] = Σ dK-contribution[q_head]` +(same for dV); dQ is unaffected. The grid stays per query head — shrinking it +by `group` would starve the GPU at small `B·H_kv`, and the kernel is +compute-bound, so K/V-load reuse from walking the group in one CTA is not +worth that trade. Instead each CTA writes its q head's dK/dV epilogue tiles to +`dk_ws`/`dv_ws`, io-dtype buffers with an `H_q`-sized head axis (one slot per +query head, so the group's partials coexist), and a lightweight `reduce` +kernel then produces dK/dV: one thread per 16-byte output vector, accumulating +the group's slices in fp32 in fixed q-head order — bandwidth-bound and bitwise +deterministic by construction. For MHA (`H_q == H_kv`) the buffers alias the +`dk`/`dv` outputs themselves — the same epilogue writes the results directly, +nothing extra is carved, and no `reduce` kernel is launched. + ### Causal and sliding-window masks Masking is applied twice, cheaply: @@ -193,8 +215,9 @@ only the unused relay operand remains in the kernel ABI. 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 +- GQA/MQA: any `H_kv` dividing `H_q` (including `H_kv == 1`) +- 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 + 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` elements + each); padded-D and non-compact layouts add staging copies diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 91ec3e4fe..24496a8bb 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -185,9 +185,9 @@ def check_support(self) -> bool: for label, val in (("B", b), ("H_q", h_q), ("H_kv", h_kv), ("S_q", s_q), ("S_kv", s_kv), ("D", d_qk)): self._value_error_if(int(val) <= 0, f"{label} must be > 0; got {val}") - self._not_implemented_error_if( - h_q != h_kv, - f"SM120 DSL SDPA backward does not implement GQA / MQA; got H_q={h_q}, H_kv={h_kv}", + self._value_error_if( + 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( @@ -290,6 +290,7 @@ def compile(self) -> None: sq=self.s_q_max, skv=self.s_k_max, d=self.head_dim_padded, + kvh=self.h_kv, ) self._logger.debug("compile completed") @@ -298,17 +299,27 @@ 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.""" + + if self.h_q == self.h_kv: + return 0 + return self.batch_size * self.s_k_max * self.h_q * self.head_dim_padded + def scratch_workspace_bytes(self) -> int: """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) + + 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.""" 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_sem_bytes = ws_align(self._dq_sem_len() * 4) + dkv_ws_bytes = 2 * ws_align(self._dkv_ws_elems() * self.dtype.itemsize) 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 + return delta_bytes + dq_accum_bytes + dq_sem_bytes + dkv_ws_bytes + staging_bytes def execute( self, @@ -382,7 +393,16 @@ def _staged_out_bshd(tensor: torch.Tensor): lse = stats_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) kernels = self._compiled_kernel - # Three-kernel chain + # dK/dV destinations for the main kernel: MHA writes the (staged) + # outputs directly; GQA stages per-q-head partials for the reduce. + 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) + + # Kernel chain (dot -> main -> [reduce] -> cvt) kernels.dot(o, do, delta, dq_accum, dq_sem, current_stream) kernels.main( q, @@ -393,12 +413,15 @@ def _staged_out_bshd(tensor: torch.Tensor): delta, dq_accum, dq_sem, - dk, - dv, + dk_ws, + dv_ws, cutlass.Float32(scale_log2), cutlass.Float32(scale_val), current_stream, ) + # GQA only + if kernels.reduce is not None: + kernels.reduce(dk_ws, dv_ws, dk, dv, 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: diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 79c318866..c1a876572 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -243,6 +243,7 @@ def _sm120_spec() -> EngineSpec: # 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}), + gqa=True, causal=True, bottom_right=True, swa=True, diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index c9ad5aa3b..c8f4ab901 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -20,7 +20,8 @@ * Supported input dtypes: Float16 and BFloat16 (output dtype matches) * 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 +* GQA/MQA: H_q must be a multiple of H_kv +* No dropout/alibi/softcap * Optional causal (top-left or bottom-right) and sliding-window masks * LSE input is the natural-log forward stats, fp32 (B, H, SQ) contiguous @@ -544,16 +545,16 @@ def load_tma_tile(self, s_dst, tma_desc, mbar, batch, head, seq, rows: cutlass.C @cute.kernel def kernel( self, - q: cute.Tensor, # [B, SQ, H, D] io dtype (BSHD) - k: cute.Tensor, # [B, SKV, H, D] - v: cute.Tensor, # [B, SKV, H, D] - do: cute.Tensor, # [B, SQ, H, D] - lse: cute.Tensor, # [B, H, SQ] fp32 (natural-log LSE) - delta: cute.Tensor, # [B, H, SQ_r128] fp32 (dot_do_o output) - dq_accum: cute.Tensor, # [B*SQ_r128*H*D] fp32 (scrambled, zeroed) - 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 + 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] + 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) tma_q_desc: cutlass.GridConstant[cuda.TensorMap], tma_k_desc: cutlass.GridConstant[cuda.TensorMap], tma_v_desc: cutlass.GridConstant[cuda.TensorMap], @@ -592,7 +593,7 @@ def kernel( VREG_PAIRS = SDP_NPER // 16 # V-in-regs frag pairs / chunk tidx, _, _ = cute.arch.thread_idx() - n_block, head, batch = cute.arch.block_idx() + n_block, q_head, batch = cute.arch.block_idx() lane = tidx % 32 warp = cute.arch.warp_idx() g_lane = lane // 4 @@ -600,16 +601,18 @@ def kernel( SQ = q.shape[1] SKV = k.shape[1] - H = q.shape[2] + HQ = q.shape[2] + HKV = k.shape[2] + GROUP = HQ // HKV # query heads per KV head (1 = plain MHA) SQ_R = ((SQ + 127) // 128) * 128 kv_base = n_block * N - row_stride = H * d # BSHD gmem row stride + q_row_stride = HQ * d # row stride of HQ-headed BSHD tensors (Q side; also dk_ws/dv_ws, whose head axis is HQ) lse_ptr = lse.iterator.raw_ptr() dd_ptr = delta.iterator.raw_ptr() dqa_ptr = dq_accum.iterator.raw_ptr() - dk_ptr = dk.iterator.raw_ptr() - dv_ptr = dv.iterator.raw_ptr() + dkws_ptr = dk_ws.iterator.raw_ptr() + dvws_ptr = dv_ws.iterator.raw_ptr() if cutlass.const_expr(self.deterministic): dqsem_ptr = dq_sem.iterator.raw_ptr() @@ -660,26 +663,24 @@ def kernel( prims.fence_mbarrier_init() prims.barrier_cta_sync(0) - # gmem tile bases for this (batch, head, n_block / m_block). - qhd_base = (batch * SQ) * row_stride + head * d - khd_base = (batch * SKV + kv_base) * row_stride + head * d - lse_base = (batch * H + head) * SQ - dd_base = (batch * H + head) * SQ_R + kv_head = q_head // GROUP + lse_base = (batch * HQ + q_head) * SQ + dd_base = (batch * HQ + q_head) * SQ_R if cutlass.const_expr(self.deterministic): - det_sem = dqsem_ptr + (batch * H + head) * ((SQ + M - 1) // M) + det_sem = dqsem_ptr + (batch * HQ + q_head) * ((SQ + M - 1) // M) if warp == self.load_warp_id: prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(v_mbar, N * d * io_dtype.bytes) prims.mbarrier_arrive_expect_tx(k_mbar, N * d * io_dtype.bytes) - self.load_tma_tile(sV, tma_v_desc, v_mbar, batch, head, kv_base, rows=N) - self.load_tma_tile(sK, tma_k_desc, k_mbar, batch, head, kv_base, rows=N) + 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) if n_iters > 0: if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(q_full, M * d * io_dtype.bytes) - self.load_tma_tile(sQ, tma_q_desc, q_full, batch, head, (m_block_max - 1) * M, rows=M) + self.load_tma_tile(sQ, tma_q_desc, q_full, batch, q_head, (m_block_max - 1) * M, rows=M) if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(do_full, M * d * io_dtype.bytes) self.load_tma_tile( @@ -687,7 +688,7 @@ def kernel( tma_do_desc, do_full, batch, - head, + q_head, (m_block_max - 1) * M, rows=M, ) @@ -722,7 +723,7 @@ def kernel( tma_q_desc, next_q_full, batch, - head, + q_head, next_m * M, rows=M, ) @@ -731,14 +732,14 @@ def kernel( 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) + self.load_tma_tile(sdO, tma_do_desc, do_full, batch, q_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) + self.load_tma_tile(sQ, tma_q_desc, q_full, batch, q_head, next_m * M, rows=M) elif warp < self.load_warp_id: prims.setmaxregister(240, prims.SetMaxRegisterAction.INCREASE) @@ -1094,10 +1095,10 @@ def kernel( 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 + dqa_base = ((batch * SQ_R + q_row0) * HQ + q_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) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, HQ, 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: @@ -1134,10 +1135,10 @@ def kernel( 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 + dqa_base = ((batch * SQ_R + q_row0) * HQ + q_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) + _bwd_dq_scatter(acc_dq, dqa_ptr, dqa_base, math_tidx, HQ, 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( @@ -1162,7 +1163,7 @@ def kernel( if cutlass.const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() - # ---- epilogue: dK/dV through smem (sdK aliases sK, sdV aliases sV) -------- + # epilogue: dK/dV through smem (sdK aliases sK, sdV aliases sV). cute.arch.barrier(barrier_id=2, number_of_threads=256) sdK = sK sdV = sV @@ -1188,17 +1189,21 @@ def kernel( ) cute.arch.barrier(barrier_id=2, number_of_threads=256) - # smem -> gmem + # 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): - g_off = khd_base + row * row_stride + col - copy16_smem_to_gmem(tile_ptr(sdK, row, col, page=PAGE, rows=N), dk_ptr + g_off) - copy16_smem_to_gmem(tile_ptr(sdV, row, col, page=PAGE, rows=N), dv_ptr + g_off) + 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) else: prims.setmaxregister(24, prims.SetMaxRegisterAction.DECREASE) @@ -1214,8 +1219,8 @@ def __call__( delta: cute.Tensor, dq_accum: cute.Tensor, dq_sem: cute.Tensor, - dk: cute.Tensor, - dv: cute.Tensor, + dk_ws: cute.Tensor, + dv_ws: cute.Tensor, softmax_scale_log2: cutlass.Float32, attn_scale: cutlass.Float32, stream: cuda_driver.CUstream, @@ -1236,8 +1241,8 @@ def __call__( delta, dq_accum, dq_sem, - dk, - dv, + dk_ws, + dv_ws, tma_q_desc, tma_k_desc, tma_v_desc, @@ -1475,6 +1480,84 @@ def _convert_dq_host( ) +# --------------------------------------------------------------------------- +# Reduce kernel: per-q-head dk_ws/dv_ws partials (io dtype) -> dK/dV over the group +# --------------------------------------------------------------------------- + + +@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) + dk: cute.Tensor, # [B, SKV, HKV, D] io dtype out + dv: cute.Tensor, # [B, SKV, HKV, D] io dtype out + d: cutlass.Constexpr[int], + group: cutlass.Constexpr[int], + io_dtype: cutlass.Constexpr[Type[cutlass.Numeric]], + use_pdl: cutlass.Constexpr[bool], +): + # one thread per 16 B output vector, serial fp32 accumulation over the group's q-head slices (fixed order -> deterministic). + if cutlass.const_expr(use_pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + bidx, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + B = dk.shape[0] + SKV = dk.shape[1] + 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) + + +@cute.jit +def _dkv_reduce_host( + dk_ws: cute.Tensor, + dv_ws: cute.Tensor, + dk: cute.Tensor, + dv: cute.Tensor, + d: 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( + grid=(cute.ceil_div(out_vecs, 256), 1, 1), + block=(256, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + @lru_cache(maxsize=None) def compile( # noqa: A001 compute_capability: tuple[int, int], @@ -1483,9 +1566,17 @@ def compile( # noqa: A001 sq: int = 128, skv: int = 128, d: int = 128, + kvh: int = 0, ) -> SimpleNamespace: - """Compile and cache the three-kernel backward chain for one compact BSHD shape.""" + """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) + if qh % kvh: + raise ValueError(f"GQA requires qh to be a multiple of kvh; got qh={qh}, kvh={kvh}") bwd = SM120FusedMultiHeadAttentionFP16Backward( in_dtype=STORAGE_DTYPE, is_causal=PARAMS.is_causal, @@ -1508,19 +1599,24 @@ def _fake(dtype, shape): ) fake_q = _fake(STORAGE_DTYPE, (b, sq, qh, d)) - fake_k = _fake(STORAGE_DTYPE, (b, skv, qh, d)) - fake_v = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_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, qh, d)) - fake_dv = _fake(STORAGE_DTYPE, (b, skv, qh, d)) + fake_dk = _fake(STORAGE_DTYPE, (b, skv, kvh, d)) + fake_dv = _fake(STORAGE_DTYPE, (b, skv, kvh, d)) fake_lse = _fake(cutlass.Float32, (b, qh, sq)) fake_delta = _fake(cutlass.Float32, (b, qh, sq_r)) fake_dq_accum = _fake(cutlass.Float32, (b * sq_r * qh * d,)) # 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 + # (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_stream = make_fake_stream(use_tvm_ffi_env_stream=False) options = "--enable-tvm-ffi" @@ -1549,8 +1645,8 @@ def _fake(dtype, shape): fake_delta, fake_dq_accum, fake_dq_sem, - fake_dk, - fake_dv, + fake_dk_ws, + fake_dv_ws, cutlass.Float32(1.0), cutlass.Float32(1.0), fake_stream, @@ -1570,4 +1666,19 @@ def _fake(dtype, shape): fake_stream, options=options, ) - return SimpleNamespace(dot=compiled_dot, main=compiled_main, cvt=compiled_cvt) + compiled_reduce = None + if has_gqa: + compiled_reduce = cute.compile( + _dkv_reduce_host, + fake_dk_ws, + fake_dv_ws, + fake_dk, + fake_dv, + d, + qh // kvh, + STORAGE_DTYPE, + bwd.use_pdl, + fake_stream, + options=options, + ) + return SimpleNamespace(dot=compiled_dot, main=compiled_main, cvt=compiled_cvt, reduce=compiled_reduce) 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 f198257d0..37d8ea994 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -85,14 +85,27 @@ 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, staged: tuple[torch.Tensor, ...] = ()) -> int: +def _expected_workspace_bytes( + batch: int, + h_q: int, + s_q: int, + head_dim: int, + staged: tuple[torch.Tensor, ...] = (), + h_kv: int | None = None, + s_kv: int | None = None, + io_itemsize: int = 2, +) -> 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 - 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) + 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()) return base + staging @@ -191,7 +204,9 @@ 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) - assert workspace_size == _expected_workspace_bytes(batch, h_q, q_gpu.shape[2], head_dim, staged) + 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() + ) workspace = torch.empty(max(workspace_size, 1), dtype=torch.uint8, device="cuda") variant_pack = { @@ -217,7 +232,8 @@ def _tolerances(dtype: torch.dtype) -> dict: def _run_case( *, batch: int = 2, - heads: int = 4, + h_q: int = 4, + h_kv: int | None = None, s_q: int = 512, s_kv: int = 512, head_dim: int = 64, @@ -232,15 +248,16 @@ def _run_case( layout: str = "bshd", grad_layout: str = "bshd", ) -> str: + h_kv = h_q if h_kv is None else h_kv scale = 1.0 / math.sqrt(head_dim) - 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) + 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) 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, layout=layout).copy_(o) + o = _bhsd(batch, h_q, s_q, head_dim, dtype, empty=True, layout=layout).copy_(o) dq, dk, dv, plan_name = _run_bwd_graph( q, k, @@ -466,6 +483,55 @@ def test_sdpa_bwd_dsl_sm120_padded_head_dim_wrapper(head_dim: int): torch.testing.assert_close(out["dv_tensor"].float(), dv_ref.float(), **tol) +@pytest.mark.L0 +@pytest.mark.parametrize( + ("h_q", "h_kv", "is_causal"), + [(8, 2, True), (8, 1, False)], + ids=["gqa_8_2_causal", "mqa_8_1_dense"], +) +@torch_fork_set_rng(seed=17) +def test_sdpa_bwd_dsl_sm120_gqa(h_q: int, h_kv: int, is_causal: bool): + """GQA / MQA head groups: the grid keeps one CTA per query head; each + head's dK/dV partial stages through dk_ws/dv_ws and the reduce kernel + sums the group per KV head.""" + + _run_case(h_q=h_q, h_kv=h_kv, s_q=1024, s_kv=1024, head_dim=64, is_causal=is_causal) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=18) +def test_sdpa_bwd_dsl_sm120_gqa_swa_bf16(): + """GQA composed with causal + sliding window, bf16, d=128.""" + + _run_case(h_q=8, h_kv=2, s_q=1024, s_kv=1024, head_dim=128, dtype=torch.bfloat16, is_causal=True, window_size_left=255) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=19) +def test_sdpa_bwd_dsl_sm120_gqa_causal_br_tails(): + """GQA with bottom-right causal, unequal seq lens, and a partial Q tile.""" + + _run_case(h_q=4, h_kv=2, s_q=193, s_kv=257, head_dim=128, is_causal=True, causal_bottom_right=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=20) +def test_sdpa_bwd_dsl_sm120_gqa_padded_head_dim(): + """GQA through the head-dim envelope (D=96 zero-pads to 128).""" + + _run_case(h_q=8, h_kv=2, s_q=512, s_kv=512, head_dim=96, is_causal=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=21) +def test_sdpa_bwd_dsl_sm120_gqa_deterministic_numeric(): + """deterministic + GQA composes: dQ uses the relay, while dK/dV come + from the fixed-order group reduce (deterministic in both modes); the + graph routes to the engine with parity.""" + + _run_case(h_q=8, h_kv=2, s_q=512, s_kv=512, head_dim=64, is_causal=True, deterministic=True) + + @pytest.mark.L0 @torch_fork_set_rng(seed=5) def test_sdpa_bwd_dsl_sm120_auto_routing(): @@ -509,10 +575,11 @@ def _run_bitwise_case(n_runs: int = 3, **case_kwargs) -> None: s_q = case_kwargs.pop("s_q", 1024) s_kv = case_kwargs.pop("s_kv", 1024) head_dim = case_kwargs.pop("head_dim", 64) + h_kv = case_kwargs.pop("h_kv", heads) 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) + k = _bhsd(batch, h_kv, s_kv, head_dim, dtype) + v = _bhsd(batch, h_kv, s_kv, head_dim, dtype) do = _bhsd(batch, heads, s_q, head_dim, dtype) o, stats, _, _, _ = _ref_bwd( q, @@ -555,6 +622,21 @@ def test_sdpa_bwd_dsl_sm120_deterministic_bitwise_tails_knobs(): _run_bitwise_case(s_q=1000, s_kv=999, head_dim=64, is_causal=True, q_tile=128, kv_tile=64) +@pytest.mark.L0 +@pytest.mark.parametrize(("mask", "h_kv"), [("dense", 2), ("causal", 1)], ids=["dense_gqa2", "causal_mqa"]) +@torch_fork_set_rng(seed=22) +def test_sdpa_bwd_dsl_sm120_gqa_deterministic_bitwise(mask: str, h_kv: int): + """Repeated deterministic GQA/MQA runs are bitwise identical: the relay + fixes dQ's fp32 add order and the reduce kernel sums the group's dK/dV + partials in fixed q-head order.""" + + _run_bitwise_case( + h_kv=h_kv, + head_dim=64, + is_causal=mask != "dense", + ) + + 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).""" diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index ce2bdceef..f18ea50e8 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -985,9 +985,13 @@ def test_bwd_probe_rejects_forward_graph(monkeypatch): assert not _eligible(_mk_bwd_graph()) -def test_bwd_probe_rejects_gqa(monkeypatch): +def test_bwd_probe_gqa(monkeypatch): monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) - assert not _bwd_eligible(_mk_bwd_graph(h_kv=H // 2)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(h_kv=H // 2)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(h_kv=1)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(h_kv=H // 2, use_deterministic_algorithm=True)) + # H_q must be a multiple of H_kv + assert not _bwd_eligible(_mk_bwd_graph(h_kv=3)) def test_bwd_probe_rejects_unsupported_head_dim(monkeypatch): @@ -1057,8 +1061,6 @@ def test_bwd_knob_domains(monkeypatch): def test_bwd_mismatch_reason_strings(monkeypatch): monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) caps = bwd_engines.ENGINE_SPECS[0].capabilities - reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(h_kv=H // 2))) - assert reason is not None and "GQA" in reason reason = bwd_engines.mismatch(caps, _facts(_mk_bwd_graph(d=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))) From 3048a00c693c3c450898a4d5861044689daab1a0 Mon Sep 17 00:00:00 2001 From: barretw Date: Tue, 11 Aug 2026 21:55:40 -0700 Subject: [PATCH 2/6] add padding mask --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 23 ++- python/cudnn/sdpa/bwd/api_dsl.py | 60 +++++++ python/cudnn/sdpa/bwd/config_sm120.py | 5 + python/cudnn/sdpa/bwd/engines.py | 12 ++ .../cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 65 +++++-- .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 160 +++++++++++++++++- .../sdpa/frost/test_sdpa_graph_analyzer.py | 22 +++ 7 files changed, 330 insertions(+), 17 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 29fa43f5f..8f1c79a67 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -41,6 +41,8 @@ grads = sdpa_bwd_wrapper_dsl_sm120( 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) + seq_q_lens=None, # (B,) int32 per-batch Q lengths (padding mask) + seq_kv_lens=None, # (B,) int32 per-batch KV lengths (padding mask) ) dq, dk, dv = grads["dq_tensor"], grads["dk_tensor"], grads["dv_tensor"] ``` @@ -174,7 +176,7 @@ deterministic by construction. For MHA (`H_q == H_kv`) the buffers alias the `dk`/`dv` outputs themselves — the same epilogue writes the results directly, nothing extra is carved, and no `reduce` kernel is launched. -### Causal and sliding-window masks +### Causal, sliding-window, and padding masks Masking is applied twice, cheaply: @@ -183,12 +185,26 @@ Masking is applied twice, cheaply: 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. + (`do_mask_causal` / `do_mask_window` / `do_mask_pad` 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. +**Padding mask** (per-batch `seq_kv_lens`, optionally `seq_q_lens`) reuses +both layers: `q_block_max` trims to `ceil(seq_q_lens[b] / tile_q)` and a KV +tile fully inside the pad drains without loads or compute, while boundary +tiles mask scores at `seq_kv_lens[b]`. With bottom-right alignment the +diagonal anchors at the **actual** lengths (`diag_off = seq_kv_lens[b] − +seq_q_lens[b]`), matching the SM120 forward kernel and cuDNN padded-graph +semantics. Q rows at or past `seq_q_lens[b]` ride on the forward's +`LSE = −inf` convention (`P = 0`), so `dQ` rows past `seq_q_lens[b]` and +`dK`/`dV` rows past `seq_kv_lens[b]` come out exactly zero. The length +tensors are `None`-specialized kernel parameters: a specialization built +without them carries neither the parameters nor any padding code, and +needs no extra workspace either way. + ### Deterministic vs. non-deterministic dQ The default path's relaxed atomics make the fp32 add order — hence the @@ -214,7 +230,8 @@ only the unused relay operand remains in the kernel ABI. - 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) + (left-window offset, with or without causal), padding (per-batch + `seq_len_kv` required, `seq_len_q` optional; composes with the other masks) - GQA/MQA: any `H_kv` dividing `H_q` (including `H_kv == 1`) - No dropout / bias / ALiBi / sinks / softcap / THD - Workspace (carved from the caller's buffer): fp32 `delta` and `dq_accum` diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 24496a8bb..5a496f164 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -73,6 +73,8 @@ def __init__( scale_softmax: Optional[float] = None, tile_m: Optional[int] = None, tile_n: Optional[int] = None, + seq_kv_lens_present: bool = False, + seq_q_lens_present: bool = False, ) -> None: super().__init__() self._warn_experimental_api() @@ -95,6 +97,8 @@ def __init__( self.scale_softmax = scale_softmax self.tile_m = None if tile_m is None else int(tile_m) self.tile_n = None if tile_n is None else int(tile_n) + self.seq_kv_lens_present = bool(seq_kv_lens_present) + self.seq_q_lens_present = bool(seq_q_lens_present) self.batch_size: Optional[int] = None self.s_q_max: Optional[int] = None @@ -129,6 +133,8 @@ def execute( scale_softmax: Optional[float] = None, workspace: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, + seq_q_lens: Optional[torch.Tensor] = None, + seq_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Execute the compiled kernel chain using the common operand set.""" @@ -238,6 +244,10 @@ def check_support(self) -> bool: self.causal_bottom_right and not self.is_causal, "causal_bottom_right requires is_causal=True", ) + self._value_error_if( + self.seq_q_lens_present and not self.seq_kv_lens_present, + "seq_q_lens_present requires seq_kv_lens_present (per-batch Q lengths are only supported as part of the padding mask)", + ) self._value_error_if( self.window_size_left is not None and self.window_size_left < 0, f"window_size_left must be non-negative, got {self.window_size_left}", @@ -281,6 +291,8 @@ def compile(self) -> None: deterministic=self.deterministic, q_tile=self.q_tile, kv_tile=self.kv_tile, + seq_kv_lens_present=self.seq_kv_lens_present, + seq_q_lens_present=self.seq_q_lens_present, ) self._k_mod = _load_sm120_kernel_module(params) self._compiled_kernel = self._k_mod.compile( @@ -307,6 +319,22 @@ def _dkv_ws_elems(self) -> int: return 0 return self.batch_size * self.s_k_max * self.h_q * self.head_dim_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).""" + self._value_error_if( + seq_lens.dtype != torch.int32, + f"{name} must be int32; got {seq_lens.dtype}", + ) + self._value_error_if( + seq_lens.numel() != self.batch_size, + f"{name} must have B = {self.batch_size} elements; got {seq_lens.numel()}", + ) + self._value_error_if( + not seq_lens.is_contiguous(), + f"{name} must be contiguous (bound to the kernel as a flat (B,) view)", + ) + 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]) + dq_sem (int32 flat [B*H*ceil(SQ/32)], deterministic relay counters) @@ -335,12 +363,31 @@ def execute( scale_softmax: Optional[float] = None, workspace: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, + seq_q_lens: Optional[torch.Tensor] = None, + seq_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Execute tensors matching the compiled specialization.""" if self._compiled_kernel is None: raise RuntimeError("SdpaBwdDslSm120 kernel is not compiled") + self._value_error_if( + self.seq_kv_lens_present and seq_kv_lens is None, + "seq_kv_lens is required by this compiled specialization", + ) + self._value_error_if( + not self.seq_kv_lens_present and seq_kv_lens is not None, + "this specialization was compiled without per-batch KV lengths; construct the API with seq_kv_lens_present=True", + ) + self._value_error_if( + self.seq_q_lens_present and seq_q_lens is None, + "seq_q_lens is required by this compiled specialization", + ) + self._value_error_if( + not self.seq_q_lens_present and seq_q_lens is not None, + "this specialization was compiled without per-batch Q lengths; construct the API with seq_q_lens_present=True", + ) + scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) scale_log2 = scale_val * math.log2(math.e) @@ -381,6 +428,9 @@ def _staged_out_bshd(tensor: torch.Tensor): b, s, h, _ = view.shape return carver.take(b * s * h * d_pad, self.dtype).view(b, s, h, d_pad), view + seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if seq_q_lens is not None else None + 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) @@ -415,6 +465,8 @@ def _staged_out_bshd(tensor: torch.Tensor): dq_sem, dk_ws, dv_ws, + seq_q_t, + seq_kv_t, cutlass.Float32(scale_log2), cutlass.Float32(scale_val), current_stream, @@ -448,6 +500,8 @@ def sdpa_bwd_wrapper_dsl_sm120( window_size_left: Optional[int] = None, deterministic: bool = False, scale_softmax: Optional[float] = None, + seq_q_lens: Optional[torch.Tensor] = None, + seq_kv_lens: Optional[torch.Tensor] = None, ) -> TupleDict: """Run SM120 SDPA backward and return ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=...)``.""" @@ -470,6 +524,8 @@ def sdpa_bwd_wrapper_dsl_sm120( window_size_left, bool(deterministic), scale_softmax, + seq_q_lens is not None, + seq_kv_lens is not None, ) api = _wrapper_api_cache.get(cache_key) if api is None: @@ -488,6 +544,8 @@ def sdpa_bwd_wrapper_dsl_sm120( window_size_left=window_size_left, deterministic=deterministic, scale_softmax=scale_softmax, + seq_kv_lens_present=seq_kv_lens is not None, + seq_q_lens_present=seq_q_lens is not None, ) api.check_support() api.compile() @@ -506,5 +564,7 @@ def sdpa_bwd_wrapper_dsl_sm120( dv_tensor=dv_tensor, scale_softmax=scale_softmax, workspace=workspace, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, ) return TupleDict(dq_tensor=dq_tensor, dk_tensor=dk_tensor, dv_tensor=dv_tensor) diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py index 93285c363..b75885b19 100644 --- a/python/cudnn/sdpa/bwd/config_sm120.py +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -38,6 +38,9 @@ class TemplateParams: use_pdl: bool = True q_tile: int = 0 kv_tile: int = 0 + # Padding mask: per-batch int32 lengths; seq_q is only valid with seq_kv. + seq_kv_lens_present: bool = False + seq_q_lens_present: bool = False def validate_params(params: TemplateParams) -> None: @@ -54,6 +57,8 @@ def validate_params(params: TemplateParams) -> None: raise ValueError("SM120 SDPA bwd: causal_top_left requires is_causal=True") if params.window_size_left is not None and params.window_size_left < 0: raise ValueError(f"SM120 SDPA bwd: window_size_left must be non-negative; got {params.window_size_left}") + if params.seq_q_lens_present and not params.seq_kv_lens_present: + raise ValueError("SM120 SDPA bwd: seq_q_lens_present requires seq_kv_lens_present (padding mask)") if params.q_tile not in (0,) + SEQ_Q_TILES: raise ValueError(f"SM120 SDPA bwd: q_tile must be one of {(0,) + SEQ_Q_TILES} (0 = per-head-dim default); got {params.q_tile}") if params.kv_tile not in (0,) + SEQ_KV_TILES: diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index c1a876572..0eb16987f 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -247,6 +247,7 @@ def _sm120_spec() -> EngineSpec: causal=True, bottom_right=True, swa=True, + padded=True, layouts=frozenset({"bshd", "dense_flex"}), deterministic=True, tile_ms=frozenset(_SM120_Q_TILES), @@ -316,6 +317,9 @@ def _desc(geom, dtype, name: str) -> "Any": name=name, ) + seq_kv_t = facts.seq_kv_t if facts.padded else None + seq_q_t = facts.seq_q_t if facts.padded else None + api = _adapter_sm120()( sample_q=_desc(q_geom, facts.dtype, "q"), sample_k=_desc(k_geom, facts.dtype, "k"), @@ -333,6 +337,8 @@ def _desc(geom, dtype, name: str) -> "Any": 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, + seq_kv_lens_present=seq_kv_t is not None, + seq_q_lens_present=seq_q_t is not None, ) api.check_support() # raises ValueError / NotImplementedError if unsupported api.compile() @@ -353,6 +359,8 @@ def _desc(geom, dtype, name: str) -> "Any": dq=facts.dq_t, dk=facts.dk_t, dv=facts.dv_t, + seq_len_kv=seq_kv_t, + seq_len_q=seq_q_t, ) def _canonical_view(buf, geom): @@ -371,6 +379,8 @@ def _canonical_view(buf, geom): def _execute(variant_pack, workspace=None, stream=None): resolved = ga.resolve_variant_pack(variant_pack, binding) + seq_kv_buf = resolved.get(id(binding.seq_len_kv)) if binding.seq_len_kv is not None else None + seq_q_buf = resolved.get(id(binding.seq_len_q)) if binding.seq_len_q is not None else None api.execute( q_tensor=_canonical_view(resolved[id(binding.q)], q_geom), k_tensor=_canonical_view(resolved[id(binding.k)], k_geom), @@ -381,6 +391,8 @@ def _execute(variant_pack, workspace=None, stream=None): 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), + seq_q_lens=seq_q_buf, + seq_kv_lens=seq_kv_buf, 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 c8f4ab901..77d200e40 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -22,7 +22,7 @@ serves any other multiple of 8 up to 256 by zero-padding D. * GQA/MQA: H_q must be a multiple of H_kv * No dropout/alibi/softcap -* Optional causal (top-left or bottom-right) and sliding-window masks +* Optional causal (top-left or bottom-right), sliding-window masks and padding masks. * LSE input is the natural-log forward stats, fp32 (B, H, SQ) contiguous One backward call is three kernel launches through the per-shape @@ -33,7 +33,7 @@ from functools import lru_cache from types import SimpleNamespace -from typing import Type +from typing import Optional, Type import cuda.bindings.driver as cuda_driver import cutlass @@ -468,12 +468,16 @@ def __init__( use_pdl: bool = True, q_tile: int = 0, kv_tile: int = 0, + seq_kv_lens_present: bool = False, + seq_q_lens_present: bool = False, ): self.in_dtype = in_dtype 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.seq_kv_lens_present = bool(seq_kv_lens_present) + self.seq_q_lens_present = bool(seq_q_lens_present) self.d = head_dim self.use_pdl = bool(use_pdl) self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim] @@ -555,6 +559,8 @@ def kernel( 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) + 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], tma_k_desc: cutlass.GridConstant[cuda.TensorMap], tma_v_desc: cutlass.GridConstant[cuda.TensorMap], @@ -608,6 +614,14 @@ def kernel( 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) + # Per-batch actual lengths (Padding mask) + seqlen_q = SQ + seqlen_kv = SKV + if cutlass.const_expr(self.seq_kv_lens_present): + seqlen_kv = cute.math.max(cutlass.Int32(0), cute.math.min(seq_kv_lens[batch], cutlass.Int32(SKV))) + if cutlass.const_expr(self.seq_q_lens_present): + seqlen_q = cute.math.max(cutlass.Int32(0), cute.math.min(seq_q_lens[batch], cutlass.Int32(SQ))) + lse_ptr = lse.iterator.raw_ptr() dd_ptr = delta.iterator.raw_ptr() dqa_ptr = dq_accum.iterator.raw_ptr() @@ -618,11 +632,18 @@ def kernel( PARTIAL_Q = (SQ % M) != 0 PARTIAL_KV = (SKV % N) != 0 + # Skip the kv >= SKV check when padding already masks kv >= seqlen_kv (<= SKV). + MASK_KV_GLOBAL = PARTIAL_KV and not self.seq_kv_lens_present + # Fully-masked rows have LSE = -inf; flip to +inf so P = exp2(finite - inf) = 0, not NaN. + FLIP_MASKED_LSE = ( + (self.is_causal and not self.causal_top_left) or self.window_size_left is not None or self.seq_kv_lens_present or self.seq_q_lens_present + ) - m_block_max = (SQ + M - 1) // M + m_block_max = (seqlen_q + M - 1) // M if cutlass.const_expr(self.is_causal or self.window_size_left is not None): if cutlass.const_expr(self.is_causal and not self.causal_top_left): - diag_off = SKV - SQ + # The diagonal anchors at the actual lengths. + diag_off = seqlen_kv - seqlen_q else: diag_off = cutlass.Int32(0) if cutlass.const_expr(self.is_causal): @@ -635,6 +656,10 @@ def kernel( rows_hi = cute.math.max(last_q_row + 1, cutlass.Int32(0)) m_block_max = cute.math.min(m_block_max, (rows_hi + M - 1) // M) n_iters = cute.math.max(m_block_max - m_block_min, cutlass.Int32(0)) + if cutlass.const_expr(self.seq_kv_lens_present): + # Fully-padded KV tile: skip the loop; the epilogue writes zero dK/dV. + if kv_base >= seqlen_kv: + 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 @@ -748,8 +773,8 @@ def kernel( math_warp = warp math_tidx = tidx m_block = m_block_max - 1 - if cutlass.const_expr(self.window_size_left is not None): - # For bottom-right SWA, a KV tile might lie completely before every query's sliding window + if cutlass.const_expr(self.window_size_left is not None or self.seq_q_lens_present): + # m_block_max can be 0 in bottom-right SWA, or seq_len_q[b] == 0 m_block = cute.math.max(m_block, cutlass.Int32(0)) wm_s = math_warp % WM_SDP wn_s = math_warp // WM_SDP @@ -775,7 +800,7 @@ def kernel( val = inf else: val = (lse_ptr + lse_base + r_abs).load() - if cutlass.const_expr((self.is_causal and not self.causal_top_left) or self.window_size_left is not None): + if cutlass.const_expr(FLIP_MASKED_LSE): # 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")): @@ -880,6 +905,8 @@ def kernel( do_mask_causal = (m_block * M) < (kv_base + N - diag_off) if cutlass.const_expr(self.window_size_left is not None): do_mask_window = kv_base < (m_block * M + M - 1 + diag_off - self.window_size_left) + if cutlass.const_expr(self.seq_kv_lens_present): + do_mask_pad = (kv_base + N) > seqlen_kv neg_inf = cutlass.Float32(float("-inf")) for rep in cutlass.range_constexpr(SDP_REPS): for nf in cutlass.range_constexpr(SDP_NF): @@ -915,7 +942,15 @@ def kernel( s2 = neg_inf if kv_a1 < lo8: s3 = neg_inf - if cutlass.const_expr(PARTIAL_KV): + if cutlass.const_expr(self.seq_kv_lens_present): + if do_mask_pad: + if kv_a0 >= seqlen_kv: + s0 = neg_inf + s2 = neg_inf + if kv_a1 >= seqlen_kv: + s1 = neg_inf + s3 = neg_inf + if cutlass.const_expr(MASK_KV_GLOBAL): if kv_a0 >= SKV: s0 = neg_inf s2 = neg_inf @@ -1089,7 +1124,7 @@ def kernel( 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): + if cutlass.const_expr(FLIP_MASKED_LSE): # 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")): @@ -1129,7 +1164,7 @@ def kernel( 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): + if cutlass.const_expr(FLIP_MASKED_LSE): # 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")): @@ -1221,6 +1256,8 @@ def __call__( dq_sem: cute.Tensor, dk_ws: cute.Tensor, dv_ws: cute.Tensor, + seq_q_lens: Optional[cute.Tensor], + seq_kv_lens: Optional[cute.Tensor], softmax_scale_log2: cutlass.Float32, attn_scale: cutlass.Float32, stream: cuda_driver.CUstream, @@ -1243,6 +1280,8 @@ def __call__( dq_sem, dk_ws, dv_ws, + seq_q_lens, + seq_kv_lens, tma_q_desc, tma_k_desc, tma_v_desc, @@ -1587,6 +1626,8 @@ def compile( # noqa: A001 use_pdl=PARAMS.use_pdl, q_tile=PARAMS.q_tile, kv_tile=PARAMS.kv_tile, + seq_kv_lens_present=PARAMS.seq_kv_lens_present, + seq_q_lens_present=PARAMS.seq_q_lens_present, ) sq_r = ceil_div(sq, 128) * 128 @@ -1617,6 +1658,8 @@ def _fake(dtype, shape): 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_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) options = "--enable-tvm-ffi" @@ -1647,6 +1690,8 @@ def _fake(dtype, shape): fake_dq_sem, fake_dk_ws, fake_dv_ws, + fake_seq_q_lens, + fake_seq_kv_lens, cutlass.Float32(1.0), cutlass.Float32(1.0), fake_stream, 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 37d8ea994..1d1c4829e 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -68,6 +68,7 @@ def _ref_bwd( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + padding: tuple[list[int], list[int]] | None = None, ): """Reference via the canonical refs (sdpa/fp16_ref.py).""" @@ -78,9 +79,11 @@ def _ref_bwd( right_bound = 0 if is_causal else None # The refs take the cuDNN window LENGTH; window_size_left is the offset W = L - 1. left_bound = None if window_size_left is None else window_size_left + 1 - o_ref, stats_ref, _, _ = compute_ref(q, k, v, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, torch_type=q.dtype) + o_ref, stats_ref, _, _ = compute_ref( + q, k, v, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, padding=padding, torch_type=q.dtype + ) dq, dk, dv, _, _ = compute_ref_backward( - q, k, v, o_ref, do, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, torch_type=q.dtype + q, k, v, o_ref, do, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, padding=padding, torch_type=q.dtype ) return o_ref.to(q.dtype), stats_ref.contiguous(), dq.to(q.dtype), dk.to(q.dtype), dv.to(q.dtype) @@ -128,6 +131,8 @@ def _run_bwd_graph( q_tile: int | None = None, kv_tile: int | None = None, grad_layout: str = "bshd", + seq_q_lens: torch.Tensor | None = None, + seq_kv_lens: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]: """Build and execute the SM120 FROST backward graph; returns (dq, dk, dv, plan_name).""" @@ -172,6 +177,12 @@ def _run_bwd_graph( bwd_kwargs["sliding_window_length"] = window_size_left + 1 if deterministic: bwd_kwargs["use_deterministic_algorithm"] = True + seq_q_t = seq_kv_t = None + if seq_q_lens is not None or seq_kv_lens is not None: + assert seq_q_lens is not None and seq_kv_lens is not None + seq_q_t = graph.tensor_like(seq_q_lens, name="seq_q") + seq_kv_t = graph.tensor_like(seq_kv_lens, name="seq_kv") + bwd_kwargs.update(use_padding_mask=True, seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) dq, dk, dv = graph.sdpa_backward(**bwd_kwargs) dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride()) @@ -220,6 +231,8 @@ def _run_bwd_graph( dk: dk_gpu, dv: dv_gpu, } + if seq_q_t is not None: + variant_pack.update({seq_q_t: seq_q_lens, seq_kv_t: seq_kv_lens}) graph.execute(variant_pack, workspace) torch.cuda.synchronize() return dq_gpu, dk_gpu, dv_gpu, plan_name @@ -247,6 +260,7 @@ def _run_case( kv_tile: int | None = None, layout: str = "bshd", grad_layout: str = "bshd", + padding: tuple[list[int], list[int]] | None = None, ) -> str: h_kv = h_q if h_kv is None else h_kv scale = 1.0 / math.sqrt(head_dim) @@ -255,9 +269,13 @@ def _run_case( v = _bhsd(batch, h_kv, s_kv, head_dim, dtype, layout=layout) do = _bhsd(batch, h_q, 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 + q, k, v, do, scale=scale, is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, padding=padding ) o = _bhsd(batch, h_q, s_q, head_dim, 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) + seq_kv_lens = torch.tensor(padding[1], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) dq, dk, dv, plan_name = _run_bwd_graph( q, k, @@ -274,11 +292,20 @@ def _run_case( q_tile=q_tile, kv_tile=kv_tile, grad_layout=grad_layout, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, ) tol = _tolerances(dtype) torch.testing.assert_close(dq.float(), dq_ref.float(), **tol) torch.testing.assert_close(dk.float(), dk_ref.float(), **tol) torch.testing.assert_close(dv.float(), dv_ref.float(), **tol) + if padding is not None: + for b, (len_q, len_kv) in enumerate(zip(*padding)): + if len_q < s_q: + assert dq[b, :, len_q:, :].abs().max().item() == 0.0, f"batch {b}: dQ padded rows must be exactly zero" + if len_kv < s_kv: + assert dk[b, :, len_kv:, :].abs().max().item() == 0.0, f"batch {b}: dK padded rows must be exactly zero" + assert dv[b, :, len_kv:, :].abs().max().item() == 0.0, f"batch {b}: dV padded rows must be exactly zero" return plan_name @@ -403,6 +430,119 @@ 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("mask", ["dense", "causal_tl", "causal_br"]) +@torch_fork_set_rng(seed=20) +def test_sdpa_bwd_dsl_sm120_padding_mask(mask: str): + """Padding mask (per-batch seq lens): full-length, tile-boundary, and + sub-tile batches; bottom-right diagonals anchor at the actual lengths.""" + + _run_case( + batch=3, + s_q=512, + s_kv=512, + head_dim=64, + is_causal=mask != "dense", + causal_bottom_right=mask == "causal_br", + padding=([512, 300, 17], [512, 128, 65]), + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=21) +def test_sdpa_bwd_dsl_sm120_padding_mask_tails(): + """Padding mask on top of non-tile-multiple global sequence tails.""" + + _run_case( + batch=2, + s_q=193, + s_kv=257, + head_dim=128, + is_causal=True, + padding=([193, 100], [200, 33]), + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=22) +def test_sdpa_bwd_dsl_sm120_padding_sliding_window(): + """Padding mask + bottom-right sliding window (the window follows the + per-batch diagonal anchor).""" + + _run_case( + batch=3, + s_q=512, + s_kv=512, + head_dim=64, + is_causal=True, + causal_bottom_right=True, + window_size_left=127, + padding=([512, 300, 65], [512, 260, 64]), + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=23) +def test_sdpa_bwd_dsl_sm120_padding_zero_lengths(): + """Zero-length batches: seq_len_kv[b] == 0 (no visible key) and + seq_len_q[b] == 0 (no query) drain to all-zero gradients.""" + + _run_case( + batch=3, + s_q=512, + s_kv=512, + head_dim=64, + padding=([512, 0, 33], [0, 512, 48]), + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=27) +def test_sdpa_bwd_dsl_sm120_padding_gqa(): + """Padding mask composes with GQA: the group-reduce sums per-q-head + partials whose padded rows are zero, so dK/dV padding stays exactly zero.""" + + _run_case( + batch=3, + h_q=4, + h_kv=2, + s_q=512, + s_kv=512, + head_dim=64, + is_causal=True, + causal_bottom_right=True, + padding=([512, 300, 17], [512, 128, 65]), + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=25) +def test_sdpa_bwd_dsl_sm120_padding_kv_only_wrapper(): + """KV-only padding through the direct wrapper: seq_q_lens omitted means + every batch runs the full S_q.""" + + _require_dsl() + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, s_q, s_kv, head_dim, dtype = 2, 4, 512, 512, 64, torch.float16 + scale = 1.0 / math.sqrt(head_dim) + kv_lens = [317, 64] + 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, padding=([s_q] * batch, kv_lens)) + 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, scale_softmax=scale, seq_kv_lens=torch.tensor(kv_lens, dtype=torch.int32, device="cuda")) + 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) + for b, len_kv in enumerate(kv_lens): + assert out["dk_tensor"][b, :, len_kv:, :].abs().max().item() == 0.0, f"batch {b}: dK padded rows must be exactly zero" + assert out["dv_tensor"][b, :, len_kv:, :].abs().max().item() == 0.0, f"batch {b}: dV padded rows must be exactly zero" + + @pytest.mark.L0 @pytest.mark.parametrize("head_dim", [192, 256]) @pytest.mark.parametrize("mask", ["dense", "causal_tl", "causal_br"]) @@ -568,7 +708,7 @@ def test_sdpa_bwd_dsl_sm120_deterministic_numeric(mask: str): ) -def _run_bitwise_case(n_runs: int = 3, **case_kwargs) -> None: +def _run_bitwise_case(n_runs: int = 3, padding: tuple[list[int], list[int]] | None = None, **case_kwargs) -> None: """Same inputs, ``n_runs`` independent graph runs: outputs must be bitwise equal.""" batch, heads, dtype = 2, 4, torch.float16 @@ -590,8 +730,12 @@ def _run_bitwise_case(n_runs: int = 3, **case_kwargs) -> None: 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"), + padding=padding, ) o = _bhsd(batch, heads, s_q, head_dim, dtype, empty=True).copy_(o) + if padding is not None: + case_kwargs["seq_q_lens"] = torch.tensor(padding[0], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) + case_kwargs["seq_kv_lens"] = torch.tensor(padding[1], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) 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): @@ -637,6 +781,14 @@ def test_sdpa_bwd_dsl_sm120_gqa_deterministic_bitwise(mask: str, h_kv: int): ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=26) +def test_sdpa_bwd_dsl_sm120_deterministic_bitwise_padding(): + """Bitwise reproducibility with a padding mask (per-batch relay trims).""" + + _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): """Deterministic run(s) through the direct wrapper (D>128 has no graph surface); returns (outputs per run, references).""" diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index f18ea50e8..1da0dceae 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -893,6 +893,7 @@ def _mk_bwd_graph( grad_strides: tuple | None = None, bias: bool = False, dbias: bool = False, + seq_lens: str | None = None, # "kv" / "both" (padding mask) or "q_only" **bwd_kwargs, ): g = _mk_graph() @@ -921,6 +922,12 @@ def _mk_bwd_graph( if dbias: dbias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="dBias") bwd_kwargs.update(dBias=dbias_t) + if seq_lens in ("kv", "both"): + seq_kv_t = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_kv") + bwd_kwargs.update(use_padding_mask=True, seq_len_kv=seq_kv_t) + if seq_lens in ("both", "q_only"): + seq_q_t = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_q") + bwd_kwargs.update(seq_len_q=seq_q_t) dq, dk, dv = g.sdpa_backward(name="sb", q=q, k=k, v=v, o=o, dO=do, stats=stats, attn_scale=0.125, **bwd_kwargs) _finish_output(dq, q_dims, grad_strides or _bshd_strides(H, s_q, d)) _finish_output(dk, (B, h_kv, s_kv, d), grad_strides or _bshd_strides(h_kv, s_kv, d)) @@ -1017,6 +1024,21 @@ def test_bwd_probe_accepts_deterministic(monkeypatch): assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(use_deterministic_algorithm=True)) +def test_bwd_probe_accepts_padding_mask(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(seq_lens="kv")) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(seq_lens="both")) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(seq_lens="both", use_causal_mask_bottom_right=True)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(seq_lens="both", use_causal_mask=True, sliding_window_length=64)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(seq_lens="both", use_deterministic_algorithm=True)) + + +def test_bwd_probe_rejects_seq_len_q_without_padding_mask(monkeypatch): + # Bare seq_len_q is per-batch Q trimming, which the kernel has no path for. + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert not _bwd_eligible(_mk_bwd_graph(seq_lens="q_only")) + + def test_bwd_probe_rejects_bias(monkeypatch): monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) assert not _bwd_eligible(_mk_bwd_graph(bias=True)) From 364d67b3627df8abe96ddc05563b19a47334543d Mon Sep 17 00:00:00 2001 From: barretw Date: Tue, 11 Aug 2026 22:32:46 -0700 Subject: [PATCH 3/6] fix --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 11 +++++----- python/cudnn/sdpa/bwd/api_dsl.py | 4 ++++ .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 20 +++++++++++++++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 8f1c79a67..6c086e722 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -77,7 +77,7 @@ One backward call is three launches (four under GQA), overlapped with programmatic dependent launch (PDL) so each kernel's prologue runs under its predecessor's tail: -``` +```text dot delta = rowsum(O ∘ dO); zeroes dq_accum (and, when deterministic, the relay counters) main the fused five-GEMM pass; writes dK/dV into dk_ws/dv_ws (aliased to the dk/dv outputs for MHA, per-q-head partial buffers for GQA); accumulates dQ into dq_accum @@ -91,7 +91,7 @@ Grid is `(num_kv_tiles, H_q, B)` — one CTA owns one KV tile of one **query** head (its KV head is `q_head // group`), loads K/V **once**, and walks every q-tile of its (batch, head) in descending order. Per q-tile iteration: -``` +```text 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) @@ -231,10 +231,11 @@ only the unused relay operand remains in the kernel ABI. 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), padding (per-batch - `seq_len_kv` required, `seq_len_q` optional; composes with the other masks) + `seq_kv_len` required, `seq_q_len` optional; composes with the other masks) - GQA/MQA: any `H_kv` dividing `H_q` (including `H_kv == 1`) - 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); GQA adds - the io-dtype `dk_ws`/`dv_ws` partials buffers (`B·S_kv·H_q·D` elements - each); padded-D and non-compact layouts add staging copies + 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 diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 5a496f164..e6de829e5 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -321,6 +321,10 @@ def _dkv_ws_elems(self) -> int: 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).""" + self._value_error_if( + seq_lens.device != self.q_desc.device, + f"{name} must be on {self.q_desc.device} (with Q); got {seq_lens.device}", + ) self._value_error_if( seq_lens.dtype != torch.int32, f"{name} must be int32; got {seq_lens.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 1d1c4829e..ac85a37dd 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -515,6 +515,26 @@ def test_sdpa_bwd_dsl_sm120_padding_gqa(): ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=29) +def test_sdpa_bwd_dsl_sm120_padding_rejects_cpu_seq_lens(): + """A CPU length tensor must be rejected up front — the kernel would + otherwise receive a host pointer (illegal access or garbage lengths).""" + + _require_dsl() + from cudnn.sdpa.bwd.api_dsl import sdpa_bwd_wrapper_dsl_sm120 + + batch, heads, s, head_dim, dtype = 2, 4, 256, 64, torch.float16 + q = _bhsd(batch, heads, s, head_dim, dtype) + k = _bhsd(batch, heads, s, head_dim, dtype) + v = _bhsd(batch, heads, s, head_dim, dtype) + do = _bhsd(batch, heads, s, head_dim, dtype) + o = torch.zeros_like(q) # never consumed: execute rejects before launching + stats = torch.zeros(batch, heads, s, 1, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="seq_kv_lens must be on"): + sdpa_bwd_wrapper_dsl_sm120(q, k, v, o, do, stats, seq_kv_lens=torch.tensor([s, s], dtype=torch.int32)) + + @pytest.mark.L0 @torch_fork_set_rng(seed=25) def test_sdpa_bwd_dsl_sm120_padding_kv_only_wrapper(): From 5c9506cb19304d355cdac6e6b87c00e43fb1c979 Mon Sep 17 00:00:00 2001 From: barretw Date: Wed, 12 Aug 2026 01:38:52 -0700 Subject: [PATCH 4/6] right band --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 18 ++++-- python/cudnn/sdpa/bwd/api_dsl.py | 14 +++++ python/cudnn/sdpa/bwd/config_sm120.py | 5 ++ python/cudnn/sdpa/bwd/engines.py | 10 ++-- .../cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 22 ++++--- .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 58 ++++++++++++++++--- .../sdpa/frost/test_sdpa_graph_analyzer.py | 9 +++ 7 files changed, 111 insertions(+), 25 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 6c086e722..4a9dbf6ad 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -39,6 +39,8 @@ grads = sdpa_bwd_wrapper_dsl_sm120( is_causal=True, causal_bottom_right=False, window_size_left=None, # W: keys with k < q + diag - W are masked + window_size_right=None, # R: widen the causal diagonal right by R keys + # (keep k <= q + diag + R deterministic=False, # ordered dQ KV-tile reduction (bitwise-reproducible) scale_softmax=None, # None -> 1/sqrt(D) seq_q_lens=None, # (B,) int32 per-batch Q lengths (padding mask) @@ -181,9 +183,11 @@ nothing extra is carved, and no `reduce` kernel is launched. 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. + (`q_block_min`, bottom-right via `diag_off = S_kv − S_q`, a right band via + the compile-time widening `q_block_min = (kv_base − diag_off − R) / tile_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` / `do_mask_pad` gates); interior tiles skip it. @@ -229,9 +233,11 @@ only the unused relay operand remains in the kernel ABI. - 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), padding (per-batch - `seq_kv_len` required, `seq_q_len` optional; composes with the other masks) +- 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 + causal), padding (per-batch `seq_kv_len` required, `seq_q_len` optional; + composes with the other masks) - GQA/MQA: any `H_kv` dividing `H_q` (including `H_kv == 1`) - No dropout / bias / ALiBi / sinks / softcap / THD - Workspace (carved from the caller's buffer): fp32 `delta` and `dq_accum` diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index e6de829e5..f73036679 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -69,6 +69,7 @@ def __init__( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, deterministic: bool = False, scale_softmax: Optional[float] = None, tile_m: Optional[int] = None, @@ -93,6 +94,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.window_size_right = None if window_size_right is None else int(window_size_right) self.deterministic = bool(deterministic) self.scale_softmax = scale_softmax self.tile_m = None if tile_m is None else int(tile_m) @@ -252,6 +254,14 @@ def check_support(self) -> bool: self.window_size_left is not None and self.window_size_left < 0, f"window_size_left must be non-negative, got {self.window_size_left}", ) + self._value_error_if( + self.window_size_right is not None and self.window_size_right < 0, + f"window_size_right must be non-negative, got {self.window_size_right}", + ) + self._value_error_if( + self.window_size_right is not None and not self.is_causal, + "window_size_right widens the causal diagonal and requires is_causal=True", + ) self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") self.compute_capability = torch.cuda.get_device_capability(self.q_desc.device) @@ -288,6 +298,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, + window_size_right=self.window_size_right, deterministic=self.deterministic, q_tile=self.q_tile, kv_tile=self.kv_tile, @@ -502,6 +513,7 @@ def sdpa_bwd_wrapper_dsl_sm120( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, deterministic: bool = False, scale_softmax: Optional[float] = None, seq_q_lens: Optional[torch.Tensor] = None, @@ -526,6 +538,7 @@ def sdpa_bwd_wrapper_dsl_sm120( bool(is_causal), bool(causal_bottom_right), window_size_left, + window_size_right, bool(deterministic), scale_softmax, seq_q_lens is not None, @@ -546,6 +559,7 @@ def sdpa_bwd_wrapper_dsl_sm120( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + window_size_right=window_size_right, deterministic=deterministic, scale_softmax=scale_softmax, seq_kv_lens_present=seq_kv_lens is not None, diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py index b75885b19..51eaa7d38 100644 --- a/python/cudnn/sdpa/bwd/config_sm120.py +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -34,6 +34,7 @@ class TemplateParams: is_causal: bool = False causal_top_left: bool = False window_size_left: int | None = None + window_size_right: int | None = None deterministic: bool = False use_pdl: bool = True q_tile: int = 0 @@ -57,6 +58,10 @@ def validate_params(params: TemplateParams) -> None: raise ValueError("SM120 SDPA bwd: causal_top_left requires is_causal=True") if params.window_size_left is not None and params.window_size_left < 0: raise ValueError(f"SM120 SDPA bwd: window_size_left must be non-negative; got {params.window_size_left}") + if params.window_size_right is not None and params.window_size_right < 0: + raise ValueError(f"SM120 SDPA bwd: window_size_right must be non-negative; got {params.window_size_right}") + if params.window_size_right is not None and not params.is_causal: + raise ValueError("SM120 SDPA bwd: window_size_right widens the causal diagonal and requires is_causal=True") if params.seq_q_lens_present and not params.seq_kv_lens_present: raise ValueError("SM120 SDPA bwd: seq_q_lens_present requires seq_kv_lens_present (padding mask)") if params.q_tile not in (0,) + SEQ_Q_TILES: diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 0eb16987f..0b2ea4aeb 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -206,9 +206,9 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", requested: if fact and not cap: return f"graph uses {label}, which this engine does not support" - if facts.bottom_right and not facts.causal: - return "bottom-right alignment requires a causal upper bound" - if facts.causal and facts.bottom_right and not capabilities.bottom_right: + if facts.bottom_right and not (facts.causal or facts.right_band_widening): + return "bottom-right alignment requires a causal upper bound (plain or right-widened)" + if facts.bottom_right and not capabilities.bottom_right: return "graph uses bottom-right causal, which this engine does not support" # The kernel consumes the forward stats as a contiguous natural-log LSE @@ -246,6 +246,7 @@ def _sm120_spec() -> EngineSpec: gqa=True, causal=True, bottom_right=True, + right_band_widening=True, swa=True, padded=True, layouts=frozenset({"bshd", "dense_flex"}), @@ -330,9 +331,10 @@ def _desc(geom, dtype, name: str) -> "Any": 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, + is_causal=facts.causal or facts.right_band_widening, causal_bottom_right=facts.bottom_right, window_size_left=facts.window_left, + window_size_right=(facts.right_bound if facts.right_band_widening else None), deterministic=facts.deterministic, scale_softmax=facts.scale, tile_m=requested.tile_m if requested is not None else None, diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index 77d200e40..6ce696ecd 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -22,7 +22,8 @@ serves any other multiple of 8 up to 256 by zero-padding D. * GQA/MQA: H_q must be a multiple of H_kv * No dropout/alibi/softcap -* Optional causal (top-left or bottom-right), sliding-window masks and padding masks. +* Optional causal (top-left or bottom-right), right-band-widened causal + (window_size_right), sliding-window masks and padding masks. * LSE input is the natural-log forward stats, fp32 (B, H, SQ) contiguous One backward call is three kernel launches through the per-shape @@ -463,6 +464,7 @@ def __init__( is_causal: bool = False, causal_top_left: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, deterministic: bool = False, head_dim: int = 128, use_pdl: bool = True, @@ -475,6 +477,8 @@ def __init__( self.is_causal = is_causal self.causal_top_left = bool(causal_top_left) self.window_size_left = window_size_left + self.window_size_right = window_size_right + self.right_slack = window_size_right if window_size_right is not None else 0 self.deterministic = bool(deterministic) self.seq_kv_lens_present = bool(seq_kv_lens_present) self.seq_q_lens_present = bool(seq_q_lens_present) @@ -647,7 +651,8 @@ def kernel( else: diag_off = cutlass.Int32(0) if cutlass.const_expr(self.is_causal): - m_block_min = cute.math.max(kv_base - diag_off, cutlass.Int32(0)) // M + # kv is visible to q when kv <= q + diag_off + right_slack + m_block_min = cute.math.max(kv_base - diag_off - self.right_slack, cutlass.Int32(0)) // M else: m_block_min = cutlass.Int32(0) if cutlass.const_expr(self.window_size_left is not None): @@ -902,7 +907,7 @@ def kernel( # Mask + softmax (scores -> P, unscaled by attn_scale) and the # P store to smem. if cutlass.const_expr(self.is_causal): - do_mask_causal = (m_block * M) < (kv_base + N - diag_off) + do_mask_causal = (m_block * M) < (kv_base + N - diag_off - self.right_slack) if cutlass.const_expr(self.window_size_left is not None): do_mask_window = kv_base < (m_block * M + M - 1 + diag_off - self.window_size_left) if cutlass.const_expr(self.seq_kv_lens_present): @@ -922,13 +927,15 @@ def kernel( s3 = acc_s[off + 3] if cutlass.const_expr(self.is_causal): if do_mask_causal: - if kv_a0 > r0 + diag_off: + hi0 = r0 + diag_off + self.right_slack + hi8 = r8 + diag_off + self.right_slack + if kv_a0 > hi0: s0 = neg_inf - if kv_a1 > r0 + diag_off: + if kv_a1 > hi0: s1 = neg_inf - if kv_a0 > r8 + diag_off: + if kv_a0 > hi8: s2 = neg_inf - if kv_a1 > r8 + diag_off: + if kv_a1 > hi8: s3 = neg_inf if cutlass.const_expr(self.window_size_left is not None): if do_mask_window: @@ -1621,6 +1628,7 @@ def compile( # noqa: A001 is_causal=PARAMS.is_causal, causal_top_left=PARAMS.causal_top_left, window_size_left=PARAMS.window_size_left, + window_size_right=PARAMS.window_size_right, deterministic=PARAMS.deterministic, head_dim=d, use_pdl=PARAMS.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 ac85a37dd..9d78da2e4 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -68,6 +68,7 @@ def _ref_bwd( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, padding: tuple[list[int], list[int]] | None = None, ): """Reference via the canonical refs (sdpa/fp16_ref.py).""" @@ -76,7 +77,7 @@ def _ref_bwd( from sdpa.fp16_ref import compute_ref, compute_ref_backward diag_align = cudnn.diagonal_alignment.BOTTOM_RIGHT if causal_bottom_right else cudnn.diagonal_alignment.TOP_LEFT - right_bound = 0 if is_causal else None + right_bound = window_size_right if window_size_right is not None else (0 if is_causal else None) # The refs take the cuDNN window LENGTH; window_size_left is the offset W = L - 1. left_bound = None if window_size_left is None else window_size_left + 1 o_ref, stats_ref, _, _ = compute_ref( @@ -126,6 +127,7 @@ def _run_bwd_graph( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, deterministic: bool = False, select: bool = True, q_tile: int | None = None, @@ -169,12 +171,18 @@ def _run_bwd_graph( "stats": stats, "attn_scale": scale, } - if causal_bottom_right: - bwd_kwargs["use_causal_mask_bottom_right"] = True - elif is_causal: - bwd_kwargs["use_causal_mask"] = True - if window_size_left is not None: - bwd_kwargs["sliding_window_length"] = window_size_left + 1 + if window_size_right is not None: + bwd_kwargs["diagonal_band_right_bound"] = window_size_right + bwd_kwargs["diagonal_alignment"] = cudnn.diagonal_alignment.BOTTOM_RIGHT if causal_bottom_right else cudnn.diagonal_alignment.TOP_LEFT + if window_size_left is not None: + bwd_kwargs["diagonal_band_left_bound"] = window_size_left + 1 + else: + if causal_bottom_right: + bwd_kwargs["use_causal_mask_bottom_right"] = True + elif is_causal: + 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 seq_q_t = seq_kv_t = None @@ -254,6 +262,7 @@ def _run_case( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, deterministic: bool = False, select: bool = True, q_tile: int | None = None, @@ -269,7 +278,16 @@ def _run_case( v = _bhsd(batch, h_kv, s_kv, head_dim, dtype, layout=layout) do = _bhsd(batch, h_q, 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, padding=padding + q, + k, + v, + do, + scale=scale, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_size_left=window_size_left, + window_size_right=window_size_right, + padding=padding, ) o = _bhsd(batch, h_q, s_q, head_dim, dtype, empty=True, layout=layout).copy_(o) seq_q_lens = seq_kv_lens = None @@ -287,6 +305,7 @@ def _run_case( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + window_size_right=window_size_right, deterministic=deterministic, select=select, q_tile=q_tile, @@ -430,6 +449,29 @@ 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 +@torch_fork_set_rng(seed=46) +def test_sdpa_bwd_dsl_sm120_right_band(): + """diagonal_band_right_bound > 0: the causal diagonal widened right by a + compile-time R (keep kv <= q + diag + R).""" + + _run_case(s_q=256, s_kv=256, head_dim=64, window_size_right=32) # top-left band + _run_case(s_q=192, s_kv=320, head_dim=64, causal_bottom_right=True, window_size_right=48) # bottom-right anchor + _run_case(s_q=256, s_kv=256, head_dim=64, window_size_left=64, window_size_right=32) # full band + _run_case(s_q=193, s_kv=257, head_dim=128, window_size_right=24) # ragged tails + _run_case(s_q=128, s_kv=128, head_dim=64, window_size_right=300) # R >= S_kv clamps to dense + _run_case(s_q=256, s_kv=256, head_dim=64, window_size_right=32, deterministic=True) # relay turns unaffected by R + _run_case( + s_q=256, + s_kv=256, + head_dim=64, + causal_bottom_right=True, + window_size_right=48, + window_size_left=96, + padding=([230, 120], [180, 240]), + ) # per-batch diagonal + R + + @pytest.mark.L0 @pytest.mark.parametrize("mask", ["dense", "causal_tl", "causal_br"]) @torch_fork_set_rng(seed=20) diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 1da0dceae..92abf4af0 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -1018,6 +1018,15 @@ 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_accepts_right_band_widening(monkeypatch): + # diagonal_band_right_bound > 0 lowers as causal with a right offset. + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + import cudnn + + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(diagonal_band_right_bound=16)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(s_q=S // 2, diagonal_band_right_bound=16, diagonal_alignment=cudnn.diagonal_alignment.BOTTOM_RIGHT)) + + 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)) From 30de4269378b1332f56fb82c3395e68294e18bfc Mon Sep 17 00:00:00 2001 From: barretw Date: Wed, 12 Aug 2026 21:04:26 -0700 Subject: [PATCH 5/6] add dsink --- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md | 11 +- python/cudnn/sdpa/bwd/api_dsl.py | 45 ++++++++ python/cudnn/sdpa/bwd/config_sm120.py | 5 + python/cudnn/sdpa/bwd/engines.py | 15 +++ .../cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 103 +++++++++++++++++- .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 83 +++++++++++--- .../sdpa/frost/test_sdpa_graph_analyzer.py | 16 +++ 7 files changed, 258 insertions(+), 20 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md index 4a9dbf6ad..6aefb562a 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_sm120.md @@ -45,6 +45,8 @@ grads = sdpa_bwd_wrapper_dsl_sm120( scale_softmax=None, # None -> 1/sqrt(D) seq_q_lens=None, # (B,) int32 per-batch Q lengths (padding mask) seq_kv_lens=None, # (B,) int32 per-batch KV lengths (padding mask) + sink_token=None, # fp32 (1, H_q, 1, 1) sink logits; adds dsink_tensor + # to the result ) dq, dk, dv = grads["dq_tensor"], grads["dk_tensor"], grads["dv_tensor"] ``` @@ -85,6 +87,8 @@ main the fused five-GEMM pass; writes dK/dV into dk_ws/dv_ws (aliased to the outputs for MHA, per-q-head partial buffers for GQA); accumulates dQ into dq_accum reduce GQA only: dK/dV = fixed-order sum of each KV head's group of q-head partials cvt dq_accum (fp32, scrambled) -> dQ (io dtype), applying attn_scale +dsink dSink_token graphs only, summing over every batch b and query row q: + dsink[h] = -sum_{b,q} exp(sink[h] - LSE[b,h,q]) * delta[b,h,q] ``` ### Main-kernel pipeline: KV-stationary, five chained GEMMs @@ -239,7 +243,12 @@ only the unused relay operand remains in the kernel ABI. causal), padding (per-batch `seq_kv_len` required, `seq_q_len` optional; composes with the other masks) - GQA/MQA: any `H_kv` dividing `H_q` (including `H_kv == 1`) -- No dropout / bias / ALiBi / sinks / softcap / THD +- Sink tokens: sink logits input and optional `dSink_token` output. dQ/dK/dV + need no sink code (the forward LSE already folds the sink into the softmax + denominator); the `dSink_token` output adds one tiny `dsink` reduce kernel + (`dsink[h] = -sum p_sink * delta`, fixed order — bitwise deterministic in + both modes) +- 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` diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index f73036679..5724a3421 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -66,6 +66,8 @@ def __init__( sample_dq: torch.Tensor | TensorDesc, sample_dk: torch.Tensor | TensorDesc, sample_dv: torch.Tensor | TensorDesc, + sample_sink: Optional[torch.Tensor | TensorDesc] = None, + sample_dsink: Optional[torch.Tensor | TensorDesc] = None, is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: Optional[int] = None, @@ -90,6 +92,8 @@ def __init__( self.dq_desc = self._make_tensor_desc(sample_dq, name="dQ") self.dk_desc = self._make_tensor_desc(sample_dk, name="dK") self.dv_desc = self._make_tensor_desc(sample_dv, name="dV") + self.sink_desc = self._make_tensor_desc(sample_sink, name="sink") if sample_sink is not None else None + self.dsink_desc = self._make_tensor_desc(sample_dsink, name="dSink") if sample_dsink is not None else None self.is_causal = bool(is_causal) self.causal_bottom_right = bool(causal_bottom_right) @@ -137,6 +141,8 @@ def execute( current_stream: Optional[cuda.CUstream] = None, seq_q_lens: Optional[torch.Tensor] = None, seq_kv_lens: Optional[torch.Tensor] = None, + sink_tensor: Optional[torch.Tensor] = None, + dsink_tensor: Optional[torch.Tensor] = None, ) -> None: """Execute the compiled kernel chain using the common operand set.""" @@ -262,6 +268,22 @@ def check_support(self) -> bool: self.window_size_right is not None and not self.is_causal, "window_size_right widens the causal diagonal and requires is_causal=True", ) + self._value_error_if( + self.dsink_desc is not None and self.sink_desc is None, + "dSink output requires a sink logits input", + ) + for desc in (self.sink_desc, self.dsink_desc): + if desc is None: + continue + self._value_error_if( + tuple(desc.shape) != (1, h_q, 1, 1), + f"{desc.name} must be (1, H_q, 1, 1) = (1, {h_q}, 1, 1); got {tuple(desc.shape)}", + ) + self._value_error_if( + not desc.is_contiguous(), + f"{desc.name} must be contiguous; got stride {desc.stride}", + ) + self._check_dtype(desc, torch.float32, name=desc.name) self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") self.compute_capability = torch.cuda.get_device_capability(self.q_desc.device) @@ -304,6 +326,8 @@ def compile(self) -> None: kv_tile=self.kv_tile, seq_kv_lens_present=self.seq_kv_lens_present, seq_q_lens_present=self.seq_q_lens_present, + sink_present=self.sink_desc is not None, + dsink_present=self.dsink_desc is not None, ) self._k_mod = _load_sm120_kernel_module(params) self._compiled_kernel = self._k_mod.compile( @@ -380,6 +404,8 @@ def execute( current_stream: Optional[cuda.CUstream] = None, seq_q_lens: Optional[torch.Tensor] = None, seq_kv_lens: Optional[torch.Tensor] = None, + sink_tensor: Optional[torch.Tensor] = None, + dsink_tensor: Optional[torch.Tensor] = None, ) -> None: """Execute tensors matching the compiled specialization.""" @@ -402,6 +428,14 @@ def execute( not self.seq_q_lens_present and seq_q_lens is not None, "this specialization was compiled without per-batch Q lengths; construct the API with seq_q_lens_present=True", ) + self._value_error_if( + self.dsink_desc is not None and (sink_tensor is None or dsink_tensor is None), + "sink_tensor and dsink_tensor are required by this compiled specialization", + ) + self._value_error_if( + self.dsink_desc is None and dsink_tensor is not None, + "this specialization was compiled without a dSink output; construct the API with sample_dsink", + ) scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) scale_log2 = scale_val * math.log2(math.e) @@ -490,6 +524,8 @@ def _staged_out_bshd(tensor: torch.Tensor): if kernels.reduce is not None: kernels.reduce(dk_ws, dv_ws, dk, dv, current_stream) 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), 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]) @@ -518,12 +554,14 @@ def sdpa_bwd_wrapper_dsl_sm120( scale_softmax: Optional[float] = None, seq_q_lens: Optional[torch.Tensor] = None, seq_kv_lens: Optional[torch.Tensor] = None, + sink_token: Optional[torch.Tensor] = None, ) -> TupleDict: """Run SM120 SDPA backward and return ``TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=...)``.""" dq_tensor = torch.empty_strided(q_tensor.shape, q_tensor.stride(), dtype=q_tensor.dtype, device=q_tensor.device) dk_tensor = torch.empty_strided(k_tensor.shape, k_tensor.stride(), dtype=k_tensor.dtype, device=k_tensor.device) dv_tensor = torch.empty_strided(v_tensor.shape, v_tensor.stride(), dtype=v_tensor.dtype, device=v_tensor.device) + dsink_tensor = torch.empty_like(sink_token) if sink_token is not None else None # check_support()/compile() run only on a miss, so the key must carry the # full signature of every operand the specialization depends on (dq/dk/dv @@ -543,6 +581,7 @@ def sdpa_bwd_wrapper_dsl_sm120( scale_softmax, seq_q_lens is not None, seq_kv_lens is not None, + sink_token is not None, ) api = _wrapper_api_cache.get(cache_key) if api is None: @@ -556,6 +595,8 @@ def sdpa_bwd_wrapper_dsl_sm120( sample_dq=dq_tensor, sample_dk=dk_tensor, sample_dv=dv_tensor, + sample_sink=sink_token, + sample_dsink=dsink_tensor, is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, @@ -584,5 +625,9 @@ def sdpa_bwd_wrapper_dsl_sm120( workspace=workspace, seq_q_lens=seq_q_lens, seq_kv_lens=seq_kv_lens, + sink_tensor=sink_token, + dsink_tensor=dsink_tensor, ) + if dsink_tensor is not None: + return TupleDict(dq_tensor=dq_tensor, dk_tensor=dk_tensor, dv_tensor=dv_tensor, dsink_tensor=dsink_tensor) return TupleDict(dq_tensor=dq_tensor, dk_tensor=dk_tensor, dv_tensor=dv_tensor) diff --git a/python/cudnn/sdpa/bwd/config_sm120.py b/python/cudnn/sdpa/bwd/config_sm120.py index 51eaa7d38..7db5a1386 100644 --- a/python/cudnn/sdpa/bwd/config_sm120.py +++ b/python/cudnn/sdpa/bwd/config_sm120.py @@ -42,6 +42,9 @@ class TemplateParams: # Padding mask: per-batch int32 lengths; seq_q is only valid with seq_kv. seq_kv_lens_present: bool = False seq_q_lens_present: bool = False + # Sink Attention + sink_present: bool = False + dsink_present: bool = False def validate_params(params: TemplateParams) -> None: @@ -64,6 +67,8 @@ def validate_params(params: TemplateParams) -> None: raise ValueError("SM120 SDPA bwd: window_size_right widens the causal diagonal and requires is_causal=True") if params.seq_q_lens_present and not params.seq_kv_lens_present: raise ValueError("SM120 SDPA bwd: seq_q_lens_present requires seq_kv_lens_present (padding mask)") + if params.dsink_present and not params.sink_present: + raise ValueError("SM120 SDPA bwd: dsink_present requires sink_present (a dSink output needs the sink logits)") if params.q_tile not in (0,) + SEQ_Q_TILES: raise ValueError(f"SM120 SDPA bwd: q_tile must be one of {(0,) + SEQ_Q_TILES} (0 = per-head-dim default); got {params.q_tile}") if params.kv_tile not in (0,) + SEQ_KV_TILES: diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 0b2ea4aeb..ba22e404f 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -206,6 +206,8 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", requested: if fact and not cap: return f"graph uses {label}, which this engine does not support" + if facts.has_dsink and not facts.has_sink: + return "dSink_token output requires a sink_token input" if facts.bottom_right and not (facts.causal or facts.right_band_widening): return "bottom-right alignment requires a causal upper bound (plain or right-widened)" if facts.bottom_right and not capabilities.bottom_right: @@ -249,6 +251,8 @@ def _sm120_spec() -> EngineSpec: right_band_widening=True, swa=True, padded=True, + sink=True, + dsink=True, layouts=frozenset({"bshd", "dense_flex"}), deterministic=True, tile_ms=frozenset(_SM120_Q_TILES), @@ -320,6 +324,9 @@ def _desc(geom, dtype, name: str) -> "Any": seq_kv_t = facts.seq_kv_t if facts.padded else None seq_q_t = facts.seq_q_t if facts.padded else None + # Sink ports: geometry straight from the IR tensors (fp32 (1, H_q, 1, 1)). + sink_geom = (tuple(facts.sink_t.get_dim()), tuple(facts.sink_t.get_stride())) if facts.has_sink else None + dsink_geom = (tuple(facts.dsink_t.get_dim()), tuple(facts.dsink_t.get_stride())) if facts.has_dsink else None api = _adapter_sm120()( sample_q=_desc(q_geom, facts.dtype, "q"), @@ -331,6 +338,8 @@ def _desc(geom, dtype, name: str) -> "Any": sample_dq=_desc(dq_geom, facts.dtype, "dQ"), sample_dk=_desc(dk_geom, facts.dtype, "dK"), sample_dv=_desc(dv_geom, facts.dtype, "dV"), + sample_sink=_desc(sink_geom, torch.float32, "sink") if sink_geom is not None else None, + sample_dsink=_desc(dsink_geom, torch.float32, "dSink") if dsink_geom is not None else None, is_causal=facts.causal or facts.right_band_widening, causal_bottom_right=facts.bottom_right, window_size_left=facts.window_left, @@ -363,6 +372,8 @@ def _desc(geom, dtype, name: str) -> "Any": dv=facts.dv_t, seq_len_kv=seq_kv_t, seq_len_q=seq_q_t, + sink_token=facts.sink_t if facts.has_sink else None, + dsink=facts.dsink_t if facts.has_dsink else None, ) def _canonical_view(buf, geom): @@ -383,6 +394,8 @@ def _execute(variant_pack, workspace=None, stream=None): resolved = ga.resolve_variant_pack(variant_pack, binding) seq_kv_buf = resolved.get(id(binding.seq_len_kv)) if binding.seq_len_kv is not None else None seq_q_buf = resolved.get(id(binding.seq_len_q)) if binding.seq_len_q is not None else None + sink_buf = _canonical_view(resolved[id(binding.sink_token)], sink_geom) if binding.sink_token is not None else None + dsink_buf = _canonical_view(resolved[id(binding.dsink)], dsink_geom) if binding.dsink is not None else None api.execute( q_tensor=_canonical_view(resolved[id(binding.q)], q_geom), k_tensor=_canonical_view(resolved[id(binding.k)], k_geom), @@ -395,6 +408,8 @@ def _execute(variant_pack, workspace=None, stream=None): dv_tensor=_canonical_view(resolved[id(binding.dv)], dv_geom), seq_q_lens=seq_q_buf, seq_kv_lens=seq_kv_buf, + sink_tensor=sink_buf, + dsink_tensor=dsink_buf, 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 6ce696ecd..ade431000 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -29,7 +29,8 @@ One backward call is three kernel launches through the per-shape ``compile()`` cache at the bottom of this module: ``dot`` (delta = rowsum(dO*O), also zeroes the dq_accum workspace), ``main`` (the fused -five-GEMM pass writing dK/dV), and ``cvt`` (dq_accum fp32 -> dQ io dtype). +five-GEMM pass writing dK/dV), and ``cvt`` (dq_accum fp32 -> dQ io dtype); +GQA adds ``reduce`` and a dSink_token output adds ``dsink``. """ from functools import lru_cache @@ -472,6 +473,7 @@ def __init__( kv_tile: int = 0, seq_kv_lens_present: bool = False, seq_q_lens_present: bool = False, + sink_present: bool = False, ): self.in_dtype = in_dtype self.is_causal = is_causal @@ -482,6 +484,8 @@ def __init__( self.deterministic = bool(deterministic) self.seq_kv_lens_present = bool(seq_kv_lens_present) 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.use_pdl = bool(use_pdl) self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim] @@ -810,6 +814,10 @@ def kernel( # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. if val == cutlass.Float32(float("-inf")): val = cutlass.Float32(float("inf")) + if cutlass.const_expr(self.trim_q_rows): + # explicit padded-row trim (sink LSE is finite there) + if r_abs >= seqlen_q: + val = cutlass.Float32(float("inf")) lse_r[rep * 2 + hf] = val * cutlass.Float32(_LOG2E) while not prims.mbarrier_try_wait_parity(v_mbar, cutlass.Int32(0)): @@ -1136,6 +1144,10 @@ def kernel( # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. if val == cutlass.Float32(float("-inf")): val = cutlass.Float32(float("inf")) + if cutlass.const_expr(self.trim_q_rows): + # explicit padded-row trim (sink LSE is finite there) + 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 if cutlass.const_expr(self.deterministic): @@ -1176,6 +1188,10 @@ def kernel( # -inf - (-inf) = NaN. Use +inf to reconstruct P = 0. if val == cutlass.Float32(float("-inf")): val = cutlass.Float32(float("inf")) + if cutlass.const_expr(self.trim_q_rows): + # explicit padded-row trim (sink LSE is finite there) + 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 if cutlass.const_expr(self.deterministic): @@ -1604,6 +1620,74 @@ def _dkv_reduce_host( ) +@cute.kernel +def _dsink_kernel( + lse: cute.Tensor, # [B, HQ, SQ] fp32 (natural-log, sink folded in by the fwd) + delta: cute.Tensor, # [B, HQ, SQ_r128] fp32 (dot_do_o output) + sink: cute.Tensor, # [HQ] fp32 sink logits + dsink: cute.Tensor, # [HQ] fp32 out + use_pdl: cutlass.Constexpr[bool], +): + """dsink[h] = -sum_{b,q} exp(sink[h] - lse[b,h,q]) * delta[b,h,q]. + + One warp per query head, fixed reduction order -> bitwise deterministic.""" + if cutlass.const_expr(use_pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + head, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + B = lse.shape[0] + HQ = lse.shape[1] + SQ = lse.shape[2] + SQ_R = delta.shape[2] + lse_ptr = lse.iterator.raw_ptr() + delta_ptr = delta.iterator.raw_ptr() + s_val = (sink.iterator.raw_ptr() + head).load() + inf = cutlass.Float32(float("inf")) + acc = cutlass.Float32(0.0) + batch = cutlass.Int32(0) + # batch loop + while batch < B: + lse_base = (batch * HQ + head) * SQ + delta_base = (batch * HQ + head) * SQ_R + q = cutlass.Int32(tidx) + while q < SQ: + lv = (lse_ptr + lse_base + q).load() + # Skip padded / trimmed rows which carry LSE = -inf and delta = 0 + if lv > -inf and lv < inf: + dd = (delta_ptr + delta_base + q).load() + acc = acc + cute.math.exp2((s_val - lv) * cutlass.Float32(_LOG2E), fastmath=True) * dd + q = q + 32 + batch = batch + 1 + for sh in cutlass.range_constexpr(5): + acc = acc + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=acc, + offset=1 << (4 - sh), + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + if tidx == 0: + (dsink.iterator.raw_ptr() + head).store(-acc) + + +@cute.jit +def _dsink_host( + lse: cute.Tensor, + delta: cute.Tensor, + sink: cute.Tensor, + dsink: cute.Tensor, + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +): + _dsink_kernel(lse, delta, sink, dsink, use_pdl).launch( + grid=(lse.shape[1], 1, 1), + block=(32, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + @lru_cache(maxsize=None) def compile( # noqa: A001 compute_capability: tuple[int, int], @@ -1636,6 +1720,7 @@ def compile( # noqa: A001 kv_tile=PARAMS.kv_tile, seq_kv_lens_present=PARAMS.seq_kv_lens_present, seq_q_lens_present=PARAMS.seq_q_lens_present, + sink_present=PARAMS.sink_present, ) sq_r = ceil_div(sq, 128) * 128 @@ -1734,4 +1819,18 @@ def _fake(dtype, shape): fake_stream, options=options, ) - return SimpleNamespace(dot=compiled_dot, main=compiled_main, cvt=compiled_cvt, reduce=compiled_reduce) + compiled_dsink = None + if PARAMS.dsink_present: + fake_sink = _fake(cutlass.Float32, (qh,)) + fake_dsink = _fake(cutlass.Float32, (qh,)) + compiled_dsink = cute.compile( + _dsink_host, + fake_lse, + fake_delta, + fake_sink, + fake_dsink, + bwd.use_pdl, + fake_stream, + options=options, + ) + return SimpleNamespace(dot=compiled_dot, main=compiled_main, cvt=compiled_cvt, reduce=compiled_reduce, dsink=compiled_dsink) 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 9d78da2e4..b07d161bc 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -70,6 +70,7 @@ def _ref_bwd( window_size_left: int | None = None, window_size_right: int | None = None, padding: tuple[list[int], list[int]] | None = None, + sink_token: torch.Tensor | None = None, ): """Reference via the canonical refs (sdpa/fp16_ref.py).""" @@ -81,12 +82,32 @@ def _ref_bwd( # The refs take the cuDNN window LENGTH; window_size_left is the offset W = L - 1. left_bound = None if window_size_left is None else window_size_left + 1 o_ref, stats_ref, _, _ = compute_ref( - q, k, v, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, padding=padding, torch_type=q.dtype + q, + k, + v, + attn_scale=scale, + diag_align=diag_align, + right_bound=right_bound, + left_bound=left_bound, + padding=padding, + sink_token=sink_token, + torch_type=q.dtype, ) - dq, dk, dv, _, _ = compute_ref_backward( - q, k, v, o_ref, do, attn_scale=scale, diag_align=diag_align, right_bound=right_bound, left_bound=left_bound, padding=padding, torch_type=q.dtype + dq, dk, dv, _, dsink = compute_ref_backward( + q, + k, + v, + o_ref, + do, + attn_scale=scale, + diag_align=diag_align, + right_bound=right_bound, + left_bound=left_bound, + padding=padding, + sink_token=sink_token, + torch_type=q.dtype, ) - return o_ref.to(q.dtype), stats_ref.contiguous(), dq.to(q.dtype), dk.to(q.dtype), dv.to(q.dtype) + return o_ref.to(q.dtype), stats_ref.contiguous(), dq.to(q.dtype), dk.to(q.dtype), dv.to(q.dtype), dsink def _expected_workspace_bytes( @@ -135,8 +156,9 @@ def _run_bwd_graph( grad_layout: str = "bshd", seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]: - """Build and execute the SM120 FROST backward graph; returns (dq, dk, dv, plan_name).""" + sink_gpu: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, "torch.Tensor | None", str]: + """Build and execute the SM120 FROST backward graph; returns (dq, dk, dv, dsink, plan_name).""" _require_dsl() import cudnn @@ -185,6 +207,12 @@ def _run_bwd_graph( bwd_kwargs["sliding_window_length"] = window_size_left + 1 if deterministic: bwd_kwargs["use_deterministic_algorithm"] = True + sink_t = dsink_t = dsink_gpu = None + if sink_gpu is not None: + sink_t = graph.tensor_like(sink_gpu, name="sink") + dsink_gpu = torch.empty_like(sink_gpu) + dsink_t = graph.tensor_like(dsink_gpu, name="dSink") + bwd_kwargs.update(sink_token=sink_t, dSink_token=dsink_t) seq_q_t = seq_kv_t = None if seq_q_lens is not None or seq_kv_lens is not None: assert seq_q_lens is not None and seq_kv_lens is not None @@ -241,9 +269,11 @@ def _run_bwd_graph( } if seq_q_t is not None: variant_pack.update({seq_q_t: seq_q_lens, seq_kv_t: seq_kv_lens}) + if sink_t is not None: + variant_pack.update({sink_t: sink_gpu, dsink_t: dsink_gpu}) graph.execute(variant_pack, workspace) torch.cuda.synchronize() - return dq_gpu, dk_gpu, dv_gpu, plan_name + return dq_gpu, dk_gpu, dv_gpu, dsink_gpu, plan_name def _tolerances(dtype: torch.dtype) -> dict: @@ -270,6 +300,7 @@ def _run_case( layout: str = "bshd", grad_layout: str = "bshd", padding: tuple[list[int], list[int]] | None = None, + sink: bool = False, ) -> str: h_kv = h_q if h_kv is None else h_kv scale = 1.0 / math.sqrt(head_dim) @@ -277,7 +308,8 @@ def _run_case( 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) - o, stats, dq_ref, dk_ref, dv_ref = _ref_bwd( + 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, k, v, @@ -288,13 +320,14 @@ def _run_case( window_size_left=window_size_left, window_size_right=window_size_right, padding=padding, + sink_token=sink_gpu, ) o = _bhsd(batch, h_q, s_q, head_dim, 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) seq_kv_lens = torch.tensor(padding[1], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) - dq, dk, dv, plan_name = _run_bwd_graph( + dq, dk, dv, dsink, plan_name = _run_bwd_graph( q, k, v, @@ -313,11 +346,14 @@ def _run_case( grad_layout=grad_layout, seq_q_lens=seq_q_lens, seq_kv_lens=seq_kv_lens, + sink_gpu=sink_gpu, ) tol = _tolerances(dtype) torch.testing.assert_close(dq.float(), dq_ref.float(), **tol) torch.testing.assert_close(dk.float(), dk_ref.float(), **tol) torch.testing.assert_close(dv.float(), dv_ref.float(), **tol) + if sink: + torch.testing.assert_close(dsink.float(), dsink_ref.float(), **tol) if padding is not None: for b, (len_q, len_kv) in enumerate(zip(*padding)): if len_q < s_q: @@ -379,7 +415,7 @@ def test_sdpa_bwd_dsl_sm120_causal_br_sq_gt_skv(s_q: int, s_kv: int): 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, causal_bottom_right=True) + o, stats, dq_ref, dk_ref, dv_ref, _ = _ref_bwd(q, k, v, do, scale=scale, is_causal=True, causal_bottom_right=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, causal_bottom_right=True, scale_softmax=scale) tol = _tolerances(dtype) @@ -472,6 +508,19 @@ def test_sdpa_bwd_dsl_sm120_right_band(): ) # per-batch diagonal + R +@pytest.mark.L0 +@torch_fork_set_rng(seed=49) +def test_sdpa_bwd_dsl_sm120_sink(): + """Sink attention""" + + _run_case(s_q=256, s_kv=256, head_dim=64, sink=True) # dense + _run_case(s_q=256, s_kv=256, head_dim=64, is_causal=True, sink=True) # causal + _run_case(h_q=8, h_kv=2, s_q=256, s_kv=256, head_dim=64, sink=True) # GQA + _run_case(s_q=256, s_kv=256, head_dim=64, sink=True, deterministic=True) # fixed-order reduce + _run_case(s_q=193, s_kv=257, head_dim=128, is_causal=True, sink=True) # ragged tails + _run_case(s_q=256, s_kv=256, head_dim=64, sink=True, padding=([230, 120], [180, 240])) # padded rows skip (LSE = -inf guard) + + @pytest.mark.L0 @pytest.mark.parametrize("mask", ["dense", "causal_tl", "causal_br"]) @torch_fork_set_rng(seed=20) @@ -593,7 +642,7 @@ def test_sdpa_bwd_dsl_sm120_padding_kv_only_wrapper(): 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, padding=([s_q] * batch, kv_lens)) + o, stats, dq_ref, dk_ref, dv_ref, _ = _ref_bwd(q, k, v, do, scale=scale, padding=([s_q] * batch, kv_lens)) 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, scale_softmax=scale, seq_kv_lens=torch.tensor(kv_lens, dtype=torch.int32, device="cuda")) tol = _tolerances(dtype) @@ -633,7 +682,7 @@ def test_sdpa_bwd_dsl_sm120_large_d_wrapper(mask: str, head_dim: int): 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, 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) @@ -676,7 +725,7 @@ def test_sdpa_bwd_dsl_sm120_padded_head_dim_wrapper(head_dim: int): 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, 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) @@ -783,7 +832,7 @@ def _run_bitwise_case(n_runs: int = 3, padding: tuple[list[int], list[int]] | No k = _bhsd(batch, h_kv, s_kv, head_dim, dtype) v = _bhsd(batch, h_kv, s_kv, head_dim, dtype) do = _bhsd(batch, heads, s_q, head_dim, dtype) - o, stats, _, _, _ = _ref_bwd( + o, stats, _, _, _, _ = _ref_bwd( q, k, v, @@ -799,8 +848,8 @@ def _run_bitwise_case(n_runs: int = 3, padding: tuple[list[int], list[int]] | No case_kwargs["seq_q_lens"] = torch.tensor(padding[0], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) case_kwargs["seq_kv_lens"] = torch.tensor(padding[1], dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) 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): + 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" @@ -863,7 +912,7 @@ def _run_wrapper_det_case(head_dim: int, *, s_q: int, s_kv: int, is_causal: bool 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, 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) diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 92abf4af0..7c959183f 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -893,6 +893,8 @@ def _mk_bwd_graph( grad_strides: tuple | None = None, bias: bool = False, dbias: bool = False, + sink: bool = False, + dsink: bool = False, seq_lens: str | None = None, # "kv" / "both" (padding mask) or "q_only" **bwd_kwargs, ): @@ -922,6 +924,12 @@ def _mk_bwd_graph( if dbias: dbias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="dBias") bwd_kwargs.update(dBias=dbias_t) + if sink: + sink_t = g.tensor(dim=(1, H, 1, 1), stride=(H, 1, 1, 1), data_type=cudnn.data_type.FLOAT, name="sink") + bwd_kwargs.update(sink_token=sink_t) + if dsink: + dsink_t = g.tensor(dim=(1, H, 1, 1), stride=(H, 1, 1, 1), data_type=cudnn.data_type.FLOAT, name="dSink") + bwd_kwargs.update(dSink_token=dsink_t) if seq_lens in ("kv", "both"): seq_kv_t = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_kv") bwd_kwargs.update(use_padding_mask=True, seq_len_kv=seq_kv_t) @@ -1048,6 +1056,14 @@ def test_bwd_probe_rejects_seq_len_q_without_padding_mask(monkeypatch): assert not _bwd_eligible(_mk_bwd_graph(seq_lens="q_only")) +def test_bwd_probe_accepts_sink(monkeypatch): + # dSink without the sink input is rejected. + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(sink=True)) + assert _BWD_ENGINE in _bwd_eligible(_mk_bwd_graph(sink=True, dsink=True, use_causal_mask=True)) + assert not _bwd_eligible(_mk_bwd_graph(dsink=True)) + + def test_bwd_probe_rejects_bias(monkeypatch): monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) assert not _bwd_eligible(_mk_bwd_graph(bias=True)) From fa63cb2ae95aa00f2bb67c6794153cad097418a9 Mon Sep 17 00:00:00 2001 From: barretw Date: Wed, 12 Aug 2026 23:55:59 -0700 Subject: [PATCH 6/6] fix --- python/cudnn/sdpa/bwd/api_dsl.py | 8 ++++++-- python/cudnn/sdpa/bwd/engines.py | 4 ++-- python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py | 12 +++++++++--- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/bwd/api_dsl.py b/python/cudnn/sdpa/bwd/api_dsl.py index 5724a3421..7702552fc 100644 --- a/python/cudnn/sdpa/bwd/api_dsl.py +++ b/python/cudnn/sdpa/bwd/api_dsl.py @@ -275,6 +275,10 @@ def check_support(self) -> bool: for desc in (self.sink_desc, self.dsink_desc): if desc is None: continue + self._value_error_if( + desc.device != self.q_desc.device, + f"{desc.name} must be on {self.q_desc.device} (with Q); got {desc.device}", + ) self._value_error_if( tuple(desc.shape) != (1, h_q, 1, 1), f"{desc.name} must be (1, H_q, 1, 1) = (1, {h_q}, 1, 1); got {tuple(desc.shape)}", @@ -525,7 +529,7 @@ def _staged_out_bshd(tensor: torch.Tensor): kernels.reduce(dk_ws, dv_ws, dk, dv, current_stream) 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), current_stream) + 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)): if user_view is not None: user_view.copy_(staged[..., : self.head_dim]) @@ -581,7 +585,7 @@ def sdpa_bwd_wrapper_dsl_sm120( scale_softmax, seq_q_lens is not None, seq_kv_lens is not None, - sink_token is not None, + _tensor_signature(sink_token) if sink_token is not None else None, ) api = _wrapper_api_cache.get(cache_key) if api is None: diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 885242e33..dc1298c77 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -355,8 +355,8 @@ def _desc(geom, dtype, name: str) -> "Any": sample_dq=_desc(dq_geom, facts.dtype, "dQ"), sample_dk=_desc(dk_geom, facts.dtype, "dK"), sample_dv=_desc(dv_geom, facts.dtype, "dV"), - sample_sink=_desc(sink_geom, torch.float32, "sink") if sink_geom is not None else None, - sample_dsink=_desc(dsink_geom, torch.float32, "dSink") if dsink_geom is not None else None, + sample_sink=_desc(sink_geom, facts.sink_t.get_data_type(), "sink") if sink_geom is not None else None, + sample_dsink=_desc(dsink_geom, facts.dsink_t.get_data_type(), "dSink") if dsink_geom is not None else None, is_causal=facts.causal or facts.right_band_widening, causal_bottom_right=facts.bottom_right, window_size_left=facts.window_left, diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py index ade431000..e26e2728d 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py @@ -1626,6 +1626,7 @@ def _dsink_kernel( delta: cute.Tensor, # [B, HQ, SQ_r128] fp32 (dot_do_o output) sink: cute.Tensor, # [HQ] fp32 sink logits dsink: cute.Tensor, # [HQ] fp32 out + seq_q_lens: Optional[cute.Tensor], # [B] int32; None unless seq_q_lens_present use_pdl: cutlass.Constexpr[bool], ): """dsink[h] = -sum_{b,q} exp(sink[h] - lse[b,h,q]) * delta[b,h,q]. @@ -1650,10 +1651,13 @@ def _dsink_kernel( while batch < B: lse_base = (batch * HQ + head) * SQ delta_base = (batch * HQ + head) * SQ_R + q_bound = SQ + if cutlass.const_expr(seq_q_lens is not None): + q_bound = cute.math.max(cutlass.Int32(0), cute.math.min(seq_q_lens[batch], cutlass.Int32(SQ))) q = cutlass.Int32(tidx) - while q < SQ: + while q < q_bound: lv = (lse_ptr + lse_base + q).load() - # Skip padded / trimmed rows which carry LSE = -inf and delta = 0 + # Padded / trimmed rows carry LSE = -inf: skip them, exp(sink + inf) * 0 = NaN if lv > -inf and lv < inf: dd = (delta_ptr + delta_base + q).load() acc = acc + cute.math.exp2((s_val - lv) * cutlass.Float32(_LOG2E), fastmath=True) * dd @@ -1677,10 +1681,11 @@ def _dsink_host( delta: cute.Tensor, sink: cute.Tensor, dsink: cute.Tensor, + seq_q_lens: Optional[cute.Tensor], use_pdl: cutlass.Constexpr[bool], stream: cuda_driver.CUstream, ): - _dsink_kernel(lse, delta, sink, dsink, use_pdl).launch( + _dsink_kernel(lse, delta, sink, dsink, seq_q_lens, use_pdl).launch( grid=(lse.shape[1], 1, 1), block=(32, 1, 1), stream=stream, @@ -1829,6 +1834,7 @@ def _fake(dtype, shape): fake_delta, fake_sink, fake_dsink, + fake_seq_q_lens, bwd.use_pdl, fake_stream, options=options,