From 01a96720c6f690ac00adeed94b762b193d8472cd Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Fri, 17 Jul 2026 16:58:11 -0700 Subject: [PATCH] [spec decoding] replace torch.multinomial with several native torch op in rejection sampling (#31620) --- python/sglang/srt/speculative/spec_utils.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 4787423bb..ad133122d 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -114,7 +114,16 @@ def resolve_num_tokens_per_req( def fast_sample(probs: torch.Tensor, num_samples: int = 1): - sample_index = torch.multinomial(probs, num_samples=num_samples) + """Gumbel-max draw: argmax(probs / Exp(1)). Distributionally equivalent to + torch.multinomial minus its device-side validity assert, which a capturing + CUDA graph would replay every step.""" + q = torch.empty_like(probs, dtype=torch.float32).exponential_(1.0) + q.clamp_min_(torch.finfo(torch.float32).tiny) + scores = probs.float() / q + if num_samples == 1: + sample_index = scores.argmax(dim=-1, keepdim=True) + else: + sample_index = scores.topk(num_samples, dim=-1).indices sample_p = probs.gather(1, sample_index) return sample_p, sample_index