[CPU][Bugfix][Speculative Decoding] Accept USE_FP64_GUMBEL in CPU recovered-tokens sampler (#46069)
Signed-off-by: hillel.darshan <hillel.darshan@intel.com> Co-authored-by: Cursor <cursoragent@cursor.com>
hillel.darshan committed
Jun 23, 2026 at 14:54 UTC
3554ada5d82c50fc55a1ca98d169fafccc435827
1 file changed
+3
-1
vllm/utils/cpu_triton_utils.py
+3
-1
@@ -300,6 +300,7 @@ def _sample_recovered_tokens_kernel_impl(
300
vocab_size,
301
BLOCK_SIZE=None,
302
NO_DRAFT_PROBS=False,
303
+ USE_FP64_GUMBEL=False,
304
):
305
# C++ reads integer tensors as int64_t*; ensure correct dtype.
306
orig_dtype = output_token_ids.dtype
@@ -310,7 +311,8 @@ def _sample_recovered_tokens_kernel_impl(
311
_ensure_int64(draft_token_ids),
312
draft_probs,
313
target_probs,
313
- inv_q,
314
+ # C++ kernel reads inv_q as float32.
315
+ inv_q.to(torch.float32),
316
vocab_size,
317
NO_DRAFT_PROBS,
318
)