From e068355831fe9c9433a4fb5e742906cf2cdb225f Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Mon, 15 Jun 2026 22:21:26 -0700 Subject: [PATCH] [spec decoding] supports step 0 in adaptive spec decoding (updating draft kv cache without draft decoding) (#23994) Co-authored-by: shuwenn <2508695655@qq.com> --- python/sglang/srt/environ.py | 4 + .../srt/speculative/adaptive_spec_params.py | 62 ++++++--- python/sglang/srt/speculative/draft_utils.py | 3 +- .../sglang/srt/speculative/eagle_worker_v2.py | 130 +++++++++++++++--- .../spec/eagle/test_adaptive_speculative.py | 115 +++++++++++++++- .../unit/spec/test_adaptive_spec_params.py | 52 ++++--- 6 files changed, 313 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 300a76bef..6d2960303 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -638,6 +638,10 @@ class Envs: # Spec Config SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) + # Skip draft_extend while adaptive spec is at steps=0 (drafting disabled). + # Saves the per-step draft forward, but the draft KV goes stale: an upshift + # back to steps>0 starts from a cold draft state (low accept until it recovers). + SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False) # Master switch for all async-asserted invariant probes (NaN, Inf, OOB, # page alignment). Off in prod; tests turn it on to fail-fast on # numerical / index violations instead of getting silent NaN cascades. diff --git a/python/sglang/srt/speculative/adaptive_spec_params.py b/python/sglang/srt/speculative/adaptive_spec_params.py index 286c7e13c..134289299 100644 --- a/python/sglang/srt/speculative/adaptive_spec_params.py +++ b/python/sglang/srt/speculative/adaptive_spec_params.py @@ -19,7 +19,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# TODO: add step=0 (nospec fallback) for BS>=8 once supported. DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = { "1": { "candidate_steps": [1, 3, 7], @@ -28,13 +27,19 @@ DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = { "ceiling_coeff": 0, }, "8": { - "candidate_steps": [1, 3], + "candidate_steps": [0, 1, 3], "up_hysteresis": 0.0, "down_hysteresis": 0.0, "ceiling_coeff": 0, }, "32": { - "candidate_steps": [1], + "candidate_steps": [0, 1], + "up_hysteresis": 0.0, + "down_hysteresis": 0.0, + "ceiling_coeff": 0, + }, + "64": { + "candidate_steps": [0], "up_hysteresis": 0.0, "down_hysteresis": 0.0, "ceiling_coeff": 0, @@ -102,10 +107,11 @@ def _load_adaptive_config( if ( not isinstance(steps, list) or not steps - or not all(isinstance(s, int) and s > 0 for s in steps) + or not all(isinstance(s, int) and s >= 0 for s in steps) ): raise ValueError( - f"BS {key}: candidate_steps must be a list of positive ints, got {steps!r}" + f"BS {key}: candidate_steps must be a list of non-negative ints, " + f"got {steps!r}" ) bs_entries[int(key)] = entry @@ -172,10 +178,13 @@ class AdaptiveStepSlot: if not num_correct_drafts_per_req: return False - batch_avg = sum(num_correct_drafts_per_req) / len(num_correct_drafts_per_req) - self.ema_accept_len = ( - 1 - self.ema_alpha - ) * self.ema_accept_len + self.ema_alpha * batch_avg + if self.current_steps > 0: + batch_avg = sum(num_correct_drafts_per_req) / len( + num_correct_drafts_per_req + ) + self.ema_accept_len = ( + 1 - self.ema_alpha + ) * self.ema_accept_len + self.ema_alpha * batch_avg self._batch_count += 1 if self._batch_count <= self.warmup_batches: @@ -190,23 +199,39 @@ class AdaptiveStepSlot: """Recompute steps from EMA. Returns True if params changed.""" old_steps = self.current_steps current_idx = self.candidate_steps.index(old_steps) + old_idx = current_idx + + # Probe the smallest positive step after a zero-step nospec interval. + if old_steps == 0: + current_idx = min(current_idx + 1, len(self.candidate_steps) - 1) + target = self.candidate_steps[current_idx] + if target > 0 and self.ema_accept_len < 0: + # A slot initialized at steps=0 has no draft acceptance history; + # start the first positive-step probe from that step's neutral EMA. + self.ema_accept_len = float(target - 1) + return self._apply_target_steps(old_steps, target) # TODO: Consider limiting step changes to avoid overshooting. while current_idx > 0: prev_step = self.candidate_steps[current_idx - 1] - drop_threshold = prev_step - 0.5 + self.down_hysteresis + # A zero-step candidate disables drafting. Treat zero accepted drafts + # as low enough to reach it when it is the floor candidate. + drop_threshold = 0.5 if prev_step == 0 else prev_step - 0.5 + drop_threshold += self.down_hysteresis if self.ema_accept_len <= drop_threshold: current_idx -= 1 else: break - while current_idx < len(self.candidate_steps) - 1: - current_step = self.candidate_steps[current_idx] - rise_threshold = current_step - 0.5 + self.up_hysteresis - if self.ema_accept_len > rise_threshold: - current_idx += 1 - else: - break + moved_down = current_idx < old_idx + if not moved_down: + while current_idx < len(self.candidate_steps) - 1: + current_step = self.candidate_steps[current_idx] + rise_threshold = current_step - 0.5 + self.up_hysteresis + if self.ema_accept_len > rise_threshold: + current_idx += 1 + else: + break target = self.candidate_steps[current_idx] # EMA ceiling: only caps downward — never blocks step-ups, so the @@ -218,6 +243,9 @@ class AdaptiveStepSlot: current_idx -= 1 target = self.candidate_steps[current_idx] + return self._apply_target_steps(old_steps, target) + + def _apply_target_steps(self, old_steps: int, target: int) -> bool: if target != old_steps: self.current_steps = target log_info_on_rank0( diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index 10bc78954..49ef47b7b 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -37,7 +37,8 @@ class DraftBackendFactory: return backend_map[backend_type]() def create_decode_backend(self): - if self.speculative_num_steps == 1: + # No multi-step draft backend for steps=0 (nospec) or steps=1. + if self.speculative_num_steps <= 1: return None backend_map = { diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index f3557a1fa..22ce04998 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1011,33 +1011,129 @@ class EAGLEWorkerV2(BaseSpecWorker): topk=self.topk, capture_hidden_mode=capture_mode, ) - with ( - self.draft_worker.draft_tp_context( - self.draft_worker.draft_runner.tp_group - ), - speculative_moe_backend_context(), - speculative_moe_a2a_backend_context(), - spec_stage_span("draft"), - ): - verify_input: EagleVerifyInput = self.draft_worker.draft(batch) + if self.speculative_num_steps == 0: + # Drafting disabled (high batch size). _draft_extend below still + # runs, keeping draft KV warm for when the batch shrinks. + verify_input = self._build_trivial_verify_input(batch) + else: + with ( + self.draft_worker.draft_tp_context( + self.draft_worker.draft_runner.tp_group + ), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + spec_stage_span("draft"), + ): + verify_input: EagleVerifyInput = self.draft_worker.draft(batch) assert verify_input.is_verify_input() batch.spec_info = verify_input batch_output = self.verify(batch) # Publish before draft_extend so the fence is at verify-end. if on_publish is not None: on_publish(batch_output.new_seq_lens) - with ( - self.draft_worker.draft_tp_context( - self.draft_worker.draft_runner.tp_group - ), - speculative_moe_backend_context(), - speculative_moe_a2a_backend_context(), - spec_stage_span("draft_extend"), + if ( + self.speculative_num_steps == 0 + and envs.SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND.get() ): - self.draft_worker._draft_extend_for_decode(batch, batch_output) + self._stub_skipped_draft_extend(batch, batch_output) + else: + with ( + self.draft_worker.draft_tp_context( + self.draft_worker.draft_runner.tp_group + ), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + spec_stage_span("draft_extend"), + ): + self.draft_worker._draft_extend_for_decode(batch, batch_output) return batch_output + def _build_trivial_verify_input(self, batch: ScheduleBatch) -> EagleVerifyInput: + """Build a 1-node EagleVerifyInput rooted at the previous bonus token. + + Used when ``speculative_num_steps == 0`` to skip drafting while still + routing through the existing TARGET_VERIFY graph captured at + ``draft_token_num=1``: the kernel always accepts the root and samples + one new bonus token from target logits -- functionally a plain decode. + """ + if batch.forward_mode.is_idle(): + return EagleVerifyInput.create_idle_input( + topk=self.topk, spec_steps=0, num_verify_tokens=1 + ) + + draft_input: EagleDraftInput = batch.spec_info + bs = batch.seq_lens.shape[0] + device = self.device + + retrieve_index = torch.arange(bs, dtype=torch.long, device=device).unsqueeze(1) + retrieve_next_token = torch.full((bs, 1), -1, dtype=torch.long, device=device) + retrieve_next_sibling = torch.full((bs, 1), -1, dtype=torch.long, device=device) + + attn_backend = self._target_worker.model_runner.attn_backend + mask_buf, position_buf = attn_backend.get_verify_buffers_to_fill_after_draft() + if mask_buf is not None: + custom_mask = mask_buf + custom_mask.fill_(True) + else: + if batch.seq_lens_sum is not None: + seq_lens_sum = batch.seq_lens_sum + elif batch.seq_lens_cpu is not None: + seq_lens_sum = int(batch.seq_lens_cpu.sum()) + else: + seq_lens_sum = bs * attn_backend.max_context_len + custom_mask = torch.ones(seq_lens_sum + bs, dtype=torch.bool, device=device) + + if position_buf is not None: + positions = position_buf + positions[:bs].copy_(batch.seq_lens) + else: + positions = batch.seq_lens.to(torch.int64) + + return EagleVerifyInput( + draft_token=draft_input.bonus_tokens, + custom_mask=custom_mask, + positions=positions, + retrieve_index=retrieve_index, + retrieve_next_token=retrieve_next_token, + retrieve_next_sibling=retrieve_next_sibling, + retrieve_cum_len=None, + spec_steps=0, + topk=self.topk, + draft_token_num=1, + capture_hidden_mode=CaptureHiddenMode.FULL, + seq_lens_sum=None, + seq_lens_cpu=None, + ) + + def _stub_skipped_draft_extend( + self, batch: ScheduleBatch, batch_output: GenerationBatchResult + ) -> None: + """Fill shape-valid stubs on next_draft_input when draft_extend is skipped. + + ``verify`` already set ``bonus_tokens`` (the only field the next steps=0 + verify reads). The overlap FutureMap still stashes topk_p/topk_index/ + hidden_states, so provide zeroed tensors of the right shape. They are never + consumed while at steps=0; an upshift to steps>0 would draft from this stale + state (cold recovery), which is the documented cost of this experimental flag. + """ + next_draft_input: EagleDraftInput = batch_output.next_draft_input + bs = batch.seq_lens.shape[0] + device = self.device + next_draft_input.topk_p = torch.zeros( + (bs, self.topk), dtype=torch.float32, device=device + ) + next_draft_input.topk_index = torch.zeros( + (bs, self.topk), dtype=torch.int64, device=device + ) + hidden_size = EagleDraftInput.hidden_size_for(self.draft_worker) + if hidden_size is not None: + next_draft_input.hidden_states = torch.zeros( + (bs, hidden_size), + dtype=EagleDraftInput.dtype_for(self.draft_worker), + device=device, + ) + def on_verify_complete_cpu( self, num_correct_drafts_per_req: list[int], batch_size: int = 0 ) -> None: diff --git a/test/registered/spec/eagle/test_adaptive_speculative.py b/test/registered/spec/eagle/test_adaptive_speculative.py index bcf17b7e0..1c33e0aa1 100644 --- a/test/registered/spec/eagle/test_adaptive_speculative.py +++ b/test/registered/spec/eagle/test_adaptive_speculative.py @@ -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() diff --git a/test/registered/unit/spec/test_adaptive_spec_params.py b/test/registered/unit/spec/test_adaptive_spec_params.py index 44734c385..6cb71f1fb 100644 --- a/test/registered/unit/spec/test_adaptive_spec_params.py +++ b/test/registered/unit/spec/test_adaptive_spec_params.py @@ -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()