diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 143249b95..776ff0f72 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -424,12 +424,17 @@ def _handle_eagle_family(server_args: ServerArgs) -> None: "--enable-deterministic-inference; the sampling kernel draws " "coins from the global RNG and is not batch-invariant." ) + from sglang.srt.arg_groups.overrides import resolved_view - if resolved_view(server_args).enable_multi_layer_eagle: - raise NotImplementedError( - "--speculative-use-rejection-sampling is not supported with " - "multi-layer EAGLE (--enable-multi-layer-eagle)." + if ( + resolved_view(server_args).enable_multi_layer_eagle + and server_args.speculative_eagle_topk != 1 + ): + raise ValueError( + "--speculative-use-rejection-sampling with multi-layer EAGLE " + "(--enable-multi-layer-eagle) requires --speculative-eagle-topk 1; " + "rejection sampling is only implemented for the linear (topk=1) chain." ) logger.info( "Rejection sampling is enabled for speculative decoding " diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 68adadee6..078581ee2 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -150,8 +150,8 @@ class EagleDraftInput(SpecInput): # shape: (b, topk) topk_p: torch.Tensor = None topk_index: torch.Tensor = None - # shape: (b, vocab) - single-step draft proposal q from draft-extend; - # only set under rejection sampling. + # Draft proposal q from draft-extend, only set under rejection sampling: + # (b, vocab) single-layer; (b, num_steps, vocab) multi-layer chain. draft_probs: torch.Tensor = None # shape: (b, hidden_size) - one hidden per req, consumed by `draft` forward. # None when the spec algorithm's draft doesn't read hidden_states diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index d65b902da..bab29bf30 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -86,6 +86,7 @@ from sglang.srt.speculative.spec_utils import ( record_stream_each, record_stream_for_v2_verify, renorm_draft_probs, + sample_draft_proposal, select_top_k_tokens, spec_stage_span, ) @@ -691,12 +692,10 @@ class EagleDraftWorker(EagleDraftWorkerBase): maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") maybe_detect_inf(logits_output.next_token_logits, f"draft_forward step {i}") if self.server_args.speculative_use_rejection_sampling: - probs = renorm_draft_probs( + probs, topk_p, topk_index = sample_draft_proposal( logits_output.next_token_logits, - forward_batch.sampling_info, - self.server_args.speculative_use_rejection_sampling, + forward_batch.sampling_info.temperatures, ) - topk_p, topk_index = fast_sample(probs, num_samples=1) draft_probs_list.append(probs) elif self.topk == 1 and not _is_hip: topk_index = torch.argmax( @@ -994,13 +993,10 @@ class EagleDraftWorker(EagleDraftWorkerBase): # The draft-extend graph only anchors full logits; selected-row topk is # owned by the worker for both graph and eager paths. if self.server_args.speculative_use_rejection_sampling: - probs = renorm_draft_probs( + ret_draft_probs, ret_topk_p, ret_topk_index = sample_draft_proposal( draft_logits_output.next_token_logits, - batch.sampling_info, - self.server_args.speculative_use_rejection_sampling, + batch.sampling_info.temperatures, ) - ret_topk_p, ret_topk_index = fast_sample(probs, num_samples=1) - ret_draft_probs = probs elif self.topk == 1 and not _is_hip: # Gated to CUDA: see #26358 — ROCm's argmax tie-break corrupts # MTP draft selection on FP8 logits. diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 5361342c4..c051ac242 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -64,6 +64,7 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, record_stream_each, record_stream_for_v2_verify, + sample_draft_proposal, select_top_k_tokens, ) from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens @@ -124,6 +125,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): self.topk = server_args.speculative_eagle_topk self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.use_rejection_sampling = server_args.speculative_use_rejection_sampling assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, ( "multi-layer EAGLE requires speculative_num_draft_tokens == " "speculative_num_steps + 1, " @@ -305,6 +307,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): capture_hidden_mode=None, seq_lens_sum=None, seq_lens_cpu=None, + draft_probs=draft_input.draft_probs, ) def draft_forward(self, forward_batch: ForwardBatch): @@ -441,6 +444,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): topk_p_list = [] topk_index_list = [] + draft_probs_list = [] for step in range(self.speculative_num_steps): output: ModelRunnerOutput = self.draft_runner_list[step].forward( forward_batch @@ -453,8 +457,16 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): output.logits_output.next_token_logits, f"draft_extend_for_prefill step {step}", ) - probs = torch.softmax(output.logits_output.next_token_logits, dim=-1) - topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + if self.use_rejection_sampling and self.topk == 1: + # Sample X ~ q and stash q for the first verify's Leviathan step. + probs, topk_p, topk_index = sample_draft_proposal( + output.logits_output.next_token_logits, + forward_batch.sampling_info.temperatures, + ) + draft_probs_list.append(probs) + else: + probs = torch.softmax(output.logits_output.next_token_logits, dim=-1) + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) topk_index_list.append(topk_index) # Chain-style: use this step's output hidden_states as next step's input @@ -484,6 +496,12 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, ) + # q [bs, num_steps, vocab] for the first verify's Leviathan step (RS only). + next_draft_input.draft_probs = ( + torch.stack(draft_probs_list, dim=1) + if self.use_rejection_sampling and draft_probs_list + else None + ) return next_draft_input @@ -526,6 +544,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): ) ret_topk_p_list = [] ret_topk_index_list = [] + ret_draft_probs_list = [] next_token_ids_backup = batch_result.next_token_ids.clone() if can_cuda_graph: @@ -534,7 +553,16 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): # against it and the chain is advanced in place between steps. cgr.prepare(forward_batch) for step in range(self.speculative_num_steps): - _, ret_topk_p, ret_topk_index = cgr.replay(step) + _out, ret_topk_p, ret_topk_index = cgr.replay(step) + if self.use_rejection_sampling and self.topk == 1: + # Re-pick X ~ q worker-side so the chain rotation carries it + # to step N+1 (per-step graph does not sample in-graph). + sel = cgr.buffers.select_index[: cgr.raw_bs] + probs, ret_topk_p, ret_topk_index = sample_draft_proposal( + _out.next_token_logits[sel], + forward_batch.sampling_info.temperatures, + ) + ret_draft_probs_list.append(probs) ret_topk_p_list.append(ret_topk_p.clone()) ret_topk_index_list.append(ret_topk_index.clone()) # Advance the draft chain by rotating the shared input_ids window @@ -572,11 +600,17 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): draft_logits_output = self.draft_runner_list[step].forward( forward_batch ) - probs = torch.softmax( - draft_logits_output.logits_output.next_token_logits[select_index], - dim=-1, - ) - ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) + logits_sel = draft_logits_output.logits_output.next_token_logits[ + select_index + ] + if self.use_rejection_sampling and self.topk == 1: + probs, ret_topk_p, ret_topk_index = sample_draft_proposal( + logits_sel, forward_batch.sampling_info.temperatures + ) + ret_draft_probs_list.append(probs) + else: + probs = torch.softmax(logits_sel, dim=-1) + ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) # Chain-style: use this step's output hidden_states as next step's input if ( self.chain_mtp_hidden_states @@ -609,6 +643,12 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): torch.cat(ret_topk_index_list, dim=1).clone(), None, ) + # q [bs, num_steps, vocab] carries the per-chain-step draft distributions + # to the next verify's Leviathan step (accept iff coin*q < p). None + # otherwise (default target-only tree sampling). + next_draft_input.draft_probs = ( + torch.stack(ret_draft_probs_list, dim=1) if ret_draft_probs_list else None + ) class MultiLayerEagleWorkerV2(BaseSpecWorker): diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index c5fc46e20..240678864 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -95,6 +95,18 @@ def renorm_draft_probs( return torch.softmax(next_token_logits / sampling_info.temperatures, dim=-1) +def sample_draft_proposal(next_token_logits: torch.Tensor, temperatures: torch.Tensor): + """Leviathan draft proposal: q = softmax(logits / T), X ~ q. + + Returns (q, q(X), X). The verify's accept test coin*q(X) < p(X) is unbiased + only if q is exactly the distribution X was drawn from, so callers must hand + the returned q (not a recomputed one) to the verify. + """ + probs = torch.softmax(next_token_logits / temperatures, dim=-1) + topk_p, topk_index = fast_sample(probs, num_samples=1) + return probs, topk_p, topk_index + + # 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()