minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
David Wang
2026-06-14 00:32:02 -07:00
committed by GitHub
parent bb48405c31
commit 8c5320b37e
6 changed files with 250 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashinfer_backend.py
+40 -15
View File
@@ -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:
python/sglang/srt/models/dflash.py
+35 -2
View File
@@ -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
+31
View File
@@ -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
+69 -4
View File
@@ -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
+3 -10
View File
@@ -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.
test/registered/attention/unittests/swa/test_flashinfer.py
+72
View File
@@ -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()