From 3554ada5d82c50fc55a1ca98d169fafccc435827 Mon Sep 17 00:00:00 2001 From: "hillel.darshan" <141740534+hillelda@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:54:07 +0300 Subject: [PATCH] [CPU][Bugfix][Speculative Decoding] Accept USE_FP64_GUMBEL in CPU recovered-tokens sampler (#46069) Signed-off-by: hillel.darshan Co-authored-by: Cursor --- vllm/utils/cpu_triton_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index ea0383a9d4..657afad838 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -300,6 +300,7 @@ def _sample_recovered_tokens_kernel_impl( vocab_size, BLOCK_SIZE=None, NO_DRAFT_PROBS=False, + USE_FP64_GUMBEL=False, ): # C++ reads integer tensors as int64_t*; ensure correct dtype. orig_dtype = output_token_ids.dtype @@ -310,7 +311,8 @@ def _sample_recovered_tokens_kernel_impl( _ensure_int64(draft_token_ids), draft_probs, target_probs, - inv_q, + # C++ kernel reads inv_q as float32. + inv_q.to(torch.float32), vocab_size, NO_DRAFT_PROBS, )