From 8c5320b37e99816ea000b3ecb34a0a89b3fa608d Mon Sep 17 00:00:00 2001 From: David Wang <21328423+dcw02@users.noreply.github.com> Date: Sun, 14 Jun 2026 03:32:02 -0400 Subject: [PATCH] dflash add sliding window attention draft layer support (#27469) --- .../layers/attention/flashinfer_backend.py | 55 ++++++++++---- python/sglang/srt/models/dflash.py | 37 +++++++++- python/sglang/srt/speculative/dflash_utils.py | 31 ++++++++ .../speculative_target_verify_runner.py | 73 ++++++++++++++++++- .../attention/unittests/swa/README.md | 13 +--- .../unittests/swa/test_flashinfer.py | 72 ++++++++++++++++++ 6 files changed, 250 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 47d8b40f3..11b917208 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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 = ( - spec_info.generate_attn_arg_prefill( - req_pool_indices, - paged_kernel_lens, - paged_kernel_lens_sum, - self.req_to_token, + if spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY: + 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, + 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: diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 58537eba7..71b56e311 100644 --- a/python/sglang/srt/models/dflash.py +++ b/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) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 13d55ebc6..e63efe4ea 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/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) diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py index 6ce5c18b1..b7b0389ff 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py +++ b/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, diff --git a/test/registered/attention/unittests/swa/README.md b/test/registered/attention/unittests/swa/README.md index 93a94e7b6..573929ba8 100644 --- a/test/registered/attention/unittests/swa/README.md +++ b/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 - (`FlashInferIndicesUpdaterPrefill.update_sliding_window`, - `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. +- **FlashInfer SWA `DRAFT_EXTEND`** — not covered here. The FlashInfer SWA + coverage added for this path is limited to DFLASH `TARGET_VERIFY`. - **`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. diff --git a/test/registered/attention/unittests/swa/test_flashinfer.py b/test/registered/attention/unittests/swa/test_flashinfer.py index 899f511fa..5635a0906 100644 --- a/test/registered/attention/unittests/swa/test_flashinfer.py +++ b/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()