dflash add sliding window attention draft layer support (#27469)

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

@@ -40,7 +40,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
-from sglang.srt.speculative.spec_info import SpecInput
+from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
 from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
     draft_kv_indices_used_len,
@@ -1306,7 +1306,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
-        prefix_lens: torch.Tensor,
+        prefix_lens: Optional[torch.Tensor],
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
@@ -1324,7 +1324,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
-        prefix_lens: torch.Tensor,
+        prefix_lens: Optional[torch.Tensor],
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
@@ -1334,6 +1334,7 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
     ):
         if use_ragged:
+            assert prefix_lens is not None
             # TODO: remove this device sync, we can use forward_batch.extend_prefix_lens_cpu
             # and forward_batch.extend_seq_lens_cpu
             paged_kernel_lens = prefix_lens
@@ -1365,7 +1366,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
-        prefix_lens: torch.Tensor,
+        prefix_lens: Optional[torch.Tensor],
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
@@ -1374,6 +1375,18 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
     ):
+        if prefix_lens is None:
+            num_accept_tokens = getattr(spec_info, "num_accept_tokens", None)
+            prefix_lens = (
+                seq_lens
+                if num_accept_tokens is None
+                else seq_lens
+                - num_accept_tokens[: seq_lens.shape[0]].to(
+                    device=seq_lens.device, dtype=seq_lens.dtype
+                )
+            )
+        sliding_window_size = self.sliding_window_size
+        assert sliding_window_size is not None
         for wrapper_id in range(2):
             swa_paged_custom_mask = None
             if wrapper_id == 0:
@@ -1382,7 +1395,7 @@ class FlashInferIndicesUpdaterPrefill:
                     # the paged wrapper sees prefix-only. Trim to the last `window` tokens
                     # (required for SWATokenToKVPoolAllocator; also keeps mask O(window)).
                     effective_start = torch.clamp(
-                        prefix_lens - self.sliding_window_size, min=0
+                        prefix_lens - sliding_window_size, min=0
                     )
                     paged_kernel_lens = prefix_lens - effective_start
                     paged_kernel_lens_sum = paged_kernel_lens.sum().item()
@@ -1394,7 +1407,7 @@ class FlashInferIndicesUpdaterPrefill:
                     # window attention use paged only
                     paged_kernel_lens = torch.minimum(
                         seq_lens,
-                        torch.tensor(self.sliding_window_size) + seq_lens - prefix_lens,
+                        sliding_window_size + seq_lens - prefix_lens,
                     )
                     paged_kernel_lens_sum = paged_kernel_lens.sum().item()
                     kv_start_idx = seq_lens - paged_kernel_lens
@@ -1473,7 +1486,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
-        prefix_lens: torch.Tensor,
+        prefix_lens: Optional[torch.Tensor],
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
@@ -1522,7 +1535,7 @@ class FlashInferIndicesUpdaterPrefill:
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
         seq_lens: torch.Tensor,
-        prefix_lens: torch.Tensor,
+        prefix_lens: Optional[torch.Tensor],
         kv_start_idx: torch.Tensor,
         kv_indptr: torch.Tensor,
         qo_indptr: torch.Tensor,
@@ -1535,6 +1548,7 @@ class FlashInferIndicesUpdaterPrefill:
     ):
         bs = len(seq_lens)
         if spec_info is None:
+            assert prefix_lens is not None
             assert len(seq_lens) == len(req_pool_indices)
             # Normal extend
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
@@ -1559,14 +1573,25 @@ class FlashInferIndicesUpdaterPrefill:
             custom_mask = cross_attention_custom_mask
         else:
             assert isinstance(spec_info, SpecInput)
-            kv_indices, kv_indptr, qo_indptr, custom_mask = (
+            if spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY:
-                spec_info.generate_attn_arg_prefill(
+                kv_indices, kv_indptr, qo_indptr, custom_mask = (
-                    req_pool_indices,
+                    spec_info.generate_attn_arg_prefill(
-                    paged_kernel_lens,
+                        req_pool_indices,
-                    paged_kernel_lens_sum,
+                        paged_kernel_lens,
-                    self.req_to_token,
+                        paged_kernel_lens_sum,
+                        self.req_to_token,
+                        kv_start_idx=kv_start_idx,
+                    )
+                )
+            else:
+                kv_indices, kv_indptr, qo_indptr, custom_mask = (
+                    spec_info.generate_attn_arg_prefill(
+                        req_pool_indices,
+                        paged_kernel_lens,
+                        paged_kernel_lens_sum,
+                        self.req_to_token,
+                    )
                 )
-            )
 
         # extend part
         if use_ragged:

python/sglang/srt/models/dflash.py

@@ -28,6 +28,8 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.speculative.dflash_utils import (
     can_dflash_slice_qkv_weight,
+    get_dflash_attention_sliding_window_size,
+    get_dflash_layer_types,
     parse_dflash_draft_config,
 )
 from sglang.srt.utils import is_npu
@@ -39,6 +41,31 @@ if _is_npu:
 logger = logging.getLogger(__name__)
 
 
+def _get_dflash_layer_attention_params(
+    config, layer_id: int
+) -> Tuple[int, AttentionType]:
+    layer_types = get_dflash_layer_types(config)
+    if layer_types is None:
+        return -1, AttentionType.ENCODER_ONLY
+    if layer_id >= len(layer_types):
+        raise ValueError(
+            "DFLASH config.layer_types must contain one entry per draft layer. "
+            f"Got {len(layer_types)} entries, layer_id={layer_id}."
+        )
+
+    layer_type = layer_types[layer_id]
+    if layer_type == "full_attention":
+        return -1, AttentionType.ENCODER_ONLY
+    if layer_type == "sliding_attention":
+        sliding_window_size = get_dflash_attention_sliding_window_size(config)
+        assert sliding_window_size is not None
+        return sliding_window_size, AttentionType.DECODER
+    raise ValueError(
+        "Unsupported DFLASH draft layer type. "
+        f"layer_types[{layer_id}]={layer_type!r}."
+    )
+
+
 class DFlashAttention(nn.Module):
     def __init__(self, config, layer_id: int) -> None:
         super().__init__()
@@ -112,14 +139,17 @@ class DFlashAttention(nn.Module):
         )
 
         self.scaling = head_dim**-0.5
-        # DFlash uses non-causal attention over the draft block.
+        self.sliding_window_size, self.attn_type = _get_dflash_layer_attention_params(
+            config, layer_id
+        )
         self.attn = RadixAttention(
             num_heads=self.num_heads,
             head_dim=head_dim,
             scaling=self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
-            attn_type=AttentionType.ENCODER_ONLY,
+            sliding_window_size=self.sliding_window_size,
+            attn_type=self.attn_type,
         )
 
     def forward_prepare_npu(self, positions, hidden_states):
@@ -319,6 +349,9 @@ class DFlashDraftModel(nn.Module):
 
         self.block_size = draft_config.resolve_block_size(default=16)
 
+    def get_attention_sliding_window_size(self) -> Optional[int]:
+        return get_dflash_attention_sliding_window_size(self.config)
+
     def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor:
         """Project concatenated target-layer hidden states into draft hidden_size."""
         expected = int(self.fc.in_features)

python/sglang/srt/speculative/dflash_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Sequence
 from dataclasses import dataclass
 from numbers import Integral
 from typing import Any, List, Optional, Tuple
@@ -325,6 +326,36 @@ def build_target_layer_ids(num_target_layers: int, num_draft_layers: int) -> Lis
     ]
 
 
+def get_dflash_layer_types(config: Any) -> Optional[Sequence[str]]:
+    text_config = _get_text_config(config)
+    layer_types = _cfg_get(text_config, "layer_types", _cfg_get(config, "layer_types"))
+    if layer_types is None:
+        return None
+    if isinstance(layer_types, str) or not isinstance(layer_types, Sequence):
+        raise ValueError(
+            "DFLASH config.layer_types must be a sequence of attention type strings."
+        )
+    return layer_types
+
+
+def get_dflash_attention_sliding_window_size(config: Any) -> Optional[int]:
+    layer_types = get_dflash_layer_types(config)
+    if layer_types is None or "sliding_attention" not in layer_types:
+        return None
+
+    text_config = _get_text_config(config)
+    sliding_window = _cfg_get(
+        text_config, "sliding_window", _cfg_get(config, "sliding_window")
+    )
+    if sliding_window is None:
+        raise ValueError(
+            "DFLASH sliding_attention layers require config.sliding_window."
+        )
+
+    # HF sliding windows include the current token; SGLang stores window_left.
+    return int(sliding_window) - 1
+
+
 def _cfg_get(config: Any, key: str, default: Any = None) -> Any:
     if isinstance(config, dict):
         return config.get(key, default)

python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from typing import Literal
 
 import torch
@@ -259,6 +260,50 @@ def _make_custom_masks(
     return masks_by_req, torch.cat(flattened_masks, dim=0)
 
 
+def _make_flashinfer_dflash_swa_builtin_masks(
+    case,
+    *,
+    device: str,
+) -> list[torch.Tensor]:
+    """Mirror FlashInfer DFLASH verify's production no-custom-mask path."""
+    draft_token_num = _check_target_verify_case(case)
+    window = int(case.sliding_window_size)
+    masks_by_req = []
+    q_idx = torch.arange(
+        draft_token_num,
+        dtype=torch.int32,
+        device=device,
+    ).unsqueeze(1)
+    for prefix_len in case.prefix_lens:
+        seq_len = prefix_len + draft_token_num
+        prefix_start = max(0, int(prefix_len) - window)
+        k_idx = torch.arange(seq_len, dtype=torch.int32, device=device).unsqueeze(0)
+        masks_by_req.append((k_idx >= prefix_start) & (k_idx <= prefix_len + q_idx))
+
+    return masks_by_req
+
+
+def _expected_case_and_masks_for_spec_verify(
+    case,
+    *,
+    topk: int,
+    spec_kind: SpecVerifyKind,
+    device: str,
+):
+    if (
+        spec_kind == "dflash"
+        and case.backend == "flashinfer"
+        and getattr(case, "sliding_window_size", None) is not None
+    ):
+        reference_case = replace(case, sliding_window_size=None)
+        return reference_case, _make_flashinfer_dflash_swa_builtin_masks(
+            case, device=device
+        )
+
+    masks_by_req, _ = _make_custom_masks(case, topk=topk, device=device)
+    return case, masks_by_req
+
+
 def _make_retrieve_tensors(
     case,
     *,
@@ -299,6 +344,14 @@ def _make_spec_verify_input(
     if spec_kind == "dflash":
         if topk != 1:
             raise ValueError("DFlash verify is linear and expects topk=1.")
+        if (
+            case.backend == "flashinfer"
+            and getattr(case, "sliding_window_size", None) is not None
+        ):
+            # Production DFLASH disables custom verify masks for FlashInfer
+            # backends. SWA metadata clips the cached prefix; the backend causal
+            # path handles the draft block.
+            custom_mask = None
         return DFlashVerifyInput(
             draft_token=batch.input_ids,
             positions=batch.positions,
@@ -375,12 +428,18 @@ def _target_verify_expected_output(
     case,
     inputs,
     topk: int,
+    spec_kind: SpecVerifyKind,
     device: str,
 ):
-    masks_by_req, _ = _make_custom_masks(case, topk=topk, device=device)
+    reference_case, masks_by_req = _expected_case_and_masks_for_spec_verify(
+        case,
+        topk=topk,
+        spec_kind=spec_kind,
+        device=device,
+    )
     return reference_fn(
         fixture.reference_module,
-        case,
+        reference_case,
         inputs["prefix_hidden"],
         inputs["input_hidden"],
         masks_by_req,
@@ -457,6 +516,7 @@ def _run_spec_verify_cuda_graph_case(
                 case=spec_case,
                 inputs=inputs,
                 topk=topk,
+                spec_kind=spec_kind,
                 device=device,
             )
         ),
@@ -498,7 +558,12 @@ def run_dense_spec_verify_case(
         device=device,
     )
     _prepare_target_verify_batch(fixture.forward_batch, case, device)
-    masks_by_req, _ = _make_custom_masks(case, topk=topk, device=device)
+    reference_case, masks_by_req = _expected_case_and_masks_for_spec_verify(
+        case,
+        topk=topk,
+        spec_kind=spec_kind,
+        device=device,
+    )
     fixture.forward_batch.spec_info = _make_spec_verify_input(
         case,
         fixture.forward_batch,
@@ -509,7 +574,7 @@ def run_dense_spec_verify_case(
     inputs = dense_fixture_inputs(fixture)
     expected = dense_attention_reference_with_custom_mask(
         fixture.reference_module,
-        case,
+        reference_case,
         inputs["prefix_hidden"],
         inputs["input_hidden"],
         masks_by_req,

test/registered/attention/unittests/swa/README.md

@@ -17,7 +17,7 @@ Columns are runner modes; rows are attention backends. Cells use:
 |---|---|---|---|---|---|---|---|---|---|---|---|---|
 | `torch_native` | ✓ no-prefix + prefix window edges, MHA + GQA decode window edges (uses explicit SDPA local-attention mask) | — (no CG hooks) | — (no CG path) | — (no CG path) | — | — | — | — | — | — | — | — |
 | `triton` | ✓ no-prefix lengths below/equal/above window + prefix lengths below/equal/above window | ✓ within-window decode (`prefix_lens=(1,2,3)`, `window=4`) + above-window decode (`prefix_lens=(7,8,9)`, `window=4`) | ✓ no-prefix window edges, prefix-within-window MHA extend | ✓ same as PCG | ✓ EAGLE chain (topk=1) + EAGLE tree (topk=2), `window=4` | ✓ EAGLE tree within-window + EAGLE chain above-window (`prefix_lens=(6,8)`, `window=4`) | — | — | — | — | — | — |
-| `flashinfer` | ✓ no-prefix lengths below/equal/above window (`head_dim=64` for SM90) | ✓ within-window decode | ✓ no-prefix window edges (MHA extend) | ✓ same as PCG | blocked: SWA prefill updater needs `prefix_lens != None`, target-verify passes `None` (`flashinfer_backend.py:1296-1344` consumed by `init_forward_metadata` at `flashinfer_backend.py:742,754`) | blocked: same prefill updater contract | — | — | — | — | — | — |
+| `flashinfer` | ✓ no-prefix lengths below/equal/above window (`head_dim=64` for SM90) | ✓ within-window decode | ✓ no-prefix window edges (MHA extend) | ✓ same as PCG | ✓ DFLASH chain (`topk=1`, `window=4`) | ✓ DFLASH chain (`topk=1`, `window=4`) | — | — | — | — | — | — |
 
 ## Input And Config Coverage
 
@@ -50,12 +50,8 @@ Columns are runner modes; rows are attention backends. Cells use:
 
 ## Production-Unsupported
 
-- **FlashInfer SWA `TARGET_VERIFY` / `DRAFT_EXTEND`** — the SWA prefill updater
+- **FlashInfer SWA `DRAFT_EXTEND`** — not covered here. The FlashInfer SWA
-  (`FlashInferIndicesUpdaterPrefill.update_sliding_window`,
+  coverage added for this path is limited to DFLASH `TARGET_VERIFY`.
-  `flashinfer_backend.py:1296-1344`) requires non-`None` `prefix_lens`. The
-  target-verify and draft-extend code paths pass `prefix_lens=None` at
-  `flashinfer_backend.py:742,754`, so the SWA prefill kernel cannot be reached
-  without a separate fix to the prefill metadata contract.
 - **`torch_native` SWA speculative / CUDA graph** — no CG hooks; all graph
   integration is structurally unsupported.
 
@@ -65,6 +61,3 @@ Columns are runner modes; rows are attention backends. Cells use:
   separately (the above-window case currently asserts within tolerance with the
   matching reference rule; if a real backend regression appears, lower the
   tolerance).
-- FlashInfer SWA verify path would need a new metadata contract that threads
-  `prefix_lens` through the target-verify replay; until that lands the fixture
-  is intentionally inactive.

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

@@ -20,6 +20,10 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
 from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
     run_dense_cuda_graph_decode_case,
 )
+from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
+    run_dense_spec_verify_case,
+    run_dense_spec_verify_cuda_graph_case,
+)
 from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
     run_dense_split_op_extend_case,
 )
@@ -89,6 +93,40 @@ class TestFlashInferSWAAttentionBackendCorrectness(CustomTestCase):
             16,
         ),
     )
+    SPEC_VERIFY_CASES = (
+        (
+            DenseAttentionCase(
+                name="runner_dflash_verify_swa_chain",
+                backend="flashinfer",
+                forward_mode=ForwardMode.TARGET_VERIFY,
+                num_heads=4,
+                num_kv_heads=4,
+                page_size=16,
+                prefix_lens=(3, 5),
+                extend_lens=(3, 3),
+                sliding_window_size=4,
+            ),
+            1,
+            "dflash",
+        ),
+    )
+    SPEC_VERIFY_CUDA_GRAPH_CASES = (
+        (
+            DenseAttentionCase(
+                name="runner_cuda_graph_dflash_verify_swa_chain",
+                backend="flashinfer",
+                forward_mode=ForwardMode.TARGET_VERIFY,
+                num_heads=4,
+                num_kv_heads=4,
+                page_size=16,
+                prefix_lens=(3, 5),
+                extend_lens=(3, 3),
+                sliding_window_size=4,
+            ),
+            1,
+            "dflash",
+        ),
+    )
 
     def test_projected_swa_attention_cases(self):
         for case in self.CASES:
@@ -171,6 +209,40 @@ class TestFlashInferSWAAttentionBackendCorrectness(CustomTestCase):
                         hidden_size=self.HIDDEN_SIZE,
                     )
 
+    def test_runner_mode_spec_verify_cases(self):
+        for case, topk, spec_kind in self.SPEC_VERIFY_CASES:
+            with self.subTest(
+                case=case.name,
+                backend=case.backend,
+                topk=topk,
+                spec_kind=spec_kind,
+            ):
+                run_dense_spec_verify_case(
+                    self,
+                    case,
+                    topk=topk,
+                    spec_kind=spec_kind,
+                    head_dim=self.HEAD_DIM,
+                    hidden_size=self.HIDDEN_SIZE,
+                )
+
+    def test_runner_mode_spec_verify_cuda_graph_cases(self):
+        for case, topk, spec_kind in self.SPEC_VERIFY_CUDA_GRAPH_CASES:
+            with self.subTest(
+                case=case.name,
+                backend=case.backend,
+                topk=topk,
+                spec_kind=spec_kind,
+            ):
+                run_dense_spec_verify_cuda_graph_case(
+                    self,
+                    case,
+                    topk=topk,
+                    spec_kind=spec_kind,
+                    head_dim=self.HEAD_DIM,
+                    hidden_size=self.HIDDEN_SIZE,
+                )
+
 
 if __name__ == "__main__":
     unittest.main()