[V1] Warm up seeded and greedy sampler paths (#54425, #54455)

Pre-compile native Gumbel and bfloat16 rejection kernels by cycling stochastic,
seeded, and greedy sampling configs during startup warmup.

Signed-off-by: Manikanta Bandham <bandhammanikanta@gmail.com>

Author: bandham-manikanta

tests/v1/worker/test_gpu_sampler_flags.py

@@ -88,3 +88,31 @@ def test_logits_processing_cache_only_checks_active_requests():
 
     assert not np.any(sampler.needs_logits_processing[sampling_only])
     assert np.any(sampler.needs_logits_processing[with_processing])
+
+
+def test_all_sampler_warmup_configs():
+    """Test that for_all_sampler_warmup_configs returns a list of configurations
+    covering:
+    - FlashInfer (unseeded stochastic),
+    - native top-k/top-p + Gumbel (seeded), and
+    - greedy (model dtype) paths."""
+    configs = SamplingParams.for_all_sampler_warmup_configs()
+    assert len(configs) >= 3
+
+    unseeded = configs[0]
+    assert unseeded.seed is None
+    assert unseeded.temperature > 0.0
+
+    seeded = configs[1]
+    assert seeded.seed is not None
+
+    greedy = configs[2]
+    assert greedy.temperature == 0.0
+
+    sampler = _make_sampler()
+    sampler.add_request(0, 1, unseeded)
+    sampler.add_request(1, 1, seeded)
+    sampler.add_request(2, 1, greedy)
+
+    assert sampler.needs_logits_processing[0]
+    assert not sampler.needs_logits_processing[2]

vllm/sampling_params.py

@@ -1270,6 +1270,15 @@ def for_sampler_warmup() -> "SamplingParams":
             prompt_logprobs=1,
         )
 
+    @classmethod
+    def for_all_sampler_warmup_configs(cls) -> list["SamplingParams"]:
+        """Return SamplingParams covering all sampler warmup configurations."""
+        return [
+            cls.for_sampler_warmup(),
+            cls(temperature=0.9, seed=42),
+            cls(temperature=0.0),
+        ]
+
 
 class BeamSearchParams(
     msgspec.Struct,

vllm/v1/worker/gpu/warmup.py

@@ -281,14 +281,14 @@ def warmup_kernels(
 
     # SamplingParams exercising all sampling features.
     if model_runner.is_pooling_model:
-        sampling_params = None
+        sampling_params_list = [None]
         pooling_task = model_runner.model_config.get_pooling_task(
             model_runner.get_supported_tasks()
         )
         pooling_params = PoolingParams(task=pooling_task)
         pooling_params.verify(model_runner.model_config)
     else:
-        sampling_params = SamplingParams.for_sampler_warmup()
+        sampling_params_list = SamplingParams.for_all_sampler_warmup_configs()
         pooling_params = None
 
     # Assign distinct block IDs per request per group. 0 null block, start from 1.
@@ -309,7 +309,7 @@ def _alloc_blocks(num_blocks: int) -> list[int]:
             Request(
                 req_ids[i],
                 prompt_token_ids,
-                sampling_params,
+                sampling_params_list[i % len(sampling_params_list)],
                 pooling_params,
                 mm_features=warmup_mm_features,
             ),
@@ -410,10 +410,20 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None:
                 # Exercise the model paths that split a batch by whether each
                 # request received draft tokens.
                 decode_steps.append(([0, 1], [False, False]))
-        if num_reqs > 1:
+
+            # Single-request steps covering unseeded, seeded, and greedy configs.
             decode_steps.append(([0], [use_spec_decode]))
             if use_spec_decode:
                 decode_steps.append(([0], [False]))
+
+            decode_steps.append(([1], [use_spec_decode]))
+            if use_spec_decode:
+                decode_steps.append(([1], [False]))
+
+            if num_reqs >= 3:
+                decode_steps.append(([2], [use_spec_decode]))
+                if use_spec_decode:
+                    decode_steps.append(([2], [False]))
         elif use_spec_decode:
             decode_steps.append(([0], [False]))
 

vllm/v1/worker/gpu_model_runner.py

@@ -6446,6 +6446,14 @@ def _dummy_sampler_run(
                 logits,
                 all_greedy_metadata,
             )
+            if self.model_config.dtype != logits.dtype:
+                model_dtype_logits = logits.to(self.model_config.dtype)
+                self.rejection_sampler(
+                    dummy_spec_decode_metadata,
+                    draft_probs,
+                    model_dtype_logits,
+                    all_greedy_metadata,
+                )
             torch.accelerator.synchronize()
         return sampler_output