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

[spec decoding] supports step 0 in adaptive spec decoding (updating draft kv cache without draft decoding) (#23994)

Browse Source 
Co-authored-by: shuwenn <2508695655@qq.com>
This commit is contained in:
Qiaolin Yu
2026-06-15 22:21:26 -07:00
committed by GitHub
co-authored by shuwenn
parent 800aaefc9e
commit e068355831
6 changed files with 313 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/spec/eagle/test_adaptive_speculative.py
+114 -1
View File
@@ -18,7 +18,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=76, stage="base-b", runner_config="1-gpu-large")
register_cuda_ci(est_time=160, stage="base-b", runner_config="1-gpu-large")
HIGH_ACCEPT_PROMPT = (
"Output exactly 128 new lines. "
@@ -198,5 +198,118 @@ class TestAdaptiveSpeculativeServer(CustomTestCase):
)
class TestAdaptiveZeroStepBatchSizeServer(CustomTestCase):
"""steps=0 (nospec) fallback triggered by batch size.
Config routes BS>=8 -> steps=0 (drafting disabled) and BS<8 -> steps=3, so the
server cycles steps=3 -> steps=0 -> steps=3 as load rises and falls.
"""
model = DEFAULT_TARGET_MODEL_EAGLE
draft_model = DEFAULT_DRAFT_MODEL_EAGLE
base_url = DEFAULT_URL_FOR_TEST
COUNT_PROMPT = "Count from 1 to 400, separated by commas. Output only the numbers."
@classmethod
def setUpClass(cls):
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
json.dump(
{
"1": {"candidate_steps": [3], "warmup_batches": 0},
"8": {"candidate_steps": [0], "warmup_batches": 0},
},
f,
)
cls.adaptive_config_path = f.name
try:
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--attention-backend",
"triton",
"--speculative-algorithm",
"EAGLE",
"--speculative-draft-model-path",
cls.draft_model,
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--speculative-adaptive",
"--speculative-adaptive-config",
cls.adaptive_config_path,
"--max-running-requests",
"32",
"--skip-server-warmup",
"--mem-fraction-static",
"0.7",
],
)
except Exception:
os.unlink(cls.adaptive_config_path)
raise
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process"):
kill_process_tree(cls.process.pid)
if os.path.exists(cls.adaptive_config_path):
os.unlink(cls.adaptive_config_path)
def _steps(self) -> int:
r = requests.get(self.base_url + "/server_info", timeout=30)
self.assertEqual(r.status_code, 200, r.text)
return r.json()["internal_states"][0]["speculative_num_steps"]
def test_batch_size_step_cycle(self):
"""The server cycles steps=3 -> steps=0 -> steps=3 as load rises and falls:
a BS=1 request drafts at steps=3; a 14-way batch (BS>=8) routes the worker
to nospec steps=0; a following BS=1 request returns to steps=3 with drafting
restored (high accept rate again)."""
one = {"temperature": 0, "max_new_tokens": 64, "ignore_eos": True}
def generate_single() -> dict:
r = requests.post(
self.base_url + "/generate",
json={"text": self.COUNT_PROMPT, "sampling_params": one},
timeout=600,
)
self.assertEqual(r.status_code, 200, r.text)
return r.json()["meta_info"]
# Phase 1: BS=1 -> steps=3, drafting active.
m1 = generate_single()
self.assertEqual(self._steps(), 3, "expected steps=3 at BS=1")
self.assertGreater(
m1["spec_accept_rate"], 0.8, f"not drafting at steps=3: {m1}"
)
# Phase 2: BS=14 -> the worker switches to nospec steps=0. Equal-length
# requests finish together, so the last decode batch (and thus the state)
# is at BS=14 -> steps=0.
full = {"temperature": 0, "max_new_tokens": 128, "ignore_eos": True}
r = requests.post(
self.base_url + "/generate",
json={"text": [self.COUNT_PROMPT] * 14, "sampling_params": [full] * 14},
timeout=600,
)
self.assertEqual(r.status_code, 200, r.text)
self.assertEqual(self._steps(), 0, "BS>=8 did not switch to steps=0")
# Phase 3: BS=1 -> steps=3 again, drafting restored.
m3 = generate_single()
self.assertEqual(self._steps(), 3, "did not reopen to steps=3")
self.assertGreater(
m3["spec_accept_rate"], 0.8, f"drafting not restored after steps=0: {m3}"
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_adaptive_spec_params.py
+35 -17
View File
@@ -198,6 +198,30 @@ class TestAdaptiveStepSlot(unittest.TestCase):
self.assertEqual(params.current_steps, 1)
self.assertEqual(params.ema_accept_len, 0.375)
def test_zero_step_mixed_slot_drops_probes_and_rechecks(self):
params = self._make_params_from_config(
3,
{
"candidate_steps": [0, 3],
"ema_alpha": 1.0,
"warmup_batches": 0,
"update_interval": 1,
"down_hysteresis": 0.0,
},
)
self.assertTrue(params.update([0, 0]))
self.assertEqual(params.current_steps, 0)
self.assertEqual(params.ema_accept_len, 0.0)
self.assertTrue(params.update([3, 3]))
self.assertEqual(params.current_steps, 3)
self.assertEqual(params.ema_accept_len, 0.0)
self.assertTrue(params.update([0, 0]))
self.assertEqual(params.current_steps, 0)
self.assertEqual(params.ema_accept_len, 0.0)
def test_ceiling_coeff_caps_steps(self):
params = self._make_params_from_config(
7,
@@ -223,10 +247,11 @@ class TestAdaptiveStepSlot(unittest.TestCase):
class TestAdaptiveSpeculativeParams(unittest.TestCase):
def test_default_config_loads(self):
params = AdaptiveSpeculativeParams(initial_steps=3)
self.assertEqual(params._bs_list, [1, 8, 32])
self.assertEqual(params._bs_list, [1, 8, 32, 64])
self.assertEqual(params._slots[1].candidate_steps, [1, 3, 7])
self.assertEqual(params._slots[8].candidate_steps, [1, 3])
self.assertEqual(params._slots[32].candidate_steps, [1])
self.assertEqual(params._slots[8].candidate_steps, [0, 1, 3])
self.assertEqual(params._slots[32].candidate_steps, [0, 1])
self.assertEqual(params._slots[64].candidate_steps, [0])
def test_config_file(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
@@ -287,13 +312,6 @@ class TestAdaptiveSpeculativeParams(unittest.TestCase):
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_zero_steps_raises(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"1": {"candidate_steps": [0]}}, f)
f.flush()
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_global_hysteresis_inherited(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
@@ -333,20 +351,20 @@ class TestBatchSizeRouting(unittest.TestCase):
# A batch maps to the largest slot BS <= batch (floor), capped at the top slot.
self.assertEqual(params._route(1).candidate_steps, [1, 3, 7])
self.assertEqual(params._route(7).candidate_steps, [1, 3, 7])
self.assertEqual(params._route(8).candidate_steps, [1, 3])
self.assertEqual(params._route(31).candidate_steps, [1, 3])
self.assertEqual(params._route(32).candidate_steps, [1])
self.assertEqual(params._route(1000).candidate_steps, [1])
self.assertEqual(params._route(8).candidate_steps, [0, 1, 3])
self.assertEqual(params._route(31).candidate_steps, [0, 1, 3])
self.assertEqual(params._route(32).candidate_steps, [0, 1])
self.assertEqual(params._route(1000).candidate_steps, [0])
def test_cuda_graph_bs_pads_batch_up_before_routing(self):
params = self._params()
params.set_cuda_graph_bs([4, 8, 16, 32])
# bs=5 pads up to the captured graph BS 8 -> slot bs=8.
self.assertEqual(params._route(5).candidate_steps, [1, 3])
self.assertEqual(params._route(5).candidate_steps, [0, 1, 3])
# bs=17 pads up to 32 -> slot bs=32.
self.assertEqual(params._route(17).candidate_steps, [1])
self.assertEqual(params._route(17).candidate_steps, [0, 1])
# A batch larger than every captured BS keeps its own value -> top slot.
self.assertEqual(params._route(100).candidate_steps, [1])
self.assertEqual(params._route(100).candidate_steps, [0])
def test_cuda_graph_bs_for_step_prunes_unreachable_graphs(self):
params = self._params()