From a01afdd52656343b7d6736961a3a220b909a3a65 Mon Sep 17 00:00:00 2001 From: weireweire Date: Wed, 1 Jul 2026 08:22:05 +0800 Subject: [PATCH] Support real draft tokens to simulated acceptance (#29645) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- python/sglang/srt/environ.py | 1 + python/sglang/srt/speculative/eagle_utils.py | 24 ++++++++++++++++++++ python/sglang/srt/speculative/spec_utils.py | 22 +++++++++++++++++- 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a9166094e..7a47e0b19 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -288,6 +288,7 @@ class Envs: SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) SGLANG_SIMULATE_ACC_METHOD = EnvStr("match-expected") + SGLANG_SIMULATE_ACC_TOKEN_MODE = EnvStr("fixed") SGLANG_SIMULATE_UNIFORM_EXPERTS = EnvBool(False) SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS = EnvBool(False) SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp") diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index b75bce92e..90f147384 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -436,6 +436,7 @@ def eagle_sample( from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_utils import ( SIMULATE_ACC_LEN, + SIMULATE_ACC_TOKEN_MODE, generate_simulated_accept_index, ) from sglang.srt.utils.async_probe import maybe_detect_nan, sanitize_nan_logits @@ -493,6 +494,7 @@ def eagle_sample( num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device) # Sample tokens + target_predict = None if sampling_info.is_all_greedy or _is_npu or _is_hip: target_predict = torch.argmax(next_token_logits, dim=-1) target_predict = target_predict.reshape(bs, verify_input.draft_token_num) @@ -607,11 +609,33 @@ def eagle_sample( # Do simulation. The helper builds (and returns) a replacement # accept_index of width spec_steps + 1, so pass max_tree_depth - 1 # to keep the simulated width identical to the real one. + if SIMULATE_ACC_TOKEN_MODE not in ("fixed", "real-draft-token"): + raise ValueError( + "Invalid SGLANG_SIMULATE_ACC_TOKEN_MODE " + f"{SIMULATE_ACC_TOKEN_MODE!r}; expected 'fixed' or " + "'real-draft-token'." + ) + + if SIMULATE_ACC_TOKEN_MODE == "real-draft-token": + if verify_input.tree_topk != 1: + raise ValueError( + "SGLANG_SIMULATE_ACC_LEN with real draft tokens currently " + "requires speculative_eagle_topk=1." + ) + + # Use target argmax as the synthetic bonus for non-greedy requests. + if target_predict is None: + target_predict = torch.argmax(next_token_logits, dim=-1).reshape( + bs, verify_input.draft_token_num + ) accept_index = generate_simulated_accept_index( accept_index=accept_index, predict=predict, # mutable num_correct_drafts=num_correct_drafts, # mutable + candidates=candidates, + target_predict=target_predict, simulate_acc_len=SIMULATE_ACC_LEN, + simulate_acc_token_mode=SIMULATE_ACC_TOKEN_MODE, bs=bs, spec_steps=verify_input.max_tree_depth - 1, ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index da3593318..c1ca85bb8 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -97,6 +97,7 @@ def renorm_draft_probs( # Simulate acceptance length for benchmarking purposes SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get() # turn off if < 0 SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get() +SIMULATE_ACC_TOKEN_MODE = envs.SGLANG_SIMULATE_ACC_TOKEN_MODE.get() TREE_TRAVERSE_TIME_THRESHOLD = 1 # TODO: set this properly TREE_SPEC_KERNEL_AVAILABLE = ( @@ -316,11 +317,16 @@ def generate_simulated_accept_index( accept_index, predict, num_correct_drafts, + candidates, + target_predict, bs, spec_steps, simulate_acc_len: float = SIMULATE_ACC_LEN, simulate_acc_method: str = SIMULATE_ACC_METHOD, + simulate_acc_token_mode: str = SIMULATE_ACC_TOKEN_MODE, ): + use_real_draft_tokens = simulate_acc_token_mode == "real-draft-token" + assert simulate_acc_len > 0.0 simulate_acc_len = _sample_simulated_acc_len( simulate_acc_len, simulate_acc_method, spec_steps + 1 @@ -334,7 +340,21 @@ def generate_simulated_accept_index( simulate_acc_len, device=accept_index.device ) num_correct_drafts.fill_(simulate_acc_len - 1) - predict.fill_(100) # some legit token id + + if not use_real_draft_tokens: + predict.fill_(100) # some legit token id + return sim_accept_index + + # Use the topk=1 draft chain for forced acceptance, then a target-derived bonus. + if simulate_acc_len > 1: + draft_node_indices = sim_accept_index[:, : simulate_acc_len - 1].long() + predict[draft_node_indices] = candidates[:, 1:simulate_acc_len].to( + dtype=predict.dtype + ) + bonus_node_indices = sim_accept_index[:, simulate_acc_len - 1].long() + predict[bonus_node_indices] = target_predict[:, simulate_acc_len - 1].to( + dtype=predict.dtype + ) return sim_accept_index