diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index b48f8249e6..6a0c59874b 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/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, diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 199ce3b54a..9d3f60d6f4 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/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]) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 5418147500..848d984295 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/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(