From 25058ec4c7543be8391d9e90f6517dae19f884d5 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Thu, 13 Aug 2026 11:16:49 -0700 Subject: [PATCH] frost(sdpa): bottom-right diagonal + sliding window on the SM100 rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The band shifts wholesale with its diagonal: under BOTTOM_RIGHT the SWA lower limit is q + (S_kv - S_q) - W, the same causal_diag offset the upper (causal) limit has always applied. Add that offset to the two SWA terms — the KV-loop tile bounds (compute_kv_loop_bounds: the swa left bound and the unmasked-lo anchor) and the per-element mask (apply_mask_chunk) — where it folds to the plain top-left math when the diagonal is zero. Lift the three gates (config validation, adapter backstop, bottom_right_with_swa on the f16 and per-tensor FP8 rows); THD gets the per-sequence diagonal for free through resolve_seqlen_q, and the widened band composes (BR + SWA + window_right). The MXFP8 row stays gated: the one mhas graph BR+SWA admits there trips the executor's pre-existing SF-size mismatch (SF-layout issue independent of the mask) — flip once that plumbing is fixed. Suite: swa_br joins the dense and THD combo axes (48 combos), the torch reference's SWA term becomes diagonal-aware, probe tests updated. mhas routing (paged/bwd_ragged excluded, pip-9.23 env): 949/2274 (41.7%), up from 38.6% — d128 +29, d192x128 +5, d256 +17, fp8 +20 graphs. Co-Authored-By: Claude Fable 5 --- python/cudnn/frost/tile_dsl/mask.py | 8 ++++- python/cudnn/sdpa/fwd/api_dsl.py | 7 ---- python/cudnn/sdpa/fwd/config_sm100.py | 2 -- python/cudnn/sdpa/fwd/engines.py | 5 +++ .../cudnn/sdpa/fwd/kernels/_common_sm100.py | 10 ++++-- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 13 +++++--- .../sdpa/frost/test_sdpa_graph_analyzer.py | 32 +++++++------------ 7 files changed, 38 insertions(+), 39 deletions(-) diff --git a/python/cudnn/frost/tile_dsl/mask.py b/python/cudnn/frost/tile_dsl/mask.py index 7e451ebf4..bd9d1bf69 100644 --- a/python/cudnn/frost/tile_dsl/mask.py +++ b/python/cudnn/frost/tile_dsl/mask.py @@ -31,7 +31,13 @@ def apply_mask_chunk( return reg_S neg_inf = cutlass.Float32(mask_value) - q_minus_w = q_abs - cutlass.Int32(window_left) if (mask_flags & MASK_SWA) else None + # The whole band shifts with the diagonal: under BOTTOM_RIGHT the SWA + # lower limit is q + (S_kv - S_q) - W — the same causal_diag offset the + # upper (causal) limit uses below. Top-left keeps the plain q - W. + q_minus_w = None + if mask_flags & MASK_SWA: + swa_base = (q_abs + causal_diag) if bottom_right else q_abs + q_minus_w = swa_base - cutlass.Int32(window_left) # window_right is the compile-time diagonal-band right bound (cuDNN # diagonal_band_right_bound): kv columns up to q + window_right (plus the # bottom-right diagonal offset) stay unmasked. 0 = plain causal. diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index e8bc28633..33da805fd 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -849,17 +849,10 @@ def check_support(self) -> bool: band_right is not None and not self.is_causal, "SM100 DSL SDPA: window_size_right widens the causal diagonal and requires is_causal=True", ) - # The kernels' bottom-right diagonal path excludes a left bound: - # bottom_right requires a right bound and rejects window_left - # (see config_sm100._validate_params). self._value_error_if( self.causal_bottom_right and not self.is_causal, "SM100 DSL SDPA: causal_bottom_right requires is_causal=True", ) - self._value_error_if( - self.causal_bottom_right and swa_left is not None, - "SM100 DSL SDPA: causal_bottom_right cannot be combined with a left sliding-window (kernel gap)", - ) # 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. diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index 1dedb6c5e..09bba508f 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -111,8 +111,6 @@ def _validate_params(flavor: str, k: TemplateParams) -> None: if k.bottom_right: if k.window_right is None: raise ValueError(f"{flavor}: bottom_right anchors the band's diagonal and requires a right bound (window_right)") - if k.window_left is not None: - raise ValueError(f"{flavor}: bottom_right + window_left (SWA) is not supported") if k.thd_varlen and not k.seq_kv_lens_present: raise ValueError(f"{flavor}: THD/varlen requires SEQ_KV_LENS_PRESENT (per-sequence padded masking)") if k.seq_q_lens_present: diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index d263c6256..5488607ec 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -404,6 +404,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), causal=True, bottom_right=True, + bottom_right_with_swa=True, right_band_widening=True, swa=True, padded=True, @@ -453,6 +454,9 @@ def _sm100_mxfp8_spec(d: int) -> EngineSpec: is_mxfp8=True, causal=True, bottom_right=True, + # BR+SWA stays off this row: the one mhas graph it admits trips the + # mxfp8 executor's SF-size mismatch (pre-existing SF-layout issue, + # independent of the mask) — flip once that plumbing is fixed. right_band_widening=True, swa=True, padded=True, @@ -491,6 +495,7 @@ def _sm100_fp8_spec(d: int) -> EngineSpec: is_fp8=True, causal=True, bottom_right=True, + bottom_right_with_swa=True, right_band_widening=True, swa=True, padded=True, diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index ff47dc4cd..2f2f7242c 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -226,8 +226,12 @@ def compute_kv_loop_bounds( right = cute.math.min(right, kv_hi_caus) if cutlass.const_expr(mask_flags & MASK_SWA): - cond = q_row_coord > cutlass.Int32(window_left) - delta = q_row_coord - cutlass.Int32(window_left) + # The whole band shifts with the diagonal: under BOTTOM_RIGHT the SWA + # lower bound is q + (S_kv - S_q) - W, same anchor the causal upper + # bound uses (causal_diag folds to 0 for top-left). + swa_base = q_row_coord + causal_diag + cond = swa_base > cutlass.Int32(window_left) + delta = swa_base - cutlass.Int32(window_left) kv_lo_swa = cutlass.Int32( arith.select( cond.ir_value(), @@ -255,7 +259,7 @@ def compute_kv_loop_bounds( unmasked_lo = left if cutlass.const_expr(mask_flags & MASK_SWA): - anchor = q_row_coord + cutlass.Int32(cga_tile_m - 1 - window_left) + anchor = q_row_coord + causal_diag + cutlass.Int32(cga_tile_m - 1 - window_left) swa_unmasked_lo = _div_up(anchor, tile_n) cond = anchor > cutlass.Int32(0) swa_unmasked_lo = cutlass.Int32( 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 1492ff16e..a779f72dc 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -128,7 +128,8 @@ def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, band_ lim = (i + (s_kv - s_q) if bottom_right else i) + band_right masked = masked | (j > lim) if swa_window is not None: - masked = masked | (j < i - swa_window) + 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) @@ -671,6 +672,7 @@ def _mask_graph_kwargs(mask): "band": dict(diagonal_band_right_bound=_COMBO_BAND_R), "band_br": dict(diagonal_band_right_bound=_COMBO_BAND_R, diagonal_alignment=cudnn.diagonal_alignment.BOTTOM_RIGHT), "band_swa": dict(diagonal_band_right_bound=_COMBO_BAND_R, diagonal_band_left_bound=_COMBO_SWA_W + 1), + "swa_br": dict(use_causal_mask_bottom_right=True, sliding_window_length=_COMBO_SWA_W + 1), }[mask] @@ -684,6 +686,7 @@ def _mask_ref_kwargs(mask): "band": dict(is_causal=True, band_right=_COMBO_BAND_R), "band_br": dict(is_causal=True, bottom_right=True, band_right=_COMBO_BAND_R), "band_swa": dict(is_causal=True, band_right=_COMBO_BAND_R, swa_window=_COMBO_SWA_W), + "swa_br": dict(is_causal=True, bottom_right=True, swa_window=_COMBO_SWA_W), }[mask] @@ -812,7 +815,7 @@ def _dense_buf(packed, s_max, t, H): def _combo_dense(d, dtype, H_q, H_kv, scale, sink_t, mask): b = 2 - s_q, s_kv = (128, 256) if mask in ("causal_br", "band_br") else (256, 256) + s_q, s_kv = (128, 256) if mask in ("causal_br", "band_br", "swa_br") else (256, 256) 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) @@ -835,7 +838,7 @@ def _combo_thd(d, dtype, H_q, H_kv, scale, sink_t, mask): dev = "cuda" seq_lens_q = [200, 150] seq_lens_kv = [180, 120] - if mask in ("causal_br", "band_br"): + if mask in ("causal_br", "band_br", "swa_br"): # Bottom-right masks: keep seq_len_kv[b] >= seq_len_q[b] so no sequence # has fully-masked rows (the torch softmax reference NaNs on those). seq_lens_q = [150, 90] @@ -1039,12 +1042,12 @@ def test_dsl_sm100_thd_cu_seq_len_zero_lens(): _COMBO_MASKS = { - "dense": ["none", "causal", "causal_br", "swa", "padded", "band", "band_br", "band_swa"], + "dense": ["none", "causal", "causal_br", "swa", "padded", "band", "band_br", "band_swa", "swa_br"], # THD forces padding internally, so its mask axis rides on top of that. # causal_br: the kernels anchor the BR diagonal at each sequence's own # (seq_len_q[b], seq_len_kv[b]) from the cu_seqlen metadata. # band/band_br: diagonal-band right-bound widening (BAND_RIGHT). - "thd": ["none", "causal", "swa", "causal_br", "band", "band_br"], + "thd": ["none", "causal", "swa", "causal_br", "band", "band_br", "swa_br"], } diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index e7774bc7a..cbc4a5204 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -477,24 +477,6 @@ def test_resolve_generate_stats(): assert cfg.stats_t is not None -def test_probe_rejects_bottom_right_plus_swa(): - # Kernel gap: CAUSAL_BOTTOM_RIGHT excludes SWA (config would assert at import). - g = _mk_graph() - q, k, v, dims, strides = _mk_qkv(g) - o, _ = g.sdpa( - name="s", - q=q, - k=k, - v=v, - attn_scale=0.1, - is_inference=True, - use_causal_mask_bottom_right=True, - sliding_window_length=128, - ) - _finish_output(o, dims, strides) - assert not _eligible(g) - - def test_probe_rejects_bottom_right_swa_only(): # Kernel gap: CAUSAL_BOTTOM_RIGHT requires MASK_CAUSAL; BOTTOM_RIGHT # alignment with only a left band has no causal bit. @@ -605,11 +587,19 @@ def test_sm120_probe_accepts_causal_swa_on_both_minors(monkeypatch): assert not any("sm100" in name for name in elig) +def test_probe_accepts_bottom_right_with_swa(): + # The band shifts wholesale with the diagonal: the SM100 kernels apply the + # same causal_diag offset to the SWA lower limit as to the causal upper one. + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g) + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_causal_mask_bottom_right=True, sliding_window_length=64) + _finish_output(o, dims, strides) + assert engines.engine_name(512) in _eligible(g) + + def test_sm120_probe_accepts_bottom_right_with_swa(monkeypatch): - # SM120-only notch: the SM100 kernels' BR diagonal excludes SWA. Facts are - # cached per graph, so each device family probes a freshly built graph. + # BR + SWA is served on both families now; this pins the SM120 row's claim. kwargs = dict(use_causal_mask_bottom_right=True, sliding_window_length=64) - assert not _eligible(_mk_sm120_graph(**kwargs)) # (10, 0): no row serves BR+SWA monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) assert _SM120 in _eligible(_mk_sm120_graph(**kwargs))