frost(sdpa): bound the SM80 backward q-loop by the sliding window (6-12x on gpt_oss-style SWA) + window-aware deterministic relay

python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py

@@ -52,7 +52,9 @@
 extra GMEM beyond the small counter; perf-insensitive (gated knob).  dK/dV are
 already deterministic (each CTA owns distinct K rows → no cross-tile atomic).
 All semaphore code folds out under ``const_expr(deterministic)`` → the default
-(non-deterministic) path is byte-identical.
+(non-deterministic) path is byte-identical.  With a sliding window the q-loop
+is also bounded from above, so a q-tile's visitors start at the window floor;
+the relay then counts turns from that first visitor (``relay_turn``).
 
 Why two sub-groups: register capacity.  dV_acc + dK_acc + frags at d=128 bust
 the 255-reg/thread Ampere cap in a single group; splitting halves the live
@@ -405,6 +407,28 @@ def _bprop_kernel(
     else:
         q_lo_tile = cutlass.Int32(0)
     n_iters = cutlass.Int32(arith.maxsi((n_q_tiles_eff - q_lo_tile).ir_value(), cutlass.Int32(0).ir_value()))
+    # ---- SWA compute-skip (UPPER bound).  With a left window W the mask keeps
+    #      kv >= anchor(q) - W (anchor = q, or q + causal_diag under bottom-right,
+    #      see _mask_p), so this kv-tile is attended by NO q with anchor(q) >
+    #      kv_base + tile_kv - 1 + W, i.e. q >= kv_base + tile_kv + W - diag.
+    #      Cap the q-loop there — the mirror of the forward's kv_left trim.
+    #      Without it every kv-tile sweeps the q-tiles to the end of the sequence
+    #      and masks them: O(S^2) work for an O(S*W) problem (gpt_oss W=128 on
+    #      A100: 5x slower than the backend at 2k, 75x at 32k).  Fully-masked
+    #      tiles (q_hi <= q_lo) run 0 iters and the epilogue stores dK = dV = 0.
+    #      Deterministic relay: a high-end cut means a q-tile's kv-tile visitors
+    #      no longer start at 0 (a lower tile may skip a q-tile this one runs),
+    #      so the relay counts turns from the FIRST visitor of the q-tile —
+    #      see relay_turn at the dQ atomicAdd (mirrors bprop_f16_sm120.py's
+    #      `relay_turn`, which inverts the same clamp).
+    if cutlass.const_expr(mask_flags & MASK_SWA):
+        _q_hi_abs = kv_base + cutlass.Int32(tile_kv + swa_window)
+        if cutlass.const_expr(causal_bottom_right):
+            _q_hi_abs = _q_hi_abs - causal_diag
+        _q_hi_abs = cutlass.Int32(arith.maxsi(_q_hi_abs.ir_value(), cutlass.Int32(0).ir_value()))
+        _q_hi_t = (_q_hi_abs + cutlass.Int32(tile_q - 1)) // cutlass.Int32(tile_q)
+        _q_hi_t = cutlass.Int32(arith.minsi(_q_hi_t.ir_value(), n_q_tiles_eff.ir_value()))
+        n_iters = cutlass.Int32(arith.maxsi((_q_hi_t - q_lo_tile).ir_value(), cutlass.Int32(0).ir_value()))
     # THD over-provisioned grid: this CTA's kv-tile may start past the packed
     # sequence's KV length (n_kv_tiles = ceil(max_skv/tile_kv) covers the longest
     # sequence).  Force 0 q-iters for such tiles → no compute, and the dV/dK
@@ -926,8 +950,28 @@ def _bprop_kernel(
         if cutlass.const_expr(deterministic):
             sem_idx = (batch * cutlass.Int32(H) + head) * sem_q_stride + q_iter
             sem_ptr = Pointer(cutlass.make_array_view(DQ_SEM).data_ptr() + sem_idx, dtype=cutlass.Int32)
+            # Relay turn of this kv-tile for this q-tile.  Without a window
+            # every kv-tile 0..kv_tile-1 visits the q-tile (the causal q_lo
+            # skip only removes HIGHER kv-tiles), so the turn is kv_tile.  With
+            # a left window the SWA q-loop upper bound (above) means the
+            # q-tile's visitors start at the window floor of its FIRST row:
+            # kv_first = max((q_row0 + diag - W) // tile_kv, 0) — the inverse
+            # of the q_hi clamp — so count turns from there (SM120 precedent:
+            # bprop_f16_sm120.py `relay_turn`).  The counter is per (seq, head,
+            # q_tile) and every visitor computes the same kv_first, so the
+            # ordering (ascending kv_tile) and the release (turn + 1) stay
+            # consistent.  Folds to kv_tile when MASK_SWA is off.
+            if cutlass.const_expr(mask_flags & MASK_SWA):
+                _q_row0 = q_iter * cutlass.Int32(tile_q)
+                if cutlass.const_expr(causal_bottom_right):
+                    _q_row0 = _q_row0 + causal_diag
+                _kv_first = (_q_row0 - cutlass.Int32(swa_window)) // cutlass.Int32(tile_kv)
+                _kv_first = cutlass.Int32(arith.maxsi(_kv_first.ir_value(), cutlass.Int32(0).ir_value()))
+                relay_turn = kv_tile - _kv_first
+            else:
+                relay_turn = kv_tile
             if tidx == cutlass.Int32(0):
-                _dq_sem_wait(sem_ptr, kv_tile)
+                _dq_sem_wait(sem_ptr, relay_turn)
             nvvm.barrier_cta_sync()
         if cutlass.const_expr(dq_smem_coalesce or has_rope):
             # COALESCED via SMEM staging (sDQ).  Stage the dq_acc C-fragment into
@@ -1015,7 +1059,7 @@ def _bprop_kernel(
             nvvm.barrier_cta_sync()
             cute.arch.fence_acq_rel_gpu()
             if tidx == cutlass.Int32(0):
-                cute.arch.atomic_exch(sem_ptr, kv_tile + cutlass.Int32(1), sem="release", scope="gpu")
+                cute.arch.atomic_exch(sem_ptr, relay_turn + cutlass.Int32(1), sem="release", scope="gpu")
 
         # ---- B3: before next Q-iter overwrites the ring stage ---------------
         nvvm.barrier_cta_sync()

test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py

@@ -47,8 +47,13 @@ def _bshd_randn(b, h, s, d, **kw):
     return torch.randn((b, s, h, d), **kw).permute(0, 2, 1, 3)
 
 
-def _ref_grads(q, k, v, do, *, is_causal, window_left, scale):
-    """fp32 autograd reference; returns (o, dq, dk, dv)."""
+def _ref_grads(q, k, v, do, *, is_causal, window_left, scale, causal_bottom_right=False):
+    """fp32 autograd reference; returns (o, dq, dk, dv).
+
+    Mask semantics mirror the kernel's ``_mask_p``: the diagonal anchor is
+    ``q`` (top-left) or ``q + (s_kv - s_q)`` (bottom-right); causal keeps
+    ``kv <= anchor``; a left window keeps ``kv >= anchor - W`` (also without
+    causal — the non-causal "swa" mask)."""
     _, h_q, s_q, _ = q.shape
     _, h_kv, s_kv, _ = k.shape
     g = h_q // h_kv
@@ -59,15 +64,15 @@ def _ref_grads(q, k, v, do, *, is_causal, window_left, scale):
     k_exp = k_ref.repeat_interleave(g, dim=1)
     v_exp = v_ref.repeat_interleave(g, dim=1)
     scores = torch.matmul(q_ref, k_exp.transpose(-1, -2)) * scale
-    if is_causal:
-        # Top-left diagonal, matching the wrapper's default mask (the tests
-        # never pass causal_bottom_right); with a bottom-right-anchored
-        # reference the two would agree only while s_q == s_kv.
+    if is_causal or window_left >= 0:
         i = torch.arange(s_q, device=q.device).view(s_q, 1)
         j = torch.arange(s_kv, device=q.device).view(1, s_kv)
-        keep = j <= i
+        anchor = i + (s_kv - s_q) if causal_bottom_right else i
+        keep = torch.ones(s_q, s_kv, dtype=torch.bool, device=q.device)
+        if is_causal:
+            keep &= j <= anchor
         if window_left >= 0:
-            keep &= (i - j) <= window_left
+            keep &= (anchor - j) <= window_left
         scores = scores.masked_fill(~keep, float("-inf"))
     probs = torch.softmax(scores, dim=-1)
     o = torch.matmul(probs, v_exp)
@@ -247,7 +252,7 @@ def test_sdpa_bwd_sm80_d64_fast_path(monkeypatch):
     torch.testing.assert_close(dv_d.float(), dv_g.float(), rtol=2e-2, atol=2e-2)
 
 
-def _thd_run_and_check(lens, h, d_qk, d_v, *, check_grads=True):
+def _thd_run_and_check(lens, h, d_qk, d_v, *, check_grads=True, window_left=-1):
     """Run the THD fwd+bwd wrappers on packed random inputs; when
     ``check_grads``, compare each sequence's dQ/dK/dV slice against the dense
     fp32 autograd reference (the existing ``_ref_grads`` pattern)."""
@@ -262,8 +267,11 @@ def _thd_run_and_check(lens, h, d_qk, d_v, *, check_grads=True):
     k = torch.randn(1, t, h, d_qk, dtype=torch.float16, device="cuda")
     v = torch.randn(1, t, h, d_v, dtype=torch.float16, device="cuda")
     do = torch.randn(1, t, h, d_v, dtype=torch.float16, device="cuda")
-    fwd = sdpa_fwd_wrapper_sm80(q, k, v, is_causal=True, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu, max_s_q=int(max(lens)))
-    out = sdpa_bwd_wrapper_sm80(q, k, v, fwd["o_tensor"], do, fwd["lse_tensor"], is_causal=True, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu)
+    window = (window_left, 0)
+    fwd = sdpa_fwd_wrapper_sm80(q, k, v, is_causal=True, window_size=window, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu, max_s_q=int(max(lens)))
+    out = sdpa_bwd_wrapper_sm80(
+        q, k, v, fwd["o_tensor"], do, fwd["lse_tensor"], is_causal=True, window_size=window, cum_seqlen_q_tensor=cu, cum_seqlen_k_tensor=cu
+    )
     if check_grads:
         for i in range(len(lens)):
             lo, hi = int(cu[i]), int(cu[i + 1])
@@ -272,7 +280,9 @@ def _thd_run_and_check(lens, h, d_qk, d_v, *, check_grads=True):
             def _bhsd(x, _lo=lo, _hi=hi):
                 return x[0, _lo:_hi].permute(1, 0, 2).unsqueeze(0)
 
-            _, dq_ref, dk_ref, dv_ref = _ref_grads(_bhsd(q), _bhsd(k), _bhsd(v), _bhsd(do), is_causal=True, window_left=-1, scale=1.0 / math.sqrt(d_qk))
+            _, dq_ref, dk_ref, dv_ref = _ref_grads(
+                _bhsd(q), _bhsd(k), _bhsd(v), _bhsd(do), is_causal=True, window_left=window_left, scale=1.0 / math.sqrt(d_qk)
+            )
             torch.testing.assert_close(_bhsd(out["dq_tensor"]).to(torch.float32), dq_ref, rtol=3e-2, atol=3e-2)
             torch.testing.assert_close(_bhsd(out["dk_tensor"]).to(torch.float32), dk_ref, rtol=3e-2, atol=3e-2)
             torch.testing.assert_close(_bhsd(out["dv_tensor"]).to(torch.float32), dv_ref, rtol=3e-2, atol=3e-2)
@@ -546,3 +556,129 @@ def run(lse):
     )
     with pytest.raises(ValueError, match="device"):
         eng.check_support()
+
+
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_sm80_bwd_swa_long_seq_smoke():
+    """One long-sequence sliding-window case at L0 (the q-loop bound is a kernel
+    change, so a single S >> W check runs by default); the three window
+    geometries sweep at L1."""
+    test_sm80_bwd_swa_long_seq(True, False, 2048, 2048)
+
+
+@pytest.mark.L1
+@pytest.mark.parametrize(
+    "is_causal,bottom_right,s_q,s_kv",
+    [(True, False, 2048, 2048), (False, False, 2048, 2048), (True, True, 1536, 2048)],
+    ids=["causal_swa_tl", "swa_only", "causal_swa_br"],
+)
+@torch_fork_set_rng(seed=0)
+def test_sm80_bwd_swa_long_seq(is_causal, bottom_right, s_q, s_kv):
+    """Sliding window at S >> W: the backward bounds each kv-tile's q-loop to
+    the window (q <= kv + W - br_diag), the mirror of the forward's kv_left
+    trim, instead of sweeping every q-tile and masking. Correctness against the
+    dense reference across the three window geometries the kernel serves:
+    top-left causal+window, window-only (non-causal), bottom-right causal+window
+    with s_q != s_kv (the anchor shifts by s_kv - s_q)."""
+    try:
+        from cudnn.sdpa.bwd import sdpa_bwd_wrapper_sm80
+        from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80
+    except ImportError as e:
+        pytest.skip(f"SM80 SDPA API not available: {e}")
+
+    b, h, d, w = 1, 4, 64, 128
+    scale = 1.0 / math.sqrt(d)
+    q = _bshd_randn(b, h, s_q, d, dtype=torch.float16, device="cuda")
+    k = _bshd_randn(b, h, s_kv, d, dtype=torch.float16, device="cuda")
+    v = _bshd_randn(b, h, s_kv, d, dtype=torch.float16, device="cuda")
+    do = _bshd_randn(b, h, s_q, d, dtype=torch.float16, device="cuda")
+    # window_size_right is a causal-band attribute: the wrappers take -1
+    # (unset) for a non-causal left window and 0 (no widening) under causal.
+    window = (w, 0 if is_causal else -1)
+    fwd = sdpa_fwd_wrapper_sm80(
+        q_tensor=q, k_tensor=k, v_tensor=v, is_causal=is_causal, window_size=window, causal_bottom_right=bottom_right, scale_softmax=scale
+    )
+    out = sdpa_bwd_wrapper_sm80(
+        q_tensor=q,
+        k_tensor=k,
+        v_tensor=v,
+        o_tensor=fwd["o_tensor"],
+        do_tensor=do,
+        lse_tensor=fwd["lse_tensor"],
+        is_causal=is_causal,
+        window_size=window,
+        causal_bottom_right=bottom_right,
+        scale_softmax=scale,
+    )
+    _, dq_ref, dk_ref, dv_ref = _ref_grads(q, k, v, do, is_causal=is_causal, window_left=w, scale=scale, causal_bottom_right=bottom_right)
+    torch.testing.assert_close(out["dq_tensor"].to(torch.float32), dq_ref, rtol=3e-2, atol=3e-2)
+    torch.testing.assert_close(out["dk_tensor"].to(torch.float32), dk_ref, rtol=3e-2, atol=3e-2)
+    torch.testing.assert_close(out["dv_tensor"].to(torch.float32), dv_ref, rtol=3e-2, atol=3e-2)
+
+
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_sm80_bwd_thd_swa():
+    """THD + sliding window: the q-loop bound is per packed sequence (in-seq
+    q/kv indices); sequences longer and shorter than the window mix."""
+    try:
+        _thd_run_and_check([96, 1024, 640], h=4, d_qk=64, d_v=64, window_left=128)
+    except ImportError as e:
+        pytest.skip(f"SM80 SDPA API not available: {e}")
+
+
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_sm80_bwd_swa_deterministic_smoke():
+    """Top-left deterministic + window at L0 (the relay's first-visitor turn is
+    the hang-risk path, so one case runs by default); bottom-right at L1."""
+    test_sm80_bwd_swa_deterministic(False, 2048, 2048)
+
+
+@pytest.mark.L1
+@pytest.mark.parametrize("bottom_right,s_q,s_kv", [(False, 2048, 2048), (True, 1536, 2048)], ids=["tl", "br"])
+@torch_fork_set_rng(seed=0)
+def test_sm80_bwd_swa_deterministic(bottom_right, s_q, s_kv):
+    """Deterministic dQ + sliding window at S >> W: the window-bounded q-loop
+    means a q-tile's kv-tile visitors start at the window floor, not at 0, so
+    the ordered relay counts turns from that first visitor (the SM120 kernel's
+    `relay_turn`). Must be bitwise repeatable AND agree with the
+    non-deterministic path — a wrong first-visitor would hang or mis-order."""
+    try:
+        from cudnn.sdpa.bwd import sdpa_bwd_wrapper_sm80
+        from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80
+    except ImportError as e:
+        pytest.skip(f"SM80 SDPA API not available: {e}")
+
+    b, h, d, w = 1, 4, 64, 128
+    scale = 1.0 / math.sqrt(d)
+    q = _bshd_randn(b, h, s_q, d, dtype=torch.float16, device="cuda")
+    k = _bshd_randn(b, h, s_kv, d, dtype=torch.float16, device="cuda")
+    v = _bshd_randn(b, h, s_kv, d, dtype=torch.float16, device="cuda")
+    do = _bshd_randn(b, h, s_q, d, dtype=torch.float16, device="cuda")
+    window = (w, 0)
+    fwd = sdpa_fwd_wrapper_sm80(q_tensor=q, k_tensor=k, v_tensor=v, is_causal=True, window_size=window, causal_bottom_right=bottom_right, scale_softmax=scale)
+
+    def run(det):
+        return sdpa_bwd_wrapper_sm80(
+            q_tensor=q,
+            k_tensor=k,
+            v_tensor=v,
+            o_tensor=fwd["o_tensor"],
+            do_tensor=do,
+            lse_tensor=fwd["lse_tensor"],
+            is_causal=True,
+            window_size=window,
+            causal_bottom_right=bottom_right,
+            scale_softmax=scale,
+            deterministic=det,
+        )
+
+    a, bb = run(True), run(True)
+    for key in ("dq_tensor", "dk_tensor", "dv_tensor"):
+        assert torch.equal(a[key], bb[key]), key
+    nd = run(False)
+    _, dq_ref, _, _ = _ref_grads(q, k, v, do, is_causal=True, window_left=w, scale=scale, causal_bottom_right=bottom_right)
+    torch.testing.assert_close(a["dq_tensor"].to(torch.float32), dq_ref, rtol=3e-2, atol=3e-2)
+    torch.testing.assert_close(a["dq_tensor"].to(torch.float32), nd["dq_tensor"].to(torch.float32), rtol=2e-2, atol=2e-2)