[Bugfix][CPU][Spec Decode] Honor use_fp64_gumbel in CPU recovered-token sampling

.buildkite/hardware_tests/cpu.yaml

@@ -36,6 +36,8 @@ steps:
   - vllm/v1/sample/ops/topk_topp_triton.py
   - vllm/v1/sample/ops/topk_topp_sampler.py
   - tests/v1/sample/test_topk_topp_sampler.py
+  - vllm/utils/cpu_triton_utils.py
+  - tests/v1/sample/test_rejection_sampler.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -57,7 +59,8 @@ steps:
       tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \
       tests/kernels/mamba/test_causal_conv1d.py \
       tests/kernels/mamba/test_mamba_ssm.py \
-      tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
+      tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp \
+      tests/v1/sample/test_rejection_sampler.py::test_cpu_sample_recovered_tokens_preserves_fp64_inv_q"
   parallelism: 2
 
 # Speculative decoding under a CPU platform. The unit tests already pass on

csrc/cpu/spec_decode_utils.cpp

@@ -427,7 +427,10 @@ void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
   }
 }
 
-void sample_recovered_tokens_kernel_impl(
+namespace {
+
+template <typename inv_q_t>
+void sample_recovered_tokens_impl(
     torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& draft_token_ids,
     const std::optional<torch::Tensor>& draft_probs,
@@ -441,7 +444,7 @@ void sample_recovered_tokens_kernel_impl(
   const float* draft_probs_ptr =
       no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
   const float* target_probs_ptr = target_probs.data_ptr<float>();
-  const float* inv_q_ptr = inv_q.data_ptr<float>();
+  const inv_q_t* inv_q_ptr = inv_q.data_ptr<inv_q_t>();
 
   const int64_t target_stride = target_probs.stride(0);
   const int64_t draft_probs_stride =
@@ -454,7 +457,7 @@ void sample_recovered_tokens_kernel_impl(
     int64_t end_idx = cu_draft_ptr[req_idx];
     int64_t num_draft_tokens = end_idx - start_idx;
 
-    const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride;
+    const inv_q_t* req_inv_q = inv_q_ptr + req_idx * inv_q_stride;
 
     for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
       int64_t token_idx = start_idx + pos;
@@ -467,7 +470,7 @@ void sample_recovered_tokens_kernel_impl(
                          : (draft_probs_ptr + token_idx * draft_probs_stride);
 
       int64_t best_id = 0;
-      float best_val = -1.0f;
+      inv_q_t best_val = -1;
 
       for (int64_t v = 0; v < vocab_size; ++v) {
         float prob = token_target_probs[v];
@@ -478,7 +481,7 @@ void sample_recovered_tokens_kernel_impl(
           prob = diff > 0.0f ? diff : 0.0f;
         }
 
-        float val = prob * req_inv_q[v];
+        inv_q_t val = static_cast<inv_q_t>(prob) * req_inv_q[v];
         if (val > best_val) {
           best_val = val;
           best_id = v;
@@ -489,4 +492,24 @@ void sample_recovered_tokens_kernel_impl(
   }
 }
 
+}  // namespace
+
+void sample_recovered_tokens_kernel_impl(
+    torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
+    const torch::Tensor& draft_token_ids,
+    const std::optional<torch::Tensor>& draft_probs,
+    const torch::Tensor& target_probs, const torch::Tensor& inv_q,
+    const int64_t vocab_size, const bool no_draft_probs) {
+  // `inv_q` carries the exponential-race noise at the precision the caller
+  // selected via `use_fp64_gumbel`. Dispatch on its dtype so fp64 noise keeps
+  // full precision through both the score and the argmax, matching the Triton
+  // kernel. The fp32 instantiation is unchanged.
+  AT_DISPATCH_FLOATING_TYPES(
+      inv_q.scalar_type(), "sample_recovered_tokens_kernel_impl", [&] {
+        sample_recovered_tokens_impl<scalar_t>(
+            output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs,
+            target_probs, inv_q, vocab_size, no_draft_probs);
+      });
+}
+
 }  // namespace cpu_utils

tests/v1/sample/test_rejection_sampler.py

@@ -10,6 +10,7 @@
 
 from tests.v1.sample.utils import create_allowed_token_ids
 from vllm.platforms import current_platform
+from vllm.utils import cpu_triton_utils
 from vllm.v1.sample.logits_processor import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.rejection_sampler import (
@@ -1000,6 +1001,50 @@ def test_sample_recovered_tokens_uses_fp64_exponential_race_when_requested():
     assert torch.equal(actual, expected)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cpu(),
+    reason="exercises the CPU fallback kernel for sample_recovered_tokens",
+)
+@pytest.mark.parametrize(
+    "inv_q_row,fp64_choice,fp32_choice",
+    [
+        # Near-tie only fp64 can resolve: both entries round to 1.0 in fp32.
+        ([1.0, 1.0 + 1e-10, 1.0], 1, 0),
+        # Lower-tail draws whose reciprocals both overflow fp32 to +inf.
+        ([1e39, 2e39, 1.0], 1, 0),
+    ],
+)
+def test_cpu_sample_recovered_tokens_preserves_fp64_inv_q(
+    inv_q_row, fp64_choice, fp32_choice
+):
+    """fp64 gumbel noise must survive the CPU kernel, as it does on Triton.
+
+    The fp32 expectations double as a guard that the fp32 path is unchanged.
+    """
+    vocab_size = 3
+    target_probs = torch.tensor([[1.0, 1.0, 0.5]], dtype=torch.float32)
+    draft_token_ids = torch.tensor([2], dtype=torch.int64)
+    cu_num_draft_tokens = torch.tensor([1], dtype=torch.int64)
+
+    def recovered_token(dtype: torch.dtype) -> int:
+        out = torch.zeros(1, dtype=torch.int64)
+        cpu_triton_utils._sample_recovered_tokens_kernel_impl(
+            out,
+            cu_num_draft_tokens,
+            draft_token_ids,
+            None,
+            target_probs,
+            torch.tensor([inv_q_row], dtype=dtype),
+            vocab_size,
+            NO_DRAFT_PROBS=True,
+            USE_FP64_GUMBEL=dtype is torch.float64,
+        )
+        return int(out.item())
+
+    assert recovered_token(torch.float64) == fp64_choice
+    assert recovered_token(torch.float32) == fp32_choice
+
+
 @pytest.mark.parametrize("no_draft_probs", [True, False])
 @pytest.mark.parametrize(
     "vocab_size",

vllm/utils/cpu_triton_utils.py

@@ -417,9 +417,9 @@ def _sample_recovered_tokens_kernel_impl(
     USE_FP64_GUMBEL=False,
 ):
     # USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller
-    # has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes
-    # `inv_q` directly, so the flag is accepted for interface parity and the
-    # value is read at its existing dtype.
+    # has already applied to `inv_q` (fp64 vs fp32). The C++ kernel dispatches on
+    # that dtype, so `inv_q` is passed through unchanged to keep fp64 noise at
+    # full precision through the argmax, matching the Triton kernel.
     # C++ reads integer tensors as int64_t*; ensure correct dtype.
     orig_dtype = output_token_ids.dtype
     output_i64 = _ensure_int64(output_token_ids)
@@ -429,8 +429,7 @@ def _sample_recovered_tokens_kernel_impl(
         _ensure_int64(draft_token_ids),
         draft_probs,
         target_probs,
-        # C++ kernel reads inv_q as float32.
-        inv_q.to(torch.float32),
+        inv_q,
         vocab_size,
         NO_DRAFT_PROBS,
     )