[Speculative] Seed rejection-sampling draft proposals for deterministic inference

benchmark/kernels/sampling/bench_deterministic_fast_sample.py

@@ -0,0 +1,89 @@
+"""Benchmark seeded versus unseeded EAGLE draft Gumbel sampling.
+
+Example:
+    python benchmark/kernels/sampling/bench_deterministic_fast_sample.py \
+        --batch-sizes 1 8 32 128 --vocab-sizes 32000 128256 163840
+"""
+
+import argparse
+
+import torch
+import triton
+
+from sglang.srt.speculative.spec_utils import fast_sample
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8, 32, 128])
+    parser.add_argument(
+        "--vocab-sizes", type=int, nargs="+", default=[32000, 128256, 163840]
+    )
+    parser.add_argument("--warmup-ms", type=int, default=100)
+    parser.add_argument("--rep-ms", type=int, default=500)
+    parser.add_argument(
+        "--modes", nargs="+", choices=("eager", "graph"), default=["graph"]
+    )
+    return parser.parse_args()
+
+
+def capture(fn):
+    fn()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    return graph.replay
+
+
+def main():
+    args = parse_args()
+    device = torch.device("cuda")
+    print(f"device={torch.cuda.get_device_name(device)}")
+    print("mode\tbatch\tvocab\tunseeded_us\tseeded_us\tslowdown")
+
+    for vocab_size in args.vocab_sizes:
+        for batch_size in args.batch_sizes:
+            logits = torch.randn(
+                (batch_size, vocab_size), device=device, dtype=torch.float32
+            )
+            probs = torch.softmax(logits, dim=-1)
+            seeds = torch.arange(batch_size, device=device, dtype=torch.int64) + 12345
+            positions = (
+                torch.arange(batch_size, device=device, dtype=torch.int64) + 4096
+            )
+
+            def unseeded():
+                return fast_sample(probs)
+
+            def seeded():
+                return fast_sample(
+                    probs,
+                    sampling_seed=seeds,
+                    positions=positions,
+                    draft_step=1,
+                )
+
+            # Compile/JIT and populate allocator caches before timing.
+            unseeded()
+            seeded()
+            torch.cuda.synchronize()
+
+            for mode in args.modes:
+                unseeded_fn = capture(unseeded) if mode == "graph" else unseeded
+                seeded_fn = capture(seeded) if mode == "graph" else seeded
+                unseeded_ms = triton.testing.do_bench(
+                    unseeded_fn, warmup=args.warmup_ms, rep=args.rep_ms
+                )
+                seeded_ms = triton.testing.do_bench(
+                    seeded_fn, warmup=args.warmup_ms, rep=args.rep_ms
+                )
+                print(
+                    f"{mode}\t{batch_size}\t{vocab_size}\t"
+                    f"{unseeded_ms * 1000:.1f}\t{seeded_ms * 1000:.1f}\t"
+                    f"{seeded_ms / unseeded_ms:.2f}x"
+                )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/srt/arg_groups/speculative_hook.py

@@ -1126,6 +1126,12 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
             raise ValueError(
                 "--speculative-use-rejection-sampling requires --speculative-eagle-topk=1."
             )
+        if cfg.enable_deterministic_inference and not cfg.device.startswith("cuda"):
+            raise ValueError(
+                "Deterministic rejection sampling is currently supported only on "
+                "CUDA-compatible devices because draft proposal sampling uses a "
+                "Triton hash kernel."
+            )
         if (
             cfg.speculative_accept_threshold_single != 1.0
             or cfg.speculative_accept_threshold_acc != 1.0
@@ -1136,13 +1142,6 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
                 "--speculative-accept-threshold-acc; rejection sampling ignores "
                 "the accept thresholds."
             )
-        if cfg.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 (
             resolved_view(server_args).enable_multi_layer_eagle
             and cfg.speculative_eagle_topk != 1

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -77,6 +77,7 @@ class EagleDraftInputBuffers(ForwardInputBuffers):
     topk_p: torch.Tensor
     topk_index: torch.Tensor
     draft_probs: Optional[torch.Tensor]
+    sampling_seed: Optional[torch.Tensor]
     hidden_states: Optional[torch.Tensor]
     global_num_tokens_gpu: Optional[torch.Tensor]
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
@@ -203,6 +204,12 @@ def __init__(
                 if get_spec().speculative_use_rejection_sampling
                 else None
             )
+            sampling_seed = (
+                torch.zeros((self.max_bs,), dtype=torch.int64)
+                if get_exec().deterministic.enable_deterministic_inference
+                and get_spec().speculative_use_rejection_sampling
+                else None
+            )
             _hidden_size, _hidden_dtype = get_draft_recurrent_hidden_state_spec(
                 model_runner
             )
@@ -271,6 +278,7 @@ def __init__(
             topk_p=topk_p,
             topk_index=topk_index,
             draft_probs=draft_probs,
+            sampling_seed=sampling_seed,
             hidden_states=hidden_states,
             global_num_tokens_gpu=global_num_tokens_gpu,
             global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
@@ -455,6 +463,11 @@ def capture_one_shape(
             need_top_k_sampling=False,
             need_min_p_sampling=False,
             vocab_size=self.model_runner.model_config.vocab_size,
+            sampling_seed=(
+                self.buffers.sampling_seed[:num_seqs]
+                if self.buffers.sampling_seed is not None
+                else None
+            ),
         )
 
         forward_batch = ForwardBatch(
@@ -667,6 +680,10 @@ def execute(self, forward_batch: ForwardBatch):
             self.temperatures[:raw_bs].copy_(
                 forward_batch.sampling_info.temperatures[:raw_bs]
             )
+            if self.buffers.sampling_seed is not None:
+                self.buffers.sampling_seed[:raw_bs].copy_(
+                    forward_batch.sampling_info.sampling_seed[:raw_bs]
+                )
             self.top_ks[:raw_bs].copy_(forward_batch.sampling_info.top_ks[:raw_bs])
 
         # TODO(ch-wan): support num_token_non_padded

python/sglang/srt/speculative/eagle_utils.py

@@ -642,11 +642,9 @@ def _seeded_verify_coins(
     per-draft rejection coins; column draft_token_num drives the final
     fallback-sampling coin.
 
-    Scope: this seeds only the verify-side RNG. With rejection sampling the
-    draft workers still pick candidates via unseeded multinomial
-    (fast_sample in eagle_worker_v2), so that mode stays non-deterministic
-    until the draft RNG is seeded in a follow-up; top-k/greedy draft
-    selection is already deterministic.
+    Draft proposal draws use a separate domain-separated seeded stream in
+    ``fast_sample``. Keeping the streams separate preserves the independence
+    required by rejection sampling.
     """
     from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -789,6 +789,9 @@ def draft_forward(self, forward_batch: ForwardBatch):
                         logits_output.next_token_logits,
                         forward_batch.sampling_info.temperatures,
                         forward_batch.sampling_info.top_ks,
+                        sampling_seed=forward_batch.sampling_info.sampling_seed,
+                        positions=forward_batch.positions + 1,
+                        draft_step=i + 1,
                     )
                     draft_probs_list.append(probs)
                     forward_batch.positions.add_(1)
@@ -998,7 +1001,12 @@ def _draft_extend_for_prefill(
             use_rejection_sampling,
         )
         if use_rejection_sampling:
-            topk_p, topk_index = fast_sample(probs, num_samples=1)
+            topk_p, topk_index = fast_sample(
+                probs,
+                num_samples=1,
+                sampling_seed=batch.sampling_info.sampling_seed,
+                positions=batch.seq_lens,
+            )
         else:
             topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
         return EagleDraftInput(
@@ -1161,6 +1169,8 @@ def _draft_extend_for_decode(
                 draft_logits_output.next_token_logits,
                 batch.sampling_info.temperatures,
                 batch.sampling_info.top_ks,
+                sampling_seed=batch.sampling_info.sampling_seed,
+                positions=batch_result.new_seq_lens,
             )
         elif self.topk == 1 and _is_hip:
             ret_topk_p, ret_topk_index = draft_topk1_argmax_only(

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -135,6 +135,7 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers):
     # of draft_probs selects the in-graph proposal branch in _run_step_body).
     temperatures: Optional[torch.Tensor]
     draft_probs: Optional[torch.Tensor]
+    sampling_seed: Optional[torch.Tensor]
 
 
 class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
@@ -392,7 +393,17 @@ def _sample_draft_proposal(self, ret, bs: int):
             buffers.temperatures[:bs],
             buffers.draft_probs[:bs, self.step],
         )
-        ret.topk_p, ret.topk_index = fast_sample(probs, num_samples=1)
+        ret.topk_p, ret.topk_index = fast_sample(
+            probs,
+            num_samples=1,
+            sampling_seed=(
+                buffers.sampling_seed[:bs]
+                if buffers.sampling_seed is not None
+                else None
+            ),
+            positions=buffers.seq_lens[:bs],
+            draft_step=self.step,
+        )
 
     def _run_step_body(self, forward_batch: ForwardBatch, num_tokens: int, bs: int):
         """One draft step's body: model forward + chain-hidden write + top-k.
@@ -678,9 +689,15 @@ def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers:
                     (max_bs, self.speculative_num_steps, vocab_size),
                     dtype=torch.float,
                 )
+                sampling_seed = (
+                    torch.zeros((max_bs,), dtype=torch.int64)
+                    if get_exec().deterministic.enable_deterministic_inference
+                    else None
+                )
             else:
                 temperatures = None
                 draft_probs = None
+                sampling_seed = None
 
             if self.require_gathered_buffer:
                 if self.require_mlp_tp_gather:
@@ -724,6 +741,7 @@ def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers:
             global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
             temperatures=temperatures,
             draft_probs=draft_probs,
+            sampling_seed=sampling_seed,
         )
 
     def _prepare_extra(self, forward_batch: ForwardBatch) -> None:
@@ -833,6 +851,11 @@ def prepare(self, forward_batch: ForwardBatch, *, staged: bool = False):
             self.seq_len_fill_value,
         )
 
+        if buffers.sampling_seed is not None:
+            buffers.sampling_seed[:raw_bs].copy_(
+                forward_batch.sampling_info.sampling_seed[:raw_bs]
+            )
+
         # Refresh the host mirror only when published; hand replay None
         # otherwise so no consumer reads a stale buffer.
         if forward_batch.seq_lens_cpu is not None:
@@ -928,10 +951,10 @@ class OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner(
     Rejection sampling is supported by sampling X ~ q inside the graph
     (_sample_draft_proposal, selected by the draft_probs buffer's presence):
     the captured rotation then carries the sampled token, prepare() stages the
-    temperatures sampling_info cannot deliver in-graph, and the worker clones
-    the per-step q off buffers.draft_probs after replay. torch.multinomial
-    draws through the graph-registered Philox generator, so each replay gets
-    fresh coins.
+    temperatures and sampling seeds cannot be delivered in-graph, and the
+    worker clones the per-step q off buffers.draft_probs after replay. Seeded
+    draws are keyed by request position and draft step; unseeded draws keep
+    using the graph-registered Philox generator.
     """
 
     rotates_in_graph = True

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -679,6 +679,9 @@ def _draft_extend_for_prefill(
                 probs, topk_p, topk_index = sample_draft_proposal(
                     output.logits_output.next_token_logits,
                     forward_batch.sampling_info.temperatures,
+                    sampling_seed=forward_batch.sampling_info.sampling_seed,
+                    positions=forward_batch.seq_lens,
+                    draft_step=step,
                 )
                 draft_probs_list.append(probs)
             else:
@@ -875,6 +878,9 @@ def _draft_extend_for_decode(
                         probs, ret_topk_p, ret_topk_index = sample_draft_proposal(
                             step_logits,
                             forward_batch.sampling_info.temperatures,
+                            sampling_seed=forward_batch.sampling_info.sampling_seed,
+                            positions=forward_batch.seq_lens,
+                            draft_step=step,
                         )
                         ret_draft_probs_list.append(probs)
                     if rotates_in_graph:
@@ -948,7 +954,11 @@ def _draft_extend_for_decode(
                     ]
                 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
+                        logits_sel,
+                        forward_batch.sampling_info.temperatures,
+                        sampling_seed=forward_batch.sampling_info.sampling_seed,
+                        positions=forward_batch.seq_lens,
+                        draft_step=step,
                     )
                     ret_draft_probs_list.append(probs)
                 else: