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

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

Browse Source
This commit is contained in:
Qiaolin Yu
2026-07-07 14:50:24 -07:00
committed by GitHub
parent 60f502a4fd
commit 801571e949
5 changed files with 76 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+9 -4
View File
@@ -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 "
python/sglang/srt/speculative/eagle_info.py
+2 -2
View File
@@ -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
python/sglang/srt/speculative/eagle_worker_v2.py
+5 -9
View File
@@ -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.
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+48 -8
View File
@@ -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):
python/sglang/srt/speculative/spec_utils.py
+12
View File
@@ -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()