[Bugfix][MRV2] Fix speculative warmup deadlock with pipeline parallelism (#62129)

Signed-off-by: chenjr97 <chenjr97@mail2.sysu.edu.cn>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/v1/worker/test_gpu_warmup_blocks.py

@@ -88,6 +88,7 @@ def _make_runner(
     """Stub model runner exposing only what the warmup entry points read."""
     return SimpleNamespace(
         num_speculative_steps=num_spec_steps,
+        speculative_config=SimpleNamespace() if num_spec_steps else None,
         adaptive_verification=None,
         rejection_sampler=None,
         decode_query_len=num_spec_steps + 1,

tests/v1/worker/test_mixed_warmup_gate.py

@@ -74,6 +74,7 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
     runner = SimpleNamespace(
         adaptive_verification=manager,
         rejection_sampler=rejection_sampler,
+        speculative_config=SimpleNamespace(),
     )
 
     def run_steps(model_runner, execute, sample, greedy=False):
@@ -91,3 +92,27 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
     assert runner.adaptive_verification is manager
     assert manager.cost_tables is None
     assert rejection_sampler.enable_adaptive_verification
+
+
+@pytest.mark.parametrize("speculative", [False, True])
+def test_kernel_warmup_runs_same_rounds_on_every_pp_stage(monkeypatch, speculative):
+    """PP receivers need a matching sender for every speculative warmup round."""
+    rounds = []
+    monkeypatch.setattr(
+        warmup,
+        "_warmup_kernels",
+        lambda runner, execute, sample, greedy=False: rounds.append(greedy),
+    )
+    for is_last_pp_rank in (False, True):
+        rounds.clear()
+        runner = SimpleNamespace(
+            adaptive_verification=None,
+            speculative_config=SimpleNamespace() if speculative else None,
+            rejection_sampler=(
+                SimpleNamespace(enable_adaptive_verification=False)
+                if is_last_pp_rank and speculative
+                else None
+            ),
+        )
+        warmup.warmup_kernels(runner, _fail, _fail)
+        assert rounds == ([False, True] if speculative else [False])

vllm/v1/worker/gpu/warmup.py

@@ -238,7 +238,8 @@ def warmup_kernels(
         rejection_sampler.enable_adaptive_verification = False
     try:
         _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)
-        if rejection_sampler is not None:
+        # All PP ranks must run the same forwards to match sends and receives.
+        if model_runner.speculative_config is not None:
             # Greedy batches skip logits processing, so rejection kernels
             # also see model-dtype logits instead of FP32.
             _warmup_kernels(