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

[Spec] Remove deprecated EAGLE v1 DRAFT_EXTEND forward mode (#28129)

Browse Source
This commit is contained in:
Cheng Wan
2026-06-13 13:00:21 -07:00
committed by GitHub
parent d988d5d681
commit bde6bccf39
56 changed files with 148 additions and 2484 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/unittests/dense/test_fa3.py
-107
View File
@@ -20,10 +20,7 @@ from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner i
run_dense_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_dense_draft_extend_cuda_graph_case,
run_dense_draft_extend_v2_cuda_graph_case,
run_dense_eagle_draft_extend_case,
run_dense_eagle_draft_extend_cuda_graph_runner_case,
run_dense_eagle_draft_extend_v2_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
@@ -64,62 +61,6 @@ class TestFA3DenseAttentionBackendCorrectness(CustomTestCase):
prefix_lens=(14, 15, 16),
),
)
DRAFT_EXTEND_CASES = (
(
DenseAttentionCase(
name="runner_fa3_eagle_draft_extend",
backend="fa3",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_fa3_frozen_kv_mtp_draft_extend",
backend="fa3",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
DRAFT_EXTEND_CUDA_GRAPH_CASES = (
(
DenseAttentionCase(
name="runner_cuda_graph_fa3_eagle_draft_extend",
backend="fa3",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_cuda_graph_fa3_frozen_kv_mtp_draft_extend",
backend="fa3",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
DRAFT_EXTEND_V2_CUDA_GRAPH_CASES = (
DenseAttentionCase(
name="runner_cuda_graph_fa3_eagle_draft_extend_v2_fixed_tokens",
@@ -247,18 +188,6 @@ class TestFA3DenseAttentionBackendCorrectness(CustomTestCase):
"ngram",
),
)
EAGLE_DRAFT_EXTEND_RUNNER_CASES = (
DenseAttentionCase(
name="runner_fa3_eagle_draft_extend_cuda_graph_runner",
backend="fa3",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
)
EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES = (
DenseAttentionCase(
name="runner_fa3_eagle_draft_extend_v2_cuda_graph_runner_fixed_tokens",
@@ -450,42 +379,6 @@ class TestFA3DenseAttentionBackendCorrectness(CustomTestCase):
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_dense_eagle_draft_extend_case(
self,
case,
spec_kind=spec_kind,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_draft_extend_cuda_graph_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CUDA_GRAPH_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_dense_draft_extend_cuda_graph_case(
self,
case,
spec_kind=spec_kind,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_eagle_draft_extend_cuda_graph_runner_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
test/registered/attention/unittests/dense/test_fa4.py
-107
View File
@@ -19,10 +19,7 @@ from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner i
run_dense_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_dense_draft_extend_cuda_graph_case,
run_dense_draft_extend_v2_cuda_graph_case,
run_dense_eagle_draft_extend_case,
run_dense_eagle_draft_extend_cuda_graph_runner_case,
run_dense_eagle_draft_extend_v2_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
@@ -59,62 +56,6 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase):
prefix_lens=(14, 15, 16),
),
)
DRAFT_EXTEND_CASES = (
(
DenseAttentionCase(
name="runner_fa4_eagle_draft_extend",
backend="fa4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_fa4_frozen_kv_mtp_draft_extend",
backend="fa4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
DRAFT_EXTEND_CUDA_GRAPH_CASES = (
(
DenseAttentionCase(
name="runner_cuda_graph_fa4_eagle_draft_extend",
backend="fa4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_cuda_graph_fa4_frozen_kv_mtp_draft_extend",
backend="fa4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
DRAFT_EXTEND_V2_CUDA_GRAPH_CASES = (
DenseAttentionCase(
name="runner_cuda_graph_fa4_eagle_draft_extend_v2_fixed_tokens",
@@ -242,18 +183,6 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase):
"ngram",
),
)
EAGLE_DRAFT_EXTEND_RUNNER_CASES = (
DenseAttentionCase(
name="runner_fa4_eagle_draft_extend_cuda_graph_runner",
backend="fa4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
)
EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES = (
DenseAttentionCase(
name="runner_fa4_eagle_draft_extend_v2_cuda_graph_runner_fixed_tokens",
@@ -439,42 +368,6 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase):
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_dense_eagle_draft_extend_case(
self,
case,
spec_kind=spec_kind,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_draft_extend_cuda_graph_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CUDA_GRAPH_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_dense_draft_extend_cuda_graph_case(
self,
case,
spec_kind=spec_kind,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_eagle_draft_extend_cuda_graph_runner_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
test/registered/attention/unittests/dense/test_flashinfer.py
-113
View File
@@ -19,11 +19,6 @@ 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_draft_extend_runner import (
run_dense_draft_extend_cuda_graph_case,
run_dense_eagle_draft_extend_case,
run_dense_eagle_draft_extend_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_dense_eagle_draft_cuda_graph_runner_case,
run_dense_frozen_kv_mtp_cuda_graph_runner_case,
@@ -251,74 +246,6 @@ class TestFlashInferDenseAttentionBackendCorrectness(CustomTestCase):
"ngram",
),
)
EAGLE_DRAFT_EXTEND_CASES = (
(
DenseAttentionCase(
name="runner_eagle_draft_extend_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(2, 5),
extend_lens=(1, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_frozen_kv_mtp_draft_extend_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(2, 5),
extend_lens=(1, 3),
),
"frozen_kv_mtp",
),
)
DRAFT_EXTEND_CUDA_GRAPH_CASES = (
(
DenseAttentionCase(
name="runner_cuda_graph_eagle_draft_extend_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(2, 5),
extend_lens=(1, 3),
),
"eagle",
),
(
DenseAttentionCase(
name="runner_cuda_graph_frozen_kv_mtp_draft_extend_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(2, 5),
extend_lens=(1, 3),
),
"frozen_kv_mtp",
),
)
EAGLE_DRAFT_EXTEND_RUNNER_CASES = (
DenseAttentionCase(
name="runner_eagle_draft_extend_cuda_graph_runner_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(2, 5),
extend_lens=(2, 4),
),
)
EAGLE_DRAFT_RUNNER_CASES = (
(
DenseAttentionCase(
@@ -475,46 +402,6 @@ class TestFlashInferDenseAttentionBackendCorrectness(CustomTestCase):
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name,
backend=case.backend,
spec_kind=spec_kind,
):
run_dense_eagle_draft_extend_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
spec_kind=spec_kind,
)
def test_runner_mode_draft_extend_cuda_graph_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CUDA_GRAPH_CASES:
with self.subTest(
case=case.name,
backend=case.backend,
spec_kind=spec_kind,
):
run_dense_draft_extend_cuda_graph_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
spec_kind=spec_kind,
)
def test_runner_mode_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_eagle_draft_extend_cuda_graph_runner_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
def test_runner_mode_eagle_draft_cuda_graph_runner_cases(self):
for case, topk, num_draft_tokens in self.EAGLE_DRAFT_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend, topk=topk):
test/registered/attention/unittests/dsa/test_dsa.py
+2 -38
View File
@@ -31,9 +31,6 @@ from sglang.test.kits.attention_unittest.attention_methods.dsa_attention import
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_dsa_sparse_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_dsa_eagle_draft_extend_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_dsa_eagle_draft_cuda_graph_runner_case,
)
@@ -217,8 +214,8 @@ class TestDSAAttentionBackendCorrectness(CustomTestCase):
self, self.DECODE_IMPL_CASE, impl
)
# Speculative forward-mode coverage. TARGET_VERIFY, DRAFT_EXTEND,
# and DRAFT_EXTEND_V2 all route through the `dsa_decode_impl`
# Speculative forward-mode coverage. TARGET_VERIFY and
# DRAFT_EXTEND_V2 both route through the `dsa_decode_impl`
# dispatcher (the same kernel selection as plain DECODE) but
# produce different `seqlens_expanded` and `cu_seqlens_q` from
# `dsa_backend.py:469-529`. `DSAMockModelRunner.__init__` derives
@@ -235,16 +232,6 @@ class TestDSAAttentionBackendCorrectness(CustomTestCase):
prefix_lens=(128,),
extend_lens=(3,),
),
DSAAttentionCase(
name="dsa_sparse_draft_extend",
backend="dsa",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=1,
page_size=DSA_PAGE_SIZE,
prefix_lens=(128,),
extend_lens=(3,),
),
DSAAttentionCase(
name="dsa_sparse_draft_extend_v2",
backend="dsa",
@@ -382,29 +369,6 @@ class TestDSAAttentionBackendCorrectness(CustomTestCase):
with self.subTest(case=case.name, backend=case.backend):
run_dsa_eagle_draft_cuda_graph_runner_case(self, case)
# EAGLE production draft-extend CUDA-graph runner. Routes through
# `DraftBackendFactory._create_dsa_prefill_backend` which returns a
# single `DeepseekSparseAttnBackend` (not multi-step); the forward
# goes through `forward_extend` with `dsa_decode_impl` selected via
# `is_draft_extend(include_v2=True)`.
EAGLE_DRAFT_EXTEND_CASES = (
DSAAttentionCase(
name="runner_eagle_draft_extend_cuda_graph_dsa",
backend="dsa",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
num_kv_heads=1,
page_size=DSA_PAGE_SIZE,
prefix_lens=(128, 192),
extend_lens=(2, 3),
),
)
def test_runner_mode_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dsa_eagle_draft_extend_cuda_graph_runner_case(self, case)
# CG decode replay with FP8 KV cache. Captures and replays through
# `flashmla_kv` (the only FP8-compatible decode kernel). The
# `_clone_dsa_sparse_cache` hook is reused as-is — it snapshots the
test/registered/attention/unittests/dsv4/test_deepseek_v4.py
-75
View File
@@ -34,16 +34,11 @@ from sglang.test.kits.attention_unittest.attention_methods.dsv4_attention import
make_dsv4_cases,
run_dsv4_attention_case,
run_dsv4_compress_attention_case,
run_dsv4_draft_extend_attention_case,
run_dsv4_target_verify_attention_case,
)
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import ( # noqa: E402
run_dsv4_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import ( # noqa: E402
run_dsv4_eagle_draft_extend_cuda_graph_case,
run_dsv4_eagle_draft_extend_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import ( # noqa: E402
run_dsv4_eagle_draft_cuda_graph_runner_case,
)
@@ -246,48 +241,6 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase):
):
run_dsv4_eagle_verify_cuda_graph_case(self, case, topk=1)
# EAGLE DRAFT_EXTEND is SWA-only for DSV4 (see runner docstring).
DRAFT_EXTEND_CASES = (
DSV4AttentionCase(
name="dsv4_swa_eagle_draft_extend",
backend="dsv4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=64,
page_size=DSV4_PAGE_SIZE,
prefix_lens=(64, 96),
extend_lens=(2, 4),
),
)
def test_eagle_draft_extend_cases(self):
for case in self.DRAFT_EXTEND_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dsv4_draft_extend_attention_case(self, case)
# CUDA-graph capture/replay for EAGLE DRAFT_EXTEND — SWA only
# (init_forward_metadata_draft_extend uses need_compress=False; see
# `Production-Unsupported` in dsv4/README.md). Uniform `extend_lens`
# because DSV4 `forward(compress_ratio=0)` asserts
# `swa_page_indices.shape[0] == q.shape[0]` and the graph metadata
# builder uses uniform `num_tokens_per_bs = max_num_tokens // max_bs`
# (see `deepseek_v4_backend.py:646-647`).
EAGLE_DRAFT_EXTEND_CUDA_GRAPH_CASES = (
DSV4AttentionCase(
name="runner_cuda_graph_dsv4_swa_eagle_draft_extend",
backend="dsv4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=64,
page_size=DSV4_PAGE_SIZE,
prefix_lens=(64, 96),
extend_lens=(4, 4),
),
)
def test_runner_mode_eagle_draft_extend_cuda_graph_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dsv4_eagle_draft_extend_cuda_graph_case(self, case)
# Production EAGLE draft graph runner (chain only, SWA only). The runner
# routes through `DeepseekV4MultiStepBackend` (one `DeepseekV4AttnBackend`
# per draft step), captures a fixed batch, and replays distinct request
@@ -309,34 +262,6 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase):
with self.subTest(case=case.name, backend=case.backend):
run_dsv4_eagle_draft_cuda_graph_runner_case(self, case)
# Production EAGLE draft-extend graph runner (SWA only). Routes through
# the prefill-side `DeepseekV4AttnBackend` (single backend, not
# multi-step); `init_forward_metadata_draft_extend` forces
# `need_compress=False` so C4/C128 is structurally unreachable for this
# path.
# Uniform `extend_lens` because the DSV4 graph contract requires
# `q.shape[0] == swa_page_indices.shape[0]` and the
# `init_forward_metadata_draft_extend` graph path uses
# `num_tokens_per_bs = max_num_tokens // max_bs` (see
# `deepseek_v4_backend.py:646-647`). Same constraint as the metadata-
# style draft_extend CG case.
PRODUCTION_EAGLE_DRAFT_EXTEND_RUNNER_CASES = (
DSV4AttentionCase(
name="runner_production_eagle_draft_extend_dsv4_swa",
backend="dsv4",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=64,
page_size=DSV4_PAGE_SIZE,
prefix_lens=(64, 96),
extend_lens=(4, 4),
),
)
def test_runner_mode_production_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.PRODUCTION_EAGLE_DRAFT_EXTEND_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dsv4_eagle_draft_extend_cuda_graph_runner_case(self, case)
class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
"""CPU-only checks for the DSV4 BCG metadata replay contract."""
test/registered/attention/unittests/gdn/test_triton.py
-45
View File
@@ -25,9 +25,6 @@ from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_gdn_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_gdn_eagle_draft_extend_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
run_gdn_eagle_verify_case,
run_gdn_eagle_verify_cuda_graph_case,
@@ -306,48 +303,6 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
self, case, topk=topk, spec_kind=spec_kind
)
# EAGLE / Frozen-KV MTP DRAFT_EXTEND eager — `HybridLinearAttnBackend`
# raises `ValueError("Invalid forward mode")` for DRAFT_EXTEND CG
# capture (`hybrid_linear_attn_backend.py:509,572`), so CG is
# structurally blocked across the family (GDN/KDA/Lightning/Mamba2).
# The EXTEND-style gated-delta recurrence reference doubles as the
# DRAFT_EXTEND reference across both spec kinds.
EAGLE_DRAFT_EXTEND_CASES = (
(
GDNAttentionCase(
name="runner_eagle_draft_extend_gdn",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
GDNAttentionCase(
name="runner_frozen_kv_mtp_draft_extend_gdn",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_gdn_eagle_draft_extend_case(self, case, spec_kind=spec_kind)
# Spy directly on each sub-backend's `init_forward_metadata*` so
# dispatch-layer slice mutations show up as a missing call, which
# forward-output assertions can miss when the fixture happens to
test/registered/attention/unittests/kda/test_triton.py
-41
View File
@@ -19,9 +19,6 @@ from sglang.test.kits.attention_unittest.attention_methods.kda_attention import
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_kda_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_kda_eagle_draft_extend_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
run_kda_eagle_verify_case,
run_kda_eagle_verify_cuda_graph_case,
@@ -239,44 +236,6 @@ class TestTritonKDABackendCorrectness(CustomTestCase):
32,
),
)
# EAGLE / Frozen-KV MTP DRAFT_EXTEND eager. CG is structurally
# blocked across the HybridLinearAttn family
# (`hybrid_linear_attn_backend.py:509,572`).
EAGLE_DRAFT_EXTEND_CASES = (
(
KDAAttentionCase(
name="runner_eagle_draft_extend_kda",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
KDAAttentionCase(
name="runner_frozen_kv_mtp_draft_extend_kda",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_kda_eagle_draft_extend_case(self, case, spec_kind=spec_kind)
@unittest.skipIf(
is_hip(),
test/registered/attention/unittests/lightning/test_triton.py
-39
View File
@@ -18,9 +18,6 @@ from sglang.test.kits.attention_unittest.attention_methods.lightning_attention i
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_lightning_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_lightning_eagle_draft_extend_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
run_lightning_eagle_verify_case,
run_lightning_eagle_verify_cuda_graph_case,
@@ -186,42 +183,6 @@ class TestTritonLightningBackendCorrectness(CustomTestCase):
):
run_lightning_eagle_verify_cuda_graph_case(self, case, topk=topk)
# EAGLE / Frozen-KV MTP DRAFT_EXTEND eager — CG is structurally
# blocked across the HybridLinearAttn family.
EAGLE_DRAFT_EXTEND_CASES = (
(
LightningAttentionCase(
name="runner_eagle_draft_extend_lightning",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"eagle",
),
(
LightningAttentionCase(
name="runner_frozen_kv_mtp_draft_extend_lightning",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=2,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
"frozen_kv_mtp",
),
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_lightning_eagle_draft_extend_case(self, case, spec_kind=spec_kind)
# PCG/BCG split-op extend is deliberately NOT covered. Lightning's
# backend `forward_extend` flattens the output via `o.view(-1,
# tp_q_head_num * v_head_dim)` (`lightning_backend.py:335`), so eager
test/registered/attention/unittests/mamba/test_mamba2.py
-37
View File
@@ -32,9 +32,6 @@ from sglang.test.kits.attention_unittest.attention_methods.mamba2_attention impo
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_mamba2_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_mamba2_eagle_draft_extend_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
run_mamba2_eagle_verify_case,
run_mamba2_eagle_verify_cuda_graph_case,
@@ -128,33 +125,6 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
1,
),
)
# EAGLE / Frozen-KV MTP DRAFT_EXTEND eager — `HybridLinearAttnBackend`
# raises `ValueError("Invalid forward mode")` for DRAFT_EXTEND CG
# capture (`hybrid_linear_attn_backend.py:509,572`), so CG is
# structurally blocked; only eager is exercised. Same EXTEND-style
# chunked-scan recurrence reference doubles as the DRAFT_EXTEND
# reference, like the verify path.
EAGLE_DRAFT_EXTEND_CASES = tuple(
(
Mamba2AttentionCase(
name=f"runner_{spec_kind}_draft_extend_mamba2",
backend="triton",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=DEFAULT_NUM_HEADS,
head_dim=DEFAULT_HEAD_DIM,
state_size=DEFAULT_STATE_SIZE,
n_groups=DEFAULT_N_GROUPS,
conv_kernel=DEFAULT_CONV_KERNEL,
mamba_chunk_size=DEFAULT_MAMBA_CHUNK_SIZE,
hidden_size=DEFAULT_HIDDEN_SIZE,
page_size=16,
prefix_lens=(4, 7),
extend_lens=(3, 3),
),
spec_kind,
)
for spec_kind in ("eagle", "frozen_kv_mtp")
)
def test_projected_mamba2_attention_cases(self):
for case in self.CASES:
@@ -194,13 +164,6 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
with self.subTest(case=case.name, backend=case.backend, topk=topk):
run_mamba2_eagle_verify_cuda_graph_case(self, case, topk=topk)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.EAGLE_DRAFT_EXTEND_CASES:
with self.subTest(
case=case.name, backend=case.backend, spec_kind=spec_kind
):
run_mamba2_eagle_draft_extend_case(self, case, spec_kind=spec_kind)
# PCG/BCG split-op extend is deliberately NOT covered. The
# `MambaMixer2.forward` asserts `num_actual_tokens ==
# projected_states.shape[0]` (`mamba.py:467`) — the projection step
test/registered/attention/unittests/mla/test_flashinfer.py
-61
View File
@@ -17,11 +17,6 @@ from sglang.test.kits.attention_unittest.attention_methods.mla_attention import
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_mla_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_mla_draft_extend_cuda_graph_case,
run_mla_eagle_draft_extend_case,
run_mla_eagle_draft_extend_cuda_graph_runner_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_mla_eagle_draft_cuda_graph_runner_case,
)
@@ -101,39 +96,6 @@ class TestFlashInferMLAAttentionBackendCorrectness(CustomTestCase):
1,
),
)
DRAFT_EXTEND_CASES = (
MLAAttentionCase(
name="runner_eagle_draft_extend_mla_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
page_size=16,
prefix_lens=(5, 8),
extend_lens=(2, 4),
),
)
DRAFT_EXTEND_CUDA_GRAPH_CASES = (
MLAAttentionCase(
name="runner_cuda_graph_eagle_draft_extend_mla_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
page_size=16,
prefix_lens=(5, 8),
extend_lens=(2, 4),
),
)
EAGLE_DRAFT_EXTEND_RUNNER_CASES = (
MLAAttentionCase(
name="runner_eagle_draft_extend_mla_cuda_graph_runner_ragged_accept",
backend="flashinfer",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
page_size=16,
prefix_lens=(5, 8),
extend_lens=(2, 4),
),
)
EAGLE_DRAFT_RUNNER_CASES = (
(
MLAAttentionCase(
@@ -259,29 +221,6 @@ class TestFlashInferMLAAttentionBackendCorrectness(CustomTestCase):
**MLA_SHAPE_KWARGS,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case in self.DRAFT_EXTEND_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_mla_eagle_draft_extend_case(self, case, **MLA_SHAPE_KWARGS)
def test_runner_mode_eagle_draft_extend_cuda_graph_cases(self):
for case in self.DRAFT_EXTEND_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_mla_draft_extend_cuda_graph_case(
self,
case,
**MLA_SHAPE_KWARGS,
)
def test_runner_mode_eagle_draft_extend_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_RUNNER_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_mla_eagle_draft_extend_cuda_graph_runner_case(
self,
case,
**MLA_SHAPE_KWARGS,
)
def test_runner_mode_eagle_draft_cuda_graph_runner_cases(self):
# Backend gate (KNOWN_FAILURES.md §3): FlashInfer MLA multi-step
# draft CG capture/replay produces numerically wrong outputs on
test/registered/attention/unittests/mla/test_flashmla.py
-19
View File
@@ -21,9 +21,6 @@ from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner i
_init_cuda_graph_replay_metadata,
run_mla_cuda_graph_decode_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import (
run_mla_eagle_draft_extend_case,
)
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_mla_eagle_draft_cuda_graph_runner_case,
)
@@ -203,17 +200,6 @@ class TestFlashMLAAttentionBackendCorrectness(CustomTestCase):
1,
),
)
DRAFT_EXTEND_CASES = (
MLAAttentionCase(
name="runner_eagle_draft_extend_mla_flashmla_ragged_accept",
backend="flashmla",
forward_mode=ForwardMode.DRAFT_EXTEND,
num_heads=4,
page_size=64,
prefix_lens=(5, 8),
extend_lens=(2, 4),
),
)
EAGLE_DRAFT_RUNNER_CASES = (
(
MLAAttentionCase(
@@ -338,11 +324,6 @@ class TestFlashMLAAttentionBackendCorrectness(CustomTestCase):
**MLA_SHAPE_KWARGS,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case in self.DRAFT_EXTEND_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_mla_eagle_draft_extend_case(self, case, **MLA_SHAPE_KWARGS)
@unittest.skipIf(_DECODE_REQUIRES_SM90A, _DECODE_SKIP_REASON)
def test_runner_mode_eagle_draft_cuda_graph_runner_cases(self):
for case, topk, num_draft_tokens in self.EAGLE_DRAFT_RUNNER_CASES:
test/registered/kv_canary/test_self_unit_token_oracle.py
+1 -1
View File
@@ -26,7 +26,7 @@ class TestTokenOracleManager(CustomTestCase):
) -> None:
"""Verify EAGLE draft extend maps one request row to every draft token."""
forward_batch = SimpleNamespace(
forward_mode=ForwardMode.DRAFT_EXTEND,
forward_mode=ForwardMode.DRAFT_EXTEND_V2,
spec_info=SimpleNamespace(num_tokens_per_req=4),
rids_int=torch.tensor([3, 7], dtype=torch.int64, device=self.device),
bootstrap_room_ids_int=None,