Reland spec v2 tree drafting (eagle topk>1) with page_size==1 (#26866) (#26997)

Co-authored-by: Alison Shao <54658187+alisonshao@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2026-06-03 15:40:05 -04:00
committed by GitHub
co-authored by Alison Shao
parent 7f706f4cfb
commit ac99794e64
12 changed files with 176 additions and 18 deletions
@@ -20,7 +20,7 @@ from sglang.test.kits.spec_server_kits import (
)
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
register_cuda_ci(est_time=780, stage="base-b", runner_config="1-gpu-large")
class TestEagle3Perf(Eagle3Base, SpecPerfKit):
@@ -40,6 +40,20 @@ class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
)
class TestEagle3Topk16V2Retract(Eagle3Base, SpecAccuracyKit, SpecFeatureKit):
"""EAGLE3 topk=16 tree on spec v2 under retract; must not leak KV. Stresses
the accepted-path KV move (move_accepted_tokens_to_target_kvcache)."""
spec_topk = 16
spec_tokens = 64
disable_overlap = False
cuda_graph_max_bs = 5
max_running_requests = 64
gsm8k_accept_len_thres = 2.4
extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract
env_overrides = ((envs.SGLANG_TEST_RETRACT, True),)
class TestEagleLlama2AbortAll(EagleLlama2Base, AbortAllMixin):
abort_all_max_new_tokens = 4000
@@ -1,7 +1,9 @@
"""topk > 1 tree drafting (EAGLE3 topk16 + EAGLE/Llama-2 topk8).
topk > 1 always routes to spec v1; flashinfer is pinned (topk > 1 can't use fa3).
Runs on the cheap (5090) runner -- functional sanity only, no perf/stress.
topk > 1 routes to spec v1, except page_size==1 which can also stay on spec v2
(overlap). flashinfer is pinned because this runs on the cheap (5090) runner,
where fa3 (Hopper-only) isn't available -- functional sanity only, no perf/stress.
(topk > 1 on fa3 is covered on the Hopper runner in test_spec_eagle_fa3.py.)
"""
import unittest
@@ -16,7 +18,7 @@ from sglang.test.kits.spec_server_kits import (
)
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
register_cuda_ci(est_time=840, stage="base-b", runner_config="1-gpu-small")
register_cuda_ci(est_time=1180, stage="base-b", runner_config="1-gpu-small")
class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit):
@@ -31,6 +33,13 @@ class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogp
gsm8k_accept_len_thres = 2.4 # EAGLE3 topk16 gsm8k accept ~2.48
class TestEagle3Topk16SpecV2(TestEagle3Topk16, SpecFeatureKit):
"""EAGLE3 topk=16 tree on spec v2 (overlap, page1): guards the v2 tree path's
accepted-path compaction, validated by logprob_spec_v2_match."""
disable_overlap = False
class TestEagleLlama2Suite(
EagleLlama2Base,
SpecCorrectnessKit,