diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 29b26d01f..87d716fa7 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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, diff --git a/test/registered/attention/unittests/swa/test_flashinfer.py b/test/registered/attention/unittests/swa/test_flashinfer.py index 5635a0906..a8659ab7e 100644 --- a/test/registered/attention/unittests/swa/test_flashinfer.py +++ b/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.