Merge branch 'main' into feat/moss-tts-nano

Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/config/test_pipeline_registry.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the central pipeline registry (2.5/N)."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import (
+    _DIFFUSION_PIPELINES,
+    _OMNI_PIPELINES,
+    _VLLM_OMNI_PIPELINES,
+)
+from vllm_omni.config.stage_config import (
+    _PIPELINE_REGISTRY,
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+    register_pipeline,
+)
+
+
+class TestCentralRegistryDeclarations:
+    """Every in-tree pipeline must be declared exactly once in the central registry."""
+
+    def test_union_contains_all_omni(self):
+        for key in _OMNI_PIPELINES:
+            assert key in _VLLM_OMNI_PIPELINES
+
+    def test_union_contains_all_diffusion(self):
+        for key in _DIFFUSION_PIPELINES:
+            assert key in _VLLM_OMNI_PIPELINES
+
+    def test_no_duplicate_model_type_between_omni_and_diffusion(self):
+        overlap = set(_OMNI_PIPELINES) & set(_DIFFUSION_PIPELINES)
+        assert not overlap, f"Duplicate model_types across omni/diffusion: {overlap}"
+
+    def test_expected_omni_pipelines_present(self):
+        # Guard against accidental removal during future refactors.
+        assert "qwen2_5_omni" in _OMNI_PIPELINES
+        assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES
+        assert "qwen3_omni_moe" in _OMNI_PIPELINES
+        assert "qwen3_tts" in _OMNI_PIPELINES
+
+
+class TestLazyLoading:
+    """Pipelines are imported only on first access."""
+
+    def test_contains_without_import(self):
+        # ``in`` hits the lazy map, not the loaded cache.
+        assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
+
+    def test_getitem_loads_correct_pipeline(self):
+        pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
+        assert pipeline.model_type == "qwen3_omni_moe"
+        assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration"
+
+    def test_unknown_model_type_returns_none_via_get(self):
+        assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None
+
+    def test_unknown_model_type_raises_keyerror_via_getitem(self):
+        with pytest.raises(KeyError):
+            _PIPELINE_REGISTRY["not_a_real_pipeline"]
+
+    def test_iteration_yields_registered_pipelines(self):
+        keys = set(_PIPELINE_REGISTRY)
+        assert "qwen2_5_omni" in keys
+        assert "qwen3_omni_moe" in keys
+
+
+class TestDynamicRegistration:
+    """``register_pipeline()`` still works for plugins and tests."""
+
+    def test_register_adds_to_registry(self):
+        custom = PipelineConfig(
+            model_type="_test_dynamic_registration",
+            model_arch="DynamicTestModel",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="test",
+                    execution_type=StageExecutionType.LLM_AR,
+                    input_sources=(),
+                    final_output=True,
+                ),
+            ),
+        )
+        register_pipeline(custom)
+        try:
+            assert "_test_dynamic_registration" in _PIPELINE_REGISTRY
+            assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom
+        finally:
+            # Don't leak the test registration into other tests.
+            if "_test_dynamic_registration" in _PIPELINE_REGISTRY:
+                del _PIPELINE_REGISTRY["_test_dynamic_registration"]
+
+    def test_dynamic_registration_overrides_lazy_entry(self):
+        # Build a substitute for qwen3_omni_moe that we can distinguish.
+        original = _PIPELINE_REGISTRY["qwen3_omni_moe"]
+        override = PipelineConfig(
+            model_type="qwen3_omni_moe",
+            model_arch="OverriddenArch",
+            stages=original.stages,
+        )
+        register_pipeline(override)
+        try:
+            assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch"
+        finally:
+            # Remove the dynamic override so later tests see the original.
+            if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded:
+                del _PIPELINE_REGISTRY["qwen3_omni_moe"]

tests/core/sched/test_omni_scheduler_mixin.py

@@ -0,0 +1,129 @@
+"""Unit tests for OmniSchedulerMixin streaming session replacement.
+
+These tests pin the behavior of `_replace_session_with_streaming_update` against
+current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading
+vLLM, failures here should highlight incompatible changes to request state or
+update payloads early.
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _SchedulerStub(OmniSchedulerMixin):
+    """Minimal scheduler surface required by OmniSchedulerMixin."""
+
+    def __init__(self, *, log_stats: bool = False) -> None:
+        self.num_waiting_for_streaming_input = 0
+        self.log_stats = log_stats
+
+
+def _make_request(**kwargs) -> Request:
+    sp = SamplingParams(max_tokens=8)
+    defaults = dict(
+        request_id="req-mixin-test",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=sp,
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    defaults.update(kwargs)
+    return Request(**defaults)
+
+
+def _make_update(**kwargs) -> StreamingUpdate:
+    sp_new = SamplingParams(max_tokens=16)
+    defaults = dict(
+        mm_features=None,
+        prompt_token_ids=[10, 20],
+        max_tokens=32,
+        arrival_time=200.0,
+        sampling_params=sp_new,
+    )
+    defaults.update(kwargs)
+    return StreamingUpdate(**defaults)
+
+
+class TestReplaceSessionWithStreamingUpdate:
+    def test_resets_tokens_and_prompt_from_update(self) -> None:
+        sched = _SchedulerStub()
+        session = _make_request()
+        session.append_output_token_ids([7, 8])
+        session.num_computed_tokens = 99
+        session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+        update = _make_update(prompt_token_ids=[40, 41, 42])
+        sched.num_waiting_for_streaming_input = 3
+        sched._replace_session_with_streaming_update(session, update)
+
+        assert session._output_token_ids == []
+        assert list(session._all_token_ids) == [40, 41, 42]
+        assert session.prompt_token_ids == [40, 41, 42]
+        assert session.num_computed_tokens == 0
+        assert session.num_prompt_tokens == 3
+        assert session.arrival_time == 200.0
+        assert session.sampling_params is update.sampling_params
+        assert session.status == RequestStatus.WAITING
+        assert sched.num_waiting_for_streaming_input == 2
+
+    def test_none_prompt_token_ids_becomes_empty(self) -> None:
+        sched = _SchedulerStub()
+        session = _make_request()
+        session.status = RequestStatus.RUNNING
+        update = _make_update(prompt_token_ids=None)
+        sched._replace_session_with_streaming_update(session, update)
+
+        assert session.prompt_token_ids == ()
+        assert list(session._all_token_ids) == []
+        assert session.num_prompt_tokens == 0
+        assert sched.num_waiting_for_streaming_input == 0
+
+    def test_additional_information_cleared_when_update_omits_it(self) -> None:
+        sched = _SchedulerStub()
+        session = _make_request()
+        if not hasattr(session, "additional_information"):
+            pytest.skip("Request has no additional_information (Omni patch inactive?)")
+        session.additional_information = {"keep": True}
+        session.status = RequestStatus.RUNNING
+
+        base = _make_update()
+        if not hasattr(base, "additional_information"):
+            pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)")
+        update = replace(base, additional_information=None)
+
+        sched._replace_session_with_streaming_update(session, update)
+        assert session.additional_information is None
+
+    def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None:
+        sched = _SchedulerStub()
+        session = _make_request()
+        session.status = RequestStatus.RUNNING
+        sched.num_waiting_for_streaming_input = 5
+        sched._replace_session_with_streaming_update(session, _make_update())
+        assert sched.num_waiting_for_streaming_input == 5
+
+    def test_records_queued_event_when_log_stats_enabled(self) -> None:
+        sched = _SchedulerStub(log_stats=True)
+        session = _make_request()
+        session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+        sched._replace_session_with_streaming_update(session, _make_update())
+
+        assert session.events
+        assert session.events[-1].type == EngineCoreEventType.QUEUED

tests/core/test_prefix_cache.py

@@ -0,0 +1,349 @@
+import pytest
+import torch
+
+from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
+
+DEFAULT_SEQ_LEN = 15
+NUM_BLOCKS = 10
+BLOCK_SIZE = 4
+HIDDEN_SIZE = 2
+DTYPE = torch.float32
+OTHER_DTYPE = torch.float16
+DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE])
+
+
+class MockInputBatch:
+    def __init__(self, num_computed_tokens_cpu):
+        self.req_ids = ["req1", "req2"]
+        self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)}
+        self.num_computed_tokens_cpu = num_computed_tokens_cpu
+
+        # Block table is only mocked for validation of length;
+        # we don't actually need to add valid values here since
+        # we patch the table when testing.
+        class _DummyBlockTable:
+            pass
+
+        self.block_table = _DummyBlockTable()
+        self.block_table.block_tables = [None]
+
+
+def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache:
+    """Build an OmniTensorPrefixCache and init mm tensors."""
+    cache = get_omni_pcache()
+    mm_outputs = get_multimodal_outputs(feat_dims, seq_len)
+    cache.maybe_init_missing_mm_cache_keys(mm_outputs, seq_len)
+    return cache
+
+
+def get_omni_pcache() -> OmniTensorPrefixCache:
+    """Build an OmniTensorPrefixCache, but don't init mm tensors."""
+    cache = OmniTensorPrefixCache(
+        num_blocks=NUM_BLOCKS,
+        block_size=BLOCK_SIZE,
+        hidden_size=HIDDEN_SIZE,
+        hs_dtype=DTYPE,
+    )
+    return cache
+
+
+def get_multimodal_outputs(feat_dims: dict[str, int], seq_len: int) -> dict[str, torch.Tensor]:
+    fake_mm_inputs = {}
+    for mm_key, feat_dim in feat_dims.items():
+        fake_mm_inputs[mm_key] = torch.rand((seq_len, feat_dim), dtype=DTYPE)
+    return fake_mm_inputs
+
+
+### Tests for initialization
+def test_initialization_simple():
+    """Check default initialization only creates the hidden states."""
+    cache = get_omni_pcache()
+    assert isinstance(cache.hidden_states_cache, torch.Tensor)
+    assert cache.hidden_states_cache.shape == DEFAULT_SHAPE
+    assert len(cache.mm_outputs_cache) == 0
+    assert len(cache.mm_cache_keys) == 0
+
+
+def test_initialization_with_multimodal():
+    """Check initialization + registration of multimodal outputs."""
+    cache = get_omni_pcache()
+    feat_dims = {"foo": 100, "bar": 50, "baz": 10}
+    mm_outputs = get_multimodal_outputs(
+        feat_dims,
+        seq_len=DEFAULT_SEQ_LEN,
+    )
+    # Cast one of the keys to a different dtype; the dtype of the tensor
+    # that is used to initialize the cache dictates the cache dtype.
+    mm_outputs["foo"] = mm_outputs["foo"].to(OTHER_DTYPE)
+
+    cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN)
+    assert len(cache.mm_cache_keys) == 3
+    assert set(cache.mm_cache_keys) == set(feat_dims.keys())
+    for mm_key in cache.mm_cache_keys:
+        cache_tensor = cache.mm_outputs_cache[mm_key]
+        assert isinstance(cache_tensor, torch.Tensor)
+        assert cache_tensor.shape[-1] == feat_dims[mm_key]
+        assert mm_outputs[mm_key].dtype == cache_tensor.dtype
+
+
+def test_init_missing_mm_cache_keys_is_idempotent():
+    """Ensure that the cache doesn't reinitialize old keys."""
+    cache = get_omni_pcache()
+    mm_key = "foo"
+    feat_dims = {mm_key: 100}
+    mm_outputs = get_multimodal_outputs(
+        feat_dims,
+        seq_len=DEFAULT_SEQ_LEN,
+    )
+    cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN)
+    assert len(cache.mm_cache_keys) == 1
+    assert mm_key in cache.mm_cache_keys
+
+    # Cache is initialized to 0 - fill it with 1s
+    cache.mm_outputs_cache[mm_key].fill_(1)
+
+    # Ensure that running another initialization
+    # doesn't zero out our cache values
+    cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN)
+    assert len(cache.mm_cache_keys) == 1
+    assert mm_key in cache.mm_cache_keys
+    assert torch.all(cache.mm_outputs_cache[mm_key] == 1)
+
+
+### Tests for Update
+def test_update_no_multimodal():
+    """Test that slot mappings act as row indices hidden states."""
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_offset = 8
+    slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
+    new_hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=new_hidden_states,
+        multimodal_outputs=None,
+        num_tokens_unpadded=num_tokens_unpadded,
+        slot_mapping=slot_mapping,
+    )
+
+    # Ensure that if we reshape our 3D cache back to 2D, we can use the
+    # indices in our slot mappings to access the hidden states as expected
+    hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE)
+    for slot_idx, new_states in zip(slot_mapping, new_hidden_states):
+        slot_states = hs_rows[slot_idx]
+        assert torch.all(slot_states == new_states)
+
+
+@pytest.mark.parametrize(
+    "feat_dims",
+    [
+        {"foo": 100, "bar": 100},
+        {"foo": 100, "bar": 50, "baz": 10},
+    ],
+)
+def test_update_with_multimodal_outputs(feat_dims):
+    """Test that slot mappings are correct for multimodal tensors."""
+    cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
+
+    num_tokens_unpadded = 8
+    slot_offset = 8
+    slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
+    feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
+    mm_outputs = {key: torch.rand((num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys}
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=None,
+        multimodal_outputs=mm_outputs,
+        num_tokens_unpadded=num_tokens_unpadded,
+        slot_mapping=slot_mapping,
+    )
+
+    for mm_key in feat_dims.keys():
+        assert mm_key in cache.mm_outputs_cache
+        key_feat_dim = feature_dims[mm_key]
+        mm_state_rows = cache.mm_outputs_cache[mm_key].view(NUM_BLOCKS * BLOCK_SIZE, key_feat_dim)
+
+        # Similar to hidden states, but for each key in the dict;
+        # Different tensors may have different feature dims
+        new_mm_outputs = mm_outputs[mm_key]
+        for slot_idx, new_output in zip(slot_mapping, new_mm_outputs):
+            slot_states = mm_state_rows[slot_idx]
+            assert torch.all(slot_states == new_output)
+
+
+### Tests for Merging
+def fake_get_cached_block_ids(self, req_idx, *args, **kwargs):
+    """Fake block table lookup.
+
+    Assumption:
+        req_idx 0 is a cache hit with slots 8, 9, ..., 15
+        req_idx 1 is a cache miss
+    """
+    assert req_idx < 2
+    if req_idx == 0:
+        # With the slot offset we provided (8), the corresponding
+        # blocks IDs are 2 & 3 because the block size is 4.
+        return torch.tensor([2, 3], dtype=torch.long)
+    return torch.tensor([], dtype=torch.long)
+
+
+@pytest.mark.parametrize("num_tokens_padded", [None, 16])
+def test_get_merged_hidden_states(num_tokens_padded, mocker):
+    """Ensure that hidden states are merged correctly."""
+    cache = get_omni_pcache()
+
+    orig_num_tokens_unpadded = 8
+    slot_offset = 8  # We'll put our states in slots 8, 9, 10, ..., 15
+    orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
+    orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=orig_hidden_states,
+        multimodal_outputs=None,
+        num_tokens_unpadded=orig_num_tokens_unpadded,
+        slot_mapping=orig_slot_mapping,
+        num_tokens_padded=num_tokens_padded,
+    )
+
+    # Say that we have two requests, but only one of them is a cache hit
+    num_new_toks_req1 = 3
+    num_new_toks_req2 = 2
+    cache.add_prefix_cached_new_req_id("req1")
+
+    num_scheduled_tokens = {
+        "req1": num_new_toks_req1,
+        "req2": num_new_toks_req2,
+    }
+    new_hidden_states = torch.rand(
+        (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE),
+        dtype=DTYPE,
+    )
+    req1_new_states = new_hidden_states[:num_new_toks_req1]
+    req2_new_states = new_hidden_states[-num_new_toks_req2:]
+
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    merged_states = cache.get_merged_hidden_states(
+        query_start_loc=[0, num_new_toks_req1],
+        input_batch=input_batch,
+        hidden_states=new_hidden_states,
+        num_scheduled_tokens=num_scheduled_tokens,
+    )
+
+    assert "req1" in merged_states and "req2" in merged_states
+    req1_merged_states = merged_states["req1"]
+    req2_merged_states = merged_states["req2"]
+
+    # First, check the cache hit case
+    assert req1_merged_states.shape == torch.Size([orig_num_tokens_unpadded + num_new_toks_req1, HIDDEN_SIZE])
+    # Ensure that the req1 merged states are the cached states + the new req1 states
+    assert torch.all(req1_merged_states[:orig_num_tokens_unpadded] == orig_hidden_states)
+    assert torch.all(req1_merged_states[-num_new_toks_req1:] == req1_new_states)
+
+    # Next, ensure that the cache miss case only has the new states
+    assert req2_merged_states.shape == torch.Size([num_new_toks_req2, HIDDEN_SIZE])
+    assert torch.all(req2_merged_states == req2_new_states)
+
+
+@pytest.mark.parametrize("num_tokens_padded", [None, 16])
+@pytest.mark.parametrize(
+    "feat_dims",
+    [
+        {"foo": 100, "bar": 100},
+        {"foo": 100, "bar": 50, "baz": 10},
+    ],
+)
+def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker):
+    cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
+
+    orig_num_tokens_unpadded = 8
+    slot_offset = 8  # We'll put our states in slots 8, 9, 10, ..., 15
+    orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
+    feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
+    orig_mm_outputs = {
+        key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
+    }
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=None,
+        multimodal_outputs=orig_mm_outputs,
+        num_tokens_unpadded=orig_num_tokens_unpadded,
+        slot_mapping=orig_slot_mapping,
+        num_tokens_padded=num_tokens_padded,
+    )
+
+    # Similar to hs test- say that we have two requests, but only one of them is a cache hit
+    num_new_toks_req1 = 3
+    num_new_toks_req2 = 2
+    cache.add_prefix_cached_new_req_id("req1")
+
+    num_scheduled_tokens = {
+        "req1": num_new_toks_req1,
+        "req2": num_new_toks_req2,
+    }
+
+    new_mm_outputs = {}
+    for mm_key in cache.mm_cache_keys:
+        new_mm_outputs[mm_key] = torch.rand(
+            (num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]),
+            dtype=DTYPE,
+        )
+    # We also want to make sure passthrough data (outside of our keys) isn't dropped
+    new_mm_outputs["passthrough_data"] = "Something else"
+    # Lists are a special case because we can't split them yet if we want to match
+    # the nonprefix cache behavior, because this runs before post process.
+    new_mm_outputs["passthrough_list"] = ["should", "not", "split"]
+
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    merged_mm_outputs = cache.get_merged_multimodal_states(
+        query_start_loc=[0, num_new_toks_req1],
+        input_batch=input_batch,
+        multimodal_outputs=new_mm_outputs,
+        num_scheduled_tokens=num_scheduled_tokens,
+    )
+
+    # Ensure the passthrough data wasn't dropped
+    assert "passthrough_data" in merged_mm_outputs
+    assert "passthrough_list" in merged_mm_outputs
+
+    for mm_key, mm_output in merged_mm_outputs.items():
+        # Ensure passthrough data is just forwarded normally and not duplicated
+        assert isinstance(mm_output, dict)
+        assert "req1" in mm_output and "req2" in mm_output
+        if mm_key == "passthrough_data":
+            assert mm_key not in cache.mm_cache_keys
+            assert new_mm_outputs[mm_key] == mm_output["req1"]
+            assert new_mm_outputs[mm_key] == mm_output["req2"]
+        elif mm_key == "passthrough_list":
+            assert mm_key not in cache.mm_cache_keys
+            assert new_mm_outputs[mm_key] == mm_output["req1"]
+            assert new_mm_outputs[mm_key] == mm_output["req2"]
+        else:
+            assert mm_key in cache.mm_cache_keys
+            curr_feat_dim = feature_dims[mm_key]
+            # Ensure that req1 (cache hit) merged the mm data
+            req1_merged_mm_outputs = mm_output["req1"]
+            req1_new_mm_outputs = new_mm_outputs[mm_key][:num_new_toks_req1]
+
+            assert req1_merged_mm_outputs.shape == torch.Size(
+                [orig_num_tokens_unpadded + num_new_toks_req1, curr_feat_dim]
+            )
+            # Ensure that the req1 merged mm data are the cached data + the new data
+            assert torch.all(req1_merged_mm_outputs[:orig_num_tokens_unpadded] == orig_mm_outputs[mm_key])
+            assert torch.all(req1_merged_mm_outputs[-num_new_toks_req1:] == req1_new_mm_outputs)
+
+            # Ensure that req2 (cache miss) only has the new mm data
+            req2_merged_mm_outputs = mm_output["req2"]
+            req2_new_mm_outputs = new_mm_outputs[mm_key][-num_new_toks_req2:]
+
+            assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
+            assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)

tests/dfx/conftest.py

@@ -2,7 +2,9 @@ import json
 from pathlib import Path
 from typing import Any
 
-from tests.conftest import modify_stage_config
+import pytest
+
+from tests.helpers.stage_config import modify_stage_config
 
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
@@ -38,22 +40,32 @@ def modify_stage(default_path, updates, deletes):
 def create_unique_server_params(
     configs: list[dict[str, Any]],
     stage_configs_dir: Path,
-) -> list[tuple[str, str, str]]:
+) -> list[tuple[str, str, str | None, str | None, tuple[str, ...]]]:
     unique_params = []
     seen = set()
     for config in configs:
         test_name = config["test_name"]
-        model = config["server_params"]["model"]
-        stage_config_name = config["server_params"].get("stage_config_name")
+        server_params = config["server_params"]
+        model = server_params["model"]
+        stage_config_name = server_params.get("stage_config_name")
         if stage_config_name:
             stage_config_path = str(stage_configs_dir / stage_config_name)
-            delete = config["server_params"].get("delete", None)
-            update = config["server_params"].get("update", None)
+            delete = server_params.get("delete", None)
+            update = server_params.get("update", None)
             stage_config_path = modify_stage(stage_config_path, update, delete)
         else:
             stage_config_path = None
 
-        server_param = (test_name, model, stage_config_path)
+        stage_overrides = server_params.get("stage_overrides")
+        stage_overrides_json = json.dumps(stage_overrides) if stage_overrides else None
+
+        # ``extra_cli_args`` passes raw CLI flags straight through to
+        # ``vllm_omni.entrypoints.cli.main serve`` — used for flags that
+        # don't map to stage-level overrides, e.g. ``--async-chunk`` /
+        # ``--no-async-chunk`` toggling the deploy-level async_chunk bool.
+        extra_cli_args = tuple(server_params.get("extra_cli_args") or ())
+
+        server_param = (test_name, model, stage_config_path, stage_overrides_json, extra_cli_args)
         if server_param not in seen:
             seen.add(server_param)
             unique_params.append(server_param)
@@ -95,3 +107,13 @@ def create_benchmark_indices(
                 indices.append((test_name, idx))
 
     return indices
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    """Register shared CLI options for DFX benchmark suites."""
+    parser.addoption(
+        "--test-config-file",
+        action="store",
+        default=None,
+        help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
+    )

tests/dfx/perf/scripts/diffusion_result_template.json

@@ -0,0 +1,86 @@
+[
+  {
+    "test_name": null,
+    "backend": null,
+    "timestamp": null,
+    "server_params": {
+      "model": null,
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": false
+      }
+    },
+    "benchmark_params": {
+      "name": null,
+      "dataset": null,
+      "task": null,
+      "width": 0,
+      "height": 0,
+      "num-inference-steps": 0,
+      "num-prompts": 0,
+      "max-concurrency": 0,
+      "num-input-images": 0,
+      "enable-negative-prompt": false,
+      "baseline": {
+        "throughput_qps": 0,
+        "latency_mean": 0,
+        "peak_memory_mb_max": 0,
+        "peak_memory_mb_mean": 0
+      }
+    },
+    "result": {
+      "duration": 0,
+      "completed_requests": 0,
+      "failed_requests": 0,
+      "throughput_qps": 0,
+      "latency_mean": 0,
+      "latency_median": 0,
+      "latency_p99": 0,
+      "latency_p95": 0,
+      "latency_p50": 0,
+      "peak_memory_mb_max": 0,
+      "peak_memory_mb_mean": 0,
+      "peak_memory_mb_median": 0,
+      "stage_durations_mean": {},
+      "stage_durations_p50": {},
+      "stage_durations_p99": {},
+      "backend": null,
+      "model": null,
+      "dataset": null,
+      "task": null
+    },
+    "log_file": null,
+    "Model": null,
+    "Framework": null,
+    "Hardware": null,
+    "Deployment": null,
+    "Task": null,
+    "Dataset": null,
+    "resolution": null,
+    "Parallelism": null,
+    "max_concurrency": 0,
+    "Cache": null,
+    "Quantization": null,
+    "offload": null,
+    "compile": null,
+    "Attn_backend": null,
+    "num_inference_steps": 0,
+    "completed": 0,
+    "failed": 0,
+    "throughput_qps": 0,
+    "latency_mean": 0,
+    "latency_median": 0,
+    "latency_p99": 0,
+    "latency_p95": 0,
+    "latency_p50": 0,
+    "peak_memory_mb_max": 0,
+    "peak_memory_mb_mean": 0,
+    "peak_memory_mb_median": 0,
+    "stage_durations_mean": {},
+    "stage_durations_p50": {},
+    "stage_durations_p99": {},
+    "commit_sha": null,
+    "build_id": null,
+    "build_url": null,
+    "source_file": null
+  }
+]

tests/dfx/perf/scripts/result_omni_template.json

@@ -0,0 +1,55 @@
+{
+  "date": null,
+  "endpoint_type": null,
+  "backend": null,
+  "label": null,
+  "model_id": null,
+  "tokenizer_id": null,
+  "num_prompts": 0,
+  "request_rate": null,
+  "burstiness": 0,
+  "max_concurrency": 0,
+  "duration": 0,
+  "completed": 0,
+  "failed": 0,
+  "total_input_tokens": 0,
+  "total_output_tokens": 0,
+  "request_throughput": 0,
+  "request_goodput": null,
+  "output_throughput": 0,
+  "total_token_throughput": 0,
+  "total_audio_duration_s": 0,
+  "total_audio_frames": 0,
+  "audio_throughput": 0,
+  "max_output_tokens_per_s": 0,
+  "max_concurrent_requests": 0,
+  "rtfx": 0,
+  "mean_ttft_ms": 0,
+  "median_ttft_ms": 0,
+  "p99_ttft_ms": 0,
+  "mean_tpot_ms": 0,
+  "median_tpot_ms": 0,
+  "p99_tpot_ms": 0,
+  "mean_itl_ms": 0,
+  "median_itl_ms": 0,
+  "p99_itl_ms": 0,
+  "mean_e2el_ms": 0,
+  "median_e2el_ms": 0,
+  "p99_e2el_ms": 0,
+  "mean_audio_rtf": 0,
+  "median_audio_rtf": 0,
+  "p99_audio_rtf": 0,
+  "mean_audio_ttfp_ms": 0,
+  "median_audio_ttfp_ms": 0,
+  "p99_audio_ttfp_ms": 0,
+  "mean_audio_duration_s": 0,
+  "median_audio_duration_s": 0,
+  "p99_audio_duration_s": 0,
+  "baseline": {
+    "mean_ttft_ms": 0,
+    "mean_audio_ttfp_ms": 0,
+    "mean_audio_rtf": 0
+  },
+  "random_input_len": 0,
+  "random_output_len": 0
+}

tests/dfx/perf/scripts/run_benchmark.py

@@ -8,7 +8,6 @@ from typing import Any
 
 import pytest
 
-from tests.conftest import OmniServer
 from tests.dfx.conftest import (
     create_benchmark_indices,
     create_test_parameter_mapping,
@@ -16,18 +15,41 @@ from tests.dfx.conftest import (
     get_benchmark_params_for_server,
     load_configs,
 )
+from tests.helpers.runtime import OmniServer
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 
-CONFIG_FILE_PATH = str(Path(__file__).parent.parent / "tests" / "test.json")
+def _get_config_file_from_argv() -> str | None:
+    """Read ``--test-config-file`` from ``sys.argv`` at import time so parametrization can use it."""
+    import sys
+
+    for i, arg in enumerate(sys.argv):
+        if arg == "--test-config-file" and i + 1 < len(sys.argv):
+            return sys.argv[i + 1]
+        if arg.startswith("--test-config-file="):
+            return arg.split("=", 1)[1]
+    return None
+
+
+_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
+_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json")
+
+CONFIG_FILE_PATH = _get_config_file_from_argv()
+if CONFIG_FILE_PATH is None:
+    print(
+        "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json "
+        "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)"
+    )
+    CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
+
 BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-STAGE_INIT_TIMEOUT = 600
+OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json"
 
 
-STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs"
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR)
+DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
+test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS)
 
 _omni_server_lock = threading.Lock()
@@ -40,13 +62,19 @@ def omni_server(request):
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        test_name, model, stage_config_path = request.param
+        test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args = ["--stage-init-timeout", str(STAGE_INIT_TIMEOUT), "--init-timeout", "900"]
+        server_args = ["--stage-init-timeout", "600", "--init-timeout", "900"]
+        # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
+        # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
         if stage_config_path:
-            server_args = ["--stage-configs-path", stage_config_path] + server_args
+            server_args = ["--deploy-config", stage_config_path] + server_args
+        if stage_overrides:
+            server_args = ["--stage-overrides", stage_overrides] + server_args
+        if extra_cli_args:
+            server_args = list(extra_cli_args) + server_args
         with OmniServer(model, server_args) as server:
             server.test_name = test_name
             print("OmniServer started successfully")
@@ -56,24 +84,47 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
+def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
+    """Make a single path segment safe for result filenames on common filesystems."""
+    if value is None:
+        return default
+    s = str(value).strip()
+    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
+        s = s.replace(bad, "_")
+    return s if s else default
+
+
 def run_benchmark(
     args: list,
     test_name: str,
     flow,
     dataset_name: str,
     num_prompt,
+    *,
+    baseline_config: dict[str, Any] | None = None,
+    sweep_index: int | None = None,
+    request_rate: Any | None = None,
+    max_concurrency: Any | None = None,
+    random_input_len: Any | None = None,
+    random_output_len: Any | None = None,
 ) -> Any:
-    """Run a single benchmark iteration and return the parsed result JSON."""
+    """Run a single benchmark iteration and return the parsed result JSON.
+
+    After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the same
+    per-metric resolved thresholds as ``assert_result`` (via ``_baseline_thresholds_for_step``).
+    When ``random_input_len`` / ``random_output_len`` are set, they are also written into the result JSON;
+    omitted keys when not configured.
+    """
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
-    result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_{current_dt}.json"
+    ri = _safe_filename_token(random_input_len)
+    ro = _safe_filename_token(random_output_len)
+    result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
     if "--result-filename" in args:
         print(f"The result file will be overwritten by {result_filename}")
     command = (
         ["vllm", "bench", "serve", "--omni"]
         + args
         + [
-            "--num-warmups",
-            "2",
             "--save-result",
             "--result-dir",
             os.environ.get("BENCHMARK_DIR", "tests"),
@@ -97,8 +148,34 @@ def run_benchmark(
     else:
         result_dir = "./"
 
-    with open(os.path.join(result_dir, result_filename), encoding="utf-8") as f:
-        result = json.load(f)
+    result_path = os.path.join(result_dir, result_filename)
+    if not os.path.exists(result_path):
+        with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
+            template_result: dict[str, Any] = json.load(f)
+        Path(result_path).parent.mkdir(parents=True, exist_ok=True)
+        with open(result_path, "w", encoding="utf-8") as f:
+            json.dump(template_result, f, ensure_ascii=False, indent=2)
+        print(f"Benchmark result file not generated, fallback to template: {result_path}")
+        result = template_result
+    else:
+        with open(result_path, encoding="utf-8") as f:
+            result = json.load(f)
+
+    if baseline_config:
+        result["baseline"] = _baseline_thresholds_for_step(
+            baseline_config,
+            sweep_index=sweep_index,
+            request_rate=request_rate,
+            max_concurrency=max_concurrency,
+        )
+    else:
+        result["baseline"] = {}
+    if random_input_len is not None:
+        result["random_input_len"] = random_input_len
+    if random_output_len is not None:
+        result["random_output_len"] = random_output_len
+    with open(result_path, "w", encoding="utf-8") as f:
+        json.dump(result, f, ensure_ascii=False, indent=2)
     return result
 
 
@@ -168,6 +245,25 @@ def _resolve_baseline_value(
     return baseline_raw
 
 
+def _baseline_thresholds_for_step(
+    baseline_data: dict[str, Any],
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> dict[str, Any]:
+    """Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``)."""
+    return {
+        metric_name: _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
+        for metric_name, baseline_raw in baseline_data.items()
+    }
+
+
 def assert_result(
     result,
     params,
@@ -258,6 +354,12 @@ def test_performance_benchmark(omni_server, benchmark_params):
             flow=qps,
             dataset_name=dataset_name,
             num_prompt=num_prompt,
+            baseline_config=params.get("baseline"),
+            sweep_index=i,
+            request_rate=qps,
+            max_concurrency=None,
+            random_input_len=params.get("random_input_len"),
+            random_output_len=params.get("random_output_len"),
         )
         assert_result(
             result,
@@ -276,6 +378,12 @@ def test_performance_benchmark(omni_server, benchmark_params):
             flow=concurrency,
             dataset_name=dataset_name,
             num_prompt=num_prompt,
+            baseline_config=params.get("baseline"),
+            sweep_index=i,
+            request_rate=None,
+            max_concurrency=concurrency,
+            random_input_len=params.get("random_input_len"),
+            random_output_len=params.get("random_output_len"),
         )
         assert_result(
             result,

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -5,8 +5,8 @@ Supports vLLM-Omni server backend:
   - vllm-omni (default): starts DiffusionServer via vllm_omni.entrypoints.cli.main,
     benchmarks with diffusion_benchmark_serving.py --backend vllm-omni
 
-A config JSON file is REQUIRED via --config-file:
-  pytest run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+A config JSON file is REQUIRED via --test-config-file:
+  pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
 JSON config entries use a "server_type" field, and this runner executes
 the vllm-omni path.
@@ -52,19 +52,20 @@ BENCHMARK_SCRIPT = str(
 _SESSION_TIMESTAMP = datetime.now().strftime("%Y%m%d-%H%M%S")
 _RESULT_LOCK = threading.Lock()
 _BRANCHPOINT_COMMIT_SHA: str | None = None
+DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json"
 
 
 def _get_config_file_from_argv() -> str | None:
-    """Read --config-file from sys.argv at import time so pytest parametrize can use it.
+    """Read --test-config-file from sys.argv at import time so pytest parametrize can use it.
 
     pytest_addoption (below) registers the same flag so pytest does not reject it.
-    Supports both ``--config-file path`` and ``--config-file=path`` forms.
+    Supports both ``--test-config-file path`` and ``--test-config-file=path`` forms.
     Returns None if the flag is not present; callers must handle the missing case.
     """
     for i, arg in enumerate(sys.argv):
-        if arg == "--config-file" and i + 1 < len(sys.argv):
+        if arg == "--test-config-file" and i + 1 < len(sys.argv):
             return sys.argv[i + 1]
-        if arg.startswith("--config-file="):
+        if arg.startswith("--test-config-file="):
             return arg.split("=", 1)[1]
     return None
 
@@ -133,19 +134,6 @@ def _append_to_aggregated_file(record: dict[str, Any]) -> None:
             json.dump(records, f, indent=2, ensure_ascii=False)
 
 
-# Register --config-file with pytest so it does not reject the argument.
-def pytest_addoption(parser: pytest.Parser) -> None:
-    parser.addoption(
-        "--config-file",
-        action="store",
-        default=None,
-        help=(
-            "Path to the benchmark config JSON file (required). "
-            "Example: --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json"
-        ),
-    )
-
-
 _server_lock = threading.Lock()
 
 # ---------------------------------------------------------------------------
@@ -576,10 +564,17 @@ def run_benchmark(
 
     if process.returncode != 0:
         tmp_result_file.unlink(missing_ok=True)
-        raise RuntimeError(f"Benchmark script exited with code {process.returncode}")
+        print(f"ERROR:Benchmark script exited with code {process.returncode}")
 
     if not tmp_result_file.exists():
-        raise FileNotFoundError(f"Benchmark result file not found: {tmp_result_file}")
+        with open(DIFFUSION_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
+            template_payload = json.load(f)
+        # Template schema is fixed and owned by this repo:
+        # ``diffusion_result_template.json`` is a one-item list and metrics live at [0]["result"].
+        template_metrics: dict[str, Any] = template_payload[0]["result"]
+        with open(tmp_result_file, "w", encoding="utf-8") as f:
+            json.dump(template_metrics, f, ensure_ascii=False, indent=2)
+        print(f"Benchmark result file not generated, fallback to template: {tmp_result_file}")
 
     try:
         with open(tmp_result_file, encoding="utf-8") as f:

tests/dfx/perf/stage_configs/qwen3_omni.yaml

@@ -1,101 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-# The following config has been verified on 2x H100-80G GPUs.
-async_chunk: false
-stage_args:
-  - stage_id: 0
-    stage_type: llm  # Use llm stage type for AR stages
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: thinker
-      max_num_seqs: 64
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: false
-      trust_remote_code: true
-      engine_output_type: latent  # Output hidden states for talker
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: llm  # Use llm stage type for AR stages
-    runtime:
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      max_num_seqs: 64
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
-      enforce_eager: false
-      trust_remote_code: true
-      engine_output_type: latent  # Output codec codes for code2wav
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      distributed_executor_backend: "mp"
-      hf_config_name: talker_config
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: False
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 2
-    stage_type: llm  # Use llm stage type for AR stages
-    runtime:
-      devices: "1"
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 64
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 100000
-      hf_config_name: thinker_config
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

tests/dfx/perf/stage_configs/qwen3_tts.yaml

@@ -1,96 +0,0 @@
-# Stage config for running Qwen3-TTS with 2-stage architecture
-# Stage 0: Talker (text -> 8-layer RVQ codec codes)
-# Stage 1: Code2Wav (codec codes -> audio waveform)
-#
-# The following config has been verified on 1x H100-80G GPU.
-async_chunk: true
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_comprehension: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 4
-      model_stage: qwen3_tts
-      model_arch: Qwen3TTSTalkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      enforce_eager: false
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      engine_output_type: latent
-      gpu_memory_utilization: 0.3
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 512
-      max_model_len: 4096
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
-    output_connectors:
-      to_stage_1: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: false
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 4
-      model_stage: code2wav
-      model_arch: Qwen3TTSCode2Wav
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      engine_output_type: audio
-      gpu_memory_utilization: 0.2
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 8192
-      max_model_len: 32768
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
-    tts_args:
-      max_instructions_length: 500
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.0
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 4
-
-  connectors:
-    connector_of_shared_memory:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536
-        codec_streaming: true
-        connector_get_sleep_s: 0.01
-        connector_get_max_wait_first_chunk: 3000
-        connector_get_max_wait: 300
-        codec_chunk_frames: 25
-        codec_left_context_frames: 72
-
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -0,0 +1,49 @@
+[
+    {
+        "test_name": "test_qwen_image_layered_single_device",
+        "description": "Single-device baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Layered",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "640x640_steps20_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 640,
+                "height": 640,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.02,
+                    "latency_mean": 40.0,
+                    "peak_memory_mb_max": 70000,
+                    "peak_memory_mb_mean": 70000
+                }
+            },
+            {
+                "name": "1024x1024_steps35_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.005,
+                    "latency_mean": 80.0,
+                    "peak_memory_mb_max": 70000,
+                    "peak_memory_mb_mean": 70000
+                }
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -44,6 +44,52 @@
             }
         ]
     },
+    {
+        "test_name": "test_qwen_image_single_device_step_execution",
+        "description": "Single-device baseline (no parallelism) with step execution",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true,
+                "step-execution": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.30,
+                    "latency_mean": 3.50,
+                    "peak_memory_mb_mean": 67000
+                }
+            },
+            {
+                "name": "1536x1536_steps35",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.037,
+                    "latency_mean": 27.0,
+                    "peak_memory_mb_mean": 74000
+                }
+            }
+        ]
+    },
     {
         "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
         "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
@@ -59,22 +105,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1,
-                    "latency_mean": 2.7,
-                    "peak_memory_mb_mean": 61000
-                }
-            },
             {
                 "name": "1536x1536_steps35",
                 "dataset": "random",

tests/dfx/perf/tests/test_qwen_omni.json

@@ -3,7 +3,7 @@
         "test_name": "test_qwen3_omni",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_config_name": "qwen3_omni.yaml"
+            "extra_cli_args": ["--no-async-chunk"]
         },
         "benchmark_params": [
             {
@@ -109,25 +109,7 @@
         "test_name": "test_qwen3_omni_chunk",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_config_name": "qwen3_omni.yaml",
-            "update": {
-                "async_chunk": true,
-                "stage_args": {
-                    "0": {
-                        "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
-                    },
-                    "1": {
-                        "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
-                    }
-                }
-            },
-            "delete": {
-                "stage_args": {
-                    "2": [
-                        "custom_process_input_func"
-                    ]
-                }
-            }
+            "extra_cli_args": ["--async-chunk"]
         },
         "benchmark_params": [
             {
@@ -329,37 +311,5 @@
                 }
             }
         ]
-    },
-    {
-        "test_name": "test_qwen3_tts",
-        "server_params": {
-            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
-        },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-audio-speech",
-                "endpoint": "/v1/audio/speech",
-                "num_prompts": [
-                    10,
-                    40
-                ],
-                "max_concurrency": [
-                    1,
-                    4
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "extra_body": {
-                    "voice": "Vivian",
-                    "language": "English"
-                },
-                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_audio_ttfp_ms": [6000, 6000],
-                    "mean_audio_rtf": [0.3, 0.3]
-                }
-            }
-        ]
     }
 ]

tests/dfx/perf/tests/test_tts.json

@@ -0,0 +1,34 @@
+[
+    {
+        "test_name": "test_qwen3_tts",
+        "server_params": {
+            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-audio-speech",
+                "endpoint": "/v1/audio/speech",
+                "num_prompts": [
+                    10,
+                    40
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "extra_body": {
+                    "voice": "Vivian",
+                    "language": "English"
+                },
+                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_audio_ttfp_ms": [6000, 6000],
+                    "mean_audio_rtf": [0.3, 0.3]
+                }
+            }
+        ]
+    }
+]

tests/dfx/stability/conftest.py

@@ -1,125 +1,17 @@
-"""
-Stability-specific conftest: when pytest is executed under this directory,
-resource monitoring is started before each test and finalized after each test,
-so each stability test case gets its own HTML report (one report per case).
-No need to wrap pytest with `bash resource_monitor.sh run -- pytest ...`.
-"""
+"""Stability pytest hooks and fixtures."""
 
-import os
 import subprocess
 import sys
 import threading
-import time
-from pathlib import Path
 
 import pytest
 
-STABILITY_DIR = Path(__file__).resolve().parent
-RESOURCE_MONITOR_SCRIPT = STABILITY_DIR / "scripts" / "resource_monitor.sh"
-REPO_ROOT = STABILITY_DIR.parent.parent.parent
-
-
-def _start_resource_monitor():
-    """Start `resource_monitor.sh start` in the background and return `Popen` or `None`."""
-    if not RESOURCE_MONITOR_SCRIPT.is_file():
-        return None
-    try:
-        proc = subprocess.Popen(
-            ["bash", str(RESOURCE_MONITOR_SCRIPT), "start", "--backend", "gpu"],
-            cwd=str(REPO_ROOT),
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.PIPE,
-            start_new_session=True,
-        )
-        try:
-            proc.wait(timeout=2)
-            if proc.returncode != 0:
-                stderr = proc.stderr.read().decode("utf-8", errors="ignore") if proc.stderr else ""
-                if stderr.strip():
-                    sys.stderr.write(f"[Stability] Resource monitor failed to start: {stderr.strip()}\n")
-                return None
-        except subprocess.TimeoutExpired:
-            pass
-        return proc
-    except (FileNotFoundError, OSError):
-        return None
-
-
-def _get_monitor_data_root() -> Path:
-    data_root = os.environ.get("RESOURCE_MONITOR_DATA_ROOT") or os.environ.get("GPU_MONITOR_DATA_ROOT")
-    if data_root:
-        return Path(data_root)
-    return STABILITY_DIR / "gpu_monitor_data"
-
-
-def _wait_for_run_dir(timeout_sec: int = 10) -> Path | None:
-    data_root = _get_monitor_data_root()
-    run_id_file = data_root / "current_run_id"
-    deadline = time.time() + timeout_sec
-    while time.time() < deadline:
-        if run_id_file.is_file():
-            run_id = run_id_file.read_text(encoding="utf-8").strip()
-            if run_id:
-                run_dir = data_root / run_id
-                if run_dir.is_dir():
-                    return run_dir
-        time.sleep(0.5)
-    return None
-
-
-def _report_latest_gpu_samples(stop_event: threading.Event) -> None:
-    """Periodically print the latest sampled GPU line."""
-    log_interval = int(
-        os.environ.get("RESOURCE_MONITOR_LOG_INTERVAL") or os.environ.get("GPU_MONITOR_LOG_INTERVAL") or "15"
-    )
-    log_interval = max(log_interval, 1)
-    last_line = ""
-
-    time.sleep(min(log_interval, 5))
-    while not stop_event.wait(log_interval):
-        run_dir = _wait_for_run_dir(timeout_sec=1)
-        if run_dir is None:
-            continue
-        csv_file = run_dir / "gpu_metrics.csv"
-        if not csv_file.is_file():
-            continue
-        try:
-            lines = csv_file.read_text(encoding="utf-8").splitlines()
-        except OSError:
-            continue
-        if len(lines) <= 1:
-            continue
-        latest = lines[-1].strip()
-        if latest and latest != last_line:
-            last_line = latest
-            sys.stderr.write(f"[GPU] {latest}\n")
-
-
-def _finalize_resource_monitor() -> str | None:
-    """
-    Run `resource_monitor.sh finalize` for the current run and generate the report.
-    Returns the bundle dir path (for this test case's report) if successful, else None.
-    """
-    if not RESOURCE_MONITOR_SCRIPT.is_file():
-        return None
-    try:
-        result = subprocess.run(
-            ["bash", str(RESOURCE_MONITOR_SCRIPT), "finalize", "--backend", "gpu"],
-            cwd=str(REPO_ROOT),
-            capture_output=True,
-            text=True,
-            timeout=60,
-            check=False,
-        )
-        if result.returncode != 0:
-            return None
-        for line in (result.stdout or "").splitlines():
-            if line.startswith("GPU_MONITOR_BUNDLE_DIR=") or line.startswith("RESOURCE_MONITOR_BUNDLE_DIR="):
-                _, _, value = line.partition("=")
-                return value.strip() if value else None
-        return None
-    except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
-        return None
+from tests.dfx.stability.helpers import (
+    finalize_resource_monitor,
+    report_latest_gpu_samples,
+    start_resource_monitor,
+    wait_for_run_dir,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -128,19 +20,19 @@ def stability_resource_monitor_per_test(request: pytest.FixtureRequest):
     For each test under this directory: start GPU monitor before the test,
     then finalize after the test so this case gets its own report.html.
     """
-    proc = _start_resource_monitor()
+    proc = start_resource_monitor()
     stop_event = threading.Event()
     reporter: threading.Thread | None = None
 
     if proc is not None:
         reporter = threading.Thread(
-            target=_report_latest_gpu_samples,
+            target=report_latest_gpu_samples,
             args=(stop_event,),
             name="stability-resource-monitor-reporter",
             daemon=True,
         )
         reporter.start()
-        run_dir = _wait_for_run_dir(timeout_sec=5)
+        run_dir = wait_for_run_dir(timeout_sec=5)
         node_name = request.node.name
         if run_dir is not None:
             sys.stderr.write(f"[Stability] Resource monitor started for test: {node_name} | run dir: {run_dir}\n")
@@ -161,7 +53,7 @@ def stability_resource_monitor_per_test(request: pytest.FixtureRequest):
             except subprocess.TimeoutExpired:
                 proc.kill()
                 proc.wait()
-        bundle_dir = _finalize_resource_monitor()
+        bundle_dir = finalize_resource_monitor()
         node_name = request.node.name
         if bundle_dir:
             sys.stderr.write(f"[Stability] Report for test «{node_name}»: {bundle_dir}/report.html\n")

tests/dfx/stability/helpers.py

@@ -0,0 +1,117 @@
+"""Stability resource monitor helpers."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+STABILITY_DIR = Path(__file__).resolve().parent
+RESOURCE_MONITOR_SCRIPT = STABILITY_DIR / "scripts" / "resource_monitor.sh"
+REPO_ROOT = STABILITY_DIR.parent.parent.parent
+
+
+def start_resource_monitor():
+    """Start `resource_monitor.sh start` in the background and return `Popen` or `None`."""
+    if not RESOURCE_MONITOR_SCRIPT.is_file():
+        return None
+    try:
+        proc = subprocess.Popen(
+            ["bash", str(RESOURCE_MONITOR_SCRIPT), "start", "--backend", "gpu"],
+            cwd=str(REPO_ROOT),
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.PIPE,
+            start_new_session=True,
+        )
+        try:
+            proc.wait(timeout=2)
+            if proc.returncode != 0:
+                stderr = proc.stderr.read().decode("utf-8", errors="ignore") if proc.stderr else ""
+                if stderr.strip():
+                    sys.stderr.write(f"[Stability] Resource monitor failed to start: {stderr.strip()}\n")
+                return None
+        except subprocess.TimeoutExpired:
+            pass
+        return proc
+    except (FileNotFoundError, OSError):
+        return None
+
+
+def get_monitor_data_root() -> Path:
+    data_root = os.environ.get("RESOURCE_MONITOR_DATA_ROOT") or os.environ.get("GPU_MONITOR_DATA_ROOT")
+    if data_root:
+        return Path(data_root)
+    return STABILITY_DIR / "gpu_monitor_data"
+
+
+def wait_for_run_dir(timeout_sec: int = 10) -> Path | None:
+    data_root = get_monitor_data_root()
+    run_id_file = data_root / "current_run_id"
+    deadline = time.time() + timeout_sec
+    while time.time() < deadline:
+        if run_id_file.is_file():
+            run_id = run_id_file.read_text(encoding="utf-8").strip()
+            if run_id:
+                run_dir = data_root / run_id
+                if run_dir.is_dir():
+                    return run_dir
+        time.sleep(0.5)
+    return None
+
+
+def report_latest_gpu_samples(stop_event: threading.Event) -> None:
+    """Periodically print the latest sampled GPU line."""
+    log_interval = int(
+        os.environ.get("RESOURCE_MONITOR_LOG_INTERVAL") or os.environ.get("GPU_MONITOR_LOG_INTERVAL") or "15"
+    )
+    log_interval = max(log_interval, 1)
+    last_line = ""
+
+    time.sleep(min(log_interval, 5))
+    while not stop_event.wait(log_interval):
+        run_dir = wait_for_run_dir(timeout_sec=1)
+        if run_dir is None:
+            continue
+        csv_file = run_dir / "gpu_metrics.csv"
+        if not csv_file.is_file():
+            continue
+        try:
+            lines = csv_file.read_text(encoding="utf-8").splitlines()
+        except OSError:
+            continue
+        if len(lines) <= 1:
+            continue
+        latest = lines[-1].strip()
+        if latest and latest != last_line:
+            last_line = latest
+            sys.stderr.write(f"[GPU] {latest}\n")
+
+
+def finalize_resource_monitor() -> str | None:
+    """
+    Run `resource_monitor.sh finalize` for the current run and generate the report.
+    Returns the bundle dir path (for this test case's report) if successful, else None.
+    """
+    if not RESOURCE_MONITOR_SCRIPT.is_file():
+        return None
+    try:
+        result = subprocess.run(
+            ["bash", str(RESOURCE_MONITOR_SCRIPT), "finalize", "--backend", "gpu"],
+            cwd=str(REPO_ROOT),
+            capture_output=True,
+            text=True,
+            timeout=60,
+            check=False,
+        )
+        if result.returncode != 0:
+            return None
+        for line in (result.stdout or "").splitlines():
+            if line.startswith("GPU_MONITOR_BUNDLE_DIR=") or line.startswith("RESOURCE_MONITOR_BUNDLE_DIR="):
+                _, _, value = line.partition("=")
+                return value.strip() if value else None
+        return None
+    except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
+        return None

tests/dfx/stability/scripts/test_benchmark_stability.py

@@ -24,7 +24,6 @@ from typing import Any
 
 import pytest
 
-from tests.conftest import OmniServer
 from tests.dfx.conftest import (
     create_benchmark_indices,
     create_test_parameter_mapping,
@@ -33,9 +32,10 @@ from tests.dfx.conftest import (
     load_configs,
 )
 from tests.dfx.perf.scripts.run_benchmark import run_benchmark
+from tests.helpers.runtime import OmniServer
 
 STABILITY_DIR = Path(__file__).resolve().parent.parent
-STAGE_CONFIGS_DIR = STABILITY_DIR / "stage_configs"
+DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
 CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test.json")
 DEFAULT_NUM_PROMPTS_PER_BATCH = 20
 
@@ -45,7 +45,7 @@ try:
 except FileNotFoundError:
     BENCHMARK_CONFIGS = []
 
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
 server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
 
 _omni_server_lock = threading.Lock()
@@ -112,6 +112,8 @@ def _run_one_benchmark_batch(
             flow=flow,
             dataset_name=dataset_name,
             num_prompt=num_prompts,
+            random_input_len=params.get("random_input_len"),
+            random_output_len=params.get("random_output_len"),
         )
         return result
     except (FileNotFoundError, OSError) as e:
@@ -217,11 +219,20 @@ def omni_server(request):
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        test_name, model, stage_config_path = request.param
+        test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+        server_args = ["--stage-init-timeout", "120"]
+        # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
+        # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
+        if stage_config_path:
+            server_args = ["--deploy-config", stage_config_path] + server_args
+        if stage_overrides:
+            server_args = ["--stage-overrides", stage_overrides] + server_args
+        if extra_cli_args:
+            server_args = list(extra_cli_args) + server_args
+        with OmniServer(model, server_args) as server:
             server.test_name = test_name
             print("OmniServer started successfully")
             yield server

tests/dfx/stability/stage_configs/qwen3_omni.yaml

@@ -1,101 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-# The following config has been verified on 2x H100-80G GPUs.
-async_chunk: false
-stage_args:
-  - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
-    runtime:
-      devices: "0"
-      max_batch_size: 64
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: false
-      trust_remote_code: true
-      engine_output_type: latent  # Output hidden states for talker
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: llm  # Use llm stage type to launch OmniLLM
-    runtime:
-      devices: "1"
-      max_batch_size: 64
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
-      enforce_eager: false
-      trust_remote_code: true
-      engine_output_type: latent  # Output codec codes for code2wav
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      distributed_executor_backend: "mp"
-      hf_config_name: talker_config
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: False
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 2
-    stage_type: llm  # Use llm stage type to launch OmniLLM
-    runtime:
-      devices: "1"
-      max_batch_size: 64
-    engine_args:
-      model_stage: code2wav
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: false
-      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-      hf_config_name: thinker_config
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

tests/dfx/stability/tests/test.json

@@ -3,7 +3,11 @@
         "test_name": "test_qwen3_omni_stability",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_config_name": "qwen3_omni.yaml"
+            "stage_overrides": {
+                "2": {
+                    "max_num_batched_tokens": 1000000
+                }
+            }
         },
         "benchmark_params": [
             {
@@ -36,25 +40,12 @@
         "test_name": "test_qwen3_omni_stability_async_chunk",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_config_name": "qwen3_omni.yaml",
-            "update": {
-                "async_chunk": true,
-                "stage_args": {
-                    "0": {
-                        "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
-                    },
-                    "1": {
-                        "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
-                    }
+            "stage_overrides": {
+                "2": {
+                    "max_num_batched_tokens": 1000000
                 }
             },
-            "delete": {
-                "stage_args": {
-                    "2": [
-                        "custom_process_input_func"
-                    ]
-                }
-            }
+            "extra_cli_args": ["--async-chunk"]
         },
         "benchmark_params": [
             {

tests/diffusion/cache/test_cache_dit.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Model specific tests for CacheDiT enablement.
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SEPARATE_CFG_ENABLERS = [
+    cd_backend.enable_cache_for_ltx2,
+    cd_backend.enable_cache_for_wan22,
+    cd_backend.enable_cache_for_longcat_image,
+]
+
+SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
+
+
+@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS)
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
+    """Ensure that custom enablers for models with separate CFG pass
+    the param through to cache_dit correctly.
+
+    Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860
+    """
+    mock_pipeline = Mock()
+    enabler(mock_pipeline, SAMPLE_CACHE_CONFIG)
+
+    mock_cache_dit.enable_cache.assert_called_once()
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    assert adapter_kwargs["has_separate_cfg"] is True

tests/diffusion/cache/test_teacache_extractors.py

@@ -21,7 +21,7 @@ from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,

tests/diffusion/distributed/test_ulysses_uaa_perf.py

@@ -17,7 +17,7 @@ import pytest
 import torch
 import torch.distributed as dist
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.attention.parallel.ulysses import (
     _all_gather_int,
     _ulysses_all_to_all_any_o,

tests/diffusion/layers/test_norm.py

@@ -0,0 +1,453 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# ── Import tests ──
+
+
+def test_layernorm_import():
+    """Verify LayerNorm can be imported from the norm module."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm  # noqa: F401
+
+
+def test_rmsnorm_import():
+    """Verify RMSNorm can be imported from the norm module."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm  # noqa: F401
+
+
+# ── LayerNorm tests ──
+
+
+def test_layernorm_forward_shape():
+    """LayerNorm produces correct output shapes."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    batch = 2
+    seq_len = 4
+    norm = LayerNorm(dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    out = norm(x)
+
+    assert out.shape == (batch, seq_len, dim)
+
+
+def test_layernorm_forward_shape_2d():
+    """LayerNorm works with 2D input tensors."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    batch = 2
+    norm = LayerNorm(dim)
+
+    x = torch.randn(batch, dim)
+    out = norm(x)
+
+    assert out.shape == (batch, dim)
+
+
+def test_layernorm_preserves_dtype_fp32():
+    """LayerNorm preserves float32 dtype."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+
+    x = torch.randn(2, 4, dim, dtype=torch.float32)
+    out = norm(x)
+
+    assert out.dtype == torch.float32
+
+
+def test_layernorm_preserves_dtype_fp16():
+    """LayerNorm preserves float16 dtype."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+
+    x = torch.randn(2, 4, dim, dtype=torch.float16)
+    out = norm(x)
+
+    assert out.dtype == torch.float16
+
+
+def test_layernorm_preserves_dtype_bf16():
+    """LayerNorm preserves bfloat16 dtype."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+
+    x = torch.randn(2, 4, dim, dtype=torch.bfloat16)
+    out = norm(x)
+
+    assert out.dtype == torch.bfloat16
+
+
+def test_layernorm_without_elementwise_affine():
+    """LayerNorm works without elementwise_affine (no learned parameters)."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim, elementwise_affine=False)
+
+    assert norm.weight is None
+    assert norm.bias is None
+
+    x = torch.randn(2, 4, dim)
+    out = norm(x)
+
+    assert out.shape == (2, 4, dim)
+
+
+def test_layernorm_custom_eps():
+    """LayerNorm accepts custom epsilon value."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    eps = 1e-5
+    norm = LayerNorm(dim, eps=eps)
+
+    assert norm.eps == eps
+
+
+def test_layernorm_has_learnable_parameters():
+    """LayerNorm has learnable weight and bias by default."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+
+    assert norm.weight is not None
+    assert norm.bias is not None
+    assert norm.weight.shape == (dim,)
+    assert norm.bias.shape == (dim,)
+
+
+def test_layernorm_matches_fp32_reference():
+    """Verify LayerNorm produces identical output to FP32 nn.LayerNorm."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    eps = 1e-6
+    torch.manual_seed(42)
+
+    ours = LayerNorm(dim, eps=eps)
+    ref = torch.nn.LayerNorm(dim, eps=eps)
+
+    # Copy weights
+    ref.weight.data.copy_(ours.weight.data)
+    ref.bias.data.copy_(ours.bias.data)
+
+    x = torch.randn(2, 4, dim)
+
+    out_ours = ours(x)
+    out_ref = ref(x.float()).to(x.dtype)
+
+    torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)
+
+
+def test_layernorm_matches_diffusers_fp32layernorm():
+    """Verify LayerNorm produces identical output to diffusers FP32LayerNorm."""
+    from diffusers.models.normalization import FP32LayerNorm
+
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    eps = 1e-6
+    torch.manual_seed(42)
+
+    ours = LayerNorm(dim, eps=eps)
+    ref = FP32LayerNorm(dim, eps=eps)
+
+    # Copy weights
+    ref.weight.data.copy_(ours.weight.data)
+    ref.bias.data.copy_(ours.bias.data)
+
+    # Test with fp16 input to verify FP32 computation
+    x = torch.randn(2, 4, dim, dtype=torch.float16)
+
+    out_ours = ours(x)
+    out_ref = ref(x)
+
+    torch.testing.assert_close(out_ours, out_ref, atol=1e-3, rtol=1e-3)
+
+
+# ── RMSNorm tests ──
+
+
+def test_rmsnorm_forward_shape():
+    """RMSNorm produces correct output shapes."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    batch = 2
+    seq_len = 4
+    norm = RMSNorm(hidden_size)
+
+    x = torch.randn(batch, seq_len, hidden_size)
+    out = norm(x)
+
+    assert out.shape == (batch, seq_len, hidden_size)
+
+
+def test_rmsnorm_forward_shape_2d():
+    """RMSNorm works with 2D input tensors."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    batch = 2
+    norm = RMSNorm(hidden_size)
+
+    x = torch.randn(batch, hidden_size)
+    out = norm(x)
+
+    assert out.shape == (batch, hidden_size)
+
+
+def test_rmsnorm_preserves_dtype_fp32():
+    """RMSNorm preserves float32 dtype."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+
+    x = torch.randn(2, 4, hidden_size, dtype=torch.float32)
+    out = norm(x)
+
+    assert out.dtype == torch.float32
+
+
+def test_rmsnorm_preserves_dtype_fp16():
+    """RMSNorm preserves float16 dtype."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+
+    x = torch.randn(2, 4, hidden_size, dtype=torch.float16)
+    out = norm(x)
+
+    assert out.dtype == torch.float16
+
+
+def test_rmsnorm_preserves_dtype_bf16():
+    """RMSNorm preserves bfloat16 dtype."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+
+    x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16)
+    out = norm(x)
+
+    assert out.dtype == torch.bfloat16
+
+
+def test_rmsnorm_custom_eps():
+    """RMSNorm accepts custom epsilon value."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    eps = 1e-5
+    norm = RMSNorm(hidden_size, eps=eps)
+
+    assert norm.variance_epsilon == eps
+
+
+def test_rmsnorm_has_weight_parameter():
+    """RMSNorm has learnable weight parameter initialized to ones."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+
+    assert norm.weight is not None
+    assert norm.weight.shape == (hidden_size,)
+    torch.testing.assert_close(norm.weight, torch.ones(hidden_size))
+
+
+def test_rmsnorm_numerical_correctness():
+    """Verify RMSNorm produces numerically correct output."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    eps = 1e-6
+    torch.manual_seed(42)
+
+    norm = RMSNorm(hidden_size, eps=eps)
+    x = torch.randn(2, 4, hidden_size)
+
+    # Compute expected output manually
+    x_fp32 = x.to(torch.float32)
+    variance = x_fp32.pow(2).mean(-1, keepdim=True)
+    expected = x_fp32 * torch.rsqrt(variance + eps)
+    expected = norm.weight.to(torch.float32) * expected
+    expected = expected.to(x.dtype)
+
+    out = norm(x)
+
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_rmsnorm_matches_reference_implementation():
+    """Verify RMSNorm matches a reference implementation."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    def reference_rmsnorm(x, weight, eps):
+        """Reference RMSNorm implementation."""
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        out = x * torch.rsqrt(variance + eps)
+        out = weight.to(torch.float32) * out
+        return out.to(input_dtype)
+
+    hidden_size = 128
+    eps = 1e-6
+    torch.manual_seed(123)
+
+    norm = RMSNorm(hidden_size, eps=eps)
+
+    # Test with various dtypes
+    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+        x = torch.randn(4, 8, hidden_size, dtype=dtype)
+        expected = reference_rmsnorm(x, norm.weight, eps)
+        out = norm(x)
+        torch.testing.assert_close(out, expected, atol=1e-3, rtol=1e-3)
+
+
+# ── CustomOp dispatch tests ──
+
+
+def test_layernorm_inherits_from_customop():
+    """LayerNorm inherits from CustomOp for platform dispatch."""
+    from vllm_omni.diffusion.layers.custom_op import CustomOp
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    norm = LayerNorm(64)
+    assert isinstance(norm, CustomOp)
+
+
+def test_rmsnorm_inherits_from_customop():
+    """RMSNorm inherits from CustomOp for platform dispatch."""
+    from vllm_omni.diffusion.layers.custom_op import CustomOp
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    norm = RMSNorm(64)
+    assert isinstance(norm, CustomOp)
+
+
+def test_layernorm_has_platform_methods():
+    """LayerNorm has forward methods for each platform."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    norm = LayerNorm(64)
+
+    assert hasattr(norm, "forward_cuda")
+    assert hasattr(norm, "forward_hip")
+    assert hasattr(norm, "forward_xpu")
+    assert hasattr(norm, "forward_npu")
+    assert hasattr(norm, "forward_native")
+
+
+def test_rmsnorm_has_platform_methods():
+    """RMSNorm has forward methods for each platform."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    norm = RMSNorm(64)
+
+    assert hasattr(norm, "forward_cuda")
+    assert hasattr(norm, "forward_hip")
+    assert hasattr(norm, "forward_xpu")
+    assert hasattr(norm, "forward_npu")
+    assert hasattr(norm, "forward_native")
+
+
+def test_layernorm_forward_native_directly():
+    """LayerNorm.forward_native can be called directly."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+    x = torch.randn(2, 4, dim)
+
+    out = norm.forward_native(x)
+
+    assert out.shape == (2, 4, dim)
+
+
+def test_rmsnorm_forward_native_directly():
+    """RMSNorm.forward_native can be called directly."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+    x = torch.randn(2, 4, hidden_size)
+
+    out = norm.forward_native(x)
+
+    assert out.shape == (2, 4, hidden_size)
+
+
+# ── Edge case tests ──
+
+
+def test_layernorm_with_large_dim():
+    """LayerNorm works with large hidden dimensions."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 4096
+    norm = LayerNorm(dim)
+    x = torch.randn(1, 16, dim)
+
+    out = norm(x)
+
+    assert out.shape == (1, 16, dim)
+
+
+def test_rmsnorm_with_large_dim():
+    """RMSNorm works with large hidden dimensions."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 4096
+    norm = RMSNorm(hidden_size)
+    x = torch.randn(1, 16, hidden_size)
+
+    out = norm(x)
+
+    assert out.shape == (1, 16, hidden_size)
+
+
+def test_layernorm_with_single_element_batch():
+    """LayerNorm works with batch size of 1."""
+    from vllm_omni.diffusion.layers.norm import LayerNorm
+
+    dim = 64
+    norm = LayerNorm(dim)
+    x = torch.randn(1, 1, dim)
+
+    out = norm(x)
+
+    assert out.shape == (1, 1, dim)
+
+
+def test_rmsnorm_with_single_element_batch():
+    """RMSNorm works with batch size of 1."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    hidden_size = 64
+    norm = RMSNorm(hidden_size)
+    x = torch.randn(1, 1, hidden_size)
+
+    out = norm(x)
+
+    assert out.shape == (1, 1, hidden_size)

tests/diffusion/layers/test_rotary_emb_equivalence.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Numerical equivalence tests for rotary embedding implementations (#2436).
+
+Verifies that the optimized stack+flatten RoPE produces bit-identical results
+to the original strided-slice implementation across various tensor shapes and
+dtypes, ensuring the refactor is safe.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+
+def _apply_rotary_emb_helios_original(
+    hidden_states: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> torch.Tensor:
+    """Original Helios RoPE using strided slice assignment (pre-#2436)."""
+    x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+    cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
+    out = torch.empty_like(hidden_states)
+    out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2]
+    out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2]
+    return out.type_as(hidden_states)
+
+
+def _apply_rotary_emb_helios_optimized(
+    hidden_states: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> torch.Tensor:
+    """Optimized Helios RoPE using stack+flatten (post-#2436)."""
+    x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+    cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
+    rotated = torch.stack(
+        (
+            x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2],
+            x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2],
+        ),
+        dim=-1,
+    )
+    return rotated.flatten(-2, -1).type_as(hidden_states)
+
+
+def _make_inputs(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generate random hidden_states and freqs_cis for testing."""
+    torch.manual_seed(42)
+    hidden_states = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
+    # freqs_cis: [B, seq, head_dim*2] — cos and sin concatenated along last dim
+    freqs_cis = torch.randn(batch, seq_len, head_dim * 2, dtype=dtype)
+    return hidden_states, freqs_cis
+
+
+class TestHeliosRoPEEquivalence:
+    """Verify optimized Helios RoPE is numerically identical to original."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_equivalence_across_dtypes(self, dtype: torch.dtype) -> None:
+        """Optimized output must be bit-identical to original across dtypes."""
+        hidden, freqs = _make_inputs(2, 16, 8, 64, dtype=dtype)
+        original = _apply_rotary_emb_helios_original(hidden, freqs)
+        optimized = _apply_rotary_emb_helios_optimized(hidden, freqs)
+        torch.testing.assert_close(optimized, original, atol=0, rtol=0)
+
+    @pytest.mark.parametrize(
+        "batch,seq_len,num_heads,head_dim",
+        [
+            (1, 8, 1, 32),  # minimal: single batch, single head
+            (2, 16, 8, 64),  # typical transformer config
+            (1, 8192, 4, 64),  # video-scale patch tokens (720p DiT)
+            (4, 32, 16, 128),  # large head_dim
+        ],
+    )
+    def test_equivalence_across_shapes(self, batch: int, seq_len: int, num_heads: int, head_dim: int) -> None:
+        """Equivalence must hold across different tensor shapes."""
+        hidden, freqs = _make_inputs(batch, seq_len, num_heads, head_dim)
+        original = _apply_rotary_emb_helios_original(hidden, freqs)
+        optimized = _apply_rotary_emb_helios_optimized(hidden, freqs)
+        torch.testing.assert_close(optimized, original, atol=0, rtol=0)
+
+    def test_output_contiguous(self) -> None:
+        """Optimized output should be contiguous in memory."""
+        hidden, freqs = _make_inputs(2, 16, 8, 64)
+        optimized = _apply_rotary_emb_helios_optimized(hidden, freqs)
+        assert optimized.is_contiguous()
+
+    def test_output_shape_preserved(self) -> None:
+        """Output shape must match input shape."""
+        hidden, freqs = _make_inputs(2, 16, 8, 64)
+        optimized = _apply_rotary_emb_helios_optimized(hidden, freqs)
+        assert optimized.shape == hidden.shape
+
+    def test_output_dtype_preserved(self) -> None:
+        """Output dtype must match input dtype."""
+        hidden, freqs = _make_inputs(2, 16, 8, 64, dtype=torch.float16)
+        optimized = _apply_rotary_emb_helios_optimized(hidden, freqs)
+        assert optimized.dtype == hidden.dtype
+
+    def test_odd_head_dim_raises(self) -> None:
+        """Odd head_dim should fail at unflatten (not a valid RoPE config)."""
+        hidden = torch.randn(1, 4, 2, 63)
+        freqs = torch.randn(1, 4, 126)
+        with pytest.raises(RuntimeError):
+            _apply_rotary_emb_helios_optimized(hidden, freqs)

tests/diffusion/lora/test_lora_manager.py

@@ -8,7 +8,7 @@ import torch
 from vllm.lora.lora_weights import LoRALayerWeights
 from vllm.lora.utils import get_supported_lora_modules
 
-from tests.diffusion.lora.conftest import (
+from tests.diffusion.lora.helpers import (
     DummyBaseLayerWithLoRA,
     FakeLinearBase,
     fake_replace_submodule,

tests/diffusion/models/bagel/test_bagel_lora.py

@@ -11,7 +11,7 @@ import pytest
 import torch
 from safetensors.torch import save_file
 
-from tests.diffusion.lora.conftest import (
+from tests.diffusion.lora.helpers import (
     DummyBaseLayerWithLoRA,
     FakeLinearBase,
     fake_replace_submodule,

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests).
+_DMD2_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+    LTX2T2VDMD2Pipeline: LTX2Pipeline,
+    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ with the base pipeline mocked out (no model weights loaded)."""
+
+    base = _DMD2_BASE[cls]
+    od_config = MagicMock()
+    od_config.model = "/nonexistent"
+
+    def _mock_base_init(self, *a, **kw):
+        self.od_config = od_config
+
+    with patch.object(base, "__init__", _mock_base_init):
+        pipeline = object.__new__(cls)
+        torch.nn.Module.__init__(pipeline)
+        cls.__init__(pipeline, od_config=od_config)
+    return pipeline
+
+
+def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
+    sp = OmniDiffusionSamplingParams(**sp_kwargs)
+    return OmniDiffusionRequest(
+        prompts=prompts or [{"prompt": "a cat dancing"}],
+        sampling_params=sp,
+    )
+
+
+@pytest.fixture(
+    params=list(_DMD2_BASE.keys()),
+    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v"],
+)
+def pipeline(request):
+    return _make_pipeline(request.param)
+
+
+# ---------------------------------------------------------------------------
+# num_inference_steps
+# ---------------------------------------------------------------------------
+
+
+def test_num_inference_steps_forced_to_dmd2_value(pipeline):
+    req = _make_request(num_inference_steps=40)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
+
+
+def test_num_inference_steps_already_correct(pipeline):
+    req = _make_request(num_inference_steps=pipeline.num_inference_steps)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
+
+
+# ---------------------------------------------------------------------------
+# guidance_scale
+# ---------------------------------------------------------------------------
+
+
+def test_guidance_scale_forced_to_one(pipeline):
+    req = _make_request(guidance_scale=5.0, guidance_scale_provided=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+    assert req.sampling_params.guidance_scale_provided is False
+
+
+def test_guidance_scale_already_correct(pipeline):
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=False)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+
+
+def test_guidance_scale_provided_flag_cleared(pipeline):
+    """guidance_scale_provided=True must be cleared even if scale is already dmd2_guidance_scale."""
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_provided is False
+
+
+def test_guidance_scale_2_cleared(pipeline):
+    req = _make_request(guidance_scale_2=3.0)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_2 is None
+
+
+def test_guidance_scale_2_unset_unchanged(pipeline):
+    req = _make_request()
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale_2 is None
+
+
+def test_true_cfg_scale_cleared(pipeline):
+    req = _make_request(true_cfg_scale=2.0)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.true_cfg_scale is None
+
+
+def test_do_classifier_free_guidance_forced_false(pipeline):
+    req = _make_request(do_classifier_free_guidance=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.do_classifier_free_guidance is False
+
+
+def test_is_cfg_negative_forced_false(pipeline):
+    req = _make_request(is_cfg_negative=True)
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.is_cfg_negative is False
+
+
+def test_negative_prompt_stripped_from_prompt_dict(pipeline):
+    req = _make_request(prompts=[{"prompt": "a cat", "negative_prompt": "blurry"}])
+    pipeline._sanitize_dmd2_request(req)
+    assert "negative_prompt" not in req.prompts[0]
+    assert req.prompts[0]["prompt"] == "a cat"
+
+
+def test_no_negative_prompt_unchanged(pipeline):
+    req = _make_request(prompts=[{"prompt": "a cat"}])
+    pipeline._sanitize_dmd2_request(req)
+    assert req.prompts[0] == {"prompt": "a cat"}
+
+
+def test_string_prompt_not_mutated(pipeline):
+    """String prompts (not dicts) must pass through unchanged."""
+    req = _make_request(prompts=["a cat dancing"])
+    pipeline._sanitize_dmd2_request(req)
+    assert req.prompts == ["a cat dancing"]
+
+
+def test_multiple_prompts_all_sanitized(pipeline):
+    req = _make_request(
+        prompts=[
+            {"prompt": "a cat", "negative_prompt": "blurry"},
+            {"prompt": "a dog", "negative_prompt": "ugly"},
+        ]
+    )
+    pipeline._sanitize_dmd2_request(req)
+    for p in req.prompts:
+        assert "negative_prompt" not in p
+
+
+# ---------------------------------------------------------------------------
+# Clean request — nothing changes
+# ---------------------------------------------------------------------------
+
+
+def test_clean_request_no_changes(pipeline):
+    req = _make_request(
+        guidance_scale=pipeline.dmd2_guidance_scale,
+        guidance_scale_provided=False,
+        do_classifier_free_guidance=False,
+        is_cfg_negative=False,
+    )
+    pipeline._sanitize_dmd2_request(req)
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
+    assert req.sampling_params.guidance_scale_provided is False
+    assert req.sampling_params.guidance_scale_2 is None
+    assert req.sampling_params.true_cfg_scale is None
+    assert req.sampling_params.do_classifier_free_guidance is False
+    assert req.sampling_params.is_cfg_negative is False

tests/diffusion/models/dmd2/test_dmd2_scheduler.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DMD2_TIMESTEPS = [999, 937, 833, 624]
+
+# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests).
+_DMD2_BASE = {
+    WanT2VDMD2Pipeline: Wan22Pipeline,
+    WanI2VDMD2Pipeline: Wan22I2VPipeline,
+    LTX2T2VDMD2Pipeline: LTX2Pipeline,
+    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+}
+
+
+def _make_pipeline(cls):
+    """Run the DMD2 __init__ (including __init_dmd2__) with the base pipeline mocked."""
+
+    base = _DMD2_BASE[cls]
+    od_config = MagicMock()
+    od_config.model = "/nonexistent"
+
+    def _mock_base_init(self, *a, **kw):
+        self.od_config = od_config  # __init_dmd2__ needs this
+
+    with patch.object(base, "__init__", _mock_base_init):
+        pipeline = object.__new__(cls)
+        torch.nn.Module.__init__(pipeline)
+        cls.__init__(pipeline, od_config=od_config)
+    return pipeline
+
+
+def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
+    sp = OmniDiffusionSamplingParams(**sp_kwargs)
+    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
+
+
+@pytest.fixture(
+    params=list(_DMD2_BASE.keys()),
+    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v"],
+)
+def pipeline(request):
+    return _make_pipeline(request.param)
+
+
+# ---------------------------------------------------------------------------
+# forward() timestep injection
+# ---------------------------------------------------------------------------
+
+
+def _fake_parent_forward(self, req, *args, num_inference_steps=40, **kwargs):
+    """Stub that calls set_timesteps as the real parent does."""
+    self.scheduler.set_timesteps(num_inference_steps, device="cpu")
+    return MagicMock()
+
+
+def test_forward_timesteps_match_dmd2_schedule(pipeline):
+    """After forward() runs, scheduler.timesteps must equal the DMD2 training schedule."""
+    parent = _DMD2_BASE[type(pipeline)]
+
+    # Baseline: calling set_timesteps(40) without the DMD2 override gives a different schedule
+    pipeline.scheduler.set_timesteps(40, device="cpu")
+    default_timesteps = pipeline.scheduler.timesteps.long().tolist()
+    assert default_timesteps == _DMD2_TIMESTEPS, (
+        "DMD2EulerScheduler should always return DMD2 timesteps regardless of num_steps"
+    )
+
+    with patch.object(parent, "forward", _fake_parent_forward):
+        pipeline.forward(_make_request())
+
+    assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS
+
+
+def test_forward_timesteps_idempotent_across_calls(pipeline):
+    """Successive forward() calls must not cause scheduler state to drift."""
+    parent = _DMD2_BASE[type(pipeline)]
+
+    with patch.object(parent, "forward", _fake_parent_forward):
+        pipeline.forward(_make_request())
+        pipeline.forward(_make_request())
+
+    assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -2,7 +2,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.models.flux2.flux2_transformer import (
     Flux2PosEmbed,
     Flux2Transformer2DModel,

tests/diffusion/models/glm_image/test_glm_image_sp.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for GLM-Image Sequence Parallelism support."""
+
+import pytest
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_sp_groups(mocker):
+    """Set up SP and TP groups for each test function."""
+    mock_get_sp_group = mocker.patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group")
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
+
+    mock_sp_group = mocker.MagicMock()
+    mock_sp_group.world_size = 4
+    mock_get_sp_group.return_value = mock_sp_group
+
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 1
+    mock_get_tp_group.return_value = mock_tp_group
+    yield
+
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_glm_image_sp_plan_defined():
+    """Test that _sp_plan is properly defined on GlmImageTransformer2DModel."""
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImageTransformer2DModel,
+    )
+
+    assert hasattr(GlmImageTransformer2DModel, "_sp_plan")
+    plan = GlmImageTransformer2DModel._sp_plan
+    assert plan is not None
+
+    # Verify plan structure
+    assert "prepare" in plan
+    assert "proj_out" in plan
+
+
+def test_glm_image_sp_plan_valid():
+    """Validate _sp_plan structure."""
+    from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImageTransformer2DModel,
+    )
+
+    plan = GlmImageTransformer2DModel._sp_plan
+    validate_sp_plan(plan)
+
+
+def test_glm_image_prepare_module_exists():
+    """Test that GlmImagePrepare module exists."""
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImagePrepare,
+    )
+
+    assert GlmImagePrepare is not None
+
+
+def test_glm_image_attention_accepts_parallel_config():
+    """Test that GlmImageAttention accepts parallel_config parameter."""
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImageAttention,
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=2,
+        ring_degree=2,
+        tensor_parallel_size=1,
+        sequence_parallel_size=4,
+    )
+
+    attn = GlmImageAttention(
+        dim=2560,
+        num_heads=64,
+        head_dim=40,
+        parallel_config=parallel_config,
+    )
+
+    assert attn.parallel_config is not None
+    assert attn.parallel_config.sequence_parallel_size == 4
+
+
+def test_glm_image_transformer_block_accepts_parallel_config():
+    """Test that GlmImageTransformerBlock accepts parallel_config parameter."""
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImageTransformerBlock,
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=2,
+        ring_degree=2,
+        tensor_parallel_size=1,
+        sequence_parallel_size=4,
+    )
+
+    block = GlmImageTransformerBlock(
+        dim=2560,
+        num_attention_heads=64,
+        attention_head_dim=40,
+        time_embed_dim=512,
+        parallel_config=parallel_config,
+    )
+
+    assert block.attn1.parallel_config is not None
+    assert block.attn1.parallel_config.sequence_parallel_size == 4
+
+
+def test_glm_image_has_sp_support():
+    """Test that GLM-Image has SP support implemented."""
+    from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+        GlmImageTransformer2DModel,
+    )
+
+    # Check that the model has parallel_config support
+    assert hasattr(GlmImageTransformer2DModel, "__init__")
+
+    # Verify the model can be instantiated with SP config
+
+    # This test just verifies the structure exists
+    # Actual SP testing requires multi-GPU setup
+
+
+@pytest.mark.cuda
+@pytest.mark.sp
+def test_glm_image_sp_inference():
+    """Test SP inference (requires multi-GPU setup)."""
+    pytest.skip("Requires multi-GPU SP setup")

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -12,7 +12,7 @@ class TestSetForwardContextNumTokens:
 
     def test_sets_num_tokens_when_context_available(self, mocker):
         """num_tokens should be set on ForwardContext when available."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         mock_ctx = mocker.MagicMock()
         del mock_ctx.in_profile_run  # simulate missing attr
@@ -26,7 +26,7 @@ class TestSetForwardContextNumTokens:
 
     def test_sets_in_profile_run_only_if_missing(self, mocker):
         """in_profile_run should not be overwritten if already set."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         mock_ctx = mocker.MagicMock()
         mock_ctx.in_profile_run = True  # already set
@@ -40,7 +40,7 @@ class TestSetForwardContextNumTokens:
 
     def test_noop_when_context_unavailable(self, mocker):
         """Should do nothing when ForwardContext is not available."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
         mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
@@ -55,11 +55,11 @@ class TestHunyuanFusedMoEPlatformDispatch:
 
     def test_default_platform_uses_default_impl_qualname(self, mocker):
         """HunyuanFusedMoE should resolve the impl class from the platform hook."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         mock_platform = mocker.MagicMock()
         mock_platform.get_diffusion_model_impl_qualname.return_value = (
-            "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
+            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
         )
 
         mocker.patch.object(
@@ -71,7 +71,7 @@ class TestHunyuanFusedMoEPlatformDispatch:
         mock_impl = mocker.MagicMock()
         mock_resolve.return_value = mock_impl
 
-        from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
+        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
             HunyuanFusedMoE,
         )
 
@@ -80,7 +80,7 @@ class TestHunyuanFusedMoEPlatformDispatch:
         mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe")
         mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
         mock_resolve.assert_called_once_with(
-            "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
+            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
         )
         mock_impl.assert_called_once_with(prefix="")
 
@@ -90,7 +90,7 @@ class TestHunyuanFusedMoEFactory:
 
     def test_new_delegates_to_impl_class(self, mocker):
         """HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         class MockImpl:
             def __init__(self, *, prefix: str = "", **kwargs):
@@ -104,7 +104,7 @@ class TestHunyuanFusedMoEFactory:
         mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1))
         mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
 
-        from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
+        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
             HunyuanFusedMoE,
         )
 
@@ -119,7 +119,7 @@ class TestHunyuanFusedMoEFactory:
 
     def test_make_expert_params_mapping_delegates_to_impl(self, mocker):
         """make_expert_params_mapping should delegate to impl class method."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
 
         expected_mapping = [("a", "b", 0, "c")]
         mock_platform = mocker.MagicMock()
@@ -130,7 +130,7 @@ class TestHunyuanFusedMoEFactory:
         mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping)
         mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
 
-        from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
+        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
             HunyuanFusedMoE,
         )
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanImage3 AR sampler logic (stage transitions,
+ratio restriction, comprehension blocking)."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Fake token IDs for testing (avoid importing the real model).
+END_OF_THINK = 100
+RECAPTION = 101
+END_OF_RECAPTION = 102
+ANSWER = 103
+BOI = 104
+SIZE_TOKEN = 105
+EOS = 106
+RATIO_START = 200
+RATIO_END = 210
+RATIO_OTHER_START = 220
+RATIO_OTHER_END = 223
+
+
+class FakeSamplerModel:
+    """Minimal stub that replicates the sampler-relevant attributes of
+    HunyuanImage3ForConditionalGeneration without loading real weights."""
+
+    def __init__(self, *, is_comprehension: bool = False):
+        self._is_comprehension = is_comprehension
+        self._eos_token_id = EOS
+        self._end_of_think_id = END_OF_THINK
+        self._recaption_id = RECAPTION
+        self._end_of_recaption_id = END_OF_RECAPTION
+        self._answer_id = ANSWER
+        self._mrope_boi_token_id = BOI
+        self._size_token_id = SIZE_TOKEN
+        self._start_ratio_id = RATIO_START
+        self._end_ratio_id = RATIO_END
+        self._ratio_other_slices = [(RATIO_OTHER_START, RATIO_OTHER_END + 1)]
+        self._all_ratio_ids = set(range(RATIO_START, RATIO_END + 1))
+        self._all_ratio_ids.update(range(RATIO_OTHER_START, RATIO_OTHER_END + 1))
+
+        self._stage_transitions: dict[int, list[int]] = {}
+        if not is_comprehension:
+            self._stage_transitions[END_OF_THINK] = [RECAPTION]
+            self._stage_transitions[END_OF_RECAPTION] = [ANSWER, BOI, SIZE_TOKEN]
+
+        self._blocked_token_ids: set[int] = set()
+        if is_comprehension:
+            self._blocked_token_ids.update([BOI, SIZE_TOKEN])
+            self._blocked_token_ids.update(self._all_ratio_ids)
+
+    # Bind the real methods from the model class.
+    from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
+        HunyuanImage3ForConditionalGeneration as _Real,
+    )
+
+    _get_forced_token = _Real._get_forced_token
+    _apply_ratio_restriction = _Real._apply_ratio_restriction
+
+
+class TestGetForcedToken:
+    """Tests for the stateless _get_forced_token method."""
+
+    def setup_method(self):
+        self.model = FakeSamplerModel(is_comprehension=False)
+
+    def test_no_trigger_returns_none(self):
+        assert self.model._get_forced_token([1, 2, 3]) is None
+
+    def test_empty_history_returns_none(self):
+        assert self.model._get_forced_token([]) is None
+
+    def test_end_of_think_forces_recaption(self):
+        assert self.model._get_forced_token([END_OF_THINK]) == RECAPTION
+
+    def test_end_of_think_completed(self):
+        assert self.model._get_forced_token([END_OF_THINK, RECAPTION]) is None
+
+    def test_end_of_recaption_forces_answer(self):
+        tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION]
+        assert self.model._get_forced_token(tokens) == ANSWER
+
+    def test_end_of_recaption_forces_boi_after_answer(self):
+        tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER]
+        assert self.model._get_forced_token(tokens) == BOI
+
+    def test_end_of_recaption_forces_size_after_boi(self):
+        tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI]
+        assert self.model._get_forced_token(tokens) == SIZE_TOKEN
+
+    def test_full_sequence_complete(self):
+        tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI, SIZE_TOKEN]
+        assert self.model._get_forced_token(tokens) is None
+
+    def test_diverged_history_returns_none(self):
+        tokens = [END_OF_RECAPTION, 999]  # 999 != ANSWER
+        assert self.model._get_forced_token(tokens) is None
+
+    def test_later_trigger_takes_precedence(self):
+        tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION]
+        assert self.model._get_forced_token(tokens) == ANSWER
+
+    def test_trigger_with_extra_tokens_before(self):
+        tokens = [1, 2, 3, END_OF_THINK]
+        assert self.model._get_forced_token(tokens) == RECAPTION
+
+
+class TestComprehensionBlocking:
+    """Tests for comprehension mode token blocking."""
+
+    def test_blocked_tokens_masked(self):
+        model = FakeSamplerModel(is_comprehension=True)
+        vocab_size = 300
+        logits = torch.zeros(1, vocab_size)
+        logits[0, BOI] = 5.0
+        logits[0, SIZE_TOKEN] = 3.0
+        logits[0, RATIO_START] = 2.0
+        min_score = torch.finfo(logits.dtype).min
+
+        for tid in model._blocked_token_ids:
+            if tid < vocab_size:
+                logits[0, tid] = min_score
+
+        assert logits[0, BOI].item() == min_score
+        assert logits[0, SIZE_TOKEN].item() == min_score
+        assert logits[0, RATIO_START].item() == min_score
+
+    def test_non_blocked_tokens_preserved(self):
+        model = FakeSamplerModel(is_comprehension=True)
+        vocab_size = 300
+        logits = torch.zeros(1, vocab_size)
+        logits[0, 50] = 7.0
+        min_score = torch.finfo(logits.dtype).min
+
+        for tid in model._blocked_token_ids:
+            if tid < vocab_size:
+                logits[0, tid] = min_score
+
+        assert logits[0, 50].item() == 7.0
+
+
+class TestRatioRestriction:
+    """Tests for _apply_ratio_restriction (greedy: only argmax ratio survives)."""
+
+    def test_greedy_selects_single_ratio_token(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        vocab_size = 300
+        logits = torch.zeros(1, vocab_size)
+        logits[0, RATIO_START + 3] = 10.0
+        logits[0, RATIO_START + 1] = 5.0
+        logits[0, 50] = 20.0  # non-ratio, should be masked
+        min_score = torch.finfo(logits.dtype).min
+
+        model._apply_ratio_restriction(logits, 0, min_score)
+
+        assert logits[0, RATIO_START + 3].item() == 0
+        assert logits[0, RATIO_START + 1].item() == min_score
+        assert logits[0, 50].item() == min_score
+
+    def test_extra_ratio_slices_considered(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        vocab_size = 300
+        logits = torch.zeros(1, vocab_size)
+        logits[0, RATIO_OTHER_START] = 15.0
+        logits[0, RATIO_START] = 5.0
+        min_score = torch.finfo(logits.dtype).min
+
+        model._apply_ratio_restriction(logits, 0, min_score)
+
+        assert logits[0, RATIO_OTHER_START].item() == 0
+        assert logits[0, RATIO_START].item() == min_score
+
+
+class TestForceEosAfterRatio:
+    """Tests that a ratio token as last_token forces EOS."""
+
+    def test_ratio_token_forces_eos(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        vocab_size = 300
+        logits = torch.randn(1, vocab_size)
+        min_score = torch.finfo(logits.dtype).min
+
+        logits[0].fill_(min_score)
+        logits[0, model._eos_token_id] = 0
+
+        assert logits[0, EOS].item() == 0
+        non_eos_max = logits[0, :EOS].max().item()
+        assert non_eos_max == min_score

tests/diffusion/models/ltx2/test_ltx2_hsdp.py

@@ -0,0 +1,25 @@
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks():
+    model = object.__new__(LTX2VideoTransformer3DModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.norm_out = nn.LayerNorm(4)
+
+    conditions = getattr(model, "_hsdp_shard_conditions", None)
+
+    assert conditions is not None
+    assert len(conditions) == 1
+
+    matched = []
+    for name, module in model.named_modules():
+        if any(cond(name, module) for cond in conditions):
+            matched.append(name)
+
+    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]

tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    get_qwen_image_edit_plus_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    # Keep the mock config intentionally minimal: this test only needs the
+    # fields touched during pre-process initialization.
+    (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16}))
+
+    pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+    image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8))
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "combine",
+                "multi_modal_data": {"image": [image, image, image, image, image]},
+            }
+        ],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"):
+        pre_process(request)

tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py

@@ -0,0 +1,260 @@
+import inspect
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
+    QwenImagePipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
+    QwenImageEditPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    QwenImageEditPlusPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import (
+    QwenImageLayeredPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _RejectingTextEncoder:
+    dtype = torch.float32
+
+    def __call__(self, *args, **kwargs):
+        raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length")
+
+
+class _FakeModelInputs:
+    def __init__(self, total_sequence_length: int):
+        attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long)
+        self.input_ids = attention_mask.clone()
+        self.attention_mask = attention_mask
+        self.pixel_values = None
+        self.image_grid_thw = None
+
+    def to(self, device):
+        return self
+
+
+class _FakeTokenizer:
+    def __init__(self, total_sequence_length: int | list[int]):
+        if isinstance(total_sequence_length, list):
+            self.total_sequence_lengths = list(total_sequence_length)
+        else:
+            self.total_sequence_lengths = [total_sequence_length]
+
+    def __call__(self, *args, **kwargs):
+        if len(self.total_sequence_lengths) > 1:
+            total_sequence_length = self.total_sequence_lengths.pop(0)
+        else:
+            total_sequence_length = self.total_sequence_lengths[0]
+        return _FakeModelInputs(total_sequence_length)
+
+
+class _FakeProcessor(_FakeTokenizer):
+    pass
+
+
+class _FakeScheduler:
+    def __init__(self):
+        self.begin_index = None
+
+    def set_begin_index(self, begin_index: int):
+        self.begin_index = begin_index
+
+
+PIPELINE_CASES = [
+    pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"),
+    pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"),
+    pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"),
+    pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"),
+]
+
+
+def _make_pipeline(
+    pipeline_class: type,
+    *,
+    total_sequence_length: int,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0])
+    if input_kind == "processor":
+        pipeline.processor = _FakeProcessor(total_sequence_length)
+    return pipeline
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=1025,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"):
+        pipeline.encode_prompt(prompt="prompt")
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=17,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"):
+        pipeline.encode_prompt(prompt="prompt", max_sequence_length=16)
+
+
+def test_prepare_encode_defaults_to_tokenizer_max_length():
+    pipeline = object.__new__(QwenImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.tokenizer_max_length = 1024
+    pipeline.vae_scale_factor = 8
+    pipeline.default_sample_size = 128
+    pipeline.scheduler = _FakeScheduler()
+    pipeline._extract_prompts = lambda prompts: (["prompt"], None)
+
+    captured = {}
+
+    def _fake_prepare_generation_context(**kwargs):
+        captured["max_sequence_length"] = kwargs["max_sequence_length"]
+        embeds = torch.ones((1, 1, 1))
+        mask = torch.ones((1, 1), dtype=torch.long)
+        return {
+            "prompt_embeds": embeds,
+            "prompt_embeds_mask": mask,
+            "negative_prompt_embeds": None,
+            "negative_prompt_embeds_mask": None,
+            "latents": embeds,
+            "timesteps": torch.tensor([1]),
+            "do_true_cfg": False,
+            "guidance": None,
+            "img_shapes": [[(1, 1, 1)]],
+            "txt_seq_lens": [1],
+            "negative_txt_seq_lens": None,
+        }
+
+    pipeline._prepare_generation_context = _fake_prepare_generation_context
+    state = SimpleNamespace(
+        prompts=["prompt"],
+        sampling=SimpleNamespace(
+            height=None,
+            width=None,
+            num_inference_steps=None,
+            sigmas=None,
+            guidance_scale_provided=False,
+            num_outputs_per_prompt=0,
+            generator=None,
+            true_cfg_scale=None,
+            max_sequence_length=None,
+        ),
+    )
+
+    pipeline.prepare_encode(state)
+
+    assert captured["max_sequence_length"] == 1024
+
+
+@pytest.mark.parametrize(
+    ("pipeline_class", "drop_idx"),
+    [
+        pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"),
+    ],
+)
+def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion(
+    pipeline_class: type,
+    drop_idx: int,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([8, 0])
+    pipeline.processor = _FakeProcessor(drop_idx + 1500)
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="short prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImagePipeline, id="qwen-image"),
+        pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"),
+    ],
+)
+def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 34
+    pipeline.tokenizer = _FakeTokenizer([1029, 5])
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="boundary prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImageEditPipeline, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"),
+    ],
+)
+def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 64
+    pipeline.tokenizer = _FakeTokenizer([30, 20])
+    pipeline.processor = _FakeProcessor(1500)
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="short prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        QwenImagePipeline,
+        QwenImageLayeredPipeline,
+        QwenImageEditPipeline,
+        QwenImageEditPlusPipeline,
+    ],
+)
+def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
+    assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024

tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py

@@ -0,0 +1,25 @@
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import StableAudioDiTModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_stable_audio_exposes_hsdp_shard_conditions_for_transformer_blocks():
+    model = object.__new__(StableAudioDiTModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.proj_out = nn.Linear(4, 4)
+
+    conditions = getattr(model, "_hsdp_shard_conditions", None)
+
+    assert conditions is not None
+    assert len(conditions) == 1
+
+    matched = []
+    for name, module in model.named_modules():
+        if any(cond(name, module) for cond in conditions):
+            matched.append(name)
+
+    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]

tests/diffusion/models/wan2_2/conftest.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+
+class StubTransformer(nn.Module):
+    def __init__(self, *, name: str = "transformer", in_channels: int = 4, out_channels: int = 4) -> None:
+        super().__init__()
+        self.name = name
+        self.config = SimpleNamespace(
+            patch_size=(1, 2, 2),
+            in_channels=in_channels,
+            out_channels=out_channels,
+            image_dim=None,
+        )
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+    def forward(self, **kwargs):
+        hidden_states = kwargs["hidden_states"]
+        return (torch.zeros_like(hidden_states[:, : self.config.out_channels]),)
+
+
+class StubScheduler:
+    def __init__(self, timesteps: list[int]) -> None:
+        self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
+        self.config = SimpleNamespace(num_train_timesteps=1000)
+        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+
+    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+        self.set_timesteps_calls.append((num_steps, device))
+
+
+class StubVAE:
+    dtype = torch.float32
+
+    def __init__(self, z_dim: int = 4) -> None:
+        self.config = SimpleNamespace(
+            z_dim=z_dim,
+            scale_factor_temporal=4,
+            scale_factor_spatial=8,
+            latents_mean=[0.0] * z_dim,
+            latents_std=[1.0] * z_dim,
+        )
+
+    def encode(self, video: torch.Tensor):
+        latent_frames = (video.shape[2] + self.config.scale_factor_temporal - 1) // self.config.scale_factor_temporal
+        latent_height = video.shape[-2] // self.config.scale_factor_spatial
+        latent_width = video.shape[-1] // self.config.scale_factor_spatial
+        latents = torch.ones(
+            video.shape[0],
+            self.config.z_dim,
+            latent_frames,
+            latent_height,
+            latent_width,
+            dtype=video.dtype,
+            device=video.device,
+        )
+        return SimpleNamespace(latents=latents)
+
+    def decode(self, latents: torch.Tensor, return_dict: bool = False):
+        del return_dict
+        return (latents,)
+
+
+@contextmanager
+def noop_progress_bar(*args, **kwargs):
+    del args, kwargs
+
+    class Bar:
+        def update(self) -> None:
+            return None
+
+    yield Bar()

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
+    Wan22I2VPipeline,
+    get_wan22_i2v_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
+    pipeline = object.__new__(Wan22I2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = StubTransformer(name="high", in_channels=8, out_channels=4)
+    pipeline.transformer_2 = StubTransformer(name="low", in_channels=8, out_channels=4)
+    pipeline.vae = StubVAE(z_dim=4)
+    pipeline.vae_scale_factor_temporal = 4
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.expand_timesteps = expand_timesteps
+    pipeline.progress_bar = noop_progress_bar
+    return pipeline
+
+
+def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+    request = SimpleNamespace(
+        prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "red")}}],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    result = preprocess(request)
+    prompt = result.prompts[0]
+
+    assert result.sampling_params.height == 432
+    assert result.sampling_params.width == 880
+    assert prompt["multi_modal_data"]["image"].size == (880, 432)
+
+    missing_image = SimpleNamespace(
+        prompts=[{"prompt": "p", "multi_modal_data": {}}],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+    with pytest.raises(ValueError, match="No image is provided"):
+        preprocess(missing_image)
+
+
+def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    latents = torch.zeros(1, 4, 2, 4, 4)
+    condition = torch.ones_like(latents)
+    first_frame_mask = torch.ones(1, 1, 2, 4, 4)
+    first_frame_mask[:, :, 0] = 0
+    timesteps = torch.tensor([900, 100])
+
+    calls = []
+
+    def fake_predict_noise_maybe_with_cfg(**kwargs):
+        positive = kwargs["positive_kwargs"]
+        calls.append(
+            {
+                "model": positive["current_model"].name,
+                "scale": kwargs["true_cfg_scale"],
+                "timestep_shape": tuple(positive["timestep"].shape),
+                "timestep_values": positive["timestep"].clone(),
+                "hidden_states": positive["hidden_states"].clone(),
+            }
+        )
+        return torch.ones_like(latents)
+
+    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
+    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=timesteps,
+        prompt_embeds=torch.zeros(1, 2, 3),
+        negative_prompt_embeds=None,
+        image_embeds=None,
+        guidance_low=1.0,
+        guidance_high=2.0,
+        boundary_timestep=500.0,
+        dtype=torch.float32,
+        attention_kwargs={},
+        condition=condition,
+        first_frame_mask=first_frame_mask,
+    )
+
+    assert [call["model"] for call in calls] == ["high", "low"]
+    assert [call["scale"] for call in calls] == [1.0, 2.0]
+    assert calls[0]["timestep_shape"] == (1, 8)
+    timestep_dtype = calls[0]["timestep_values"].dtype
+    torch.testing.assert_close(calls[0]["timestep_values"][0, :4], torch.zeros(4, dtype=timestep_dtype))
+    torch.testing.assert_close(calls[0]["timestep_values"][0, 4:], torch.full((4,), 900, dtype=timestep_dtype))
+    torch.testing.assert_close(calls[0]["hidden_states"][:, :, 0], torch.ones(1, 4, 4, 4))
+    torch.testing.assert_close(result, torch.full_like(latents, 2.0))
+
+
+def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    latents, condition, first_frame_mask = pipeline.prepare_latents(
+        image=torch.zeros(1, 3, 16, 16),
+        batch_size=1,
+        num_channels_latents=4,
+        height=16,
+        width=16,
+        num_frames=5,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator(device="cpu").manual_seed(0),
+    )
+
+    assert latents.shape == (1, 4, 2, 2, 2)
+    assert condition.shape == (1, 4, 1, 2, 2)
+    assert first_frame_mask.shape == (1, 1, 2, 2, 2)
+    assert first_frame_mask[:, :, 0].sum() == 0
+    assert first_frame_mask[:, :, 1].sum() == 4

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _StubTransformer(nn.Module):
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+
+class _StubScheduler:
+    def __init__(self, timesteps: list[int]) -> None:
+        self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
+        self.config = SimpleNamespace(num_train_timesteps=1000)
+        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+
+    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+        self.set_timesteps_calls.append((num_steps, device))
+
+
+@contextmanager
+def _noop_progress_bar(*args, **kwargs):
+    del args, kwargs
+
+    class _Bar:
+        def update(self) -> None:
+            return None
+
+    yield _Bar()
+
+
+def _make_pipeline() -> Wan22Pipeline:
+    pipeline = object.__new__(Wan22Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = _StubTransformer()
+    pipeline.transformer_2 = None
+    pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
+    pipeline.scheduler = _StubScheduler([9, 5])
+    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 5.0
+    pipeline.vae_scale_factor_temporal = 4
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.boundary_ratio = 0.875
+    pipeline.expand_timesteps = False
+    pipeline._guidance_scale = None
+    pipeline._guidance_scale_2 = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32)
+    pipeline.progress_bar = _noop_progress_bar
+    return pipeline
+
+
+def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
+    pipeline = _make_pipeline()
+
+    prompt_embeds = torch.randn(1, 8)
+    captured: dict[str, object] = {}
+
+    def _fake_diffuse(**kwargs):
+        captured.update(kwargs)
+        return kwargs["latents"] + 1
+
+    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
+
+    req = SimpleNamespace(
+        prompts=["prompt"],
+        sampling_params=SimpleNamespace(
+            height=None,
+            width=None,
+            num_frames=1,
+            num_inference_steps=2,
+            guidance_scale_provided=False,
+            guidance_scale=None,
+            guidance_scale_2=None,
+            boundary_ratio=None,
+            generator=None,
+            seed=None,
+            num_outputs_per_prompt=1,
+            max_sequence_length=32,
+            latents=None,
+            extra_args={},
+        ),
+    )
+
+    output = pipeline.forward(req, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0)
+
+    assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
+    assert captured["guidance_low"] == 1.0
+    assert captured["guidance_high"] == 1.0
+    assert captured["boundary_timestep"] == pytest.approx(875.0)
+    assert captured["latent_condition"] is None
+    assert captured["first_frame_mask"] is None
+    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
+
+
+def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
+    pipeline = _make_pipeline()
+    latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
+    timesteps = torch.tensor([7, 3], dtype=torch.int64)
+    prompt_embeds = torch.randn(1, 8)
+
+    predict_calls: list[dict[str, object]] = []
+    scheduler_calls: list[tuple[float, int, float, bool]] = []
+
+    def _fake_predict_noise_maybe_with_cfg(**kwargs):
+        predict_calls.append(kwargs)
+        timestep = kwargs["positive_kwargs"]["timestep"]
+        assert isinstance(timestep, torch.Tensor)
+        return torch.full_like(latents, float(timestep[0].item()))
+
+    def _fake_scheduler_step_maybe_with_cfg(noise_pred, t, current_latents, do_true_cfg):
+        scheduler_calls.append(
+            (float(noise_pred[0, 0, 0, 0, 0]), int(t.item()), float(current_latents.sum()), do_true_cfg)
+        )
+        return current_latents + noise_pred
+
+    pipeline.predict_noise_maybe_with_cfg = _fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
+    pipeline.scheduler_step_maybe_with_cfg = _fake_scheduler_step_maybe_with_cfg  # type: ignore[method-assign]
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=timesteps,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=None,
+        guidance_low=1.0,
+        guidance_high=2.0,
+        boundary_timestep=5.0,
+        dtype=torch.float32,
+        attention_kwargs={},
+    )
+
+    assert len(predict_calls) == 2
+    assert predict_calls[0]["true_cfg_scale"] == 1.0
+    assert predict_calls[1]["true_cfg_scale"] == 2.0
+    assert scheduler_calls == [
+        (7.0, 7, 0.0, False),
+        (3.0, 3, 28.0, False),
+    ]
+    assert torch.equal(result, torch.full_like(latents, 10.0))

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    create_transformer_from_config,
+    load_transformer_config,
+    retrieve_latents,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _LatentDist:
+    def sample(self, generator):
+        assert isinstance(generator, torch.Generator)
+        return torch.tensor([1.0])
+
+    def mode(self):
+        return torch.tensor([2.0])
+
+
+def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
+    generator = torch.Generator(device="cpu")
+
+    assert retrieve_latents(SimpleNamespace(latent_dist=_LatentDist()), generator).item() == 1.0
+    assert retrieve_latents(SimpleNamespace(latent_dist=_LatentDist()), sample_mode="argmax").item() == 2.0
+    torch.testing.assert_close(retrieve_latents(SimpleNamespace(latents=torch.tensor([3.0]))), torch.tensor([3.0]))
+
+
+def test_retrieve_latents_rejects_unknown_encoder_output() -> None:
+    with pytest.raises(AttributeError, match="Could not access latents"):
+        retrieve_latents(SimpleNamespace())
+
+
+def test_load_transformer_config_reads_local_subfolder_config(tmp_path) -> None:
+    config_dir = tmp_path / "transformer_2"
+    config_dir.mkdir(parents=True)
+    (config_dir / "config.json").write_text(json.dumps({"patch_size": [1, 2, 2], "num_layers": 2}))
+
+    assert load_transformer_config(str(tmp_path), "transformer_2") == {"patch_size": [1, 2, 2], "num_layers": 2}
+    assert load_transformer_config(str(tmp_path), "missing") == {}
+
+
+def test_create_transformer_from_config_maps_supported_keys(monkeypatch) -> None:
+    captured = {}
+
+    class FakeTransformer:
+        def __init__(self, **kwargs) -> None:
+            captured.update(kwargs)
+
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+
+    transformer = create_transformer_from_config(
+        {
+            "patch_size": [1, 2, 2],
+            "num_attention_heads": 8,
+            "attention_head_dim": 128,
+            "in_channels": 16,
+            "out_channels": 16,
+            "text_dim": 4096,
+            "vace_layers": [0],
+            "ignored": "value",
+        }
+    )
+
+    assert isinstance(transformer, FakeTransformer)
+    assert captured == {
+        "patch_size": (1, 2, 2),
+        "num_attention_heads": 8,
+        "attention_head_dim": 128,
+        "in_channels": 16,
+        "out_channels": 16,
+        "text_dim": 4096,
+    }

tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
+    Wan22TI2VPipeline,
+    get_wan22_ti2v_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
+    pipeline = object.__new__(Wan22TI2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
+    pipeline.vae = StubVAE(z_dim=4)
+    pipeline.vae_scale_factor_temporal = 4
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.progress_bar = noop_progress_bar
+    return pipeline
+
+
+def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
+    preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
+    request = SimpleNamespace(
+        prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    result = preprocess(request)
+
+    assert result.sampling_params.height == 672
+    assert result.sampling_params.width == 1344
+    assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
+
+
+def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
+    pipeline = _make_ti2v_pipeline()
+    latents = torch.zeros(1, 4, 2, 4, 4)
+    calls = []
+
+    def fake_predict_noise_maybe_with_cfg(**kwargs):
+        calls.append(kwargs)
+        return torch.ones_like(latents)
+
+    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
+    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        prompt_embeds=torch.zeros(1, 2, 3),
+        negative_prompt_embeds=torch.zeros(1, 2, 3),
+        guidance_scale=3.0,
+        dtype=torch.float32,
+        attention_kwargs={"a": "b"},
+        num_latent_frames=2,
+        latent_height=4,
+        latent_width=4,
+    )
+
+    positive = calls[0]["positive_kwargs"]
+    assert calls[0]["do_true_cfg"] is True
+    assert positive["timestep"].shape == (1, 8)
+    torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
+    torch.testing.assert_close(positive["hidden_states"], latents)
+    torch.testing.assert_close(result, torch.ones_like(latents))
+
+
+def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
+    pipeline = _make_ti2v_pipeline()
+    latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
+        image=torch.zeros(1, 3, 16, 16),
+        batch_size=1,
+        num_channels_latents=4,
+        height=16,
+        width=16,
+        num_frames=5,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=None,
+        latents=torch.zeros(1, 4, 2, 2, 2),
+    )
+
+    torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
+    assert latent_condition.shape == (1, 4, 1, 2, 2)
+    assert first_frame_mask[:, :, 0].sum() == 0
+    assert first_frame_mask[:, :, 1].sum() == 4

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
+    Wan22VACEPipeline,
+    create_vace_transformer_from_config,
+    get_wan22_vace_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_vace_pipeline() -> Wan22VACEPipeline:
+    pipeline = object.__new__(Wan22VACEPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
+    pipeline.transformer_config = pipeline.transformer.config
+    pipeline.vae = StubVAE(z_dim=4)
+    pipeline.vae_scale_factor_temporal = 4
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.progress_bar = noop_progress_bar
+    return pipeline
+
+
+def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
+    preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
+    ref = Image.new("RGB", (320, 160), "green")
+    frame = Image.new("RGB", (64, 64), "black")
+    mask = Image.new("L", (64, 64), 255)
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "p",
+                "multi_modal_data": {
+                    "image": ref,
+                    "video": [frame],
+                    "mask": mask,
+                },
+            }
+        ],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    result = preprocess(request)
+    additional_info = result.prompts[0]["additional_information"]
+
+    assert result.sampling_params.height == 432
+    assert result.sampling_params.width == 880
+    assert additional_info["reference_images"] == [ref]
+    assert additional_info["source_video"] == [frame]
+    assert additional_info["mask"] == [mask]
+
+
+def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
+    captured = {}
+
+    class FakeVACETransformer:
+        def __init__(self, **kwargs) -> None:
+            captured.update(kwargs)
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.WanVACETransformer3DModel",
+        FakeVACETransformer,
+    )
+
+    transformer = create_vace_transformer_from_config(
+        {
+            "patch_size": [1, 2, 2],
+            "in_channels": 96,
+            "out_channels": 16,
+            "vace_layers": [0, 1, 2],
+            "vace_in_channels": 132,
+            "unknown": "ignored",
+        }
+    )
+
+    assert isinstance(transformer, FakeVACETransformer)
+    assert captured == {
+        "patch_size": (1, 2, 2),
+        "in_channels": 96,
+        "out_channels": 16,
+        "vace_layers": [0, 1, 2],
+        "vace_in_channels": 132,
+    }
+
+
+def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None:
+    pipeline = _make_vace_pipeline()
+    mask = torch.ones(1, 3, 5, 16, 16)
+    reference_images = [[torch.zeros(3, 16, 16), torch.zeros(3, 16, 16)]]
+
+    encoded = pipeline.prepare_masks(mask, reference_images)
+
+    assert encoded.shape == (1, 64, 4, 2, 2)
+    torch.testing.assert_close(encoded[:, :, :2], torch.zeros(1, 64, 2, 2, 2))
+    torch.testing.assert_close(encoded[:, :, 2:], torch.ones(1, 64, 2, 2, 2))
+
+
+def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
+    pipeline = _make_vace_pipeline()
+    latents = torch.zeros(1, 4, 1, 2, 2)
+    vace_context = torch.ones(1, 12, 1, 2, 2)
+    calls = []
+
+    def fake_predict_noise_maybe_with_cfg(**kwargs):
+        calls.append(kwargs)
+        return torch.ones_like(latents)
+
+    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
+    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=torch.tensor([5]),
+        prompt_embeds=torch.zeros(1, 2, 3),
+        negative_prompt_embeds=torch.zeros(1, 2, 3),
+        guidance_scale=4.0,
+        dtype=torch.float32,
+        attention_kwargs={},
+        vace_context=vace_context,
+        vace_context_scale=0.75,
+    )
+
+    assert calls[0]["do_true_cfg"] is True
+    assert calls[0]["true_cfg_scale"] == 4.0
+    assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
+    assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
+    torch.testing.assert_close(result, torch.ones_like(latents))

tests/diffusion/quantization/test_quantization_quality.py

@@ -32,7 +32,7 @@ import numpy as np
 import pytest
 import torch
 
-from tests.utils import hardware_marks
+from tests.helpers.mark import hardware_marks
 
 # ---------------------------------------------------------------------------
 # Configuration — add new quantization methods / models here

tests/diffusion/test_diffusion_model_runner.py

@@ -8,7 +8,7 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 
 pytestmark = [pytest.mark.diffusion]

tests/diffusion/test_diffusion_step_pipeline.py

@@ -13,7 +13,7 @@ import torch
 from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin

tests/diffusion/test_diffusion_worker.py

@@ -16,7 +16,7 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
 
 @pytest.fixture
@@ -81,17 +81,31 @@ class TestDiffusionWorkerLoadWeights:
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
+    @pytest.fixture(autouse=True)
+    def setup_allocator(self, mocker: MockerFixture):
+        """
+        Unified interception of Allocators, and provision of default security values.
+        """
+        self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        self.mock_allocator = mocker.Mock()
+        self.mock_allocator_class.get_instance.return_value = self.mock_allocator
+        self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
+        self.mock_allocator.sleep = mocker.Mock()
+
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Setup process-scoped memory mocks
         # Before sleep: 3GB used
         # After sleep: 1GB used (freed 2GB)
+        initial_usage = 3 * 1024**3
         mock_get_process_memory.side_effect = [
-            3 * 1024**3,
+            initial_usage,
             1 * 1024**3,
         ]
 
@@ -99,25 +113,29 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # Call sleep with level 1
         result = mock_gpu_worker.sleep(level=1)
 
         # Verify sleep was called with correct tags
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
-        assert result is True
+        assert bool(result) is True
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
     def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Setup process-scoped memory mocks
+        initial_usage = 5 * 1024**3
         mock_get_process_memory.side_effect = [
-            5 * 1024**3,  # Before sleep
+            initial_usage,  # Before sleep
             1 * 1024**3,  # After sleep (freed 4GB)
         ]
 
@@ -125,6 +143,7 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # Mock pipeline buffers
         mock_buffer1 = torch.randn(10, 10)
@@ -141,7 +160,7 @@ class TestDiffusionWorkerSleep:
 
         # Verify sleep was called with empty tags (offload all)
         mock_allocator.sleep.assert_called_once_with(offload_tags=tuple())
-        assert result is True
+        assert bool(result) is True
 
         # Verify buffers were saved
         assert len(mock_gpu_worker._sleep_saved_buffers) == 2
@@ -151,22 +170,26 @@ class TestDiffusionWorkerSleep:
     def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.return_value = 10 * 1024**3
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Simulate process memory increase (should trigger assertion error)
+        initial_usage = 1 * 1024**3
         mock_get_process_memory.side_effect = [
-            1 * 1024**3,  # Before sleep: 1GB used
+            initial_usage,  # Before sleep: 1GB used
             3 * 1024**3,  # After sleep: 3GB used (negative freed)
         ]
 
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # This should raise an assertion error
-        with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
-            mock_gpu_worker.sleep(level=1)
+        result = mock_gpu_worker.sleep(level=1)
+        assert result == initial_usage
 
     def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep uses device-scoped fallback when NVML is unavailable."""
@@ -184,11 +207,12 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 2 * 1024**3
 
         result = mock_gpu_worker.sleep(level=1)
 
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
-        assert result is True
+        assert bool(result) is True
 
 
 class TestDiffusionWorkerWakeUp:
@@ -202,6 +226,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Ensure no saved buffers
         mock_gpu_worker._sleep_saved_buffers = {}
@@ -211,7 +236,7 @@ class TestDiffusionWorkerWakeUp:
 
         # Verify allocator.wake_up was called
         mock_allocator.wake_up.assert_called_once_with(["weights"])
-        assert result is True
+        assert bool(result) is True
 
     def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up with saved buffers (level 2 sleep)."""
@@ -221,6 +246,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Create saved buffers
         saved_buffer1 = torch.randn(10, 10)
@@ -255,7 +281,7 @@ class TestDiffusionWorkerWakeUp:
 
         # Verify saved buffers were cleared
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
-        assert result is True
+        assert bool(result) is True
 
     def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up only restores buffers that were saved."""
@@ -265,6 +291,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Only save buffer1, not buffer2
         saved_buffer1 = torch.randn(10, 10)
@@ -293,4 +320,4 @@ class TestDiffusionWorkerWakeUp:
         # buffer2 should NOT be restored since it wasn't saved
         mock_buffer2.data.copy_.assert_not_called()
 
-        assert result is True
+        assert bool(result) is True

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_engine():
+    with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
+        engine_instance = MagicMock()
+        mock.make_engine.return_value = engine_instance
+        yield engine_instance
+
+
+@pytest.fixture
+def client(mock_engine):
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_output_type="image",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source="prompt",
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        yield c
+        c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_success(client, mock_engine):
+    # Setup mock engine step to return a successful result
+    mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
+    mock_engine.step.return_value = [mock_result]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "A test prompt", sampling_params)
+
+    # Wait for the task to be processed
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-1"
+    mock_engine.step.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_error(client, mock_engine):
+    # Setup mock engine step to raise an exception
+    mock_engine.step.side_effect = RuntimeError("Engine failure")
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-err", "A test prompt", sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-err"
+    assert output.error == "Engine failure"
+    assert not output.images
+
+
+def test_inline_shutdown(client, mock_engine):
+    assert not client._shutting_down
+
+    # Shutting down should cleanly cancel anything queued and close engine
+    client.shutdown()
+
+    assert client._shutting_down
+    mock_engine.close.assert_called_once()

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1,17 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import multiprocessing as mp
 import queue
 import threading
+import time
 from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import torch
+import zmq
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.sched import RequestScheduler
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
@@ -51,6 +59,8 @@ def _make_executor(num_gpus: int = 1):
     executor._result_mq = mock_rmq
     executor._closed = False
     executor._processes = []
+    executor.is_failed = False
+    executor._failure_callbacks = []
     return executor, req_q, res_q
 
 
@@ -334,8 +344,9 @@ class TestSerialEngineOperations:
 
 class TestCollectiveRpcTimeoutWhileLockHeld:
     """``collective_rpc(timeout=...)`` must honour its timeout even when
-    another thread holds ``engine._rpc_lock`` indefinitely (e.g. a stalled
-    ``add_req`` waiting on an unresponsive worker).
+    another thread holds ``engine._rpc_lock`` indefinitely (e.g. request
+    execution stalled on ``add_req_and_wait_for_response`` → ``execute_fn``
+    → ``collective_rpc`` while blocked on an unresponsive worker).
     """
 
     def test_rpc_times_out_when_lock_held_directly(self):
@@ -359,10 +370,10 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         with pytest.raises(TimeoutError):
             engine.collective_rpc("health", timeout=0.5)
 
-    def test_rpc_times_out_when_add_req_stalled_on_worker(self):
+    def test_rpc_times_out_when_request_execution_stalled_on_worker(self):
         """Real-world scenario the bot flagged:
 
-        ``add_req`` holds ``_rpc_lock`` while blocked on
+        The scheduler/execute path holds ``_rpc_lock`` while blocked on
         ``executor._result_mq.dequeue()`` because the worker never replies.
         A concurrent ``collective_rpc(timeout=...)`` must still time out
         instead of hanging forever waiting for the lock.
@@ -428,3 +439,353 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         t.join(5)
 
         assert result.error == "ok"
+
+
+# ───────── error handling: EngineDeadError propagation through layers ─────
+
+
+class TestMultiprocExecutorRaisesEngineDeadError:
+    """``collective_rpc`` raises ``EngineDeadError`` when the engine is failed."""
+
+    def test_collective_rpc_raises_when_is_failed(self):
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor._broadcast_mq = MagicMock()
+        executor._result_mq = MagicMock()
+        executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
+        executor.is_failed = True
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+    def test_collective_rpc_raises_mid_dequeue_when_is_failed(self):
+        """Worker dies while we are polling the dequeue loop."""
+        executor, _, res_q = _make_executor()
+
+        call_count = 0
+        orig_dequeue = executor._result_mq.dequeue
+
+        def _dying_dequeue(timeout=None):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                executor.is_failed = True
+                raise TimeoutError
+            return orig_dequeue(timeout=timeout)
+
+        executor._result_mq.dequeue = _dying_dequeue
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+
+class TestDiffusionEngineDeadErrorPassthrough:
+    """``DiffusionEngine.add_req_and_wait_for_response`` re-raises
+    ``EngineDeadError`` from executor and wraps other errors."""
+
+    def test_engine_dead_error_propagates(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=EngineDeadError())
+
+        with pytest.raises(EngineDeadError):
+            engine.add_req_and_wait_for_response(_mock_request("dead"))
+
+    def test_runtime_error_wrapped_in_output(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault"))
+
+        out = engine.add_req_and_wait_for_response(_mock_request("fault"))
+        assert isinstance(out, DiffusionOutput)
+        assert "gpu fault" in out.error
+
+
+class TestStageDiffusionClientErrorPropagation:
+    """Error surface behaviour of ``StageDiffusionClient``.
+
+    Uses ``object.__new__`` to construct a client without spawning a real
+    subprocess, then manually sets the fields needed for each test.
+    """
+
+    def _make_client(self, *, engine_dead=False, proc_alive=True):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client.final_output = True
+        client.final_output_type = "image"
+        client.default_sampling_params = None
+        client.custom_process_input_func = None
+        client.engine_input_source = None
+
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._tasks = {}
+        client._shutting_down = False
+        client._engine_dead = engine_dead
+        client._owns_process = True
+        client._proc = MagicMock(
+            is_alive=MagicMock(return_value=proc_alive),
+            exitcode=1,
+        )
+        client._request_socket = MagicMock()
+        client._response_socket = MagicMock()
+        client._encoder = MagicMock()
+        client._decoder = MagicMock()
+
+        return client
+
+    @pytest.mark.asyncio
+    async def test_add_request_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            await client.add_request_async("req-3", "test prompt", None)
+
+    def test_check_health_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            client.check_health()
+
+    def test_check_health_ok_when_alive(self):
+        client = self._make_client()
+        client.check_health()
+
+    def test_get_output_raises_engine_dead_when_dead(self):
+        """When ``_engine_dead`` is True and the output queue is empty,
+        ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
+        client = self._make_client(engine_dead=True)
+        # Simulate _drain_responses as a no-op (no ZMQ socket)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError):
+            client.get_diffusion_output_nowait()
+
+    def test_get_output_returns_none_when_alive_and_empty(self):
+        """When the engine is alive and the queue is empty, return None."""
+        client = self._make_client()
+        client._response_socket.recv.side_effect = zmq.Again
+
+        assert client.get_diffusion_output_nowait() is None
+
+    def test_check_health_raises_when_proc_dead(self):
+        """``check_health`` detects a dead subprocess via ``_proc.is_alive()``
+        and raises ``EngineDeadError``, setting ``_engine_dead`` as a
+        side effect."""
+        client = self._make_client(proc_alive=False)
+
+        with pytest.raises(EngineDeadError, match="not alive"):
+            client.check_health()
+
+        assert client._engine_dead is True
+
+    def test_get_output_raises_when_proc_dead(self):
+        """When the subprocess has died (non-signal exit) and the output
+        queue is empty, ``get_diffusion_output_nowait`` must raise
+        ``EngineDeadError`` with the exit code."""
+        client = self._make_client(proc_alive=False)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError, match="exit code"):
+            client.get_diffusion_output_nowait()
+
+        assert client._engine_dead is True
+
+    def test_get_output_returns_none_on_signal_death(self):
+        """When the subprocess was killed by a signal (exit code > 128),
+        ``get_diffusion_output_nowait`` returns ``None`` and sets
+        ``_shutting_down`` instead of raising."""
+        client = self._make_client(proc_alive=False)
+        client._proc.exitcode = 137  # SIGKILL (128 + 9)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        result = client.get_diffusion_output_nowait()
+
+        assert result is None
+        assert client._shutting_down is True
+        assert client._engine_dead is True
+
+
+# ───────── monitor thread & death sentinel integration tests ─────────
+
+
+def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool:
+    """Poll until ``get_flag()`` returns True or *timeout* elapses."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if get_flag():
+            return True
+        time.sleep(interval)
+    return False
+
+
+def _make_short_lived_process() -> mp.Process:
+    """Spawn a real subprocess that exits immediately.
+
+    The process must be started with ``"fork"`` (or the platform default)
+    so that it can use a plain ``lambda`` as its target — ``"spawn"`` would
+    fail to pickle it.
+    """
+    ctx = mp.get_context("fork")
+    p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0")
+    p.start()
+    return p
+
+
+class TestMultiprocExecutorWorkerMonitor:
+    """Integration tests for ``start_worker_monitor``.
+
+    Uses real short-lived subprocesses so that OS-level sentinel fd
+    readiness is exercised end-to-end.
+    """
+
+    def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
+        """When a worker process dies, the monitor thread must:
+        1. Set ``is_failed = True``
+        2. Call ``shutdown()`` (which sets ``_closed = True``)
+        3. Invoke all registered failure callbacks
+        """
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        # Use a no-op so shutdown() doesn't crash on None resources.
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        callback_called = threading.Event()
+        executor.register_failure_callback(callback_called.set)
+
+        executor.start_worker_monitor()
+
+        # Wait for the process to exit and the monitor to react.
+        proc.join(5)
+        assert _poll_flag(lambda: executor.is_failed), "is_failed was not set"
+        assert executor._closed, "shutdown() was not called"
+        assert callback_called.wait(timeout=2), "failure callback was not invoked"
+
+    def test_worker_monitor_noop_when_already_closed(self):
+        """If ``_closed`` is already True when the process dies (orderly
+        shutdown), the monitor must *not* set ``is_failed``."""
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = True  # already shut down
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        executor.start_worker_monitor()
+        proc.join(5)
+
+        # Give the monitor thread a chance to run (it should early-return).
+        time.sleep(0.3)
+        assert not executor.is_failed, "is_failed should remain False on orderly shutdown"
+
+
+class TestStageDiffusionClientProcMonitor:
+    """Integration test for ``StageDiffusionClient._start_proc_monitor``.
+
+    Uses a real short-lived subprocess to verify the sentinel-based
+    detection pipeline.
+    """
+
+    def test_proc_monitor_sets_engine_dead_on_process_death(self):
+        """When the subprocess dies, the monitor thread must set
+        ``_engine_dead = True``."""
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._shutting_down = False
+        client._engine_dead = False
+
+        proc = _make_short_lived_process()
+        client._proc = proc
+
+        client._start_proc_monitor()
+        proc.join(5)
+
+        assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set"
+
+
+class TestDrainResponsesDeathSentinel:
+    """Tests for death sentinel and error routing in
+    ``StageDiffusionClient._drain_responses()``.
+    """
+
+    def _make_client(self):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._engine_dead = False
+        client._shutting_down = False
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._response_socket = MagicMock()
+        client._decoder = MagicMock()
+        return client
+
+    def test_drain_responses_sets_engine_dead_on_death_sentinel(self):
+        """When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD``
+        sentinel, it must set ``_engine_dead = True`` and stop draining
+        (decoder is never called)."""
+        client = self._make_client()
+
+        # First recv returns the death sentinel, second would be a normal
+        # message but should never be reached.
+        client._response_socket.recv.side_effect = [
+            StageDiffusionProc.DIFFUSION_PROC_DEAD,
+            b"should-not-be-reached",
+        ]
+
+        client._drain_responses()
+
+        assert client._engine_dead is True
+        client._decoder.decode.assert_not_called()
+
+    def test_drain_responses_routes_error_as_omni_request_output(self):
+        """When ``_drain_responses`` receives a ``{"type": "error"}`` message
+        with a ``request_id``, it must place an ``OmniRequestOutput`` with
+        the error on ``_output_queue``."""
+        client = self._make_client()
+
+        error_msg = {
+            "type": "error",
+            "request_id": "req-fail",
+            "error": "gpu fault",
+        }
+        # First recv returns the encoded error, second raises zmq.Again.
+        client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again]
+        client._decoder.decode.return_value = error_msg
+
+        client._drain_responses()
+
+        assert not client._output_queue.empty()
+        output = client._output_queue.get_nowait()
+        assert isinstance(output, OmniRequestOutput)
+        assert output.request_id == "req-fail"
+        assert output.error == "gpu fault"
+        assert output.finished is True

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -4,10 +4,12 @@
 import threading
 from collections import deque
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -335,6 +337,27 @@ def test_cleanup_after_poll_flow(build_adapter):
     assert "ext-flow" not in adapter.request_payload
 
 
+def test_finish_requests_restores_status(build_adapter):
+    """Abort path must pop ``requests_origin_status`` and restore pre-wait status.
+
+    While ``process_pending_chunks`` holds a request off the scheduler queues, the
+    adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must
+    put that status back on the live ``Request`` so base ``Scheduler.finish_requests``
+    can finish bookkeeping without inconsistent state / crashes.
+    """
+    adapter, _ = build_adapter(stage_id=1)
+    req_id = "req-abort-during-chunk"
+    prior = RequestStatus.RUNNING
+    request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK)
+    adapter.requests_origin_status[req_id] = prior
+    requests_map = {req_id: request}
+
+    adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map)
+
+    assert request.status == prior
+    assert req_id not in adapter.requests_origin_status
+
+
 # ---------------------------------------------------------------
 # Scheduler trigger tests
 # ---------------------------------------------------------------
@@ -508,3 +531,31 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
 
     assert len(cleanup_calls) == 0
     assert len(save_calls) == 1
+
+
+def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
+    """``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base."""
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+    order: list[str] = []
+
+    adapter = mocker.MagicMock()
+
+    def _adapter_finish(request_ids, finished_status, requests):
+        order.append("adapter")
+        return []
+
+    adapter.finish_requests.side_effect = _adapter_finish
+
+    def _super_finish(_self, request_ids, finished_status):
+        order.append("super")
+        return []
+
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched.chunk_transfer_adapter = adapter
+    sched.requests = {}
+
+    with patch.object(VLLMScheduler, "finish_requests", _super_finish):
+        OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
+
+    assert order == ["adapter", "super"]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -0,0 +1,716 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for rank-aware KV transfer (TP > 1) and heterogeneous TP support.
+
+Covers:
+- _build_rank_aware_send_keys / _build_rank_aware_recv_keys
+- _get_kv_source_ranks / _get_kv_target_ranks / get_kv_connector_key
+- update_sender_info storing base host/port
+- receive path constructing per-rank metadata for connector.get()
+- Mooncake connector _query_metadata_at and partial-metadata get() path
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    KVCacheTransferData,
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+)
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    KV_RANK_PORT_STRIDE,
+)
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import (
+    KVTPTopology,
+    build_rank_aware_recv_keys,
+    build_rank_aware_send_keys,
+    get_kv_connector_key,
+    get_kv_source_ranks,
+    get_kv_target_ranks,
+    merge_received_rank_shards,
+    slice_received_rank_shard,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_manager(
+    from_tp: int = 1,
+    to_tp: int = 1,
+    local_rank: int = 0,
+    from_stage: str = "stage0",
+    to_stage: str = "stage1",
+    stage_id: str = "stage1",
+    need_recv: bool = True,
+    need_send: bool = False,
+    recv_timeout: float = 0.3,
+) -> OmniKVTransferManager:
+    """Build a manager with TP params injected, bypassing torch.distributed."""
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage=from_stage,
+        to_stage=to_stage,
+        stage_id=stage_id,
+        need_recv_cache=need_recv,
+        need_send_cache=need_send,
+        recv_timeout=recv_timeout,
+        from_tp=from_tp,
+        to_tp=to_tp,
+    )
+    with (
+        patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=local_rank),
+        patch(
+            "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size",
+            return_value=max(from_tp, to_tp),
+        ),
+    ):
+        mgr = OmniKVTransferManager(config)
+    return mgr
+
+
+def _make_payload(head_values: list[float], request_id: str = "req-1") -> dict:
+    head_tensor = torch.tensor(head_values, dtype=torch.float32).view(1, len(head_values), 1).repeat(2, 1, 1)
+    return {
+        "request_id": request_id,
+        "layer_blocks": {
+            "key_cache": [head_tensor.clone()],
+            "value_cache": [(head_tensor + 100).clone()],
+        },
+        "block_ids": [0],
+        "metadata": {"seq_len": 2},
+    }
+
+
+def _make_transfer_data(head_values: list[float], request_id: str = "req-1") -> KVCacheTransferData:
+    payload = _make_payload(head_values, request_id=request_id)
+    return KVCacheTransferData(
+        request_id=request_id,
+        layer_blocks=payload["layer_blocks"],
+        block_ids=payload["block_ids"],
+        metadata=payload["metadata"],
+    )
+
+
+# ── Key format helper ────────────────────────────────────────────────
+
+
+class TestConnectorKeyFormat:
+    def test_key_format_matches_pr2677(self):
+        key = get_kv_connector_key("req-1", "stage0", 0, 1, 2)
+        assert key == "req-1_stage0_0_1_2"
+
+    def test_key_fields_are_positional(self):
+        key = get_kv_connector_key("r", "s", 5, 3, 7)
+        parts = key.split("_")
+        assert parts == ["r", "s", "5", "3", "7"]
+
+
+# ── Source / target rank mapping ─────────────────────────────────────
+
+
+class TestRankMapping:
+    """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs."""
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        assert get_kv_target_ranks(topo) == [0]
+        assert get_kv_source_ranks(topo) == [0]
+
+    def test_homogeneous_tp2_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
+        assert get_kv_target_ranks(topo) == [1]
+        assert get_kv_source_ranks(topo) == [1]
+
+    def test_homogeneous_tp4_rank3(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3)
+        assert get_kv_target_ranks(topo) == [3]
+        assert get_kv_source_ranks(topo) == [3]
+
+    def test_sender_gt_receiver_tp4_to_tp2_rank0(self):
+        """Receiver rank 0 should receive from sender rank 0 and 1."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0)
+        assert get_kv_source_ranks(topo) == [0, 1]
+
+    def test_sender_gt_receiver_tp4_to_tp2_rank1(self):
+        """Receiver rank 1 should receive from sender rank 2 and 3."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        assert get_kv_source_ranks(topo) == [2, 3]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank0(self):
+        """Sender rank 0 should send to receiver ranks 0 and 1."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        assert get_kv_target_ranks(topo) == [0, 1]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        assert get_kv_target_ranks(topo) == [2, 3]
+
+    def test_receiver_lt_sender_source_ranks(self):
+        """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        assert get_kv_source_ranks(topo) == [0]
+
+    def test_invalid_topology_raises(self):
+        topo = KVTPTopology(source_tp_size=3, target_tp_size=2, local_rank=0)
+        with pytest.raises(ValueError, match="divisible"):
+            get_kv_source_ranks(topo)
+
+
+# ── _build_rank_aware_recv_keys ──────────────────────────────────────
+
+
+class TestBuildRankAwareRecvKeys:
+    """Verify build_rank_aware_recv_keys returns (key, from_rank) tuples."""
+
+    def test_tp1_returns_legacy_key_with_none_rank(self):
+        topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "omni_stage0_to_stage1_kv_cache_req-1"
+        assert rank is None
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "req-1_stage0_0_0_0"
+        assert rank == 0
+
+    def test_homogeneous_tp2_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "req-1_stage0_0_1_1"
+        assert rank == 1
+
+    def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self):
+        """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 2
+
+        keys = [k for k, _ in pairs]
+        ranks = [r for _, r in pairs]
+        assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_1_0"]
+        assert ranks == [0, 1]
+
+    def test_heterogeneous_tp4_to_tp2_rank1_gets_two_keys(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 2
+
+        ranks = [r for _, r in pairs]
+        assert ranks == [2, 3]
+
+    def test_heterogeneous_tp2_to_tp4_rank2_gets_one_key(self):
+        """Receiver rank 2 with source_tp=2, target_tp=4 should get 1 key from sender rank 1."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=2)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert rank == 1
+        assert key == "req-1_stage0_0_1_2"
+
+
+# ── _build_rank_aware_send_keys ──────────────────────────────────────
+
+
+class TestBuildRankAwareSendKeys:
+    def test_tp1_returns_legacy_key(self):
+        topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert keys == ["omni_stage0_to_stage1_kv_cache_req-1"]
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert keys == ["req-1_stage0_0_0_0"]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank0_sends_two_keys(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert len(keys) == 2
+        assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"]
+
+
+# ── update_sender_info stores base host/port ─────────────────────────
+
+
+class TestUpdateSenderInfoBase:
+    def test_stores_base_host_and_port(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        assert mgr._sender_base_host == "10.0.0.1"
+        assert mgr._sender_base_zmq_port == 50151
+
+    def test_rank1_adjusts_default_port_but_preserves_base(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        assert mgr._sender_base_host == "10.0.0.1"
+        assert mgr._sender_base_zmq_port == 50151
+        expected_adjusted = 50151 + 1 * KV_RANK_PORT_STRIDE
+        assert mgr.config.connector_config["sender_zmq_port"] == expected_adjusted
+
+    def test_nested_sender_info_resolves_correctly(self):
+        """Nested sender_info keyed by integer stage id should resolve
+        using recv_stages (engine_input_source → recv_from)."""
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            stage_id=2,
+            engine_input_source=[1],
+            need_recv_cache=True,
+            from_tp=2,
+            to_tp=2,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=2),
+        ):
+            mgr = OmniKVTransferManager(config)
+        mgr.update_sender_info(
+            {
+                0: {"host": "10.0.0.1", "zmq_port": 50151},
+                1: {"host": "10.0.0.2", "zmq_port": 50152},
+            }
+        )
+        assert mgr._sender_base_host == "10.0.0.2"
+        assert mgr._sender_base_zmq_port == 50152
+
+
+# ── receive path constructs per-rank metadata ────────────────────────
+
+
+class TestReceiveConstructsMetadata:
+    """Verify that receive_kv_cache_for_request passes metadata with
+    correct (host, port) to connector.get() for heterogeneous TP."""
+
+    def test_tp1_no_metadata_passed(self):
+        """TP=1: connector.get() should be called WITHOUT metadata."""
+        mgr = _make_manager(from_tp=1, to_tp=1, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        assert len(calls) > 0
+        assert calls[0]["metadata"] is None
+
+    def test_homogeneous_tp2_rank0_passes_metadata(self):
+        """TP=2 rank 0: metadata should point to sender rank 0's port."""
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        assert len(calls) > 0
+        meta = calls[0]["metadata"]
+        assert meta is not None
+        assert meta["source_host"] == "10.0.0.1"
+        assert meta["source_port"] == 50151 + 0 * KV_RANK_PORT_STRIDE
+
+    def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        meta = calls[0]["metadata"]
+        assert meta["source_port"] == 50151 + 1 * KV_RANK_PORT_STRIDE
+
+    def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self):
+        """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with
+        two different metadata entries for sender ranks 0 and 1."""
+        mgr = _make_manager(from_tp=4, to_tp=2, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        seen_ports = set()
+        for c in calls:
+            if c["metadata"]:
+                seen_ports.add(c["metadata"]["source_port"])
+        expected_ports = {
+            50151 + 0 * KV_RANK_PORT_STRIDE,
+            50151 + 1 * KV_RANK_PORT_STRIDE,
+        }
+        assert expected_ports.issubset(seen_ports)
+
+
+# ── Mooncake connector _query_metadata_at ────────────────────────────
+
+
+class TestMooncakeQueryMetadataAt:
+    """Test the connector's _query_metadata_at method and partial-metadata
+    path in get() without requiring real RDMA/Mooncake."""
+
+    def test_query_metadata_at_returns_full_metadata(self):
+        """Mock the ZMQ interaction to verify _query_metadata_at returns
+        complete metadata including data_size."""
+
+        try:
+            from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+                MooncakeTransferEngineConnector,
+                QueryResponse,
+            )
+        except ImportError:
+            pytest.skip("Mooncake not available")
+
+        import msgspec
+
+        connector = MagicMock(spec=MooncakeTransferEngineConnector)
+        connector._get_req_socket = MagicMock()
+
+        mock_socket = MagicMock()
+        resp = QueryResponse(request_id="test_key@s0_s1", data_size=4096, is_fast_path=True)
+        mock_socket.recv.return_value = msgspec.msgpack.encode(resp)
+        connector._get_req_socket.return_value = mock_socket
+
+        result = MooncakeTransferEngineConnector._query_metadata_at(
+            connector,
+            "test_key@s0_s1",
+            "10.0.0.1",
+            50151,
+        )
+
+        assert result is not None
+        assert result["source_host"] == "10.0.0.1"
+        assert result["source_port"] == 50151
+        assert result["data_size"] == 4096
+        assert result["is_fast_path"] is True
+
+    def test_query_metadata_at_returns_none_on_not_found(self):
+        try:
+            from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+                INFO_NOT_FOUND,
+                MooncakeTransferEngineConnector,
+            )
+        except ImportError:
+            pytest.skip("Mooncake not available")
+
+        connector = MagicMock(spec=MooncakeTransferEngineConnector)
+        mock_socket = MagicMock()
+        mock_socket.recv.return_value = INFO_NOT_FOUND
+        connector._get_req_socket.return_value = mock_socket
+
+        result = MooncakeTransferEngineConnector._query_metadata_at(
+            connector,
+            "test_key@s0_s1",
+            "10.0.0.1",
+            50151,
+        )
+        assert result is None
+
+
+# ── Merge / slice hooks ──────────────────────────────────────────────
+
+
+class TestMergeSliceHooks:
+    def test_single_shard_passes_through(self):
+        payload = {"layer_blocks": {"key_cache": [1]}}
+        assert merge_received_rank_shards([payload]) == payload
+
+    def test_default_merger_concats_head_dim(self):
+        p0 = _make_payload([0.0])
+        p1 = _make_payload([1.0])
+        result = merge_received_rank_shards([p0, p1])
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        value_cache = result["layer_blocks"]["value_cache"][0]
+        assert key_cache.shape == (2, 2, 1)
+        assert value_cache.shape == (2, 2, 1)
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[0.0, 1.0], [0.0, 1.0]]))
+        assert torch.equal(value_cache[:, :, 0], torch.tensor([[100.0, 101.0], [100.0, 101.0]]))
+
+    def test_custom_merger_hook_called(self):
+        merged = {"merged": True}
+        assert merge_received_rank_shards([{}, {}], merger=lambda payloads: merged) == merged
+
+    def test_slicer_hook_called(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        sliced = {"sliced": True}
+        assert slice_received_rank_shard({"full": True}, topo, slicer=lambda payload: sliced) == sliced
+
+    def test_default_slicer_extracts_rank_local_heads(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        payload = _make_payload([0.0, 1.0])
+        result = slice_received_rank_shard(payload, topo)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        value_cache = result["layer_blocks"]["value_cache"][0]
+        assert key_cache.shape == (2, 1, 1)
+        assert value_cache.shape == (2, 1, 1)
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[1.0], [1.0]]))
+        assert torch.equal(value_cache[:, :, 0], torch.tensor([[101.0], [101.0]]))
+
+    def test_presliced_payload_is_not_sliced_twice(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        payload = _make_payload([1.0])
+        payload["metadata"]["tp_head_slice"] = {"applied": True, "target_rank": 1}
+        result = slice_received_rank_shard(payload, topo)
+        assert result is payload
+
+    def test_round_trip_merge_from_tp4_to_tp2(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        source_ranks = get_kv_source_ranks(topo)
+        payloads = [_make_payload([float(rank)]) for rank in source_ranks]
+        result = merge_received_rank_shards(payloads)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[2.0, 3.0], [2.0, 3.0]]))
+
+    def test_round_trip_slice_from_tp2_to_tp4(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=3)
+        payload = _make_payload([2.0, 3.0])
+        result = slice_received_rank_shard(payload, topo)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[3.0], [3.0]]))
+
+
+class TestSenderSideSlicing:
+    def test_transfer_slices_before_sending_to_multiple_targets(self):
+        mgr = _make_manager(
+            from_tp=2,
+            to_tp=4,
+            local_rank=0,
+            need_send=True,
+            need_recv=False,
+        )
+        sent_payloads = []
+
+        class _Connector:
+            supports_raw_data = False
+
+            def put(self, from_stage, to_stage, put_key, data):
+                sent_payloads.append((put_key, KVCacheTransferData.from_bytes(data)))
+                return True, len(data), {}
+
+        mgr._connector = _Connector()
+        mgr._transfer_kv_cache(_make_transfer_data([0.0, 1.0]), "req-1")
+
+        assert [key for key, _ in sent_payloads] == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"]
+        assert sent_payloads[0][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1)
+        assert sent_payloads[1][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1)
+        assert torch.equal(
+            sent_payloads[0][1]["layer_blocks"]["key_cache"][0][:, :, 0],
+            torch.tensor([[0.0], [0.0]]),
+        )
+        assert torch.equal(
+            sent_payloads[1][1]["layer_blocks"]["key_cache"][0][:, :, 0],
+            torch.tensor([[1.0], [1.0]]),
+        )
+        assert sent_payloads[0][1]["metadata"]["tp_head_slice"]["target_rank"] == 0
+        assert sent_payloads[1][1]["metadata"]["tp_head_slice"]["target_rank"] == 1
+
+
+class _MockBroadcastGroup:
+    def __init__(self, world_size: int, rank_in_group: int, broadcast_value=None, recv_value=None):
+        self.world_size = world_size
+        self.rank_in_group = rank_in_group
+        self.broadcast_value = broadcast_value
+        self.recv_value = recv_value
+        self.broadcast_calls = []
+        self.send_calls = []
+        self.recv_calls = []
+        self.shm_broadcaster = None
+
+    def broadcast_object(self, obj=None, src: int = 0):
+        self.broadcast_calls.append((obj, src))
+        return self.broadcast_value if self.broadcast_value is not None else obj
+
+    def send_object(self, obj, dst: int):
+        self.send_calls.append((dst, obj))
+
+    def recv_object(self, src: int):
+        self.recv_calls.append(src)
+        return self.recv_value
+
+
+class TestDistributedReceive:
+    def test_tp_cfg_leader_receives_then_sends_branch_local_payloads(self):
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4, rank_in_group=2)
+        cfg_group = _MockBroadcastGroup(world_size=3, rank_in_group=0)
+
+        def _receive(req_obj, cfg_func, target_device):
+            req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])])
+            req_obj.kv_metadata = {"source": "leader"}
+            req_obj.sampling_params.past_key_values = req_obj.past_key_values
+            req_obj.sampling_params.kv_metadata = req_obj.kv_metadata
+            req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])])
+            req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"}
+            req_obj.sampling_params.cfg_img_past_key_values = SimpleNamespace(key_cache=[torch.tensor([3.0])])
+            req_obj.sampling_params.cfg_img_kv_metadata = {"source": "cfg_img"}
+            return True
+
+        mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=3,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        mgr.receive_multi_kv_cache.assert_called_once()
+        assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
+        assert req.kv_metadata == {"source": "leader"}
+        assert cfg_group.broadcast_calls == []
+        assert [dst for dst, _ in cfg_group.send_calls] == [1, 2]
+        rank1_payload = cfg_group.send_calls[0][1]
+        rank2_payload = cfg_group.send_calls[1][1]
+        assert torch.equal(rank1_payload["past_key_values"].key_cache[0], torch.tensor([1.0]))
+        assert torch.equal(rank2_payload["past_key_values"].key_cache[0], torch.tensor([1.0]))
+        assert rank1_payload["sp.cfg_active_branch"] == "cfg_text"
+        assert rank2_payload["sp.cfg_active_branch"] == "cfg_img"
+        assert rank1_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"]
+        assert rank2_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"]
+        assert "sp.cfg_branch_past_key_values" in rank1_payload
+        assert "sp.cfg_branch_past_key_values" in rank2_payload
+        assert list(rank1_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_text"]
+        assert list(rank2_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_img"]
+        assert "sp.cfg_text_past_key_values" in rank1_payload
+        assert "sp.cfg_img_past_key_values" not in rank1_payload
+        assert "sp.cfg_img_past_key_values" in rank2_payload
+        assert "sp.cfg_text_past_key_values" not in rank2_payload
+
+    def test_tp_cfg_follower_receives_local_payload_without_receiving(self):
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=1)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4, rank_in_group=3)
+        cfg_payload = {
+            "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "kv_metadata": {"source": "main"},
+            "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "sp.kv_metadata": {"source": "main"},
+            "sp.cfg_active_branch": "cfg_text",
+            "sp.cfg_branch_roles": ["cfg_text", "cfg_img"],
+            "sp.cfg_branch_past_key_values": {
+                "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+            },
+            "sp.cfg_branch_kv_metadata": {"cfg_text": {"source": "cfg-text"}},
+            "sp.cfg_text_past_key_values": SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+        }
+        cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, recv_value=cfg_payload)
+
+        mgr.receive_multi_kv_cache = MagicMock(return_value=True)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=1,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        mgr.receive_multi_kv_cache.assert_not_called()
+        assert req.kv_metadata == {"source": "main"}
+        assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0]))
+        assert torch.equal(req.sampling_params.past_key_values.key_cache[0], torch.tensor([1.0]))
+        assert req.sampling_params.cfg_active_branch == "cfg_text"
+        assert req.sampling_params.cfg_branch_roles == ["cfg_text", "cfg_img"]
+        assert torch.equal(
+            req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0],
+            torch.tensor([2.0]),
+        )
+        assert req.sampling_params.cfg_branch_kv_metadata == {"cfg_text": {"source": "cfg-text"}}
+        assert torch.equal(req.sampling_params.cfg_text_past_key_values.key_cache[0], torch.tensor([2.0]))
+        assert cfg_group.broadcast_calls == []
+        assert cfg_group.recv_calls == [0]
+
+    def test_tp_without_cfg_keeps_independent_receive_path(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
+        mgr.receive_multi_kv_cache = MagicMock(return_value=True)
+
+        with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group):
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
+
+        mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
+
+
+# ── TP auto-detect ───────────────────────────────────────────────────
+
+
+class TestAutoDetectTP:
+    def test_auto_detect_when_config_defaults(self):
+        """When config from_tp/to_tp == 1 (default), manager should auto-detect."""
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="s0",
+            stage_id="s1",
+            need_recv_cache=True,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=4),
+        ):
+            mgr = OmniKVTransferManager(config)
+        assert mgr._tp_topo.source_tp_size == 4
+        assert mgr._tp_topo.target_tp_size == 4
+
+    def test_explicit_tp_overrides_auto_detect(self):
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="s0",
+            stage_id="s1",
+            need_recv_cache=True,
+            from_tp=2,
+            to_tp=4,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=8),
+        ):
+            mgr = OmniKVTransferManager(config)
+        assert mgr._tp_topo.source_tp_size == 2
+        assert mgr._tp_topo.target_tp_size == 4

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -7,7 +7,9 @@ import pytest
 
 from vllm_omni.distributed.omni_coordinator import (
     InstanceInfo,
+    LeastQueueLengthBalancer,
     RandomBalancer,
+    RoundRobinBalancer,
     StageStatus,
 )
 
@@ -60,3 +62,173 @@ def test_load_balancer_select_returns_valid_index():
 
     assert isinstance(index, int)
     assert 0 <= index < len(instances)
+
+
+def test_round_robin_balancer_cycles_instances():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=2,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=1,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+
+    balancer = RoundRobinBalancer()
+    results = [balancer.select({}, instances) for _ in range(5)]
+
+    # Default start_index=0 => 0,1,2,0,1
+    assert results == [0, 1, 2, 0, 1]
+
+
+def test_round_robin_balancer_empty_instances_raises():
+    with pytest.raises(ValueError, match="instances must not be empty"):
+        RoundRobinBalancer().select({}, [])
+
+
+def test_round_robin_balancer_after_large_index_and_shorter_list():
+    """Large start_index % len(instances) then counter wraps with shorter list."""
+    now = time()
+    two = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    balancer = RoundRobinBalancer(start_index=7)
+    assert balancer.select({}, two) == 1  # 7 % 2
+    assert balancer.select({}, two) == 0  # next index wrapped to 0
+
+
+def test_least_queue_length_balancer_picks_min_queue():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=2,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=5,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+
+    balancer = LeastQueueLengthBalancer()
+    index = balancer.select({}, instances)
+    assert index == 1
+
+
+def test_least_queue_length_balancer_empty_instances_raises():
+    with pytest.raises(ValueError, match="instances must not be empty"):
+        LeastQueueLengthBalancer().select({}, [])
+
+
+def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    balancer = LeastQueueLengthBalancer()
+    mocker.patch(
+        "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
+        return_value=2,
+    )
+    assert balancer.select({}, instances) == 2
+
+
+def test_least_queue_length_balancer_negative_queue_raises():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=-1,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    with pytest.raises(ValueError, match="queue_length must be non-negative"):
+        LeastQueueLengthBalancer().select({}, instances)

tests/e2e/accuracy/conftest.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 import os
-import shutil
 import subprocess
 from contextlib import contextmanager
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
 
 import pytest
+import requests
 import torch
+from PIL import Image
 
-from tests.conftest import OmniServer, OmniServerParams
+from tests.helpers.runtime import OmniServer, OmniServerParams
 
 
 def pytest_addoption(parser):
@@ -183,16 +185,26 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
-def reset_artifact_dir(path: Path) -> Path:
-    if path.exists():
-        shutil.rmtree(path)
-    path.mkdir(parents=True, exist_ok=True)
-    return path
+@pytest.fixture(scope="session")
+def qwen_bear_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
+    QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+    response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "qwen_bear.png")
+    return image
 
 
-def infer_model_label(model: str) -> str:
-    label = Path(model.rstrip("/\\")).name or "model"
-    return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label)
+@pytest.fixture(scope="session")
+def rabbit_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the rabbit image from the URL and save it to the accuracy artifact root."""
+    RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
+    response = requests.get(RABBIT_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "rabbit.png")
+    return image
 
 
 def _build_accuracy_server_config(

tests/e2e/accuracy/helpers.py

@@ -0,0 +1,115 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+
+
+def reset_artifact_dir(path: Path) -> Path:
+    import shutil
+
+    if path.exists():
+        shutil.rmtree(path)
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def infer_model_label(model: str) -> str:
+    label = Path(model.rstrip("/\\")).name or "model"
+    return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label)
+
+
+def model_output_dir(parent_dir: Path, model: str) -> Path:
+    safe_model_name = model.split("/")[-1].replace(".", "_")
+    path = parent_dir / safe_model_name
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def assert_similarity(
+    *,
+    model_name: str,
+    vllm_image: Image.Image,
+    diffusers_image: Image.Image,
+    ssim_threshold: float,
+    psnr_threshold: float,
+    width: int | None = None,
+    height: int | None = None,
+    compare_mode: str = "RGB",
+) -> None:
+    requested_size = (width, height) if width is not None and height is not None else None
+    if requested_size is not None and diffusers_image.size != requested_size:
+        pytest.skip(
+            "Skipping as diffusers baseline output is corrupt and not comparable: "
+            f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}."
+        )
+
+    assert vllm_image.size == diffusers_image.size, (
+        f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}"
+    )
+
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=vllm_image,
+        reference=diffusers_image,
+        compare_mode=compare_mode,
+    )
+    print(f"{model_name} similarity metrics:")
+    print(f"  SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True")
+    print(
+        f"  PSNR: value={psnr_score:.6f} dB, threshold>={psnr_threshold:.6f} dB, range=[0, +inf), higher_is_better=True"
+    )
+
+    assert ssim_score >= ssim_threshold, (
+        f"SSIM below threshold for {model_name}: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}."
+    )
+    assert psnr_score >= psnr_threshold, (
+        f"PSNR below threshold for {model_name}: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}."
+    )
+
+
+def assert_image_sequence_similarity(
+    *,
+    model_name: str,
+    vllm_images: list[Image.Image],
+    diffusers_images: list[Image.Image],
+    ssim_threshold: float,
+    psnr_threshold: float,
+    compare_mode: str = "RGB",
+) -> None:
+    assert len(vllm_images) == len(diffusers_images), (
+        f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}"
+    )
+    for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1):
+        assert_similarity(
+            model_name=f"{model_name}[layer={index}]",
+            vllm_image=vllm_image,
+            diffusers_image=diffusers_image,
+            ssim_threshold=ssim_threshold,
+            psnr_threshold=psnr_threshold,
+            compare_mode=compare_mode,
+        )
+
+
+def compute_image_ssim_psnr(
+    *,
+    prediction: Image.Image,
+    reference: Image.Image,
+    compare_mode: str = "RGB",
+) -> tuple[float, float]:
+    pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode)
+    ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode)
+
+    ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0)
+    psnr_metric = PeakSignalNoiseRatio(data_range=1.0)
+
+    ssim_value = float(ssim_metric(pred_tensor, ref_tensor).item())
+    psnr_value = float(psnr_metric(pred_tensor, ref_tensor).item())
+    return ssim_value, psnr_value
+
+
+def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor:
+    array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor

tests/e2e/accuracy/test_gebench_h100_smoke.py

@@ -6,8 +6,8 @@ from pathlib import Path
 import pytest
 
 from benchmarks.accuracy.text_to_image.gbench import main as gbench_main
-from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir
-from tests.utils import hardware_test
+from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
+from tests.helpers.mark import hardware_test
 
 
 @pytest.mark.advanced_model

tests/e2e/accuracy/test_gedit_bench_h100_smoke.py

@@ -7,8 +7,8 @@ import pytest
 
 from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS
 from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main
-from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir
-from tests.utils import hardware_test
+from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
+from tests.helpers.mark import hardware_test
 
 
 @pytest.mark.advanced_model

tests/e2e/accuracy/test_qwen_image.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+MODEL_ID = "Qwen/Qwen-Image"
+MODEL_ENV_VAR = "QWEN_IMAGE_MODEL"
+PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
+NEGATIVE_PROMPT = "blurry, low quality"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
+            json={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": f"{WIDTH}x{HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+            },
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            width=WIDTH,
+            height=HEIGHT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+
+    vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png")
+    diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png")
+
+    assert_similarity(
+        model_name=MODEL_ID,
+        vllm_image=vllm_output,
+        diffusers_image=diffusers_output,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import gc
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
+from PIL import Image
+
+from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
+MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
+PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
+NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
+SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+
+
+def _run_vllm_omni_image_edit(
+    *,
+    omni_server: OmniServer,
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    response = requests.post(
+        f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+        data={
+            "model": omni_server.model,
+            "prompt": prompt,
+            "size": f"{WIDTH}x{HEIGHT}",
+            "n": 1,
+            "response_format": "b64_json",
+            "negative_prompt": NEGATIVE_PROMPT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "true_cfg_scale": TRUE_CFG_SCALE,
+            "seed": SEED,
+        },
+        files=[
+            ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
+            for index, image in enumerate(input_images)
+        ],
+        timeout=600,
+    )
+    response.raise_for_status()
+    payload = response.json()
+    assert len(payload["data"]) == 1
+    image = decode_base64_image(payload["data"][0]["b64_json"])
+    image.load()
+    image.save(output_path)
+    return image
+
+
+def _run_diffusers_image_edit(
+    *,
+    model: str,
+    pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline],
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
+    device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
+    try:
+        images = input_images[0] if len(input_images) == 1 else input_images
+        pipe = pipeline_class.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        ).to(device)
+        pipe.set_progress_bar_config(disable=False)
+        generator = torch.Generator(device=device).manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=prompt,
+            image=images,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            width=WIDTH,
+            height=HEIGHT,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")  # pyright: ignore[reportAttributeAccessIssue]
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+def _vllm_omni_output_single_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "vllm_omni_single.png"
+    with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_SINGLE_IMAGE,
+            input_images=[qwen_bear_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "diffusers_single.png"
+    return _run_diffusers_image_edit(
+        model=SINGLE_MODEL,
+        pipeline_class=QwenImageEditPipeline,
+        prompt=PROMPT_SINGLE_IMAGE,
+        input_images=[qwen_bear_image],
+        output_path=output_path,
+    )
+
+
+def _vllm_omni_output_multiple_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "vllm_omni_multiple.png"
+    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_MULTIPLE_IMAGE,
+            input_images=[qwen_bear_image, rabbit_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_multiple_image(
+    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "diffusers_multiple.png"
+    return _run_diffusers_image_edit(
+        model=MULTIPLE_MODEL,
+        pipeline_class=QwenImageEditPlusPipeline,
+        prompt=PROMPT_MULTIPLE_IMAGE,
+        input_images=[qwen_bear_image, rabbit_image],
+        output_path=output_path,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_edit_single_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    diffusers_image = _diffusers_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    assert_similarity(
+        model_name=SINGLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(
+    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
+)
+def test_qwen_image_edit_multiple_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    diffusers_image = _diffusers_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    assert_similarity(
+        model_name=MULTIPLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_qwen_image_layered.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_image_sequence_similarity, model_output_dir
+from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+MODEL_ID = "Qwen/Qwen-Image-Layered"
+MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL"
+PROMPT = "decompose into layers"
+NEGATIVE_PROMPT = " "
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 777
+LAYERS = 3
+RESOLUTION = 640
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _normalize_layered_images(images: object) -> list[Image.Image]:
+    if not isinstance(images, list) or not images:
+        raise AssertionError(f"Unexpected layered output container: {type(images).__name__}")
+
+    first_item = images[0]
+    if isinstance(first_item, Image.Image):
+        return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)]
+    if isinstance(first_item, (list, tuple)):
+        return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)]
+    raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}")
+
+
+def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    input_image.save(output_dir / "input.png")
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        buffer = io.BytesIO()
+        input_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+            data={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": "auto",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+                "layers": LAYERS,
+                "resolution": RESOLUTION,
+            },
+            files=[("image", ("input.png", buffer, "image/png"))],
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == LAYERS
+        output_images = []
+        for item in payload["data"]:
+            image_bytes = base64.b64decode(item["b64_json"])
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
+            image.load()
+            output_images.append(image)
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"vllm_omni_layer_{index}.png")
+        return output_images
+
+
+def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            image=input_image,
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+            num_images_per_prompt=1,
+            layers=LAYERS,
+            resolution=RESOLUTION,
+        )
+        output_images = _normalize_layered_images(result.images)
+        assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}"
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"diffusers_layer_{index}.png")
+        return output_images
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+    input_image = qwen_bear_image.convert("RGBA")
+
+    vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+    diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+
+    assert_image_sequence_similarity(
+        model_name=MODEL_ID,
+        vllm_images=vllm_outputs,
+        diffusers_images=diffusers_outputs,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        compare_mode="RGBA",
+    )

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -22,7 +22,6 @@ import torch
 from diffusers import UniPCMultistepScheduler
 from PIL import Image
 
-from tests.conftest import OmniServerParams
 from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
     _configure_scheduler,
     _ensure_wan_ftfy_fallback,
@@ -48,7 +47,8 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
     SSIM_THRESHOLD,
     WIDTH,
 )
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
 
 
 def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:

tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py

@@ -26,7 +26,7 @@ from contextlib import ExitStack
 
 import pytest
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -19,7 +19,7 @@ import numpy as np
 import pytest
 from transformers import AutoTokenizer
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -10,7 +10,7 @@ import pytest
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -64,9 +64,6 @@ stage_args:
 # Top-level runtime config with Mooncake connector
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
   connectors:
     mooncake_connector:
       name: MooncakeConnector
@@ -80,4 +77,3 @@ runtime:
   edges:
     - from: 0
       to: 1
-      window_size: -1

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -62,10 +62,6 @@ stage_args:
 # Runtime edges
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
-
   # Distributed connectors configuration (optional)
   # More connectors will be supported in the future.
   connectors:
@@ -78,4 +74,3 @@ runtime:
   edges:
     - from: 0
       to: 1
-      window_size: -1

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -1,103 +0,0 @@
-# stage config for running qwen2.5-omni for multi-stage omni runtime.
-
-# This config is optimized for CI e2e tests.
-stage_args:
-  - stage_id: 0
-    runtime:
-      process: true            # Run this stage in a separate process
-      devices: "0"
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true  # Now we only support eager mode
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-      mm_processor_cache_gb: 0
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "1"
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
-    default_sampling_params:
-      temperature: 0.9
-      top_p: 0.8
-      top_k: 40
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-      stop_token_ids: [8294]
-  - stage_id: 2
-    runtime:
-      process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-    engine_args:
-      model_stage: code2wav
-      max_num_seqs: 1
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.15
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-    engine_input_source: [1]
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1             # Simplified: trigger downstream only after full upstream completion
-    max_inflight: 1             # Simplified: process serially within each stage
-  edges:
-    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
-      to: 1
-      window_size: -1
-    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
-      to: 2
-      window_size: -1

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -22,9 +22,10 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.platforms import current_omni_platform
 
 # Reference pixel data extracted from the known-good output image
@@ -32,30 +33,30 @@ from vllm_omni.platforms import current_omni_platform
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (157, 172, 217)},
-    {"position": (400, 50), "rgb": (105, 144, 218)},
-    {"position": (700, 100), "rgb": (118, 159, 233)},
-    {"position": (150, 400), "rgb": (195, 34, 60)},
-    {"position": (512, 336), "rgb": (222, 214, 193)},
-    {"position": (700, 400), "rgb": (197, 15, 43)},
-    {"position": (100, 600), "rgb": (105, 13, 18)},
-    {"position": (400, 600), "rgb": (169, 33, 44)},
-    {"position": (700, 600), "rgb": (101, 86, 93)},
-    {"position": (256, 256), "rgb": (181, 202, 222)},
+    {"position": (100, 100), "rgb": (156, 172, 217)},
+    {"position": (400, 50), "rgb": (105, 144, 217)},
+    {"position": (700, 100), "rgb": (118, 159, 232)},
+    {"position": (150, 400), "rgb": (180, 22, 52)},
+    {"position": (512, 336), "rgb": (221, 211, 194)},
+    {"position": (700, 400), "rgb": (192, 10, 46)},
+    {"position": (100, 600), "rgb": (102, 12, 22)},
+    {"position": (400, 600), "rgb": (161, 28, 47)},
+    {"position": (700, 600), "rgb": (100, 87, 94)},
+    {"position": (256, 256), "rgb": (181, 201, 221)},
 ]
 
 if current_omni_platform.is_rocm():
     REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (156, 172, 215)},
-        {"position": (400, 50), "rgb": (106, 144, 216)},
-        {"position": (700, 100), "rgb": (118, 158, 231)},
-        {"position": (150, 400), "rgb": (183, 23, 48)},
-        {"position": (512, 336), "rgb": (218, 215, 191)},
-        {"position": (700, 400), "rgb": (194, 14, 42)},
-        {"position": (100, 600), "rgb": (105, 10, 16)},
-        {"position": (400, 600), "rgb": (167, 33, 46)},
-        {"position": (700, 600), "rgb": (102, 86, 92)},
-        {"position": (256, 256), "rgb": (181, 201, 220)},
+        {"position": (100, 100), "rgb": (156, 172, 217)},
+        {"position": (400, 50), "rgb": (105, 144, 217)},
+        {"position": (700, 100), "rgb": (118, 159, 232)},
+        {"position": (150, 400), "rgb": (180, 22, 52)},
+        {"position": (512, 336), "rgb": (221, 211, 194)},
+        {"position": (700, 400), "rgb": (192, 10, 46)},
+        {"position": (100, 600), "rgb": (102, 12, 22)},
+        {"position": (400, 600), "rgb": (161, 28, 47)},
+        {"position": (700, 600), "rgb": (100, 87, 94)},
+        {"position": (256, 256), "rgb": (181, 201, 221)},
     ]
 
 PIXEL_TOLERANCE = 10

tests/e2e/offline_inference/test_bagel_lora.py

@@ -31,9 +31,10 @@ import torch
 from PIL import Image
 from safetensors.torch import save_file
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -27,9 +27,10 @@ from typing import Any
 import pytest
 from PIL import Image
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
-from vllm_omni import Omni
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.platforms import current_omni_platform
 
 # Reference pixel data extracted from the known-good output image
@@ -37,30 +38,30 @@ from vllm_omni.platforms import current_omni_platform
 # "Generated with seed=52, num_inference_steps=15,
 # prompt='A futuristic city skyline at twilight, cyberpunk style'"
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (121, 118, 100)},
-    {"position": (400, 50), "rgb": (163, 162, 143)},
-    {"position": (700, 100), "rgb": (170, 156, 127)},
-    {"position": (150, 400), "rgb": (129, 127, 112)},
-    {"position": (512, 512), "rgb": (135, 61, 59)},
-    {"position": (700, 400), "rgb": (205, 107, 43)},
-    {"position": (100, 700), "rgb": (197, 177, 157)},
-    {"position": (400, 700), "rgb": (139, 107, 86)},
-    {"position": (700, 700), "rgb": (247, 205, 146)},
-    {"position": (256, 256), "rgb": (171, 160, 153)},
+    {"position": (100, 100), "rgb": (115, 113, 94)},
+    {"position": (400, 50), "rgb": (159, 160, 144)},
+    {"position": (700, 100), "rgb": (164, 151, 123)},
+    {"position": (150, 400), "rgb": (120, 121, 107)},
+    {"position": (512, 512), "rgb": (165, 133, 127)},
+    {"position": (700, 400), "rgb": (217, 130, 66)},
+    {"position": (100, 700), "rgb": (191, 168, 152)},
+    {"position": (400, 700), "rgb": (130, 96, 77)},
+    {"position": (700, 700), "rgb": (247, 203, 140)},
+    {"position": (256, 256), "rgb": (167, 156, 150)},
 ]
 
 if current_omni_platform.is_rocm():
     REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (123, 119, 100)},
-        {"position": (400, 50), "rgb": (162, 161, 142)},
-        {"position": (700, 100), "rgb": (171, 156, 127)},
-        {"position": (150, 400), "rgb": (131, 128, 112)},
-        {"position": (512, 512), "rgb": (134, 61, 59)},
-        {"position": (700, 400), "rgb": (204, 107, 43)},
-        {"position": (100, 700), "rgb": (201, 180, 165)},
-        {"position": (400, 700), "rgb": (140, 108, 87)},
-        {"position": (700, 700), "rgb": (247, 205, 145)},
-        {"position": (256, 256), "rgb": (171, 160, 153)},
+        {"position": (100, 100), "rgb": (115, 113, 94)},
+        {"position": (400, 50), "rgb": (159, 160, 144)},
+        {"position": (700, 100), "rgb": (164, 151, 123)},
+        {"position": (150, 400), "rgb": (120, 121, 107)},
+        {"position": (512, 512), "rgb": (165, 133, 127)},
+        {"position": (700, 400), "rgb": (217, 130, 66)},
+        {"position": (100, 700), "rgb": (191, 168, 152)},
+        {"position": (400, 700), "rgb": (130, 96, 77)},
+        {"position": (700, 700), "rgb": (247, 203, 140)},
+        {"position": (256, 256), "rgb": (167, 156, 150)},
     ]
 
 # Maximum allowed difference per color channel

tests/e2e/offline_inference/test_bagel_understanding.py

@@ -26,8 +26,9 @@ from pathlib import Path
 import pytest
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner, modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import modify_stage_config
 
 MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT"
 STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")

tests/e2e/offline_inference/test_cache_dit.py

@@ -11,8 +11,8 @@ It uses minimal settings to keep test time short for CI.
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -15,19 +15,18 @@ from __future__ import annotations
 import functools
 import io
 import os
-import tempfile
 from pathlib import Path
 from urllib.request import urlopen
 
 import numpy as np
 import pytest
 import soundfile as sf
-import yaml
 from huggingface_hub import snapshot_download
 from vllm.sampling_params import SamplingParams
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 
@@ -44,13 +43,9 @@ REFERENCE_STAGE0_TEMPERATURE = 1.0
 REFERENCE_STAGE0_REPETITION_PENALTY = 2.0
 
 
-def _stage_config(name: str) -> str:
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
-STAGE_CONFIGS = [
-    _stage_config("cosyvoice3.yaml"),
-    _stage_config("cosyvoice3_async_chunk.yaml"),
+ASYNC_CHUNK_MODES = [
+    pytest.param(False, id="sync"),
+    pytest.param(True, id="async_chunk"),
 ]
 
 
@@ -138,19 +133,6 @@ def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int):
     raise AttributeError("Unable to locate stage outputs on Omni runner")
 
 
-def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str:
-    cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8"))
-    tokenizer_path = str(model_dir / "CosyVoice-BlankEN")
-    for stage in cfg.get("stage_args", []):
-        engine_args = stage.setdefault("engine_args", {})
-        engine_args["tokenizer"] = tokenizer_path
-        engine_args["enforce_eager"] = True
-        engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]}
-    out_path = tmp_dir / Path(base_stage_config).name
-    out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
-    return str(out_path)
-
-
 def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]:
     return [
         {
@@ -165,54 +147,54 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
 @pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS)
-def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None:
+@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES)
+def test_cosyvoice3_offline_reference_zero_shot(async_chunk: bool) -> None:
     """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio."""
     prompt_audio, prompt_sr = _load_reference_prompt_wav()
     model_dir = _resolve_model_dir()
     expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"])
 
-    with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp:
-        stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp))
-        with OmniRunner(
-            str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300
-        ) as omni_runner:
-            sampling_params_list = omni_runner.get_default_sampling_params_list()
-            sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
+    with OmniRunner(
+        str(model_dir),
+        seed=42,
+        stage_configs_path=get_deploy_config_path("cosyvoice3.yaml"),
+        async_chunk=async_chunk,
+        stage_init_timeout=300,
+    ) as omni_runner:
+        sampling_params_list = omni_runner.get_default_sampling_params_list()
+        sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
 
-            outputs = omni_runner.omni.generate(
-                _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list
+        outputs = omni_runner.omni.generate(_build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list)
+
+        assert outputs, "No outputs returned"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, "No audio output found"
+
+        audio = _concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        sr_val = audio_mm.get("sr", 24000)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        sr = int(sr_val)
+        assert sr == 24000, f"Unexpected sample_rate={sr}"
+
+        duration_s = audio.size / sr
+        assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
+
+        stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
+        if stage0_outputs:
+            completion = stage0_outputs[0].outputs[0]
+            finish_reason = getattr(completion, "finish_reason", None)
+            stop_reason = getattr(completion, "stop_reason", None)
+            num_tokens = len(getattr(completion, "token_ids", []) or [])
+
+            assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
+            assert int(stop_reason) == expected_stop_token, (
+                f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
             )
-
-            assert outputs, "No outputs returned"
-            audio_mm = outputs[0].multimodal_output
-            assert "audio" in audio_mm, "No audio output found"
-
-            audio = _concat_audio(audio_mm["audio"])
-            assert audio.size > 0, "Generated audio is empty"
-
-            sr_val = audio_mm.get("sr", 24000)
-            if isinstance(sr_val, list) and sr_val:
-                sr_val = sr_val[-1]
-            if hasattr(sr_val, "item"):
-                sr_val = sr_val.item()
-            sr = int(sr_val)
-            assert sr == 24000, f"Unexpected sample_rate={sr}"
-
-            duration_s = audio.size / sr
-            assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
-
-            stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
-            if stage0_outputs:
-                completion = stage0_outputs[0].outputs[0]
-                finish_reason = getattr(completion, "finish_reason", None)
-                stop_reason = getattr(completion, "stop_reason", None)
-                num_tokens = len(getattr(completion, "token_ids", []) or [])
-
-                assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
-                assert int(stop_reason) == expected_stop_token, (
-                    f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
-                )
-                assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
-            else:
-                assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs"
+            assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
+        else:
+            assert async_chunk, "Stage-0 produced no engine outputs"

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -4,8 +4,9 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -2,8 +2,8 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 from safetensors.torch import save_file
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_dynin_omni.py

@@ -18,7 +18,7 @@ import pytest
 import torch
 from transformers import AutoTokenizer
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

tests/e2e/offline_inference/test_expert_parallel.py

@@ -18,8 +18,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_flux.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for Flux1 Schnell."""
+
+import pytest
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "black-forest-labs/FLUX.1-schnell"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_flux_schnell_text_to_image():
+    """Test FLUX.1-schnell text-to-image generation."""
+    omni = Omni(model=MODEL)
+
+    omni_outputs = list(
+        omni.generate(
+            prompts=["A photo of a cat sitting on a laptop"],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=512,
+                width=512,
+                num_inference_steps=2,
+                seed=42,
+            ),
+        )
+    )
+
+    assert len(omni_outputs) > 0
+    images = omni_outputs[0].images
+    assert len(images) == 1
+    assert isinstance(images[0], Image.Image)
+    assert images[0].size == (512, 512)

tests/e2e/offline_inference/test_flux2_klein_inpaint.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for Flux2 Klein inpainting.
+
+"""
+
+# ruff: noqa: E402
+
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from PIL import Image, ImageDraw
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+_HEIGHT = 512
+_WIDTH = 512
+_NUM_INFERENCE_STEPS = 4
+
+
+def _create_test_image(width: int = _WIDTH, height: int = _HEIGHT, color: tuple = (128, 128, 128)) -> Image.Image:
+    return Image.new("RGB", (width, height), color)
+
+
+def _create_test_mask(width: int = _WIDTH, height: int = _HEIGHT) -> Image.Image:
+    mask = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
+    return mask
+
+
+def _create_test_inputs(color: tuple = (100, 150, 200)):
+    return _create_test_image(_WIDTH, _HEIGHT, color), _create_test_mask(_WIDTH, _HEIGHT)
+
+
+def _extract_images_from_output(outputs: list) -> list[Image.Image]:
+    images = []
+    for req_output in outputs:
+        if hasattr(req_output, "images") and req_output.images:
+            images.extend(req_output.images)
+        elif hasattr(req_output, "request_output") and req_output.request_output:
+            stage_out = req_output.request_output
+            if isinstance(stage_out, OmniRequestOutput) and hasattr(stage_out, "images"):
+                images.extend(stage_out.images)
+            elif isinstance(stage_out, list):
+                for s in stage_out:
+                    if hasattr(s, "images") and s.images:
+                        images.extend(s.images)
+    return images
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_basic():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+
+        outputs = m.generate(
+            prompts=[
+                {
+                    "prompt": "Fill in the masked area with a beautiful garden",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images = _extract_images_from_output(list(outputs))
+        assert len(images) == 1
+        assert images[0].size == (_WIDTH, _HEIGHT)
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()
+
+
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_deterministic():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+        seed = 12345
+
+        gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
+        gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
+
+        outputs1 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A red flower in a field",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen1,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        outputs2 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A red flower in a field",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen2,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images1 = _extract_images_from_output(list(outputs1))
+        images2 = _extract_images_from_output(list(outputs2))
+
+        assert len(images1) == 1
+        assert len(images2) == 1
+
+        assert list(images1[0].getdata()) == list(images2[0].getdata()), (
+            "Same input with same seed should produce identical output. "
+            "This is critical for offline/online consistency."
+        )
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()
+
+
+@pytest.mark.diffusion
+def test_flux2_klein_inpaint_different_seeds_different_output():
+    m = None
+    try:
+        m = Omni(model=MODEL)
+        input_image, mask_image = _create_test_inputs()
+
+        gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42)
+        gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999)
+
+        outputs1 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A beautiful landscape",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen1,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        outputs2 = m.generate(
+            prompts=[
+                {
+                    "prompt": "A beautiful landscape",
+                    "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+                }
+            ],
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=gen2,
+                num_outputs_per_prompt=1,
+            ),
+        )
+
+        images1 = _extract_images_from_output(list(outputs1))
+        images2 = _extract_images_from_output(list(outputs2))
+
+        assert len(images1) == 1
+        assert len(images2) == 1
+
+        different_pixel_count = sum(1 for p1, p2 in zip(images1[0].getdata(), images2[0].getdata()) if p1 != p2)
+        assert different_pixel_count > 0, "Different seeds should produce different outputs"
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()

tests/e2e/offline_inference/test_flux_autoround_w4a16.py

@@ -14,8 +14,9 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import OmniRunner
-from tests.utils import DeviceMemoryMonitor, hardware_test
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_flux_kontext.py

@@ -13,7 +13,7 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 

tests/e2e/offline_inference/test_hunyuanimage3_text2img.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
 
-from tests.conftest import OmniRunner
+from tests.helpers.runtime import OmniRunner
 from vllm_omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -17,7 +17,7 @@ PROMPT = "A brown and white dog is running on the grass"
 MODEL_NAME = "tencent/HunyuanImage-3.0"
 LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
 REPO_ROOT = Path(__file__).resolve().parents[3]
-STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml"
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_t2i.yaml"
 
 pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
 

tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py

@@ -11,7 +11,7 @@ import numpy as np
 import pytest
 from PIL import Image
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 
 REPO_ROOT = Path(__file__).resolve().parents[3]
 T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"

tests/e2e/offline_inference/test_magi_human.py

@@ -8,8 +8,8 @@ import av
 import numpy as np
 import pytest
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 

tests/e2e/offline_inference/test_mammoth_moda2.py

@@ -23,8 +23,10 @@ import pytest
 import torch
 from vllm.sampling_params import SamplingParams
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 # ---------------------------------------------------------------------------
 # Constants

tests/e2e/offline_inference/test_ming_flash_omni.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
+    generate_synthetic_audio,
+    generate_synthetic_image,
+    generate_synthetic_video,
+)
+from tests.helpers.stage_config import modify_stage_config
+
+models = ["Jonathan1909/Ming-flash-omni-2.0"]
+
+# Ming-specific
+SYSTEM_PROMPT = "你是一个友好的AI助手。\n\ndetailed thinking off"
+EOS_TOKEN = "<|role_end|>"
+IMAGE_TOKEN = "<IMAGE>"
+VIDEO_TOKEN = "<VIDEO>"
+AUDIO_TOKEN = "<AUDIO>"
+
+
+def build_prompt(user_text: str) -> str:
+    """Build a Ming chat prompt."""
+    return (
+        f"<role>SYSTEM</role>{SYSTEM_PROMPT}{EOS_TOKEN}<role>HUMAN</role>{user_text}{EOS_TOKEN}<role>ASSISTANT</role>"
+    )
+
+
+def get_eager_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+            },
+        },
+    )
+    return path
+
+
+stage_configs = [get_eager_config()]
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_text_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test text-only input processing and text output generation.
+    Input Modal: text
+    Output Modal: text
+    """
+    prompt = build_prompt("请详细介绍鹦鹉的生活习性。")
+    request_config = {"prompts": prompt, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_image_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test image understanding with text output.
+    Input Modal: image + text
+    Output Modal: text
+    """
+    image = generate_synthetic_image(224, 224)["np_array"]
+    prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.")
+    request_config = {"prompts": prompt, "images": image, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test audio understanding with text output.
+    Input Modal: audio + text
+    Output Modal: text
+    """
+    audio = generate_synthetic_audio(2, 1, 16000)["np_array"]
+    if len(audio.shape) == 2:
+        audio = audio.squeeze()
+    prompt = build_prompt(f"{AUDIO_TOKEN}Please recognize the language of this speech and transcribe it. Format: oral.")
+    request_config = {"prompts": prompt, "audios": audio, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_video_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test video understanding with text output.
+    Input Modal: video + text
+    Output Modal: text
+    """
+    video = generate_synthetic_video(224, 224, 30)["np_array"]
+    prompt = build_prompt(f"{VIDEO_TOKEN}Describe what is happening in this video.")
+    request_config = {"prompts": prompt, "videos": video, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
+    """
+    Test mixed modality input (image + audio) with text output.
+    Input Modal: image + audio + text
+    Output Modal: text
+    """
+    image = generate_synthetic_image(224, 224)["np_array"]
+    audio = generate_synthetic_audio(2, 1, 16000)["np_array"]
+    if len(audio.shape) == 2:
+        audio = audio.squeeze()
+    prompt = build_prompt(f"{IMAGE_TOKEN}{AUDIO_TOKEN}Describe the image and transcribe the audio.")
+    request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]}
+
+    omni_runner_handler.send_request(request_config)

tests/e2e/offline_inference/test_omni_sleep_mode.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Omni Sleep Mode across various model architectures.
+"""
+
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+MODEL_DIFF = "riverclouds/qwen_image_random"
+
+
+def get_ack_info(ack, key, default=None):
+    if hasattr(ack, key):
+        return getattr(ack, key)
+    if isinstance(ack, dict):
+        return ack.get(key, default)
+    return default
+
+
+def get_dynamic_devices(stage_idx, num_stages, tp_size):
+    total_gpus = torch.cuda.device_count()
+    gpus_per_stage = tp_size
+    start_idx = stage_idx * gpus_per_stage
+    if start_idx + gpus_per_stage > total_gpus:
+        start_idx = start_idx % total_gpus
+    device_ids = [str(start_idx + i) for i in range(gpus_per_stage)]
+    return ",".join(device_ids)
+
+
+# Test 1: Diffusion Model (2-Stage BAGEL)
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_diffusion_model_sleep_tp(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size:
+        pytest.skip(f"Skipping TP={tp_size}")
+
+    engine_args = {
+        "model": MODEL,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "trust_remote_code": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        # BAGEL requires 2 params
+        diff_sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256)
+        llm_sp = SamplingParams()
+
+        # Warmup
+        async for _ in engine.generate("test", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        # Sleep all
+        acks = await engine.sleep(level=2)
+        statuses = [get_ack_info(ack, "status") for ack in acks]
+        assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
+
+        # Wakeup & Verify
+        await engine.wake_up()
+        async for _ in engine.generate("verify", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        print(f"Diffusion TP={tp_size} Lifecycle OK")
+    finally:
+        engine.shutdown()
+
+
+# Test 2: Multi-stage Manual Config
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_multistage_sleep_h100(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size * 2:
+        pytest.skip("Not enough GPUs")
+
+    stages = []
+    for i in range(2):
+        devs = get_dynamic_devices(i, 2, tp_size)
+        stages.append(
+            {
+                "stage_id": i,
+                "stage_type": "llm" if i == 0 else "diffusion",
+                "runtime": {"process": True, "devices": devs},
+                "engine_args": {
+                    "model": MODEL,
+                    "model_stage": "thinker" if i == 0 else "base",
+                    "tensor_parallel_size": tp_size,
+                    "gpu_memory_utilization": 0.4,
+                    "dtype": "bfloat16",
+                    "enable_sleep_mode": True,
+                    "trust_remote_code": True,
+                },
+            }
+        )
+
+    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
+
+    engine = AsyncOmni(
+        model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200
+    )
+    try:
+        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+        async for _ in engine.generate("warmup", sampling_params=[SamplingParams(), sp]):
+            pass
+
+        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        assert len(acks) == 2 * tp_size
+
+        await engine.wake_up(stage_ids=[0, 1])
+        async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]):
+            pass
+    finally:
+        engine.shutdown()
+
+
+# Test 3: Pure Diffusion Single-Stage
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_pure_diffusion_scenario(tp_size):
+    engine_args = {
+        "model": MODEL_DIFF,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        await engine.sleep(level=1)
+        await engine.wake_up()
+        async for _ in engine.generate("test", sampling_params=[SamplingParams()]):
+            pass
+        print("Pure Diffusion OK")
+    finally:
+        engine.shutdown()

tests/e2e/offline_inference/test_omnivoice.py

@@ -16,8 +16,8 @@ from pathlib import Path
 import numpy as np
 import pytest
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 
 MODEL = "k2-fsa/OmniVoice"
 

tests/e2e/offline_inference/test_ovis_image.py

@@ -16,7 +16,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 
 # Mock the OvisImageTransformer2DModel to avoid complex init if needed,

tests/e2e/offline_inference/test_quantization_fp8.py

@@ -36,8 +36,8 @@ from typing import Any
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -2,46 +2,39 @@
 E2E tests for Qwen2.5-Omni model with mixed modality inputs, audio and text output.
 """
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import (
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import (
     generate_synthetic_audio,
     generate_synthetic_image,
     generate_synthetic_video,
-    modify_stage_config,
 )
-from tests.utils import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 
+# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
+# platforms: section. NPU still uses the legacy per-platform YAML until it
+# also migrates to the new schema.
+_CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml")
+
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml"),
+    return modify_stage_config(
+        _CI_DEPLOY,
         updates={
-            "stage_args": {
-                0: {
-                    "engine_args.enforce_eager": "true",
-                },
-                1: {"engine_args.enforce_eager": "true"},
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
             },
         },
     )
-    return path
 
 
-# CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
-if current_omni_platform.is_npu():
-    stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
-elif current_omni_platform.is_xpu():
-    # Intel XPU stage config optimized for B60 GPU
-    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")
+if current_omni_platform.is_rocm() or current_omni_platform.is_xpu() or current_omni_platform.is_npu():
+    stage_config = _CI_DEPLOY
 else:
     stage_config = get_cuda_graph_config()
 

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -7,41 +7,35 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import (
-    generate_synthetic_video,
-    modify_stage_config,
-)
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_video
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 
+# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
+# platforms: section. Only CUDA needs an extra enforce_eager tweak.
+_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+
+
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
+    return modify_stage_config(
+        _CI_DEPLOY,
         updates={
-            "stage_args": {
-                0: {
-                    "engine_args.enforce_eager": "true",
-                },
-                1: {"engine_args.enforce_eager": "true"},
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
             },
         },
     )
-    return path
 
 
-# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
-elif current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
+if current_omni_platform.is_rocm() or current_omni_platform.is_xpu():
+    stage_configs = [_CI_DEPLOY]
 else:
     stage_configs = [get_cuda_graph_config()]
 

tests/e2e/offline_inference/test_qwen3_tts_base.py

@@ -13,12 +13,10 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
@@ -26,23 +24,31 @@ REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know wh
 
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        get_stage_config(),
+    """Build a temp deploy yaml mirroring the deleted qwen3_tts_no_async_chunk.yaml.
+
+    Composes the synchronous (no-async-chunk) variant on top of the bundled
+    qwen3_tts.yaml prod default, with cudagraphs disabled. Replaces the deleted
+    standalone variant yaml; same effective config, no checked-in file needed.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_tts.yaml"),
         updates={
-            "stage_args": {
+            "async_chunk": False,
+            "stages": {
                 0: {
-                    "engine_args.enforce_eager": "true",
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
+                },
+                1: {
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
                 },
-                1: {"engine_args.enforce_eager": "true"},
             },
         },
     )
-    return path
-
-
-def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
-    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
 # Same structure as test_qwen3_omni: models, stage_configs, test_params

tests/e2e/offline_inference/test_qwen3_tts_customvoice.py

@@ -13,34 +13,40 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
-from pathlib import Path
-
 import pytest
 
-from tests.conftest import modify_stage_config
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
 
 def get_cuda_graph_config():
-    path = modify_stage_config(
-        get_stage_config(),
+    """Build a temp deploy yaml mirroring the deleted qwen3_tts_no_async_chunk.yaml.
+
+    Composes the synchronous (no-async-chunk) variant on top of the bundled
+    qwen3_tts.yaml prod default, with cudagraphs disabled. Replaces the deleted
+    standalone variant yaml; same effective config, no checked-in file needed.
+    """
+    return modify_stage_config(
+        get_deploy_config_path("qwen3_tts.yaml"),
         updates={
-            "stage_args": {
+            "async_chunk": False,
+            "stages": {
                 0: {
-                    "engine_args.enforce_eager": "true",
+                    "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
+                },
+                1: {
+                    "gpu_memory_utilization": 0.2,
+                    "enforce_eager": True,
+                    "async_scheduling": False,
                 },
-                1: {"engine_args.enforce_eager": "true"},
             },
         },
     )
-    return path
-
-
-def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
-    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
 # Same structure as test_qwen3_omni: models, stage_configs, test_params

tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py

@@ -36,8 +36,8 @@ from pathlib import Path
 import pytest
 import torch
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -20,8 +20,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image
 
-from tests.conftest import OmniRunner
-from tests.utils import hardware_test
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform