[flashinfer] Pass window_left at plan time for the SWA paged prefill wrapper (#31501)

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -1894,7 +1894,8 @@ class FlashInferIndicesUpdaterPrefill:
                         prefix_lens, seq_lens, effective_start
                     )
                 else:
-                    # window attention use paged only
+                    # window attention use paged only; the trim below is
+                    # request-granular, exactness comes from plan-time window_left
                     paged_kernel_lens = torch.minimum(
                         seq_lens,
                         sliding_window_size + seq_lens - prefix_lens,
@@ -1927,6 +1928,13 @@ class FlashInferIndicesUpdaterPrefill:
                 fixed_split_size=fixed_split_size,
                 multi_item_params=multi_item_params,
                 cross_attention_custom_mask=swa_paged_custom_mask,
+                # paged-only SWA path only; ragged keeps its custom prefix
+                # mask, spec-verify keeps its tree mask
+                window_left=(
+                    sliding_window_size
+                    if (wrapper_id == 0 and not use_ragged and spec_info is None)
+                    else -1
+                ),
             )
 
     def _build_swa_prefix_custom_mask(
@@ -2043,6 +2051,7 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         seq_lens_cpu: Optional[torch.Tensor] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
+        window_left: int = -1,
     ):
         bs = len(seq_lens)
         if spec_info is None:
@@ -2178,6 +2187,10 @@ class FlashInferIndicesUpdaterPrefill:
                 max_kv_len=int(seq_lens_cpu_i32.max()),
             )
 
+        if window_left >= 0:
+            # selects the module with the per-element window mask compiled in
+            paged_plan_kwargs["window_left"] = window_left
+
         wrapper_paged.begin_forward(
             qo_indptr,
             kv_indptr,

test/registered/attention/unittests/swa/test_flashinfer.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 import torch
 
+from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.utils import is_flashinfer_available
 from sglang.test.test_utils import CustomTestCase
@@ -44,6 +45,21 @@ class TestFlashInferSWAAttentionBackendCorrectness(CustomTestCase):
     CASES = make_swa_no_prefix_input_config_cases(
         "flashinfer"
     ) + make_swa_prefix_input_config_cases("flashinfer")
+    # Paged-only prefill has no ragged pass / custom prefix mask, so the kernel
+    # must enforce the window; the long case puts tokens past the window.
+    PAGED_MODE_CASES = CASES + (
+        DenseAttentionCase(
+            name="swa_extend_no_prefix_above_window_long",
+            backend="flashinfer",
+            forward_mode=ForwardMode.EXTEND,
+            num_heads=4,
+            num_kv_heads=4,
+            page_size=16,
+            prefix_lens=(0, 0, 0),
+            extend_lens=(6, 8, 12),
+            sliding_window_size=4,
+        ),
+    )
     # Above-window decode case requires the `extend_window` reference rule
     # (window+1 keys), not the `min_seq_len_window` rule — FlashInfer's
     # decode metadata uses `clamp(seq_lens, max=window+1)` per
@@ -138,6 +154,17 @@ class TestFlashInferSWAAttentionBackendCorrectness(CustomTestCase):
                     hidden_size=self.HIDDEN_SIZE,
                 )
 
+    def test_projected_swa_attention_cases_paged_mode(self):
+        for case in self.PAGED_MODE_CASES:
+            with self.subTest(case=case.name, backend=case.backend, mode="paged"):
+                with envs.SGLANG_FLASHINFER_USE_PAGED.override(True):
+                    run_dense_attention_case(
+                        self,
+                        case,
+                        head_dim=self.HEAD_DIM,
+                        hidden_size=self.HIDDEN_SIZE,
+                    )
+
     # Layout-robustness. See dense/test_triton.py for the full rationale.
     # The default `shuffled_pages` is already exercised by
     # test_projected_swa_attention_cases on the existing case list.