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
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()