diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py new file mode 100644 index 0000000000..6941df6773 --- /dev/null +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup.""" + +from types import SimpleNamespace + +import pytest + +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + + +def _fail(*args, **kwargs): + raise AssertionError("worker callback must not run when warmup is skipped") + + +@pytest.mark.parametrize("max_num_reqs", [1, 0]) +def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): + """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 + the warmup must be skipped without touching the worker callbacks.""" + runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) + + assert ( + run_mixed_prefill_decode_warmup( + runner, + worker_execute_model=_fail, + worker_sample_tokens=_fail, + num_tokens=128, + ) + is False + ) diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 44be769e24..80dfd51996 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -130,7 +130,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with torch.inference_mode(): warmup_executed = True if is_leader: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -144,7 +144,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with flashinfer_autotune(True, cache=str(cache_path)): runner._dummy_run(**dummy_run_kwargs) else: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -236,7 +236,7 @@ def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None: ) mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens) if not mixed_warmup_done: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) run_mixed_prefill_decode_warmup( v2_runner, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index ff9a75f0ef..de47cb8880 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -34,7 +34,7 @@ def run_mixed_prefill_decode_warmup( req_id_prefix: str = "_v2_mixed_warmup", ) -> bool: """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" - if model_runner.is_pooling_model or num_tokens < 3: + if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3: return False decode_req_id = f"{req_id_prefix}_decode_"