[Test] Initialize tensor-parallel groups in native SANA accuracy fixtures

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/e2e/accuracy/sana_video/conftest.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+
+@pytest.fixture
+def sana_model_parallel(tmp_path):
+    """Construct native parallel layers with real single-rank GPU groups."""
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_distributed_env,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
+    )
+    try:
+        initialize_model_parallel()
+        yield
+    finally:
+        destroy_distributed_env()

tests/e2e/accuracy/sana_video/test_sana_video_golden.py

@@ -202,6 +202,7 @@ def test_sana_video_transformer_matches_frozen_golden(
     variant: str,
     _hardware,
     accuracy_artifact_root: Path,
+    sana_model_parallel,
 ) -> None:
     del _hardware
     from vllm.config import LoadConfig

tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py

@@ -189,6 +189,7 @@ def _cfg_prediction(
 def test_sana_video_quick_pipeline_stages_match_diffusers(
     variant: str,
     _hardware,
+    sana_model_parallel,
 ) -> None:
     """Compare the first denoising step without replacing native attention.
 

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -13,7 +13,6 @@ from vllm import SamplingParams
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_omni.config.model import OmniModelConfig
@@ -164,28 +163,6 @@ def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect
         assert state is upstream.return_value
 
 
-@pytest.mark.parametrize("omni_stage", [False, True])
-@pytest.mark.parametrize("specialized", [False, True])
-def test_auxiliary_default_state_has_omni_lifecycle_but_preserves_upstream_states(monkeypatch, omni_stage, specialized):
-    class SpecializedState(DefaultModelState):
-        pass
-
-    state_cls = SpecializedState if specialized else DefaultModelState
-    upstream_state = object.__new__(state_cls)
-    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", lambda *_args: upstream_state)
-    monkeypatch.setattr(OmniModelState, "__init__", lambda *_args: None)
-    model_config = object.__new__(OmniModelConfig) if omni_stage else SimpleNamespace()
-    state = init_omni_model_state(
-        SimpleNamespace(model_config=model_config), SimpleNamespace(), None, torch.device("cpu")
-    )
-    if omni_stage and not specialized:
-        assert isinstance(state, OmniModelState)
-        assert callable(state.run_preprocess)
-        assert callable(state.postprocess_model_output)
-    else:
-        assert state is upstream_state
-
-
 def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
     runner = _make_runner()
     calls = []

vllm_omni/worker_v2/model_states/__init__.py

@@ -4,8 +4,7 @@
 """Omni-aware ``init_model_state`` factory.
 
 Extends the upstream v2 factory with Omni architecture dispatch.
-Non-Omni architectures use upstream dispatch; its default state acquires
-Omni lifecycle support when owned by an Omni stage.
+Non-Omni architectures fall through to the upstream ``init_model_state``.
 """
 
 from __future__ import annotations
@@ -17,11 +16,8 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states import (
     init_model_state as _upstream_init_model_state,
 )
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 
-from vllm_omni.config.model import OmniModelConfig
-
 
 def init_omni_model_state(
     vllm_config: VllmConfig,
@@ -32,8 +28,8 @@ def init_omni_model_state(
     """Create the appropriate ``ModelState`` for *model*.
 
     Returns an ``OmniModelState`` when the model declares Omni lifecycle
-    capabilities. Otherwise delegates to the upstream v2 factory, wrapping
-    its default state for auxiliary Omni stages.
+    capabilities; otherwise
+    delegates to the upstream v2 factory.
     """
     factory = getattr(model, "create_omni_model_state", None)
     if factory is not None:
@@ -53,12 +49,4 @@ def init_omni_model_state(
         )
 
         return OmniModelState(vllm_config, model, encoder_cache, device)
-    state = _upstream_init_model_state(vllm_config, model, encoder_cache, device)
-    if type(state) is DefaultModelState and isinstance(getattr(vllm_config, "model_config", None), OmniModelConfig):
-        # Auxiliary stages such as the forced aligner have no model-owned
-        # hooks, but Omni's runner still needs its request/output lifecycle.
-        # Keep specialized upstream states and plain vLLM configs intact.
-        from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
-
-        return OmniModelState(vllm_config, model, encoder_cache, device)
-    return state
+    return _upstream_init_model_state(vllm_config, model, encoder_cache, device)