From 1132df761374addfec39f2b96573a8d182285dd4 Mon Sep 17 00:00:00 2001 From: Kolade Fajimi Date: Wed, 2 Sep 2026 08:03:19 +0100 Subject: [PATCH] [Bugfix][CPU][Spec Decode] Honor use_fp64_gumbel in CPU recovered-token sampling Rebase onto upstream/main (triton kernel dispatcher refactor, #43048) and fix the new CPU test to call cpu_triton_utils._sample_recovered_tokens_kernel_impl directly, since the module-level _FuncWrapper alias it used no longer exists. Co-authored-by: Claude Fable 5.1 Signed-off-by: Kolade Fajimi Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01MQjaTsvMtqQhkkjQ6wXER6 --- .buildkite/hardware_tests/cpu.yaml | 5 ++- csrc/cpu/spec_decode_utils.cpp | 33 ++++++++++++++--- tests/v1/sample/test_rejection_sampler.py | 45 +++++++++++++++++++++++ vllm/utils/cpu_triton_utils.py | 9 ++--- 4 files changed, 81 insertions(+), 11 deletions(-) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index d54197f83678..8dbd2dbe62ac 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -34,6 +34,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 " @@ -53,7 +55,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 diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp index 30192196b959..97b3a0eea8df 100644 --- a/csrc/cpu/spec_decode_utils.cpp +++ b/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 +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& 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(); const float* target_probs_ptr = target_probs.data_ptr(); - const float* inv_q_ptr = inv_q.data_ptr(); + const inv_q_t* inv_q_ptr = inv_q.data_ptr(); 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(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& 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( + output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, + target_probs, inv_q, vocab_size, no_draft_probs); + }); +} + } // namespace cpu_utils diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index 231f6adbb2d7..a97f3496ad58 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/tests/v1/sample/test_rejection_sampler.py @@ -9,6 +9,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 ( @@ -999,6 +1000,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", diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index d8e4cd389d30..84a04ffec4f2 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -421,9 +421,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) @@ -433,8 +433,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, )