From 11ff5d794c6a3dbf4306f57d0de14da5cedd6293 Mon Sep 17 00:00:00 2001 From: Haobin Guo Date: Thu, 20 Aug 2026 09:20:30 -0700 Subject: [PATCH] Retile Capabilities.bottom_right_padded_seq_q Signed-off-by: Haobin Guo --- python/cudnn/frost/README.md | 5 +- python/cudnn/sdpa/fwd/api_dsl.py | 21 --- python/cudnn/sdpa/fwd/engines.py | 15 -- .../cudnn/sdpa/fwd/kernels/_common_sm100.py | 27 ++-- .../fwd/kernels/prefill_d128_f16_sm100.py | 37 ++--- .../kernels/prefill_d192_d128_f16_sm100.py | 16 +-- .../fwd/kernels/prefill_d256_f16_sm100.py | 37 ++--- .../fwd/kernels/prefill_d512_f16_sm100.py | 37 ++--- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 128 +++++++++++++++--- .../frost/test_sdpa_fwd_split_kv_sm100.py | 17 ++- .../sdpa/frost/test_sdpa_graph_analyzer.py | 10 +- 11 files changed, 219 insertions(+), 131 deletions(-) diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 9ba0dbc7f..193193f2e 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -506,8 +506,9 @@ expressible, with one discipline separating them: - The box is pure data: per-axis fields on `Capabilities`. Covers most of the surface; adding an engine is writing a row, not logic. - A notch is a rule in `mismatch()` gated by a conjunction flag on the row - (e.g. `bottom_right_padded_seq_q: bool`). The matcher encodes the SHAPE of the - interaction once; each engine's row supplies the VERDICT. When a future + (e.g. `padded_stats: bool` — padding mask + generate_stats needs the + per-batch LSE trim). The matcher encodes the SHAPE of the interaction once; + each engine's row supplies the VERDICT. When a future kernel supports the conjunction, flip its flag -- never edit the matcher. This is what keeps interaction checks from regressing into a per-engine if-ladder: shared code may know about kinds of interactions, never about diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 04919736e..eb5380b05 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -847,15 +847,6 @@ def check_support(self) -> bool: self.causal_bottom_right and not self.is_causal, "SM100 DSL SDPA: causal_bottom_right requires is_causal=True", ) - # Backstop for the engines.bottom_right_padded_seq_q gate: with dense - # per-batch Q lengths the kernel's BR diagonal (anchored at the global - # S_q) is wrong for any batch with seq_len_q[b] < S_q. - self._value_error_if( - self.causal_bottom_right and self.seq_q_lens_present, - "SM100 DSL SDPA: causal_bottom_right with per-batch seq_len_q is not " - "supported (kernel anchors the BR diagonal at the global S_q, not " - "seq_len_q[b])", - ) if self.thd: self.seq_kv_lens_present = True self._not_implemented_error_if( @@ -884,18 +875,6 @@ def check_support(self) -> bool: self.seq_q_lens_present and self._fp8, "seq_q_lens_present (dense padded-Q LSE trim) is not plumbed for the FP8/MXFP8 kernels", ) - # Dense BR + per-batch Q lengths: the kernels anchor the bottom-right - # diagonal with the GLOBAL S_q (compute_kv_loop_bounds: causal_diag = - # seq_kv_len - seqlen_q with the scalar S_q), but cuDNN semantics - # anchor it at the per-batch (seq_len_q[b], seq_len_kv[b]) corner — - # batches with seq_len_q[b] < S_q get over-masked. KV-only padding is - # exact (actual Q length == S_q), so only this combination is gated. - self._value_error_if( - self.causal_bottom_right and self.seq_q_lens_present, - "SM100 DSL SDPA: causal_bottom_right with per-batch seq_len_q (dense " - "padded-Q trim) is not supported — the kernel anchors the BR diagonal " - "at the global S_q, not seq_len_q[b]/seq_len_kv[b]", - ) # KV-tail correctness: the kernel zero-fills the last KV tile via TMA # OOB but only *masks* those columns on the padded / causal paths. A # ragged S_kv is safe when a padding mask carries the real lengths, or diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 966adbc07..cb94c1579 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -140,13 +140,6 @@ class Capabilities: causal: bool = False bottom_right: bool = False - # Kernel gap (pre-existing): for a DENSE padded graph the BR diagonal is - # computed as seq_len_kv[b] - GLOBAL S_q, but cuDNN semantics for a dense - # padded graph carrying per-batch seq_len_q anchor it at - # (seq_len_q[b], seq_len_kv[b]) — any batch with seq_len_q[b] < S_q gets a - # wrongly shifted diagonal (extra zeroed rows at the top). KV-only padding - # (no seq_len_q tensor) is unaffected: there the actual Q length IS S_q. - bottom_right_padded_seq_q: bool = False swa: bool = False padded: bool = False sink: bool = False @@ -372,11 +365,6 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti return "bottom-right alignment requires a causal upper bound (plain or right-widened)" if not capabilities.bottom_right: return "graph uses bottom-right causal, which this kernel does not support" - if facts.padded and not facts.thd and facts.seq_q_t is not None and not capabilities.bottom_right_padded_seq_q: - return ( - "bottom-right causal with a dense padding mask carrying per-batch seq_len_q is not " - "supported (kernel anchors the BR diagonal at the global S_q, not seq_len_q[b])" - ) if facts.padded and facts.wants_stats and not facts.thd and not capabilities.padded_stats: return "padding mask with generate_stats is not supported yet (per-batch seq_len_q LSE trim not plumbed)" @@ -578,7 +566,6 @@ def _sm80_spec() -> EngineSpec: right_band_widening=True, causal=True, bottom_right=True, - bottom_right_padded_seq_q=True, swa=True, padded=True, sink=True, @@ -673,7 +660,6 @@ def _sm120_spec() -> EngineSpec: dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), causal=True, bottom_right=True, - bottom_right_padded_seq_q=True, swa=True, right_band_widening=True, padded=True, @@ -1014,7 +1000,6 @@ def _sm120_fp8_spec() -> EngineSpec: is_fp8=True, causal=True, bottom_right=True, - bottom_right_padded_seq_q=True, swa=True, right_band_widening=True, padded=True, diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index 352e4be46..95cbe07ce 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -702,20 +702,27 @@ def _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, scalar_seqlen_kv): return scalar_seqlen_kv @cute.jit - def _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, scalar_seqlen_q, n_batch): - """Per-sequence Q length for the bottom-right causal diagonal. - - THD anchors the BR diagonal at the per-sequence corner - (seq_len_q[b], seq_len_kv[b]): the actual Q length is the cu_seqlen_q - difference from the packed [kv_lens | cu_q | cu_kv] metadata buffer - (same layout _thd_decode / _thd_tma_offsets read). Dense graphs (and - every non-BR mask, where seqlen_q only feeds the unused diagonal) - keep the scalar S_q, so this folds out unless THD_VARLEN and - CAUSAL_BOTTOM_RIGHT are both set.""" + def _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, scalar_seqlen_q, n_batch, seq_q_lens_tensor=None): + """Per-batch Q length for the bottom-right causal diagonal. + + Bottom-right anchors the diagonal at the per-batch corner + (seq_len_q[b], seq_len_kv[b]). THD reads the actual Q length as the + cu_seqlen_q difference from the packed [kv_lens | cu_q | cu_kv] metadata + buffer (same layout _thd_decode / _thd_tma_offsets read); dense padded + graphs carrying per-batch Q lengths read the SEPARATE (B,)-int32 + ``seq_q_lens_tensor`` (cuDNN SEQLEN_Q style), clamped to [0, S_q]. + KV-only padding (and every non-BR mask, where seqlen_q only feeds the + unused diagonal) keeps the scalar S_q, so both reads fold out unless + CAUSAL_BOTTOM_RIGHT is set together with THD_VARLEN or + SEQ_Q_LENS_PRESENT (mutually exclusive by _validate_params). + """ if cutlass.const_expr(int(getattr(CFG, "THD_VARLEN", 0)) == 1 and int(CFG.BOTTOM_RIGHT) == 1): cu = cutlass.make_array_view(seq_kv_lens_tensor) q0 = n_batch return cutlass.Int32(cu[q0 + batch_idx + cutlass.Int32(1)]) - cutlass.Int32(cu[q0 + batch_idx]) + if cutlass.const_expr(int(getattr(CFG, "SEQ_Q_LENS_PRESENT", 0)) == 1 and int(CFG.BOTTOM_RIGHT) == 1): + arr = cutlass.make_array_view(seq_q_lens_tensor) + return cute.math.max(cutlass.Int32(0), cute.math.min(cutlass.Int32(arr[batch_idx]), scalar_seqlen_q)) return scalar_seqlen_q _thd_on = int(getattr(CFG, "THD_VARLEN", 0)) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 9cfdd3321..96baafc08 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -692,7 +692,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -834,7 +834,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1103,7 +1103,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1311,7 +1311,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1593,7 +1593,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE @@ -1725,7 +1725,7 @@ def _softmax_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @@ -1797,7 +1797,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) while is_valid_tile > cutlass.Int32(0): @@ -1962,15 +1962,20 @@ def _correction_warp_group( sO_sub_base = sO[qs].base for chunk_idx in cutlass.range_constexpr(N_CHUNKS_O): - o_addr = tmem_base_epi + cutlass.Int32(tmem_O_off + chunk_idx * O_CHUNK) - o_chunk = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr(o_addr, cutlass.Float32), - num=O_CHUNK, + o_fp16 = cutlass.Vector.from_elements( + tuple(STORAGE_DTYPE(0.0) for _ in range(O_CHUNK)), + STORAGE_DTYPE, ) - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - o_scaled = o_chunk * inv_sum - o_fp16 = o_scaled.to(STORAGE_DTYPE) + if cutlass.const_expr(not MAY_BE_EMPTY) or (bounds.right > bounds.left): + o_addr = tmem_base_epi + cutlass.Int32(tmem_O_off + chunk_idx * O_CHUNK) + o_chunk = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(o_addr, cutlass.Float32), + num=O_CHUNK, + ) + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + o_scaled = o_chunk * inv_sum + o_fp16 = o_scaled.to(STORAGE_DTYPE) col_offset_const = (chunk_idx * O_CHUNK) % D_BLOCK_SIZE block_idx_const = (chunk_idx * O_CHUNK) // D_BLOCK_SIZE @@ -2015,7 +2020,7 @@ def _correction_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) # End-of-warp tmem_dealloc: under cga2 each corr lane ALSO DSMEM-arrives diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index d0b4fb547..8c702a7ad 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -728,7 +728,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -886,7 +886,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1167,7 +1167,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1374,7 +1374,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1681,7 +1681,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE @@ -1816,7 +1816,7 @@ def _softmax_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @@ -1888,7 +1888,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) while is_valid_tile > cutlass.Int32(0): @@ -2115,7 +2115,7 @@ def _correction_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) # End-of-warp tmem_dealloc: under cga2 each corr lane ALSO DSMEM-arrives diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index a3ddef5a3..b7d624851 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -527,7 +527,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -618,7 +618,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -823,7 +823,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -971,7 +971,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1025,7 +1025,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(CFG.SOFTMAX_WG0_BASE * 32) @@ -1390,7 +1390,7 @@ def _softmax_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @@ -1441,7 +1441,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) O_CHUNK = 16 @@ -1614,15 +1614,20 @@ def _correction_warp_group( for block_idx in cutlass.range_constexpr(N_BLOCKS_EPI): for sub in cutlass.range_constexpr(CHUNKS_PER_BLK): chunk_idx_total = block_idx * CHUNKS_PER_BLK + sub - o_addr = tmem_base_epi + cutlass.Int32(LAYOUT.O_OFF + chunk_idx_total * O_CHUNK) - o_chunk = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr(o_addr, cutlass.Float32), - num=O_CHUNK, + o_out = cutlass.Vector.from_elements( + tuple(OUT_STORAGE_DTYPE(0.0) for _ in range(O_CHUNK)), + OUT_STORAGE_DTYPE, ) - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - o_scaled = o_chunk * inv_sum - o_out = o_scaled.to(OUT_STORAGE_DTYPE) + if cutlass.const_expr(not MAY_BE_EMPTY) or (bounds.right > bounds.left): + o_addr = tmem_base_epi + cutlass.Int32(LAYOUT.O_OFF + chunk_idx_total * O_CHUNK) + o_chunk = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(o_addr, cutlass.Float32), + num=O_CHUNK, + ) + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + o_scaled = o_chunk * inv_sum + o_out = o_scaled.to(OUT_STORAGE_DTYPE) col_offset_const = (chunk_idx_total * O_CHUNK) % D_BLOCK_SIZE block_offset_const = ((chunk_idx_total * O_CHUNK) // D_BLOCK_SIZE) * TMA_O_GRANU_ELEMS_LOCAL @@ -1658,7 +1663,7 @@ def _correction_warp_group( is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) if cutlass.const_expr(CFG.CTA_MMA == 2): diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index 78c914d7d..292df7846 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -860,7 +860,7 @@ def _compute_warp_group( eff_seqlen_kv = seqlen_kv else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_unmasked_lo = bounds_init.unmasked_lo @@ -1162,15 +1162,20 @@ def _compute_warp_group( tmem_sub = ((b_sub & 1) << 1) | ((b_sub & 2) >> 1) tmem_block = tmem_sub * O_BLOCKS_PER_SUB + b_intra - o_addr = tmem_O_base + cutlass.Int32(tmem_block * O_EPI_BLOCK_SIZE) - o_fp32 = nvvm.tcgen05_ld( - "32x32b", - nvvm.make_tmem_ptr(o_addr, cutlass.Float32), - num=O_EPI_BLOCK_SIZE, + o_half = cutlass.Vector.from_elements( + tuple(OUT_STORAGE_DTYPE(0.0) for _ in range(O_EPI_BLOCK_SIZE)), + OUT_STORAGE_DTYPE, ) - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - o_scaled = o_fp32 * beta - o_half = o_scaled.to(OUT_STORAGE_DTYPE) + if cutlass.const_expr(not MAY_BE_EMPTY) or (kv_right > kv_left): + o_addr = tmem_O_base + cutlass.Int32(tmem_block * O_EPI_BLOCK_SIZE) + o_fp32 = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(o_addr, cutlass.Float32), + num=O_EPI_BLOCK_SIZE, + ) + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + o_scaled = o_fp32 * beta + o_half = o_scaled.to(OUT_STORAGE_DTYPE) col_offset_const = (b * O_EPI_BLOCK_SIZE) % O_D_BLOCK block_idx_const = (b * O_EPI_BLOCK_SIZE) // O_D_BLOCK @@ -1233,7 +1238,7 @@ def _compute_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_unmasked_lo = bounds_next.unmasked_lo @@ -1348,7 +1353,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1499,7 +1504,7 @@ def _mma_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1562,7 +1567,7 @@ def _mma_warp_non_leader( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1606,7 +1611,7 @@ def _mma_warp_non_leader( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1670,7 +1675,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1768,7 +1773,7 @@ def _tmaldg_warp_group( kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch, seq_q_lens_tensor) bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 383dc36f7..92996e21b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -107,12 +107,16 @@ def test_sdpa_fwd_dsl_sm100_graph_api(dtype, is_causal, d): _THD_SENTINEL = 2048.0 -def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, band_right=0, swa_window=None, seq_kv_lens=None, sinks=None, return_stats=False): +def _ref_sdpa_full( + q, k, v, *, scale, is_causal=False, bottom_right=False, band_right=0, swa_window=None, seq_q_lens=None, seq_kv_lens=None, sinks=None, return_stats=False +): """fp32 reference matching the SM100 DSL kernel's mask + sink semantics. - q/k/v are BHSD; GQA (h_q > h_kv) is handled by expanding K/V. With - ``return_stats`` also returns the (B, H_q, S_q) LSE — logsumexp over the - masked scores (the sink joins as one extra column; fully-masked rows are - -inf without one).""" + q/k/v are BHSD; GQA (h_q > h_kv) is handled by expanding K/V. Bottom-right + anchors the diagonal at the per-batch (seq_len_q[b], seq_len_kv[b]) corner + — the global (S_q, S_kv) when no lengths are given. With ``return_stats`` + also returns the (B, H_q, S_q) LSE — logsumexp over the masked scores + (the sink joins as one extra column; fully-masked rows are -inf without + one; rows at/past seq_len_q[b] are -inf even with one).""" b, h_q, s_q, _ = q.shape _, h_kv, s_kv, _ = v.shape dev = q.device @@ -121,18 +125,17 @@ def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, band_ v_ref = v.repeat_interleave(g, dim=1).float() scores = torch.matmul(q.float(), k_ref.transpose(-1, -2)) * scale + q_lens = seq_q_lens.flatten().to(device=dev, dtype=torch.int64) if seq_q_lens is not None else torch.full((b,), s_q, dtype=torch.int64, device=dev) + kv_lens = seq_kv_lens.flatten().to(device=dev, dtype=torch.int64) if seq_kv_lens is not None else torch.full((b,), s_kv, dtype=torch.int64, device=dev) i = torch.arange(s_q, device=dev).view(1, 1, s_q, 1) j = torch.arange(s_kv, device=dev).view(1, 1, 1, s_kv) - masked = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device=dev) + diag = i + (kv_lens - q_lens).view(b, 1, 1, 1) if bottom_right else i + dead_q = i >= q_lens.view(b, 1, 1, 1) + masked = dead_q | (j >= kv_lens.view(b, 1, 1, 1)) if is_causal: - lim = (i + (s_kv - s_q) if bottom_right else i) + band_right - masked = masked | (j > lim) + masked = masked | (j > diag + band_right) if swa_window is not None: - swa_base = i + (s_kv - s_q) if bottom_right else i - masked = masked | (j < swa_base - swa_window) - if seq_kv_lens is not None: - lens = seq_kv_lens.view(b, 1, 1, 1).to(dev) - masked = masked | (j >= lens) + masked = masked | (j < diag - swa_window) scores = scores.masked_fill(masked, float("-inf")) if sinks is not None: @@ -146,6 +149,8 @@ def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, band_ if not return_stats: return o.to(q.dtype) lse = torch.logsumexp(full_scores, dim=-1) # fully-masked rows -> -inf (sink-less) + # Rows at/past seq_len_q[b] trim to -inf even with a sink. + lse = lse.masked_fill(dead_q.squeeze(-1), float("-inf")) return o.to(q.dtype), lse @@ -164,7 +169,7 @@ def _bhsd(b, h, s, d, dtype, device="cuda"): return torch.randn(b, s, h, d, device=device, dtype=dtype).transpose(1, 2) -def _run_dsl_graph(q_gpu, k_gpu, v_gpu, *, scale, dtype, sdpa_kwargs, seq_len_kv=None, sink=None): +def _run_dsl_graph(q_gpu, k_gpu, v_gpu, *, scale, dtype, sdpa_kwargs, seq_len_kv=None, seq_len_q=None, sink=None): """Build the graph, opt into the matching FROST DSL engine, execute, return O (BHSD).""" import cudnn @@ -183,9 +188,10 @@ def _run_dsl_graph(q_gpu, k_gpu, v_gpu, *, scale, dtype, sdpa_kwargs, seq_len_kv kw["seq_len_kv"] = slk kw["use_padding_mask"] = True vp[slk] = seq_len_kv - # padding_mask requires a seq_len_q companion; the kernel trims only KV, so - # Q is full and seq_len_q is accepted but unused. - seq_len_q = torch.full((b, 1, 1, 1), s_q, dtype=torch.int32, device="cuda") + # padding_mask requires a seq_len_q companion; when the caller passes + # none, synthesize full lengths (KV-only trim, every Q row live). + if seq_len_q is None: + seq_len_q = torch.full((b, 1, 1, 1), s_q, dtype=torch.int32, device="cuda") slq = g.tensor_like(seq_len_q) kw["seq_len_q"] = slq vp[slq] = seq_len_q @@ -356,6 +362,94 @@ def test_dsl_sm100_padded(dtype, d): torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_graph_api_padded_bottom_right_gqa(): + """Dense padded graph carrying SHORT per-batch seq_len_q + bottom-right + causal + SWA + GQA, served end to end through the graph probe (the + conjunction the retired bottom_right_padded_seq_q capability used to + decline while the kernels anchored the BR diagonal at the global S_q).""" + _require_dsl() + d = 128 + b, h_q, h_kv, s_q, s_kv, W = 2, 8, 2, 128, 256, 47 + scale = 1.0 / math.sqrt(d) + dtype = torch.float16 + q = _bhsd(b, h_q, s_q, d, dtype) + k = _bhsd(b, h_kv, s_kv, d, dtype) + v = _bhsd(b, h_kv, s_kv, d, dtype) + seq_len_q = torch.tensor([93, 51], dtype=torch.int32, device="cuda").view(b, 1, 1, 1) + seq_len_kv = torch.tensor([137, 79], dtype=torch.int32, device="cuda").view(b, 1, 1, 1) + o = _run_dsl_graph( + q, + k, + v, + scale=scale, + dtype=dtype, + sdpa_kwargs=dict(use_causal_mask_bottom_right=True, sliding_window_length=W + 1), + seq_len_kv=seq_len_kv, + seq_len_q=seq_len_q, + ) + o_ref = _ref_sdpa_full(q, k, v, scale=scale, is_causal=True, bottom_right=True, swa_window=W, seq_q_lens=seq_len_q, seq_kv_lens=seq_len_kv) + rows = torch.arange(s_q, device="cuda").view(1, 1, s_q, 1) + dead = rows >= seq_len_q.view(b, 1, 1, 1) + assert o[dead.expand_as(o)].abs().max().item() == 0.0, "trimmed Q rows are not zero" + torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) + + +@pytest.mark.L0 +@pytest.mark.parametrize("d_qk,d_v", [(128, 128), (192, 128), (256, 256), (512, 512)], ids=["llama_d128", "mla_d192_d128", "qwen_d256", "dsv4_d512"]) +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_bottom_right_padded_seq_q(d_qk, d_v): + """Dense bottom-right causal with per-batch seq_len_q: the diagonal (and + the SWA lower bound riding it) must anchor at the per-batch + (seq_len_q[b], seq_len_kv[b]) corner, not the global S_q. GQA + LSE; + batches one full-length / one trimming mid-tile / one zero-length, S_q + spanning several CGA tiles so both the dead-tile collapse and a mid-tile + diagonal are hit. Keep seq_len_kv[b] >= seq_len_q[b] — the suite-wide BR + convention (no sequence starts with rows that see no keys).""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + dtype = torch.float16 + b, h_q, h_kv, W = 3, 8, 4, 200 + s_q = s_kv = 1024 + scale = 1.0 / math.sqrt(d_qk) + q = _bhsd(b, h_q, s_q, d_qk, dtype) + k = _bhsd(b, h_kv, s_kv, d_qk, dtype) + v = _bhsd(b, h_kv, s_kv, d_v, dtype) + o = torch.full((b, s_q, h_q, d_v), _THD_SENTINEL, device="cuda", dtype=dtype).transpose(1, 2) + lse = torch.full((b, h_q, s_q), _THD_SENTINEL, dtype=torch.float32, device="cuda") + seq_q_lens = torch.tensor([1024, 517, 0], dtype=torch.int32, device="cuda") + seq_kv_lens = torch.tensor([1024, 800, 640], dtype=torch.int32, device="cuda") + + api = SdpaFwdDslSm100( + sample_q=q, + sample_k=k, + sample_v=v, + sample_o=o, + sample_lse=lse, + is_causal=True, + causal_bottom_right=True, + window_size_left=W, + scale_softmax=scale, + seq_kv_lens_present=True, + seq_q_lens_present=True, + ) + assert api.check_support() + api.compile() + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, seq_q_lens=seq_q_lens, seq_kv_lens=seq_kv_lens) + torch.cuda.synchronize() + + o_ref, lse_ref = _ref_sdpa_full( + q, k, v, scale=scale, is_causal=True, bottom_right=True, swa_window=W, seq_q_lens=seq_q_lens, seq_kv_lens=seq_kv_lens, return_stats=True + ) + rows = torch.arange(s_q, device="cuda").view(1, 1, s_q, 1) + dead = (rows >= seq_q_lens.view(b, 1, 1, 1)).expand_as(o) + assert o[dead].abs().max().item() == 0.0, "trimmed Q rows are not zero" + torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) + torch.testing.assert_close(lse, lse_ref, atol=2e-2, rtol=2e-2) + + @pytest.mark.L0 @pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS) @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py index 62fad4b62..76a1805f7 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py @@ -777,23 +777,28 @@ def test_split_kv_causal_other_flavors(flavor): @pytest.mark.L0 @pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") -def test_split_kv_padded_q_trim(splits): +@pytest.mark.parametrize("flavor", ["d128", "d512"]) +def test_split_kv_padded_q_trim(flavor, splits): """Dense padded-Q trim (seq_q_lens) + split. Q rows at or past the batch's actual length must come back O := 0 (cuDNN >= 9.14 convention). Under split that has to hold for EVERY split's partial, or the combine mixes live and dead rows -- the trim is applied per split in the epilogue, after which the dead row's lse = -inf makes it drop out. + """ from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full - B, H, SQ, SKV = 2, 4, 256, 2048 + kmod, d_qk, d_v = _F16_FLAVORS[flavor] + # SQ spans several CGA tiles (ROWS_PER_CLUSTER = 512 on d128) so the short + # batch has both a mid-tile trim AND fully collapsed tiles past its length. + B, H, SQ, SKV = 2, 4, 1024, 2048 kv_lens = torch.tensor([2048, 2048], dtype=torch.int32, device="cuda") - q_lens = torch.tensor([256, 137], dtype=torch.int32, device="cuda") # 2nd trims mid-tile + q_lens = torch.tensor([1024, 137], dtype=torch.int32, device="cuda") # 2nd trims mid-tile got, q, k, v, scale = _run_masked( - "prefill_d128_f16_sm100.py", - 128, - 128, + kmod, + d_qk, + d_v, splits, B=B, H=H, diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index fbfd75dc0..42edc0261 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -276,9 +276,11 @@ def test_probe_rejects_non_int32_seq_len(): assert not _eligible(g) -def test_probe_rejects_bottom_right_with_padded_seq_len_q(): - # Kernel gap (pre-existing): the BR diagonal is anchored at the global S_q, - # so dense padding with per-batch seq_len_q would shift it wrongly. +def test_probe_accepts_bottom_right_with_padded_seq_len_q(): + # The kernels anchor the BR diagonal at the per-batch + # (seq_len_q[b], seq_len_kv[b]) corner, so dense padding with per-batch + # seq_len_q is served (it used to be gated while the diagonal was anchored + # at the global S_q). g = _mk_graph() q, k, v, dims, strides = _mk_qkv(g) seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_kv") @@ -296,7 +298,7 @@ def test_probe_rejects_bottom_right_with_padded_seq_len_q(): seq_len_q=seq_q, ) _finish_output(o, dims, strides) - assert not _eligible(g) + assert engines.engine_name(512) in _eligible(g) def test_probe_rejects_seq_len_q_without_padding_mask():