| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.v1.worker.gpu import warmup |
| 10 | from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup |
| 11 | |
| 12 | |
| 13 | def _fail(*args, **kwargs): |
| 14 | raise AssertionError("worker callback must not run when warmup is skipped") |
| 15 | |
| 16 | |
| 17 | @pytest.mark.parametrize("max_num_reqs", [1, 0]) |
| 18 | def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): |
| 19 | """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 |
| 20 | the warmup must be skipped without touching the worker callbacks.""" |
| 21 | runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) |
| 22 | |
| 23 | assert ( |
| 24 | run_mixed_prefill_decode_warmup( |
| 25 | runner, |
| 26 | worker_execute_model=_fail, |
| 27 | worker_sample_tokens=_fail, |
| 28 | num_tokens=128, |
| 29 | ) |
| 30 | is False |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | @pytest.mark.parametrize("fail_warmup", [False, True]) |
| 35 | def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup): |
| 36 | """Startup must warm fixed drafts before calibration and retain its manager.""" |
| 37 | manager = SimpleNamespace(cost_tables=None) |
| 38 | rejection_sampler = SimpleNamespace(enable_adaptive_verification=True) |
| 39 | runner = SimpleNamespace( |
| 40 | adaptive_verification=manager, |
| 41 | rejection_sampler=rejection_sampler, |
| 42 | ) |
| 43 | |
| 44 | def run_steps(model_runner, execute, sample): |
| 45 | assert model_runner.adaptive_verification is None |
| 46 | assert not model_runner.rejection_sampler.enable_adaptive_verification |
| 47 | if fail_warmup: |
| 48 | raise RuntimeError("warmup failed") |
| 49 | |
| 50 | monkeypatch.setattr(warmup, "_warmup_kernels", run_steps) |
| 51 | if fail_warmup: |
| 52 | with pytest.raises(RuntimeError, match="warmup failed"): |
| 53 | warmup.warmup_kernels(runner, _fail, _fail) |
| 54 | else: |
| 55 | warmup.warmup_kernels(runner, _fail, _fail) |
| 56 | assert runner.adaptive_verification is manager |
| 57 | assert manager.cost_tables is None |
| 58 | assert rejection_sampler.enable_adaptive_verification |