[spec decoding] support rejection sampling in multi layer eagle (#30303)

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:
+        if (
-            raise NotImplementedError(
+            resolved_view(server_args).enable_multi_layer_eagle
-                "--speculative-use-rejection-sampling is not supported with "
+            and server_args.speculative_eagle_topk != 1
-                "multi-layer EAGLE (--enable-multi-layer-eagle)."
+        ):
+            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 "

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;
+    # Draft proposal q from draft-extend, only set under rejection sampling:
-    # 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

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,
+                    forward_batch.sampling_info.temperatures,
-                    self.server_args.speculative_use_rejection_sampling,
                 )
-                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,
+                batch.sampling_info.temperatures,
-                self.server_args.speculative_use_rejection_sampling,
             )
-            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.

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,6 +457,14 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 output.logits_output.next_token_logits,
                 f"draft_extend_for_prefill step {step}",
             )
+            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)
@@ -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,10 +600,16 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 draft_logits_output = self.draft_runner_list[step].forward(
                     forward_batch
                 )
-                probs = torch.softmax(
+                logits_sel = draft_logits_output.logits_output.next_token_logits[
-                    draft_logits_output.logits_output.next_token_logits[select_index],
+                    select_index
-                    dim=-1,
+                ]
+                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 (
@@ -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):

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()