From 66a9b0c84fa4e56ab0e6d3c7925dd6a68b0ae82d Mon Sep 17 00:00:00 2001
From: Zhou Taichang <tzhouam@connect.ust.hk>
Date: Tue, 28 Jul 2026 17:14:26 +0800
Subject: [PATCH] [Bugfix] Select the AR-Diffusion engine in the DreamZero
 CFG-parallel deploy config (#5439)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 tests/dreamzero/test_deploy_configs.py   | 57 ++++++++++++++++++++++++
 vllm_omni/deploy/dreamzero_tp1_cfg2.yaml |  2 +
 2 files changed, 59 insertions(+)
 create mode 100644 tests/dreamzero/test_deploy_configs.py

diff --git a/tests/dreamzero/test_deploy_configs.py b/tests/dreamzero/test_deploy_configs.py
new file mode 100644
index 000000000..8f9770add
--- /dev/null
+++ b/tests/dreamzero/test_deploy_configs.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Every shipped DreamZero deploy config must select the AR-Diffusion engine.
+
+``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing
+``engine_backend`` fails at worker startup rather than at request time. That
+only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert
+it here on CPU instead.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import yaml
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Discover siblings of a known config so a future dreamzero*.yaml is covered
+# without editing this test.
+DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent
+
+
+def _dreamzero_stages() -> list[tuple[str, int, dict]]:
+    stages = []
+    for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")):
+        config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text())
+        for index, stage in enumerate(config.get("stages") or []):
+            stages.append((name, index, stage))
+    return stages
+
+
+def test_dreamzero_deploy_configs_exist() -> None:
+    assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}"
+
+
+@pytest.mark.parametrize(
+    ("config_name", "stage_index", "stage"),
+    _dreamzero_stages(),
+    ids=lambda value: value if isinstance(value, str) else None,
+)
+def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None:
+    backend = stage.get("engine_backend")
+    assert backend, (
+        f"{config_name} stage {stage_index} does not set engine_backend; "
+        "DreamZeroPipeline requires the AR-Diffusion engine"
+    )
+    # Resolving also proves the import path is real, not just a lookalike string.
+    engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend))
+    assert issubclass(engine_class, ARDiffusionEngine), (
+        f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine"
+    )
diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
index a453d43f3..0bf3d9517 100644
--- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
+++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
@@ -11,6 +11,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     parallel_config:
       tensor_parallel_size: 1
       cfg_parallel_size: 2