From f42ec350b431d0305d34d6c70ea45fdfcd29dcad Mon Sep 17 00:00:00 2001 From: Yucheng Li Date: Sun, 21 Jun 2026 06:10:42 +0800 Subject: [PATCH] [mtp] add rejection sampling for speculative decoding (#26312) Co-authored-by: lyc508653 Co-authored-by: Qiaolin-Yu Co-authored-by: Huiqiang Jiang <30883354+iofu728@users.noreply.github.com> Co-authored-by: Yi Zhang <25844240+yizhang2077@users.noreply.github.com> Co-authored-by: Yizhong Cao <114661107+cao1zhg@users.noreply.github.com> --- .../sglang/srt/arg_groups/speculative_hook.py | 41 ++++ python/sglang/srt/managers/overlap_utils.py | 13 ++ .../srt/model_executor/forward_batch_info.py | 4 + python/sglang/srt/server_args.py | 7 + .../eagle_draft_cuda_graph_runner.py | 52 ++++- python/sglang/srt/speculative/eagle_info.py | 20 ++ python/sglang/srt/speculative/eagle_utils.py | 32 ++- .../sglang/srt/speculative/eagle_worker_v2.py | 100 +++++++-- .../sglang/srt/speculative/reject_sampling.py | 204 ++++++++++++++++++ python/sglang/srt/speculative/spec_utils.py | 22 ++ .../spec/eagle/test_eagle_reject_sampling.py | 91 ++++++++ 11 files changed, 566 insertions(+), 20 deletions(-) create mode 100644 python/sglang/srt/speculative/reject_sampling.py create mode 100644 test/registered/spec/eagle/test_eagle_reject_sampling.py diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 148c7c5f8..efef2e903 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -338,6 +338,47 @@ def _handle_eagle_family(server_args: ServerArgs) -> None: "trtllm_mha backend only supports topk = 1 for speculative decoding." ) + if server_args.speculative_use_rejection_sampling: + # Resolved alias by now: NEXTN -> EAGLE, Gemma4 draft -> FROZEN_KV_MTP. + # Only the EAGLE/EAGLE3 draft workers emit a target-vocab proposal that + # the rejection-sampling kernel consumes; everything else (STANDALONE, + # FROZEN_KV_MTP, NGRAM, DFLASH) is unsupported. + if server_args.speculative_algorithm not in ("EAGLE", "EAGLE3"): + raise NotImplementedError( + "--speculative-use-rejection-sampling is only supported for " + "EAGLE / EAGLE3 / NEXTN, not " + f"speculative_algorithm={server_args.speculative_algorithm}." + ) + if server_args.speculative_eagle_topk != 1: + raise ValueError( + "--speculative-use-rejection-sampling requires --speculative-eagle-topk=1." + ) + if ( + server_args.speculative_accept_threshold_single != 1.0 + or server_args.speculative_accept_threshold_acc != 1.0 + ): + raise ValueError( + "--speculative-use-rejection-sampling is incompatible with " + "--speculative-accept-threshold-single / " + "--speculative-accept-threshold-acc; rejection sampling ignores " + "the accept thresholds." + ) + if server_args.enable_deterministic_inference: + raise ValueError( + "--speculative-use-rejection-sampling is incompatible with " + "--enable-deterministic-inference; the sampling kernel draws " + "coins from the global RNG and is not batch-invariant." + ) + if server_args.enable_multi_layer_eagle: + raise NotImplementedError( + "--speculative-use-rejection-sampling is not supported with " + "multi-layer EAGLE (--enable-multi-layer-eagle)." + ) + logger.info( + "Rejection sampling is enabled for speculative decoding " + "(speculative_use_rejection_sampling=True)." + ) + if ( server_args.speculative_eagle_topk == 1 and server_args.speculative_num_draft_tokens diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 3af719f9d..fc6ef4a77 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -192,6 +192,15 @@ class FutureMap: device=self.device, ) + self.draft_probs_buf = None + if getattr(draft_input, "draft_probs", None) is not None: + draft_probs0 = draft_input.draft_probs[0] + self.draft_probs_buf = torch.empty( + (self.req_pool_size, *draft_probs0.shape), + dtype=draft_probs0.dtype, + device=self.device, + ) + def _resolve_spec_extras(self, batch: ScheduleBatch) -> None: if self.spec_algo.is_ngram(): # FIXME: remove once precomputed draft is supported. @@ -232,6 +241,8 @@ class FutureMap: draft_input.bonus_tokens = bonus_tokens if hidden_states is not None: draft_input.hidden_states = hidden_states + if self.draft_probs_buf is not None and draft_input.draft_probs is not None: + draft_input.draft_probs = self.draft_probs_buf[indices] elif self.need_bonus_tokens: draft_input.bonus_tokens = self.output_tokens_buf[indices] if self.need_hidden_states and not self.need_topk: @@ -352,3 +363,5 @@ class FutureMap: self.hidden_states_buf[indices] = draft_input.hidden_states.to( self.hidden_states_buf.dtype ) + if self.draft_probs_buf is not None and draft_input.draft_probs is not None: + self.draft_probs_buf[indices] = draft_input.draft_probs diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 686f1265a..0feeca831 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1336,6 +1336,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): spec_info.topk_index = self._pad_tensor_to_size( spec_info.topk_index, bs ) + if getattr(spec_info, "draft_probs", None) is not None: + spec_info.draft_probs = self._pad_tensor_to_size( + spec_info.draft_probs, bs + ) if getattr(spec_info, "num_correct_drafts", None) is not None: spec_info.num_correct_drafts = self._pad_tensor_to_size( spec_info.num_correct_drafts, bs diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 21420e265..606df4f6c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -620,6 +620,7 @@ class ServerArgs: speculative_dflash_block_size: Optional[int] = None speculative_accept_threshold_single: float = 1.0 speculative_accept_threshold_acc: float = 1.0 + speculative_use_rejection_sampling: bool = False speculative_token_map: Optional[str] = None speculative_attention_mode: str = "prefill" speculative_draft_attention_backend: Optional[str] = None @@ -6127,6 +6128,12 @@ class ServerArgs: help="The accept probability of a draft token is raised from its target probability p to min(1, p / threshold_acc).", default=ServerArgs.speculative_accept_threshold_acc, ) + parser.add_argument( + "--speculative-use-rejection-sampling", + action="store_true", + help="Use rejection sampling for speculative decoding (requires topk=1).", + default=ServerArgs.speculative_use_rejection_sampling, + ) parser.add_argument( "--speculative-token-map", type=str, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index fbbafdce1..d195fac71 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -32,6 +32,7 @@ from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backen from sglang.srt.model_executor.runner_backend_utils import ( CUDA_GRAPH_CAPTURE_FAILED_MSG, ) +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.utils import ( require_attn_tp_gather, @@ -59,6 +60,7 @@ class EagleDraftInputBuffers(ForwardInputBuffers): extend_seq_lens: torch.Tensor topk_p: torch.Tensor topk_index: torch.Tensor + draft_probs: Optional[torch.Tensor] hidden_states: Optional[torch.Tensor] global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] @@ -177,6 +179,14 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32) topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32) topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64) + draft_probs = ( + torch.zeros( + (self.max_bs, self.model_runner.model_config.vocab_size), + dtype=torch.float32, + ) + if self.model_runner.server_args.speculative_use_rejection_sampling + else None + ) _hidden_size = EagleDraftInput.hidden_size_for(self.eagle_worker) hidden_states = ( torch.zeros( @@ -187,6 +197,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): else None ) + self.temperatures = torch.ones((self.max_bs, 1), dtype=torch.float) + if self.require_gathered_buffer: if self.require_mlp_tp_gather: global_num_tokens_gpu = torch.zeros( @@ -222,6 +234,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): extend_seq_lens=extend_seq_lens, topk_p=topk_p, topk_index=topk_index, + draft_probs=draft_probs, hidden_states=hidden_states, global_num_tokens_gpu=global_num_tokens_gpu, global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, @@ -313,6 +326,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): ) topk_p = buffers.topk_p[:num_seqs] topk_index = buffers.topk_index[:num_seqs] + draft_probs = ( + buffers.draft_probs[:num_seqs] if buffers.draft_probs is not None else None + ) if self.require_mlp_tp_gather: global_num_tokens_cpu = [num_tokens] * self.dp_size @@ -345,10 +361,23 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): spec_info = EagleDraftInput( topk_p=topk_p, topk_index=topk_index, + draft_probs=draft_probs, hidden_states=hidden_states, capture_hidden_mode=capture_mode, ) + sampling_info = SamplingBatchInfo( + temperatures=self.temperatures[:num_seqs], + top_ps=torch.ones((num_seqs,), dtype=torch.float), + top_ks=torch.full((num_seqs,), -1, dtype=torch.int32), + min_ps=torch.zeros((num_seqs,), dtype=torch.float), + is_all_greedy=False, + need_top_p_sampling=False, + need_top_k_sampling=False, + need_min_p_sampling=False, + vocab_size=self.model_runner.model_config.vocab_size, + ) + forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, batch_size=num_seqs, @@ -369,6 +398,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): global_dp_buffer_len=global_dp_buffer_len, spec_algorithm=self.model_runner.spec_algorithm, spec_info=spec_info, + sampling_info=sampling_info, rids_int=rids_int, bootstrap_room_ids_int=bootstrap_room_ids_int, capture_hidden_mode=( @@ -417,8 +447,10 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): ) def _postprocess_output_to_raw_bs(self, out, raw_bs): - parent_list, top_scores_index, draft_tokens = (t[:raw_bs] for t in out) - return parent_list, top_scores_index, draft_tokens + parent_list, top_scores_index, draft_tokens, draft_probs = ( + t[:raw_bs] if t is not None else None for t in out + ) + return parent_list, top_scores_index, draft_tokens, draft_probs # ----------------------------------------------------------------- # Replay @@ -454,6 +486,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): buffers.bootstrap_room_ids_int.fill_(-1) buffers.topk_p.zero_() buffers.topk_index.zero_() + if buffers.draft_probs is not None: + buffers.draft_probs.zero_() if buffers.hidden_states is not None: buffers.hidden_states.zero_() buffers.req_pool_indices.zero_() @@ -505,11 +539,25 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): # hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync # DMA engine; foreach would force the ~3x slower compute-kernel copy. + if ( + buffers.draft_probs is not None + and forward_batch.spec_info.draft_probs is not None + ): + buffers.draft_probs[:raw_bs].copy_(forward_batch.spec_info.draft_probs) if ( buffers.hidden_states is not None and forward_batch.spec_info.hidden_states is not None ): buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states) + # Only rejection sampling reads temperatures (renorm_draft_probs); skip + # the copy otherwise to keep the non-RS path free of extra work. + if ( + self.model_runner.server_args.speculative_use_rejection_sampling + and forward_batch.sampling_info is not None + ): + self.temperatures[:raw_bs].copy_( + forward_batch.sampling_info.temperatures[:raw_bs] + ) # TODO(ch-wan): support num_token_non_padded if self.require_gathered_buffer: diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 22b783dff..4c7eb2b6e 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -8,6 +8,7 @@ from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject from sglang.srt.environ import envs from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.eagle_info_v2 import EagleDraftInputV2Mixin from sglang.srt.speculative.spec_info import SpecInput, SpecInputType @@ -42,6 +43,9 @@ class EagleVerifyInput(SpecInput): seq_lens_sum: int seq_lens_cpu: torch.Tensor grammar: BaseGrammarObject = None + # Stacked per-step draft proposal distribution q, shape (bs, num_steps, + # vocab); only set under rejection sampling. Consumed by the verify kernel. + draft_probs: torch.Tensor = None # Shape info for padding num_tokens_per_req: int = -1 # -1 auto-fills from draft_token_num. @@ -159,6 +163,9 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): # 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_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 # (e.g., STANDALONE — vanilla LLM draft). @@ -209,6 +216,7 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): dtype: Optional[torch.dtype], topk: int, capture_hidden_mode: CaptureHiddenMode, + vocab_size: int = 0, ): return cls( bonus_tokens=torch.empty((0,), device=device, dtype=torch.int32), @@ -219,6 +227,11 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): ), topk_p=torch.empty((0, topk), device=device, dtype=torch.float32), topk_index=torch.empty((0, topk), device=device, dtype=torch.int64), + draft_probs=( + torch.empty((0, vocab_size), device=device, dtype=torch.float32) + if get_global_server_args().speculative_use_rejection_sampling + else None + ), capture_hidden_mode=capture_hidden_mode, ) @@ -240,6 +253,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): self.topk_p = self.topk_p[: len(new_indices)] self.topk_index = self.topk_index[: len(new_indices)] + if self.draft_probs is not None: + self.draft_probs = self.draft_probs[: len(new_indices)] if self.hidden_states is not None: self.hidden_states = self.hidden_states[: len(new_indices)] self.bonus_tokens = self.bonus_tokens[: len(new_indices)] @@ -247,6 +262,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): # in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index` self.topk_p = self.topk_p[new_indices] self.topk_index = self.topk_index[new_indices] + if self.draft_probs is not None: + self.draft_probs = self.draft_probs[new_indices] if self.hidden_states is not None: self.hidden_states = self.hidden_states[new_indices] self.bonus_tokens = self.bonus_tokens[new_indices] @@ -267,6 +284,7 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): self.bonus_tokens = spec_info.bonus_tokens self.topk_p = spec_info.topk_p self.topk_index = spec_info.topk_index + self.draft_probs = spec_info.draft_probs return if len(spec_info.topk_index) == 0: return @@ -279,6 +297,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): ) self.topk_p = torch.cat([self.topk_p, spec_info.topk_p]) self.topk_index = torch.cat([self.topk_index, spec_info.topk_index]) + if self.draft_probs is not None and spec_info.draft_probs is not None: + self.draft_probs = torch.cat([self.draft_probs, spec_info.draft_probs]) @dataclass diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index c5a7c6963..21de68eb8 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -455,6 +455,14 @@ def eagle_sample( tree_speculative_sampling_target_only, ) + from sglang.srt.speculative.reject_sampling import ( + chain_speculative_sampling_triton, + ) + + use_rejection_sampling = ( + get_global_server_args().speculative_use_rejection_sampling + ) + # Apply temperature and get target probs expanded_temperature = torch.repeat_interleave( sampling_info.temperatures, verify_input.draft_token_num, dim=0 @@ -479,14 +487,34 @@ def eagle_sample( ) maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1) - draft_probs = torch.zeros_like(target_probs) + draft_probs = ( + verify_input.draft_probs + if use_rejection_sampling + else torch.zeros_like(target_probs) + ) + # Defense-in-depth behind the spec_hook startup allowlist: validate the + # actual kernel inputs (catches draft_probs plumbing regressions or a + # startup guard bypassed by a worker subclass) before the Triton kernel. + if use_rejection_sampling and ( + draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1] + ): + raise ValueError( + "Rejection sampling requires a target-vocab draft proposal " + "distribution; the current speculative algorithm/draft worker " + "does not produce one (draft_probs missing or vocab-mismatched)." + ) # coins for rejection sampling coins = torch.rand_like(candidates, dtype=torch.float32, device=device) # coins for final sampling coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device) - tree_speculative_sampling_target_only( + sampling_fn = ( + chain_speculative_sampling_triton + if use_rejection_sampling + else tree_speculative_sampling_target_only + ) + sampling_fn( predicts=predict, # mutable accept_index=accept_index, # mutable accept_token_num=num_correct_drafts, # mutable diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 232c1c858..dfc059d8a 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -74,11 +74,13 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( commit_mamba_states_after_verify, draft_tp_context, + fast_sample, generate_token_bitmask, load_token_map, move_accept_tokens_to_target_kvcache, record_stream_each, record_stream_for_v2_verify, + renorm_draft_probs, select_top_k_tokens, spec_stage_span, ) @@ -146,6 +148,8 @@ class EagleDraftWorker(EagleDraftWorkerBase): # Args for easy access self.device = server_args.device self.topk = server_args.speculative_eagle_topk + if self.server_args.speculative_use_rejection_sampling: + assert self.topk == 1, "Chain speculative sampling supports only topk=1" self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens self.speculative_algorithm = SpeculativeAlgorithm.from_string( @@ -222,6 +226,22 @@ class EagleDraftWorker(EagleDraftWorkerBase): self.init_token_map() self.init_lm_head() + if self.server_args.speculative_use_rejection_sampling: + target_vocab_size = self.target_worker.model_config.vocab_size + draft_vocab_size = ( + self.hot_token_id.shape[0] + if self.hot_token_id is not None + else target_vocab_size + ) + # FIXME: support reduced (hot) draft vocab by scattering draft probs + # into the target vocab via the d2t map before the sampling kernel. + if draft_vocab_size != target_vocab_size: + raise ValueError( + "--speculative-use-rejection-sampling requires the draft and " + f"target to share one vocab, but the draft vocab " + f"({draft_vocab_size}) != target vocab ({target_vocab_size})." + ) + def init_backends(self): with self.draft_tp_context( self.draft_runner.tp_group @@ -443,7 +463,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): with canary_outside_ctx: # Run draft if can_cuda_graph: - parent_list, top_scores_index, draft_tokens = ( + parent_list, top_scores_index, draft_tokens, draft_probs = ( self.cuda_graph_runner.execute(forward_batch) ) else: @@ -455,8 +475,8 @@ class EagleDraftWorker(EagleDraftWorkerBase): # `draft_forward` only does sample in this case. self.draft_attn_backend.init_forward_metadata(forward_batch) forward_batch.mark_forward_metadata_ready() - parent_list, top_scores_index, draft_tokens = self.draft_forward( - forward_batch + parent_list, top_scores_index, draft_tokens, draft_probs = ( + self.draft_forward(forward_batch) ) if batch.forward_mode.is_idle(): @@ -521,6 +541,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): capture_hidden_mode=None, seq_lens_sum=None, seq_lens_cpu=None, + draft_probs=draft_probs, ) def draft_forward(self, forward_batch: ForwardBatch): @@ -549,6 +570,8 @@ class EagleDraftWorker(EagleDraftWorkerBase): score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] parents_list: List[torch.Tensor] = [] + if self.server_args.speculative_use_rejection_sampling: + draft_probs_list: List[torch.Tensor] = [spec_info.draft_probs] # Forward multiple steps scores = None @@ -598,18 +621,25 @@ class EagleDraftWorker(EagleDraftWorkerBase): logits_output = self.draft_runner.forward(forward_batch).logits_output 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.topk == 1 and not _is_hip: - # topk=1 → degenerate single-path tree; `topk_p` is unused - # downstream, so skip softmax and just argmax over logits. - # Gated to CUDA: on ROCm the argmax tie-break diverges from - # the softmax+max path on FP8 logits and corrupts MTP draft - # selection (DSV3.2 MTP GSM8K, see #26358). + if self.server_args.speculative_use_rejection_sampling: + probs = renorm_draft_probs( + logits_output.next_token_logits, + forward_batch.sampling_info, + 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( logits_output.next_token_logits, dim=-1, keepdim=True ) topk_p = torch.ones_like(topk_index, dtype=torch.float32) else: - probs = torch.softmax(logits_output.next_token_logits, dim=-1) + probs = renorm_draft_probs( + logits_output.next_token_logits, + forward_batch.sampling_info, + self.server_args.speculative_use_rejection_sampling, + ) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) maybe_detect_oob( topk_index, @@ -640,12 +670,24 @@ class EagleDraftWorker(EagleDraftWorkerBase): draft_tokens = torch.cat(token_list, dim=1) top_scores_index = self._topk1_score_indices_prealloc[:bs] parent_list = self._topk1_parents_prealloc[:bs] - return parent_list, top_scores_index, draft_tokens + draft_probs = ( + torch.stack(draft_probs_list, dim=1) + if self.server_args.speculative_use_rejection_sampling + else None + ) + return parent_list, top_scores_index, draft_tokens, draft_probs - return organize_draft_results( + parent_list, top_scores_index, draft_tokens = organize_draft_results( score_list, token_list, parents_list, self.speculative_num_draft_tokens ) + draft_probs = ( + torch.stack(draft_probs_list, dim=1) + if self.server_args.speculative_use_rejection_sampling + else None + ) + return parent_list, top_scores_index, draft_tokens, draft_probs + def draft_extend(self): pass @@ -716,11 +758,20 @@ class EagleDraftWorker(EagleDraftWorkerBase): maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill") # Assemble the next-iter draft spec_info from the extend output. - probs = torch.softmax(logits_output.next_token_logits, dim=-1) - topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + use_rejection_sampling = self.server_args.speculative_use_rejection_sampling + probs = renorm_draft_probs( + logits_output.next_token_logits, + batch.sampling_info, + use_rejection_sampling, + ) + if use_rejection_sampling: + topk_p, topk_index = fast_sample(probs, num_samples=1) + else: + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) return EagleDraftInput( topk_p=topk_p, topk_index=topk_index, + draft_probs=probs if use_rejection_sampling else None, hidden_states=logits_output.hidden_states, bonus_tokens=next_token_ids, num_tokens_per_req=1, @@ -810,16 +861,30 @@ 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.topk == 1 and not _is_hip: + if self.server_args.speculative_use_rejection_sampling: + probs = renorm_draft_probs( + draft_logits_output.next_token_logits, + batch.sampling_info, + 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. ret_topk_index = torch.argmax( draft_logits_output.next_token_logits, dim=-1, keepdim=True ) ret_topk_p = torch.ones_like(ret_topk_index, dtype=torch.float32) + ret_draft_probs = None else: - probs = torch.softmax(draft_logits_output.next_token_logits, dim=-1) + probs = renorm_draft_probs( + draft_logits_output.next_token_logits, + batch.sampling_info, + self.server_args.speculative_use_rejection_sampling, + ) ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) + ret_draft_probs = None ret_hidden_states = draft_logits_output.hidden_states # Construct the return values @@ -833,6 +898,8 @@ class EagleDraftWorker(EagleDraftWorkerBase): ret_topk_index, ret_hidden_states, ) + if self.server_args.speculative_use_rejection_sampling: + next_draft_input.draft_probs = ret_draft_probs class EAGLEWorkerV2(BaseSpecWorker): @@ -1010,6 +1077,7 @@ class EAGLEWorkerV2(BaseSpecWorker): dtype=EagleDraftInput.dtype_for(self.draft_worker), topk=self.topk, capture_hidden_mode=capture_mode, + vocab_size=self.target_worker.model_config.vocab_size, ) if self.speculative_num_steps == 0: # Drafting disabled (high batch size). _draft_extend below still diff --git a/python/sglang/srt/speculative/reject_sampling.py b/python/sglang/srt/speculative/reject_sampling.py new file mode 100644 index 000000000..7eb4d65ca --- /dev/null +++ b/python/sglang/srt/speculative/reject_sampling.py @@ -0,0 +1,204 @@ +import triton +import triton.language as tl + + +@triton.jit +def speculative_sampling_classic_kernel( + # Pointers + Predicts, + AcceptIndex, + AcceptTokenNum, + Candidates, + RetriveIndex, + UniformSamples, + UniformSamplesFinal, + TargetProbs, + DraftProbs, + # Strides + stride_cand_b, + stride_cand_s, + stride_idx_b, + stride_idx_s, + stride_uni_b, + stride_uni_s, + stride_tp_b, + stride_tp_s, + stride_tp_v, + stride_dp_b, + stride_dp_s, + stride_dp_v, + # Constants + NUM_SLOTS: tl.constexpr, + VOCAB_SIZE: tl.constexpr, + BLOCK_V: tl.constexpr, +): + pid = tl.program_id(0) + cur_prob_row = 0 + + cand_ptr_base = Candidates + pid * stride_cand_b + idx_ptr_base = RetriveIndex + pid * stride_idx_b + uni_ptr_base = UniformSamples + pid * stride_uni_b + + root_global_idx = tl.load(idx_ptr_base + 0 * stride_idx_s) + tl.store(AcceptIndex + pid * stride_idx_b + 0 * stride_idx_s, root_global_idx) + last_accepted_global_idx = root_global_idx + + num_accept = 0 + + # Verification Loop + step = 1 + continue_verifying = 1 + + while (step < NUM_SLOTS) and (continue_verifying == 1): + draft_token = tl.load(cand_ptr_base + step * stride_cand_s) + + offset_prob = ( + (pid * stride_tp_b) + + (cur_prob_row * stride_tp_s) + + (draft_token * stride_tp_v) + ) + offset_draft = ( + (pid * stride_dp_b) + + (cur_prob_row * stride_dp_s) + + (draft_token * stride_dp_v) + ) + + p = tl.load(TargetProbs + offset_prob) + q = tl.load(DraftProbs + offset_draft) + + coin = tl.load(uni_ptr_base + (step - 1) * stride_uni_s) + + if coin * q < p: + num_accept += 1 + cur_prob_row = step + tl.store(Predicts + last_accepted_global_idx, draft_token) + + curr_global_idx = tl.load(idx_ptr_base + step * stride_idx_s) + tl.store( + AcceptIndex + pid * stride_idx_b + num_accept * stride_idx_s, + curr_global_idx, + ) + last_accepted_global_idx = curr_global_idx + + step += 1 + else: + continue_verifying = 0 + + tl.store(AcceptTokenNum + pid, num_accept) + + # Final Sampling + all_drafts_accepted = continue_verifying + coin_final = tl.load(UniformSamplesFinal + pid) + norm_sum = 0.0 + + tp_base_ptr = TargetProbs + (pid * stride_tp_b) + (cur_prob_row * stride_tp_s) + # DraftProbs has only num_steps rows (TargetProbs has num_steps + 1). When + # all drafts are accepted cur_prob_row == num_steps is out of bounds for + # DraftProbs, but the all-accepted branch samples pure target p and never + # dereferences this pointer; on rejection cur_prob_row <= num_steps - 1. + dp_base_ptr_safe = DraftProbs + (pid * stride_dp_b) + (cur_prob_row * stride_dp_s) + + # Pass 1: Sum + for v_start in range(0, VOCAB_SIZE, BLOCK_V): + v_offsets = v_start + tl.arange(0, BLOCK_V) + mask = v_offsets < VOCAB_SIZE + + p_ptr = tp_base_ptr + v_offsets * stride_tp_v + p_val = tl.load(p_ptr, mask=mask, other=0.0) + + if all_drafts_accepted: + val = p_val + else: + q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v + q_val = tl.load(q_ptr, mask=mask, other=0.0) + diff = p_val - q_val + val = tl.where(diff > 0.0, diff, 0.0) + + norm_sum += tl.sum(val) + + # Pass 2: CDF. Degenerate residual (norm_sum == 0, i.e. p == q everywhere on + # rejection) leaves the cumsum at 0 <= target_u, so final_token falls back to + # VOCAB_SIZE - 1; acceptable since this case is numerically near-impossible. + target_u = coin_final * norm_sum + cum_sum = 0.0 + final_token = VOCAB_SIZE - 1 + found = 0 + + for v_start in range(0, VOCAB_SIZE, BLOCK_V): + if found == 0: + v_offsets = v_start + tl.arange(0, BLOCK_V) + mask = v_offsets < VOCAB_SIZE + + p_ptr = tp_base_ptr + v_offsets * stride_tp_v + p_val = tl.load(p_ptr, mask=mask, other=0.0) + + if all_drafts_accepted: + val = p_val + else: + q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v + q_val = tl.load(q_ptr, mask=mask, other=0.0) + diff = p_val - q_val + val = tl.where(diff > 0.0, diff, 0.0) + + block_cumsum = tl.cumsum(val, axis=0) + total_cumsum = cum_sum + block_cumsum + + candidates_mask = total_cumsum > target_u + has_match = tl.max(candidates_mask, axis=0) + + if has_match: + match_idx = tl.argmax(candidates_mask.to(tl.int32), axis=0) + final_token = v_start + match_idx + found = 1 + + cum_sum += tl.sum(val) + + tl.store(Predicts + last_accepted_global_idx, final_token) + + +def chain_speculative_sampling_triton( + predicts, + accept_index, + accept_token_num, + candidates, + retrive_index, + retrive_next_token, + retrive_next_sibling, # not used in chain verification + uniform_samples, + uniform_samples_for_final_sampling, + target_probs, + draft_probs, + threshold_single, + threshold_acc, + deterministic, # not used +): + batch_size, num_slots = candidates.shape + vocab_size = target_probs.shape[-1] + + grid = (batch_size,) + speculative_sampling_classic_kernel[grid]( + predicts, + accept_index, + accept_token_num, + candidates, + retrive_index, + uniform_samples, + uniform_samples_for_final_sampling, + target_probs, + draft_probs, + candidates.stride(0), + candidates.stride(1), + retrive_index.stride(0), + retrive_index.stride(1), + uniform_samples.stride(0), + uniform_samples.stride(1), + target_probs.stride(0), + target_probs.stride(1), + target_probs.stride(2), + draft_probs.stride(0), + draft_probs.stride(1), + draft_probs.stride(2), + NUM_SLOTS=num_slots, + VOCAB_SIZE=vocab_size, + BLOCK_V=4096, + ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index eb7c410f1..312f366f9 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -72,6 +72,28 @@ else: logger = logging.getLogger(__name__) +def fast_sample(probs: torch.Tensor, num_samples: int = 1): + sample_index = torch.multinomial(probs, num_samples=num_samples) + sample_p = probs.gather(1, sample_index) + return sample_p, sample_index + + +def renorm_draft_probs( + next_token_logits: torch.Tensor, + sampling_info, + use_rejection_sampling: bool, +) -> torch.Tensor: + """Draft-side next-token distribution. + + Plain softmax, except under rejection sampling where logits are + temperature-scaled so the draft proposal q tracks the target sampling + temperature (higher acceptance; correctness holds for any q). + """ + if not use_rejection_sampling or not next_token_logits.size(0): + return torch.softmax(next_token_logits, dim=-1) + return torch.softmax(next_token_logits / sampling_info.temperatures, dim=-1) + + # 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() diff --git a/test/registered/spec/eagle/test_eagle_reject_sampling.py b/test/registered/spec/eagle/test_eagle_reject_sampling.py new file mode 100644 index 000000000..99fb23100 --- /dev/null +++ b/test/registered/spec/eagle/test_eagle_reject_sampling.py @@ -0,0 +1,91 @@ +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_cuda_ci(est_time=300, stage="base-b", runner_config="2-gpu-large") + +QWEN35_MODEL = "Qwen/Qwen3.5-9B" +SERVER_LAUNCH_TIMEOUT = 600 + + +class TestQwen35EagleRS(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN35_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--tp", + "2", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--speculative-use-rejection-sampling", + "--mem-fraction-static", + "0.8", + "--disable-radix-cache", + ] + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k(self): + requests.get(self.base_url + "/flush_cache") + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"{metrics=}") + + server_info = requests.get(self.base_url + "/server_info").json() + avg_spec_accept_length = server_info["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + + if is_in_ci(): + write_github_step_summary( + f"### test_gsm8k (qwen3.5-9b nextn rs)\n" + f'{metrics["score"]=:.3f}\n' + f"{avg_spec_accept_length=:.2f}\n" + ) + + self.assertGreater(metrics["score"], 0.8) + self.assertGreater(avg_spec_accept_length, 2.5) + + +if __name__ == "__main__": + unittest.main()