[Bugfix][Model Runner V2] Honor use_fp64_gumbel when FlashInfer sampling is available

tests/v1/worker/test_gpu_sampler_flags.py

@@ -13,6 +13,7 @@
     pytest.skip("CUDA required for sampler flag tests", allow_module_level=True)
 
 from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample import sampler as sampler_module
 from vllm.v1.worker.gpu.sample.logits_processor import (
     LogitsContext,
     LogitsProcessor,
@@ -30,7 +31,10 @@ class MockReasoningConfig:
     natural_reasoning_end_token_ids = [91]
 
 
-def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> Sampler:
+def _make_sampler(
+    custom_logits_processors: Sequence[LogitsProcessor] = (),
+    use_fp64_gumbel: bool = False,
+) -> Sampler:
     req_states = RequestState(
         max_num_reqs=4,
         max_model_len=64,
@@ -45,6 +49,7 @@ def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> S
         vocab_size=VOCAB_SIZE,
         device=DEVICE,
         req_states=req_states,
+        use_fp64_gumbel=use_fp64_gumbel,
         custom_logits_processors=custom_logits_processors,
     )
 
@@ -133,6 +138,35 @@ def has_topp_only_rows(idx: list[int], vocab_size: int = VOCAB_SIZE) -> bool:
     assert has_topp_only_rows([0], vocab_size=10)
 
 
+def test_fp64_gumbel_bypasses_flashinfer_sampler(monkeypatch: pytest.MonkeyPatch):
+    """FlashInfer draws fp32 noise, so an unseeded top-p request, which would
+    otherwise use the fused FlashInfer sampler, must use fp64 Gumbel sampling."""
+
+    def flashinfer_sample(*args, **kwargs):
+        raise AssertionError("FlashInfer sampler used despite use_fp64_gumbel")
+
+    monkeypatch.setattr(sampler_module, "flashinfer_sampler_supported", lambda: True)
+    monkeypatch.setattr(sampler_module, "flashinfer_sample", flashinfer_sample)
+    sampler = _make_sampler(use_fp64_gumbel=True)
+    sampler.add_request(0, SamplingParams(top_p=0.9))
+    sampler.apply_staged_writes()
+
+    idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    zeros = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    sampled, _ = sampler.sample(
+        torch.randn(1, VOCAB_SIZE, device=DEVICE),
+        expanded_idx_mapping=idx_mapping,
+        idx_mapping=idx_mapping,
+        idx_mapping_np=np.zeros(1, dtype=np.intp),
+        pos=zeros,
+        input_ids=zeros,
+        expanded_local_pos=zeros,
+        seq_lens_upper_bound_np=np.ones(1, dtype=np.int64),
+    )
+
+    assert sampled.shape == (1,)
+
+
 def test_logits_processing_cache_only_checks_active_requests():
     sampler = _make_sampler()
     sampler.add_request(0, SamplingParams(temperature=0.0))

vllm/v1/worker/gpu/sample/sampler.py

@@ -85,10 +85,12 @@ def __init__(
         self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool)
         self.num_speculative_tokens = num_speculative_tokens
         self.return_sampling_mask = return_sampling_mask
+        # The FlashInfer and XPU kernels draw fp32 noise, so they can't honor fp64.
         self.use_flashinfer = (
-            not return_sampling_mask and flashinfer_sampler_supported()
+            not return_sampling_mask
+            and not use_fp64_gumbel
+            and flashinfer_sampler_supported()
         )
-        # The XPU kernel draws fp32 exponential noise, so it can't honor fp64.
         self.use_xpu_sampler = (
             not return_sampling_mask and not use_fp64_gumbel and xpu_sampler_supported()
         )