[Tests] Select V2 for diffusion scheduler unit tests (#58272)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 23, 2026 at 02:04 UTC
0f2a15c9277f34c9afe141cf578d1f02b3bebdfe
1 file changed
+12
tests/v1/core/test_scheduler.py
+12
@@ -6896,6 +6896,13 @@ def _diffusion_request(req_id: str, extra_args: dict) -> Request:
6896
return request
6897
6898
6899
+@pytest.fixture
6900
+def diffusion_model_runner(monkeypatch):
6901
+ # These CPU tests only exercise scheduling, not Triton kernels.
6902
+ monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
6903
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
6904
+
6905
+
6906
def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
6907
scheduler = create_scheduler(
6908
async_scheduling=True,
@@ -6908,6 +6915,7 @@ def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
6915
6916
6917
@pytest.mark.parametrize("async_scheduling", [False, True])
6918
+@pytest.mark.usefixtures("diffusion_model_runner")
6919
def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
6920
config = create_scheduler(
6921
async_scheduling=async_scheduling, diffusion_canvas_length=8
@@ -6917,6 +6925,7 @@ def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
6925
)
6926
6927
6928
+@pytest.mark.usefixtures("diffusion_model_runner")
6929
def test_diffusion_scheduler_narrows_the_canvas_per_request():
6930
scheduler = _diffusion_scheduler()
6931
wide = _diffusion_request("wide", {})
@@ -6934,6 +6943,7 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request():
6943
6944
6945
@pytest.mark.parametrize("structured", [False, True])
6946
+@pytest.mark.usefixtures("diffusion_model_runner")
6947
def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
6948
"""Padded worker drafts must be narrowed before scheduling or grammar validation."""
6949
scheduler = _diffusion_scheduler()
@@ -6964,6 +6974,7 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
6974
assert narrow.structured_output_request.grammar.seen == [tokens[:4]]
6975
6976
6977
+@pytest.mark.usefixtures("diffusion_model_runner")
6978
def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
6979
scheduler = _diffusion_scheduler()
6980
one = _diffusion_request(
@@ -6986,6 +6997,7 @@ def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
6997
assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"}
6998
6999
7000
+@pytest.mark.usefixtures("diffusion_model_runner")
7001
def test_diffusion_read_deferral_keeps_a_longer_pp_wait():
7002
scheduler = _diffusion_scheduler(pipeline_parallel_size=3, use_v2_model_runner=True)
7003
read = _diffusion_request(