Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions tests/v1/worker/test_mixed_warmup_gate.py
Original file line number Diff line number Diff line change
@@ -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
)
6 changes: 3 additions & 3 deletions vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion vllm/v1/worker/gpu/warmup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_"
Expand Down
Loading