Merge branch 'main' into fix/pr4709-review-followup

tests/ar_diffusion/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
+
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+)
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
+
+
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
+
+
+def test_ar_diffusion_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
+
+
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)
+
+
+def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
+    """Regression: the worker passes ``kv_prefetch_jobs`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_jobs" in params, cls

tests/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
+
+BLOCK = 16
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    spec = make_spec()
+    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+    assert skip % 16 == 0 and skip == 16
+
+
+def test_sink_chunks_protected():
+    # sink = 1 chunk (16 tokens) is never skipped.
+    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
+    assert skip == 32
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- ARDiffusionKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    k_pools[0][BLOCK + 3].fill_(7)
+    v_pools[0][2 * BLOCK + 5].fill_(11)
+    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        cross_attn_length=L,
+        device=torch.device("cpu"),
+    )
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
+
+
+def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        local_branches=local_branches,
+        num_frame_per_block=num_frame_per_block,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
+    a single-process run (both branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids(False, 0, cap + 1)
+
+
+def test_invalid_local_branches_rejected():
+    with pytest.raises(ValueError, match="local_branches"):
+        _make_kv(local_branches=3)

tests/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        used = {int(s) // kv.block_size for s in slots}
+        assert kv.null_block_id not in used
+        assert used <= set(block_table)
+        # Resident window stays bounded by window + the in-flight chunk.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    is_negative: bool,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(is_negative)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.pos.completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 1
+    assert st._committed[False] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            is_negative=False,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.pos.completed_chunks
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        is_negative=False,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    assert torch.equal(paged, direct)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(False)
+        return out
+
+    torch._dynamo.reset()
+    from torch._dynamo.testing import CompileCounter
+
+    counter = CompileCounter()
+
+    def fn(inputs, q, k, v):
+        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+    compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+    run_one_forward(commit=True)  # history grows between calls ->
+    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+    run_one_forward(commit=False)
+
+    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"

tests/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        cross_attn_length=cross_attn_length,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(is_negative)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for neg in (False, True):
+        contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
+
+
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.pos.completed_chunks == 0
+    assert kv.window_block_ids(st.pos) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 2
+    assert st._committed[False] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.pos)) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+    assert kv.window_block_ids(st.pos) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, False, 3)
+    assert len(kv.window_block_ids(st.pos)) == 3
+
+    _prepare_and_commit(st, False, 2)
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
+
+
+def test_branches_are_independent():
+    kv, st = make_state()
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)
+    _prepare_and_commit(st, True, 1)
+    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {False: 0, True: 0}
+    assert st._paged_pending == {False: None, True: None}
+    assert kv.window_block_ids(st.pos) == []
+    assert kv.window_block_ids(st.neg) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+
+
+def test_get_cross_kv_caches_raises_before_populate():
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
+    _, st = make_state()  # cross_attn_length == 0, nothing populated
+    for neg in (False, True):
+        with pytest.raises(RuntimeError, match="cross-attn read before"):
+            st.get_cross_kv_caches(neg)
+    # Guards the full AND: the populated flag alone is not enough without a pool.
+    st._cross_text_populated[False] = True
+    with pytest.raises(RuntimeError, match="cross-attn read before"):
+        st.get_cross_kv_caches(False)
+
+
+def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for i in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        kv.write_cross_kv(i, False, k, v)
+        written.append((k, v))
+    st._cross_text_populated[False] = True
+
+    out = st.get_cross_kv_caches(False)
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
+
+
+def test_kv_create_is_noop_engine_owns_allocation():
+    from unittest.mock import MagicMock
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._ar_diffusion_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
+
+
+def test_get_kv_cache_requires_frame_aligned_seqlen():
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_failed_forward_tears_down_session(monkeypatch):
+    """Review follow-up (hsliuustc0106): a forward that dies partway must not
+    leave allocated-but-uncommitted KV behind — the session is torn down (pool
+    blocks freed, model-local state dropped) and the error propagates."""
+    from types import SimpleNamespace
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)  # session owns pool blocks
+    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
+
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.kv_cache = kv
+    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
+    runner._max_ar_diffusion_states = 4
+    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
+    runner.device = None
+    runner._perf_e2e_times = []
+
+    def boom(self, req, kv_prefetch_jobs=None):
+        raise RuntimeError("layer 17 exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    req = SimpleNamespace(
+        request_id="r0",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
+    )
+
+    with pytest.raises(RuntimeError, match="layer 17 exploded"):
+        ARDiffusionModelRunner.execute_model(runner, req)
+
+    assert "s1" not in runner._ar_diffusion_states
+    assert "s1" not in runner.pipeline._states
+    assert runner.pipeline._ar_diffusion_kv_state is None
+    # close() freed the session's resident blocks.
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure

tests/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size

tests/config/composable_parallel/test_apply.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for applying strategy specs onto merged stage configs."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.composable_parallel import (
+    Broadcast,
+    FanInByStage,
+    MeshAxisSpec,
+    RouteByStage,
+    StrategyApplyError,
+    StrategySpec,
+    TakeRank,
+    apply_strategy_specs,
+)
+from vllm_omni.config.stage_config import StageConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _tp(size: int) -> StrategySpec:
+    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
+
+
+def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
+    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
+
+
+def _stage(stage_id: int, model_stage: str, engine_args=None, runtime=None) -> StageConfig:
+    return StageConfig(
+        stage_id=stage_id,
+        model_stage=model_stage,
+        yaml_engine_args=dict(engine_args or {}),
+        yaml_runtime=dict(runtime or {"num_replicas": 1}),
+    )
+
+
+def _qwen_stages() -> list[StageConfig]:
+    return [
+        _stage(0, "thinker"),
+        _stage(1, "talker"),
+        _stage(2, "code2wav"),
+    ]
+
+
+def test_apply_tp_by_role():
+    stages = _qwen_stages()
+    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
+    # untouched roles keep their config
+    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
+
+
+def test_apply_by_model_stage():
+    stages = _qwen_stages()
+    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
+
+
+def test_apply_stage_replica_sets_num_replicas_and_surfaces_lb():
+    stages = _qwen_stages()
+    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
+    assert stages[1].yaml_runtime["num_replicas"] == 2
+    assert result.omni_lb_policy == "round-robin"
+
+
+def test_only_declared_axes_are_written():
+    stages = _qwen_stages()
+    # strategy declares only stage_replica -> tp must not be forced.
+    apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
+    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
+
+
+def test_conflict_on_explicit_tp():
+    stages = _qwen_stages()
+    stages[0].yaml_engine_args["tensor_parallel_size"] = 4
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+
+
+def test_equal_explicit_value_is_noop():
+    stages = _qwen_stages()
+    stages[0].yaml_engine_args["tensor_parallel_size"] = 2
+    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
+
+
+def test_explicit_none_conflicts_with_derived_value():
+    # An explicit YAML ``null`` (``tensor_parallel_size: null``) is a *present*
+    # value, not a missing key, so a strategy deriving a non-None size must raise
+    # rather than silently clobber the explicit None.
+    stages = _qwen_stages()
+    stages[0].yaml_engine_args["tensor_parallel_size"] = None
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+
+
+def test_missing_key_is_filled_without_conflict():
+    # A genuinely absent key (never set in the YAML) is filled by the strategy
+    # and must NOT raise — the contrast case to an explicit None.
+    stages = _qwen_stages()
+    assert "tensor_parallel_size" not in stages[0].yaml_engine_args
+    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
+
+
+def test_num_replicas_conflict():
+    stages = _qwen_stages()
+    stages[1].yaml_runtime["num_replicas"] = 3
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
+
+
+def test_device_count_ok_template():
+    stages = _qwen_stages()
+    stages[0].yaml_runtime["devices"] = "0,1"
+    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
+
+
+def test_device_count_ok_pool():
+    # tp=2 -> world=2; 2 replicas -> pool of 4 device ids is valid.
+    stages = _qwen_stages()
+    stages[1].yaml_runtime["devices"] = "0,1,2,3"
+    apply_strategy_specs(stages, {"talker": [_tp(2), _stage_replica(2)]})
+    assert stages[1].yaml_runtime["num_replicas"] == 2
+
+
+def test_device_count_mismatch():
+    stages = _qwen_stages()
+    stages[0].yaml_runtime["devices"] = "0,1,2"
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+
+
+def test_unknown_role_raises():
+    stages = _qwen_stages()
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"nonexistent": [_tp(2)]})
+
+
+def test_conflicting_lb_policy_across_roles():
+    stages = _qwen_stages()
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(
+            stages,
+            {
+                "talker": [_stage_replica(2, "round_robin")],
+                "code2wav": [_stage_replica(2, "least_queue")],
+            },
+        )

tests/config/composable_parallel/test_factory_integration.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.composable_parallel import (
+    Broadcast,
+    FanInByStage,
+    MeshAxisSpec,
+    RouteByStage,
+    StrategyApplyError,
+    StrategySpec,
+    TakeRank,
+    apply_strategy_specs,
+)
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DEPLOY = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "qwen2_5_omni.yaml"
+
+
+def _tp(size: int) -> StrategySpec:
+    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
+
+
+def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
+    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
+
+
+def _qwen_stages():
+    pipeline = OMNI_PIPELINES["qwen2_5_omni"]
+    deploy = load_deploy_config(_DEPLOY)
+    return merge_pipeline_deploy(pipeline, deploy)
+
+
+def _stage(stages, role):
+    return next(s for s in stages if s.model_stage == role)
+
+
+def test_overlay_tp_on_thinker():
+    stages = _qwen_stages()
+    # The bundled deploy pins the thinker to one GPU; a TP=2 strategy needs a
+    # matching 2-GPU layout (mirrors what a TP2 deploy would declare).
+    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
+    result = apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+    assert _stage(result.stages, "thinker").yaml_engine_args["tensor_parallel_size"] == 2
+
+
+def test_device_guard_rejects_tp_on_single_gpu_deploy():
+    # The bundled deploy pins the thinker to a single GPU, so the pre-spawn
+    # device check must refuse a TP=2 strategy on it.
+    stages = _qwen_stages()
+    with pytest.raises(StrategyApplyError):
+        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
+
+
+def test_overlay_stage_replica_on_talker():
+    stages = _qwen_stages()
+    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
+    assert _stage(result.stages, "talker").yaml_runtime["num_replicas"] == 2
+    assert result.omni_lb_policy == "round-robin"
+
+
+def test_overlay_mixed_roles():
+    stages = _qwen_stages()
+    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
+    result = apply_strategy_specs(
+        stages,
+        {
+            "thinker": [_tp(2)],
+            "talker": [_stage_replica(2, "round_robin")],
+            "code2wav": [_stage_replica(2, "round_robin")],
+        },
+    )
+    by_role = {s.model_stage: s for s in result.stages}
+    assert by_role["thinker"].yaml_engine_args["tensor_parallel_size"] == 2
+    assert by_role["talker"].yaml_runtime["num_replicas"] == 2
+    assert by_role["code2wav"].yaml_runtime["num_replicas"] == 2
+    assert result.omni_lb_policy == "round-robin"
+
+
+def _resolved(stages, role):
+    """Return a role's fully-resolved (post-CLI) OmegaConf stage config."""
+    stage = next(s for s in stages if s.model_stage == role)
+    return stage.to_omegaconf()
+
+
+def test_device_check_survives_cli_override():
+    # Strategy replicates the talker (1-GPU template -> valid at apply time),
+    # but a CLI --stage_1_devices with 3 ids must NOT slip past the device
+    # guard: effective world=1, replicas=2 admits only 1 or 2 device ids.
+    pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
+    with pytest.raises(StrategyApplyError):
+        StageConfigFactory._create_from_registry(
+            "qwen2_5_omni",
+            pipeline_cfg,
+            cli_overrides={"stage_1_devices": "0,1,2"},
+            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
+        )
+
+
+def test_cli_overrides_strategy_with_warning():
+    # Strategy derives num_replicas=2 for the talker; a CLI override to 3 wins
+    # (it is the most explicit user action) but must be surfaced loudly rather
+    # than silently. The resulting layout (1 template device) stays valid.
+    #
+    # vLLM's logger sets propagate=False, so attach a handler directly to it
+    # rather than relying on pytest's caplog (which listens on the root logger).
+    messages: list[str] = []
+
+    class _Capture(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            messages.append(record.getMessage())
+
+    log = logging.getLogger("vllm_omni.config.config_factory")
+    handler = _Capture(level=logging.WARNING)
+    log.addHandler(handler)
+    try:
+        pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
+        stages, _ = StageConfigFactory._create_from_registry(
+            "qwen2_5_omni",
+            pipeline_cfg,
+            cli_overrides={"stage_1_num_replicas": 3},
+            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
+        )
+    finally:
+        log.removeHandler(handler)
+
+    # CLI value wins in the resolved config.
+    assert _resolved(stages, "talker").runtime.num_replicas == 3
+    # ...and the override was warned about, naming the conflicting field.
+    assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)

tests/config/composable_parallel/test_spec_types.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the vendored spec/pattern data types."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.composable_parallel.aggregation import AggregationPattern, TakeRank
+from vllm_omni.config.composable_parallel.routing import RoutingPattern
+from vllm_omni.config.composable_parallel.spec import MeshAxisSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_mesh_axis_rejects_bad_kind():
+    with pytest.raises(ValueError):
+        MeshAxisSpec(kind="not_a_kind", size=2)  # type: ignore[arg-type]
+
+
+def test_mesh_axis_rejects_nonpositive_size():
+    with pytest.raises(ValueError):
+        MeshAxisSpec(kind="tp", size=0)
+
+
+def test_routing_pattern_is_closed():
+    with pytest.raises(TypeError):
+
+        class RogueRouting(RoutingPattern):  # noqa: F811
+            pass
+
+
+def test_aggregation_pattern_is_closed():
+    with pytest.raises(TypeError):
+
+        class RogueAggregation(AggregationPattern):  # noqa: F811
+            pass
+
+
+def test_take_rank_type_guard():
+    with pytest.raises(TypeError):
+        TakeRank(rank=True)  # bool is not a valid rank

tests/config/composable_parallel/test_strategy_loader.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for parsing/loading strategy files."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.composable_parallel.routing import Broadcast, RouteByStage
+from vllm_omni.config.composable_parallel.strategy_loader import (
+    StrategyLoadError,
+    load_strategy_specs,
+    parse_strategy_specs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_parse_basic():
+    data = {
+        "strategies": {
+            "thinker": [{"axis": "tp", "size": 2}],
+            "talker": [{"axis": "stage_replica", "size": 2, "routing": "round_robin"}],
+        }
+    }
+    specs = parse_strategy_specs(data)
+    assert set(specs) == {"thinker", "talker"}
+
+    (tp_spec,) = specs["thinker"]
+    assert tp_spec.mesh_axis.kind == "tp"
+    assert tp_spec.mesh_axis.size == 2
+    assert isinstance(tp_spec.routing, Broadcast)
+
+    (sr_spec,) = specs["talker"]
+    assert sr_spec.mesh_axis.kind == "stage_replica"
+    assert isinstance(sr_spec.routing, RouteByStage)
+    assert sr_spec.routing.routing_policy == "round_robin"
+
+
+def test_parse_without_strategies_key():
+    # A bare role mapping (no top-level "strategies") is also accepted.
+    specs = parse_strategy_specs({"thinker": [{"axis": "tp", "size": 4}]})
+    assert specs["thinker"][0].mesh_axis.size == 4
+
+
+def test_l1_owner_goes_to_shard_extension():
+    specs = parse_strategy_specs({"talker": [{"axis": "stage_replica", "size": 2, "l1_owner": "delegated"}]})
+    assert specs["talker"][0].shard_extension["l1_owner"] == "delegated"
+
+
+def test_missing_axis_raises():
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": [{"size": 2}]})
+
+
+def test_missing_size_raises():
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": [{"axis": "tp"}]})
+
+
+def test_routing_on_non_policy_axis_raises():
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": [{"axis": "tp", "size": 2, "routing": "random"}]})
+
+
+def test_bad_size_raises():
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": [{"axis": "tp", "size": "two"}]})
+
+
+def test_entries_must_be_list():
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": {"axis": "tp", "size": 2}})
+
+
+def test_non_mapping_entry_raises():
+    # A list whose elements are not mappings (e.g. a bare string) must raise a
+    # StrategyLoadError, not an opaque TypeError from dict(entry).
+    with pytest.raises(StrategyLoadError):
+        parse_strategy_specs({"thinker": ["tp"]})
+
+
+def test_load_from_file(tmp_path):
+    path = tmp_path / "strategy.yaml"
+    path.write_text(
+        "strategies:\n"
+        "  thinker:\n"
+        "    - axis: tp\n"
+        "      size: 2\n"
+        "  talker:\n"
+        "    - axis: stage_replica\n"
+        "      size: 2\n"
+        "      routing: least_queue\n"
+    )
+    specs = load_strategy_specs(str(path))
+    assert specs["thinker"][0].mesh_axis.size == 2
+    assert specs["talker"][0].routing.routing_policy == "least_queue"

tests/config/composable_parallel/test_translator.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for StrategySpec -> OmniParallelConfig translation."""
+
+from __future__ import annotations
+
+import dataclasses
+
+import pytest
+
+from vllm_omni.config.composable_parallel import (
+    AxisTranslationError,
+    Broadcast,
+    FanInByStage,
+    MeshAxisSpec,
+    PartitionByHash,
+    PipelineMicrobatch,
+    RouteByStage,
+    StrategySpec,
+    TakeRank,
+    Union,
+    translate_strategy_stack,
+)
+from vllm_omni.config.composable_parallel.aggregation import StitchPipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _tp(size: int) -> StrategySpec:
+    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
+
+
+def _dp(size: int, policy: str = "round_robin") -> StrategySpec:
+    return StrategySpec("dp", MeshAxisSpec("dp", size), RouteByStage(policy), Union())
+
+
+def _pp(size: int) -> StrategySpec:
+    return StrategySpec("pp", MeshAxisSpec("pp", size), PipelineMicrobatch(), StitchPipeline())
+
+
+def _ep(size: int) -> StrategySpec:
+    return StrategySpec("ep", MeshAxisSpec("ep", size), Broadcast(), Union())
+
+
+def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
+    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
+
+
+def test_tp_only():
+    cfg = translate_strategy_stack([_tp(2)])
+    assert cfg.tensor_parallel_size == 2
+    assert cfg.data_parallel_size == 1
+    assert cfg.world_size == 2
+    assert cfg.stage_replica_size == 1
+    assert cfg.l1_owners == {"tp": "engine"}
+
+
+def test_dp_is_engine_owned():
+    cfg = translate_strategy_stack([_dp(2)])
+    assert cfg.data_parallel_size == 2
+    assert cfg.world_size == 2
+    # engine DP does not emit an omni LB policy (vLLM's internal DP LB is used).
+    assert cfg.omni_lb_policy is None
+    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
+    assert cfg.l1_owners["dp"] == "engine"
+
+
+def test_tp_times_dp_world_size():
+    cfg = translate_strategy_stack([_tp(2), _dp(2)])
+    assert cfg.world_size == 4
+    assert cfg.tensor_parallel_size == 2
+    assert cfg.data_parallel_size == 2
+
+
+def test_stage_replica_is_delegated():
+    cfg = translate_strategy_stack([_stage_replica(3, "least_queue")])
+    # stage_replica is NOT a world dimension.
+    assert cfg.world_size == 1
+    assert cfg.stage_replica_size == 3
+    assert cfg.omni_lb_policy == "least-queue-length"
+    assert cfg.l1_owners["stage_replica"] == "delegated"
+    # stage_replica_size is not an engine kwarg; it's a per-stage deploy knob.
+    assert "stage_replica_size" not in cfg.as_engine_kwargs()
+
+
+def test_stage_replica_lb_policy_only_when_replicated():
+    cfg = translate_strategy_stack([_stage_replica(1)])
+    assert cfg.stage_replica_size == 1
+    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
+
+
+def test_tp_with_stage_replica():
+    cfg = translate_strategy_stack([_tp(2), _stage_replica(2)])
+    assert cfg.world_size == 2  # tp only
+    assert cfg.tensor_parallel_size == 2
+    assert cfg.stage_replica_size == 2
+    assert cfg.omni_lb_policy == "round-robin"
+
+
+def test_ep_must_match_tp_dp_product():
+    cfg = translate_strategy_stack([_tp(2), _ep(2)])
+    assert cfg.enable_expert_parallel is True
+    assert cfg.as_engine_kwargs()["enable_expert_parallel"] is True
+
+
+def test_ep_mismatch_raises():
+    with pytest.raises(Exception) as exc:
+        translate_strategy_stack([_tp(2), _ep(4)])
+    assert "must equal" in str(exc.value)
+
+
+def test_duplicate_kind_rejected():
+    with pytest.raises(AxisTranslationError):
+        translate_strategy_stack([_tp(2), _tp(2)])
+
+
+def test_unsupported_kind_rejected():
+    spec = StrategySpec("sp", MeshAxisSpec("sp_ulysses", 2), Broadcast(), TakeRank())
+    with pytest.raises(AxisTranslationError):
+        translate_strategy_stack([spec])
+
+
+def test_dp_hash_routing_rejected():
+    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), PartitionByHash(), Union())
+    with pytest.raises(NotImplementedError):
+        translate_strategy_stack([spec])
+
+
+def test_dp_invalid_routing_policy_rejected():
+    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), RouteByStage("bogus"), Union())
+    with pytest.raises(AxisTranslationError):
+        translate_strategy_stack([spec])
+
+
+def test_lb_policy_never_in_engine_kwargs_even_when_replicated():
+    # omni_lb_policy is pipeline-wide (applied at engine construction), not a
+    # per-stage engine arg, so it must never leak into as_engine_kwargs().
+    cfg = translate_strategy_stack([_tp(2), _stage_replica(2, "least_queue")])
+    assert cfg.omni_lb_policy == "least-queue-length"
+    assert "omni_lb_policy" not in cfg.as_engine_kwargs()
+
+
+def test_stage_replica_hash_routing_rejected():
+    spec = StrategySpec("sr", MeshAxisSpec("stage_replica", 2), RouteByStage("hash"), FanInByStage())
+    with pytest.raises(NotImplementedError):
+        translate_strategy_stack([spec])
+
+
+def test_dp_wrong_owner_rejected():
+    spec = StrategySpec(
+        "dp",
+        MeshAxisSpec("dp", 2),
+        RouteByStage("round_robin"),
+        Union(),
+        shard_extension={"l1_owner": "delegated"},
+    )
+    with pytest.raises(AxisTranslationError):
+        translate_strategy_stack([spec])
+
+
+def test_stage_replica_wrong_owner_rejected():
+    spec = StrategySpec(
+        "sr",
+        MeshAxisSpec("stage_replica", 2),
+        RouteByStage("round_robin"),
+        FanInByStage(),
+        shard_extension={"l1_owner": "engine"},
+    )
+    with pytest.raises(AxisTranslationError):
+        translate_strategy_stack([spec])
+
+
+@pytest.mark.parametrize(
+    "specs",
+    [
+        [_tp(2)],
+        [_dp(2)],
+        [_tp(2), _dp(2)],
+        [_tp(2), _ep(2)],
+        [_tp(2), _stage_replica(2)],
+        [_pp(2)],
+    ],
+)
+def test_as_engine_kwargs_are_valid_omni_engine_args(specs):
+    # Guard against desync: every key emitted by as_engine_kwargs() must be a
+    # real OmniEngineArgs field, otherwise constructing the per-stage engine args
+    # would explode at runtime. Keep this CPU-only by checking field names rather
+    # than instantiating the engine.
+    from vllm_omni.engine.arg_utils import OmniEngineArgs
+
+    valid_fields = {f.name for f in dataclasses.fields(OmniEngineArgs)}
+    cfg = translate_strategy_stack(specs)
+    kwargs = cfg.as_engine_kwargs()
+    unknown = set(kwargs) - valid_fields
+    assert not unknown, f"as_engine_kwargs emitted non-OmniEngineArgs fields: {unknown}"

tests/config/test_endpoint_policy.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for endpoint restrictions logic."""
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.config.endpoint_policy import (
+    EndpointRestriction,
+    OmniServingCapability,
+    shutdown_unsupported_routes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SERVER_ROUTE = "/v1/completions"
+REJECTION_REASON = "This route is banned!"
+
+
+def _make_app_with_server_route():
+    """Create a FastAPI app with a route to reject."""
+    app = FastAPI()
+
+    @app.post(SERVER_ROUTE)
+    async def existing_handler():
+        """Returns a response; we want to disable this."""
+        return {"ok": True}
+
+    return app
+
+
+def test_restricted_completions_returns_400():
+    """Ensure that simple route rejection works."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_unrestricted_completions_not_blocked():
+    """Ensure that if we don't shutdown any routes, the route stays."""
+    app = _make_app_with_server_route()
+    shutdown_unsupported_routes(app, ())
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE)
+    assert resp.status_code == 200
+    assert resp.json() == {"ok": True}
+
+
+def test_route_rejection_is_idempotent():
+    """Ensure that calling rejection twice doesn't break things."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_invalid_route_rejection():
+    """Ensure that a restriction won't crash if a route doesn't exist."""
+    # NOTE: This test is for safety since the route organization is still messy.
+    # Since the restriction needs a valid OmniServingCapability, its implied the route
+    # is a valid vLLM route, but was already removed from the app by earlier processing.
+    app = FastAPI()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"

tests/config/test_omni_config.py

@@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "distributed_executor_backend",
     } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
     assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
+
+
+def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
+    quantization_config = {
+        "method": "example_quant",
+        "weights": "weights.bin",
+    }
+
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
+
+    assert cfg.quantization_config == quantization_config

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
+from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _MockQueue:
+    def __init__(self, items: list | None = None) -> None:
+        self._items = list(items or [])
+
+    def __bool__(self) -> bool:
+        return bool(self._items)
+
+    def __iter__(self):
+        return iter(self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def add_request(self, request) -> None:
+        self._items.append(request)
+
+    def prepend_requests(self, requests) -> None:
+        self._items = list(requests) + self._items
+
+
+class _MockRequest:
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        status: RequestStatus = RequestStatus.RUNNING,
+        num_prompt_tokens: int = 4,
+        num_computed_tokens: int = 4,
+        num_output_placeholders: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.num_prompt_tokens = num_prompt_tokens
+        self.num_computed_tokens = num_computed_tokens
+        self.num_output_placeholders = num_output_placeholders
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+@pytest.fixture(autouse=True)
+def _mock_cuda_graph_platform(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
+
+
+def _make_scheduler(
+    *,
+    enable_unified_decode_graph: bool | None = True,
+    deterministic_cfm_noise: bool = False,
+    runtime_config_on_model_config: bool = False,
+) -> VoxCPM2OmniARAsyncScheduler:
+    sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
+    hf_config = SimpleNamespace()
+    model_config = SimpleNamespace(hf_config=hf_config)
+    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
+        runtime_config = SimpleNamespace(
+            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
+            deterministic_cfm_noise=deterministic_cfm_noise,
+        )
+        if runtime_config_on_model_config:
+            model_config.voxcpm2_runtime_config = runtime_config
+        else:
+            hf_config.voxcpm2_runtime_config = runtime_config
+    sched.vllm_config = SimpleNamespace(model_config=model_config)
+    return sched
+
+
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("prefill-running", num_prompt_tokens=8, num_computed_tokens=4)]
+    scheduler.waiting = _MockQueue([_MockRequest("waiting", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=False)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=None)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None:
+    scheduler = _make_scheduler(deterministic_cfm_noise=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
+    deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
+    original_waiting = _MockQueue([original_waiting_req])
+    scheduler.waiting = original_waiting
+    scheduler.policy = "fcfs"
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler._consume_pending_connector_output = lambda model_mode: None
+
+    monkeypatch.setattr(scheduler_mod, "create_request_queue", lambda _policy: _MockQueue())
+
+    def fake_upstream_schedule(self, throttle_prefills: bool = False):
+        assert self.waiting is not original_waiting
+        assert not self.waiting
+        self.waiting.add_request(deferred_by_upstream)
+        raise RuntimeError("stop before output wrapping")
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
+
+    with pytest.raises(RuntimeError, match="stop before output wrapping"):
+        scheduler.schedule()
+
+    assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]

tests/dfx/conftest.py

@@ -11,11 +11,180 @@ from typing import Any
 import pytest
 
 from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
+from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 
+def _named_pytest_marks(names: list[str]) -> list[pytest.MarkDecorator]:
+    marks: list[pytest.MarkDecorator] = []
+    for name in names:
+        name = name.strip()
+        if not name:
+            raise ValueError("mark name must be a non-empty string")
+        marks.append(getattr(pytest.mark, name))
+    return marks
+
+
+def _hardware_marks_from_dict(hw: Any) -> list[pytest.MarkDecorator]:
+    if not isinstance(hw, dict):
+        raise ValueError(f"mark.hardware_marks must be a dict, got {type(hw).__name__}")
+    res = hw.get("res")
+    if not isinstance(res, dict):
+        raise ValueError(f"mark.hardware_marks.res must be a dict, got {type(res).__name__}")
+    num_cards = hw.get("num_cards", 1)
+    return list(hardware_marks(res=res, num_cards=num_cards))
+
+
+def resolve_pytest_marks(mark_field: Any) -> list[pytest.MarkDecorator]:
+    """Convert a JSON ``mark`` field into pytest mark decorators.
+
+    Supported form (per test-case object in perf/stability JSON)::
+
+        "mark": [
+            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+            "full_model",
+            "diffusion"
+        ]
+
+    Exactly one ``hardware_marks`` object is required when ``mark`` is present.
+    ``hardware_marks`` delegates to :func:`tests.helpers.mark.hardware_marks`
+    (same shape as ``@hardware_test``). Additional array entries are registered
+    pytest marker names (strings).
+    """
+    if mark_field is None:
+        return []
+
+    if isinstance(mark_field, list):
+        marks: list[pytest.MarkDecorator] = []
+        hw_seen = False
+        for item in mark_field:
+            if isinstance(item, dict) and "hardware_marks" in item:
+                if hw_seen:
+                    raise ValueError("mark array must contain at most one hardware_marks object")
+                hw_seen = True
+                marks.extend(_hardware_marks_from_dict(item["hardware_marks"]))
+                unknown_keys = set(item) - {"hardware_marks"}
+                if unknown_keys:
+                    raise ValueError(
+                        f"mark hardware_marks object only allows hardware_marks; unknown keys: {sorted(unknown_keys)}"
+                    )
+            elif isinstance(item, str):
+                item = item.strip()
+                if not item:
+                    raise ValueError("mark name must be a non-empty string")
+                marks.extend(_named_pytest_marks([item]))
+            else:
+                raise ValueError(
+                    f"mark array entries must be hardware_marks objects or marker name strings; got {type(item).__name__}"
+                )
+        if not hw_seen:
+            raise ValueError("mark array must contain a hardware_marks object")
+        return marks
+
+    raise ValueError(f"mark must be a list; got {type(mark_field).__name__}")
+
+
+def _mark_names(mark_field: Any) -> set[str]:
+    if isinstance(mark_field, list):
+        return {str(item) for item in mark_field if isinstance(item, str)}
+    return set()
+
+
+def is_diffusion_perf_config(cfg: dict[str, Any]) -> bool:
+    """True for perf JSON cases intended for ``run_diffusion_benchmark.py``."""
+    if cfg.get("server_type") is not None:
+        return True
+    return "diffusion" in _mark_names(cfg.get("mark"))
+
+
+def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
+    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
+
+
+def create_unique_server_pytest_params(
+    configs: list[dict[str, Any]],
+    stage_configs_dir: Path,
+) -> list[Any]:
+    """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
+    marks_by_name = _marks_by_test_name(configs)
+    return [
+        pytest.param(
+            row,
+            marks=marks_by_name.get(row[0], []),
+            id=row[0],
+        )
+        for row in create_unique_server_params(configs, stage_configs_dir)
+    ]
+
+
+def create_benchmark_pytest_params(
+    benchmark_configs: list[dict[str, Any]],
+    server_to_benchmark_mapping: dict[str, dict],
+) -> list[Any]:
+    """Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
+    marks_by_name = _marks_by_test_name(benchmark_configs)
+    params: list[Any] = []
+    seen: set[str] = set()
+    for config in benchmark_configs:
+        test_name = config["test_name"]
+        if test_name in seen:
+            continue
+        seen.add(test_name)
+        params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
+        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
+        for idx, id_suffix in enumerate(id_suffixes):
+            params.append(
+                pytest.param(
+                    (test_name, idx),
+                    marks=marks_by_name.get(test_name, []),
+                    id=f"{test_name}-{id_suffix}",
+                )
+            )
+    return params
+
+
+def create_paired_benchmark_pytest_params(
+    server_entries: list[tuple[Any, str]],
+    params_by_test_name: dict[str, list[Any]],
+    marks_by_name: dict[str, list[pytest.MarkDecorator]],
+) -> list[Any]:
+    """One ``pytest.param`` per ``(server entry, benchmark index)`` pair.
+
+    Pass ``server_entry`` and ``(test_name, idx)`` as separate ``pytest.param``
+    arguments for ``@pytest.mark.parametrize("server_fixture,benchmark_params", ...)``.
+    """
+    pairs: list[Any] = []
+    for server_entry, test_name in server_entries:
+        params_list = params_by_test_name.get(test_name, [])
+        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
+        for idx, id_suffix in enumerate(id_suffixes):
+            pairs.append(
+                pytest.param(
+                    server_entry,
+                    (test_name, idx),
+                    marks=marks_by_name.get(test_name, []),
+                    id=f"{test_name}-{id_suffix}",
+                )
+            )
+    return pairs
+
+
+def create_paired_omni_benchmark_pytest_params(
+    configs: list[dict[str, Any]],
+    stage_configs_dir: Path,
+) -> list[Any]:
+    """Paired params for ``run_benchmark.py`` (omni/tts)."""
+    mapping = create_test_parameter_mapping(configs)
+    marks_by_name = _marks_by_test_name(configs)
+    server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)]
+    params_by_test_name = {
+        test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries
+    }
+    return create_paired_benchmark_pytest_params(server_entries, params_by_test_name, marks_by_name)
+
+
 def load_configs(config_path: str) -> list[dict[str, Any]]:
     try:
         abs_path = Path(config_path).resolve()
@@ -32,6 +201,29 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
+def load_benchmark_configs(
+    config_path: str | None = None,
+    *,
+    config_dir: Path | None = None,
+) -> list[dict[str, Any]]:
+    """Load one benchmark JSON file, or merge all ``*.json`` under *config_dir*.
+
+    When *config_path* is omitted, every ``*.json`` in *config_dir* is loaded and
+    concatenated. Pytest ``-m`` expressions then select cases via each entry's
+    ``mark`` field (e.g. ``-m tts`` after bulk load).
+    """
+    if config_path is not None:
+        return load_configs(config_path)
+    if config_dir is None:
+        raise ValueError("load_benchmark_configs requires config_path or config_dir")
+    configs: list[dict[str, Any]] = []
+    for path in sorted(config_dir.glob("*.json")):
+        configs.extend(load_configs(str(path)))
+    if not configs:
+        raise ValueError(f"No benchmark JSON files found under {config_dir}")
+    return configs
+
+
 def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str:
     kwargs: dict[str, Any] = {}
     if updates is not None:
@@ -293,6 +485,158 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
+def _benchmark_param_id_suffix(param: dict[str, Any], *, idx: int) -> str:
+    """Derive a readable pytest id suffix from one ``benchmark_params`` entry."""
+    name = param.get("name")
+    if isinstance(name, str) and name.strip():
+        return _safe_filename_token(name.strip())
+
+    parts: list[str] = []
+    for key in ("task", "eval_phase", "dataset_name", "dataset"):
+        value = param.get(key)
+        if value is None or value == "":
+            continue
+        token = _safe_filename_token(str(value))
+        if token != "na":
+            parts.append(token)
+    if parts:
+        return "_".join(parts)
+
+    return f"case{idx}"
+
+
+def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> list[str]:
+    """Return unique pytest id suffixes for a server's benchmark param list."""
+    raw = [_benchmark_param_id_suffix(param, idx=idx) for idx, param in enumerate(params_list)]
+    seen: dict[str, int] = {}
+    unique: list[str] = []
+    for suffix in raw:
+        count = seen.get(suffix, 0)
+        seen[suffix] = count + 1
+        if count == 0:
+            unique.append(suffix)
+        else:
+            unique.append(f"{suffix}_{count}")
+    return unique
+
+
+def extract_mark_resource_label(mark_field: Any) -> str:
+    """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values.
+
+    Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``.
+
+    Prefer :func:`get_runtime_resource_label` for perf result filenames so labels
+    reflect the machine that actually ran the benchmark.
+    """
+    if isinstance(mark_field, list):
+        for item in mark_field:
+            if isinstance(item, dict) and "hardware_marks" in item:
+                return extract_mark_resource_label(item)
+        return "na"
+    if not isinstance(mark_field, dict):
+        return "na"
+    hw = mark_field.get("hardware_marks")
+    if not isinstance(hw, dict):
+        return "na"
+    res = hw.get("res")
+    if not isinstance(res, dict) or not res:
+        return "na"
+    labels = [_safe_filename_token(value) for value in res.values()]
+    return "-".join(labels) if labels else "na"
+
+
+_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = (
+    "H100",
+    "H800",
+    "H200",
+    "H20",
+    "L40S",
+    "L40",
+    "L4",
+    "A100",
+    "A800",
+    "A10G",
+    "A10",
+    "A30",
+    "MI325",
+    "MI300",
+    "MI250",
+    "B60",
+    "S5000",
+    "910B4",
+    "910B",
+    "910",
+    "310P",
+    "A2",
+    "A3",
+)
+_RUNTIME_RESOURCE_LABEL: str | None = None
+
+
+def _normalize_runtime_device_label(raw: str) -> str:
+    """Map a platform device name to a short filename-safe resource token."""
+    if not raw or not str(raw).strip():
+        return "na"
+    upper = str(raw).upper()
+    for token in _KNOWN_RUNTIME_RESOURCE_TOKENS:
+        if token.upper() in upper:
+            return _safe_filename_token(token)
+    compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw))
+    for prefix in ("NVIDIA", "AMD", "ASCEND", "HUAWEI"):
+        if compact.upper().startswith(prefix):
+            compact = compact[len(prefix) :]
+            break
+    return _safe_filename_token(compact[:48]) if compact else "na"
+
+
+def _read_runtime_device_name(*, device_id: int = 0) -> str | None:
+    """Device name from the active Omni platform."""
+    if current_omni_platform.device_count() <= device_id:
+        return None
+    get_name = getattr(current_omni_platform, "get_device_name", None)
+    if not callable(get_name):
+        return None
+    raw = get_name(device_id)
+    if isinstance(raw, str) and raw.strip():
+        return raw.strip()
+    return None
+
+
+def get_runtime_resource_label(*, device_id: int = 0, refresh: bool = False) -> str:
+    """Return a filename-safe hardware label detected on the running machine."""
+    global _RUNTIME_RESOURCE_LABEL
+    if not refresh and _RUNTIME_RESOURCE_LABEL is not None:
+        return _RUNTIME_RESOURCE_LABEL
+    raw = _read_runtime_device_name(device_id=device_id)
+    label = _normalize_runtime_device_label(raw) if raw else "na"
+    if not refresh:
+        _RUNTIME_RESOURCE_LABEL = label
+    return label
+
+
+_FILENAME_OMIT_RESOURCE_LABELS = frozenset({"H100"})
+
+
+def hardware_json_value(resource_label: str | None) -> str:
+    """Hardware token stored in perf result JSON (empty when unknown)."""
+    token = _safe_filename_token(resource_label)
+    return "" if token == "na" else token
+
+
+def resource_label_for_filename(resource_label: str | None) -> str:
+    """Hardware token embedded in result filenames (H100 omitted on default CI pool)."""
+    token = _safe_filename_token(resource_label)
+    if token in _FILENAME_OMIT_RESOURCE_LABELS:
+        return ""
+    return token
+
+
+def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
+    """Return runtime hardware label for perf result filenames."""
+    del configs
+    return get_runtime_resource_label()
+
+
 def resolve_baseline_value(
     baseline_raw: Any,
     *,
@@ -357,6 +701,7 @@ def run_benchmark(
     max_concurrency: Any | None = None,
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
+    resource_label: str | None = None,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
@@ -367,7 +712,11 @@ def run_benchmark(
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     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"
+    hw = resource_label_for_filename(resource_label)
+    if hw:
+        result_filename = f"result_{test_name}_{hw}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
+    else:
+        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 = (
@@ -437,6 +786,7 @@ def run_benchmark(
         result["random_input_len"] = random_input_len
     if random_output_len is not None:
         result["random_output_len"] = random_output_len
+    result["Hardware"] = hardware_json_value(resource_label)
     with open(result_path, "w", encoding="utf-8") as f:
         json.dump(result, f, ensure_ascii=False, indent=2)
     return result

tests/dfx/perf/scripts/run_benchmark.py

@@ -7,21 +7,24 @@ from typing import Any
 import pytest
 
 from tests.dfx.conftest import (
-    create_benchmark_indices,
+    create_paired_omni_benchmark_pytest_params,
     create_test_parameter_mapping,
-    create_unique_server_params,
     get_benchmark_params_for_server,
-    load_configs,
+    get_runtime_resource_label,
+    is_diffusion_perf_config,
+    load_benchmark_configs,
     resolve_baseline_value,
     run_benchmark,
 )
 from tests.helpers.runtime import OmniServer
 
-pytestmark = [pytest.mark.full_model]
-
 # Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
 # (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
 # defined in the same module.
+#
+# Optional JSON field ``mark`` is applied as pytest marks via
+# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
+# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
 
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -40,22 +43,22 @@ def _get_config_file_from_argv() -> str | 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:
+    _all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR)
+    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)]
     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)"
+        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts case(s) from "
+        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; "
+        f"use -m to filter, e.g. -m tts)"
     )
-    CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
-
-BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-
+else:
+    BENCHMARK_CONFIGS = load_benchmark_configs(CONFIG_FILE_PATH)
 
 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)
+paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 
 _omni_server_lock = threading.Lock()
 
@@ -91,17 +94,11 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.fixture(params=benchmark_indices)
-def benchmark_params(request, omni_server):
-    """Benchmark parameters fixture with proper parametrization"""
+@pytest.fixture
+def benchmark_params(request):
+    """Benchmark parameters fixture; paired with ``omni_server`` via parametrization."""
     test_name, param_index = request.param
 
-    if test_name != omni_server.test_name:
-        pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}")
-
     all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
 
     if not all_params:
@@ -153,8 +150,11 @@ def assert_result(
 
 
 @pytest.mark.benchmark
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
+@pytest.mark.parametrize(
+    "omni_server,benchmark_params",
+    paired_benchmark_params,
+    indirect=["omni_server", "benchmark_params"],
+)
 def test_performance_benchmark(omni_server, benchmark_params, request):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
@@ -168,6 +168,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
     print(f"Benchmark parameters: {benchmark_params}")
 
     assert_baseline = request.config.getoption("--assert-baseline", default=False)
+    resource_label = get_runtime_resource_label()
 
     def to_list(value, default=None):
         if value is None:
@@ -239,6 +240,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=None,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
+            resource_label=resource_label,
         )
         assert_result(
             result,
@@ -264,6 +266,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=concurrency,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
+            resource_label=resource_label,
         )
         assert_result(
             result,

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -8,15 +8,21 @@ This runner separates two concepts:
 2. ``benchmark_endpoint``: which serving API the benchmark client calls.
    Examples: ``/v1/chat/completions`` and ``/v1/videos``.
 
-A config JSON file is REQUIRED via --test-config-file:
+A config JSON file may be passed via --test-config-file. If omitted, every ``*.json`` under
+``tests/dfx/perf/tests/`` is loaded and pytest ``-m`` filters by each case's ``mark``:
+  pytest run_diffusion_benchmark.py -m "diffusion"
   pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
 Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
 
-All benchmark results for a session are consolidated into a single JSON file under
-BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
-Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
-timestamp) together with the raw metrics returned by the benchmark script.
+Optional JSON field ``mark`` is applied as pytest marks on that case via
+``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
+
+All benchmark results are written under BENCHMARK_RESULT_DIR (override via the
+DIFFUSION_BENCHMARK_DIR environment variable). Each source JSON file gets one
+aggregated ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` (JSON array
+of all runs from cases in that file). Bulk load without ``--test-config-file`` uses
+the same per-file aggregation; ``-m`` only selects which cases run.
 """
 
 import json
@@ -35,8 +41,15 @@ import psutil
 import pytest
 
 from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+from tests.dfx.conftest import (
+    create_paired_benchmark_pytest_params,
+    get_runtime_resource_label,
+    hardware_json_value,
+    is_diffusion_perf_config,
+    resolve_pytest_marks,
+    resource_label_for_filename,
+)
+from tests.helpers.runtime import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -126,6 +139,10 @@ _BRANCHPOINT_COMMIT_SHA: str | None = None
 DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json"
 
 
+_DIFFUSION_SOURCE_CONFIG_KEY = "_source_config_file"
+_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
+
+
 def _get_config_file_from_argv() -> str | None:
     """Read --test-config-file from sys.argv at import time so pytest parametrize can use it.
 
@@ -142,9 +159,6 @@ def _get_config_file_from_argv() -> str | None:
 
 
 CONFIG_FILE_PATH = _get_config_file_from_argv()
-if CONFIG_FILE_PATH is None:
-    print("No config file provided, using default config file: tests/dfx/perf/tests/test_qwen_image_vllm_omni.json")
-    CONFIG_FILE_PATH = "tests/dfx/perf/tests/test_qwen_image_vllm_omni.json"
 
 # ---------------------------------------------------------------------------
 # Config loading
@@ -184,28 +198,85 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-
-_config_stem = Path(CONFIG_FILE_PATH).stem  # e.g. "test_qwen_image_vllm_omni"
-AGGREGATED_RESULT_FILE = BENCHMARK_RESULT_DIR / f"diffusion_result_{_config_stem}_{_SESSION_TIMESTAMP}.json"
+def load_diffusion_benchmark_configs(
+    config_path: str | None = None,
+    *,
+    config_dir: Path | None = None,
+) -> list[dict[str, Any]]:
+    """Load one diffusion benchmark JSON, or merge all ``*.json`` under *config_dir*."""
+    if config_path is not None:
+        configs = load_configs(config_path)
+        source = str(Path(config_path).resolve())
+        for cfg in configs:
+            cfg.setdefault(_DIFFUSION_SOURCE_CONFIG_KEY, source)
+        return configs
+    if config_dir is None:
+        raise ValueError("load_diffusion_benchmark_configs requires config_path or config_dir")
+    configs: list[dict[str, Any]] = []
+    for path in sorted(config_dir.glob("*.json")):
+        source = str(path.resolve())
+        for cfg in load_configs(str(path)):
+            cfg[_DIFFUSION_SOURCE_CONFIG_KEY] = source
+            configs.append(cfg)
+    if not configs:
+        raise ValueError(f"No benchmark JSON files found under {config_dir}")
+    return configs
 
 
-def _append_to_aggregated_file(record: dict[str, Any]) -> None:
-    """Thread-safe append of *record* to the session-level aggregated JSON file.
+if CONFIG_FILE_PATH is None:
+    _all_configs = load_diffusion_benchmark_configs(config_dir=_PERF_TESTS_DIR)
+    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if is_diffusion_perf_config(cfg)]
+    print(
+        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} diffusion case(s) from "
+        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} omni/tts; "
+        f"use -m to filter, e.g. -m diffusion)"
+    )
+else:
+    BENCHMARK_CONFIGS = load_diffusion_benchmark_configs(CONFIG_FILE_PATH)
 
-    The file contains a JSON array; each call loads the existing array (or
-    starts a new one), appends the record, and writes the file back atomically.
-    """
+_AGGREGATED_RESULT_FILES_BY_SOURCE: dict[str, Path] = {}
+
+
+def _normalized_source_path(source_file: str) -> str:
+    return str(Path(source_file).resolve())
+
+
+def _aggregated_result_file_for_source(source_file: str) -> Path:
+    """One session aggregate per source JSON (same naming as single ``--test-config-file``)."""
+    key = _normalized_source_path(source_file)
+    if key not in _AGGREGATED_RESULT_FILES_BY_SOURCE:
+        stem = Path(key).stem
+        resource = resource_label_for_filename(get_runtime_resource_label())
+        if resource:
+            result_name = f"diffusion_result_{stem}_{resource}_{_SESSION_TIMESTAMP}.json"
+        else:
+            result_name = f"diffusion_result_{stem}_{_SESSION_TIMESTAMP}.json"
+        _AGGREGATED_RESULT_FILES_BY_SOURCE[key] = BENCHMARK_RESULT_DIR / result_name
+    return _AGGREGATED_RESULT_FILES_BY_SOURCE[key]
+
+
+def _write_result_record(record: dict[str, Any]) -> Path:
+    """Append one benchmark record to the aggregate file for its source JSON."""
+    source_file = record.get("source_file") or CONFIG_FILE_PATH
+    if not source_file:
+        raise ValueError("benchmark record missing source_file")
+    target = _aggregated_result_file_for_source(str(source_file))
     with _RESULT_LOCK:
         BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True)
-        if AGGREGATED_RESULT_FILE.exists():
-            with open(AGGREGATED_RESULT_FILE, encoding="utf-8") as f:
+        if target.exists():
+            with open(target, encoding="utf-8") as f:
                 records: list[dict] = json.load(f)
         else:
             records = []
         records.append(record)
-        with open(AGGREGATED_RESULT_FILE, "w", encoding="utf-8") as f:
+        with open(target, "w", encoding="utf-8") as f:
             json.dump(records, f, indent=2, ensure_ascii=False)
+    return target
+
+
+def _append_to_aggregated_file(record: dict[str, Any]) -> None:
+    """Backward-compatible wrapper; prefer :func:`_write_result_record`."""
+    _write_result_record(record)
 
 
 _server_lock = threading.Lock()
@@ -215,14 +286,6 @@ _server_lock = threading.Lock()
 # ---------------------------------------------------------------------------
 
 
-def _get_open_port() -> int:
-    """Return an available TCP port on localhost."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return s.getsockname()[1]
-
-
 def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
     """Block until the given host:port accepts connections or timeout expires.
 
@@ -312,7 +375,7 @@ class DiffusionServer:
         self.model = server_cfg["model"]
         self.serve_args = server_cfg["serve_args"]
         self.host = "127.0.0.1"
-        self.port = port if port is not None else _get_open_port()
+        self.port = port if port is not None else get_open_port(self.host)
         self.proc: subprocess.Popen | None = None
         self.test_name: str = ""
 
@@ -482,6 +545,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
                 "serve_args": _build_serve_args(serve_args_dict),
                 "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                 "server_params": cfg["server_params"],
+                "mark": cfg.get("mark"),
             }
         )
     return result
@@ -496,19 +560,29 @@ def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]:
     return mapping
 
 
-def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
-    """Factory: return a vLLM-Omni diffusion server instance for the config."""
-    return DiffusionServer(server_cfg=server_cfg)
+def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
+    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
+
+
+def _paired_diffusion_benchmark_pytest_params(configs: list[dict[str, Any]]) -> list[Any]:
+    """Paired params for ``run_diffusion_benchmark.py``; same shape as omni runner."""
+    test_param_map = _test_param_mapping(configs)
+    server_entries = [(cfg, cfg["test_name"]) for cfg in _unique_server_params(configs)]
+    return create_paired_benchmark_pytest_params(server_entries, test_param_map, _marks_by_test_name(configs))
 
 
 # ---------------------------------------------------------------------------
 # Parametrize data
 # ---------------------------------------------------------------------------
 
-server_params = _unique_server_params(BENCHMARK_CONFIGS)
 test_param_map = _test_param_mapping(BENCHMARK_CONFIGS)
+paired_benchmark_params = _paired_diffusion_benchmark_pytest_params(BENCHMARK_CONFIGS)
+
+
+def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
+    """Factory: return a vLLM-Omni diffusion server instance for the config."""
+    return DiffusionServer(server_cfg=server_cfg)
 
-benchmark_indices: list[int] = list(range(max(len(v) for v in test_param_map.values())))
 
 # ---------------------------------------------------------------------------
 # Pytest fixtures
@@ -534,16 +608,13 @@ def diffusion_server(request):
 
 
 @pytest.fixture
-def benchmark_params(request, diffusion_server):
-    """Yield the benchmark params dict for the current (server, index) pair."""
-    param_index: int = request.param
-    test_name = diffusion_server.test_name
+def benchmark_params(request):
+    """Benchmark params for the paired server/index parametrization."""
+    test_name, param_index = request.param
 
     params_list = test_param_map.get(test_name, [])
     if not params_list:
         raise ValueError(f"No benchmark params for test: {test_name}")
-    if param_index >= len(params_list):
-        pytest.skip(f"Param index {param_index} out of range for {test_name} (has {len(params_list)} params)")
 
     current = param_index + 1
     total = len(params_list)
@@ -575,8 +646,8 @@ def run_benchmark(
     The raw metrics are written to a temporary file by the subprocess.  After
     the run completes the metrics are merged with full metadata (test_name,
     endpoint, benchmark_params, timestamp, flat reporting fields) and appended
-    to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE).  The
-    temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
+    to ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` for the
+    source JSON file. The temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
     to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
     the record.
     """
@@ -586,7 +657,12 @@ def run_benchmark(
     log_dir = BENCHMARK_RESULT_DIR / "logs"
     log_dir.mkdir(parents=True, exist_ok=True)
     endpoint_label = endpoint_filename_token(endpoint)
-    log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
+    resource_label = get_runtime_resource_label()
+    hw_for_filename = resource_label_for_filename(resource_label)
+    if hw_for_filename:
+        log_file = log_dir / f"{test_name}_{hw_for_filename}_{endpoint_label}_{timestamp}.log"
+    else:
+        log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
 
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
         tmp_result_file = Path(tmp.name)
@@ -696,7 +772,7 @@ def run_benchmark(
         "Model": model,
         "Framework": server_type,
         "API Endpoint": endpoint,
-        "Hardware": "",
+        "Hardware": hardware_json_value(resource_label),
         "Deployment": "",
         "Task": params.get("task", "t2i"),
         "Dataset": params.get("dataset", "random"),
@@ -728,8 +804,8 @@ def run_benchmark(
         "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
         "source_file": source_file,
     }
-    _append_to_aggregated_file(record)
-    print(f"\n  Result appended to: {AGGREGATED_RESULT_FILE}")
+    result_path = _write_result_record(record)
+    print(f"\n  Result saved to: {result_path}")
     print(f"  Log saved to:       {log_file}")
 
     return metrics
@@ -807,7 +883,7 @@ def assert_result(
 
 def _default_benchmark_endpoint_for_task(task: str) -> str:
     """Return the default client-side benchmark endpoint for a diffusion task."""
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
         return "/v1/chat/completions"
@@ -932,12 +1008,10 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
 # ---------------------------------------------------------------------------
 @pytest.mark.benchmark
 @pytest.mark.parametrize(
-    "diffusion_server",
-    server_params,
-    ids=[p["test_name"] for p in server_params],
-    indirect=True,
+    "diffusion_server,benchmark_params",
+    paired_benchmark_params,
+    indirect=["diffusion_server", "benchmark_params"],
 )
-@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
 def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
     """Run the diffusion performance benchmark and verify request completion.
 
@@ -962,7 +1036,10 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             test_name=test_name,
             endpoint=endpoint,
             server_cfg=server_cfg,
-            source_file=cast(str, CONFIG_FILE_PATH),
+            source_file=server_cfg.get(
+                _DIFFUSION_SOURCE_CONFIG_KEY,
+                CONFIG_FILE_PATH or f"{_PERF_TESTS_DIR}/*.json",
+            ),
         )
 
         print(f"\n{'=' * 60}")
@@ -980,7 +1057,9 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             if key in result:
                 print(f"  {key}: {result[key]:.4f}")
 
-        print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
+        source = server_cfg.get(_DIFFUSION_SOURCE_CONFIG_KEY) or CONFIG_FILE_PATH
+        if source:
+            print(f"\n  Aggregated results: {_aggregated_result_file_for_source(str(source))}")
         print("=" * 60)
 
         assert_baseline = request.config.getoption("--assert-baseline", default=False)

tests/dfx/perf/tests/test_bagel_vllm_omni.json

@@ -1,6 +1,18 @@
 [
   {
     "test_name": "test_bagel_single_device_single_stage_t2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -38,6 +50,18 @@
   },
   {
     "test_name": "test_bagel_single_device_single_stage_i2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -76,6 +100,18 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_t2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -113,6 +149,18 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_i2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
     "server_type": "vllm-omni",
     "server_params": {

tests/dfx/perf/tests/test_cosmos3_vllm_omni.json

@@ -0,0 +1,261 @@
+[
+  {
+    "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/images/generations",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps4",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 4,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "extra-body": {
+          "guidance_scale": 7.0,
+          "flow_shift": 3.0,
+          "guardrails": false
+        },
+        "baseline": {
+          "throughput_qps": 1.107,
+          "latency_mean": 0.8929,
+          "peak_memory_mb_mean": 80000.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_t2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "A robot arm is cleaning a plate in the kitchen.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false,
+          "use_resolution_template": false,
+          "use_duration_template": false
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 36.6824,
+          "peak_memory_mb_mean": 30155.4
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_i2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "The scene comes to life with smooth, natural motion.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 36.968,
+          "peak_memory_mb_mean": 30443.6
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_v2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "v2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "Continue the same scene with smooth natural motion.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false,
+          "condition_frame_indexes_vision": [
+            0,
+            1
+          ],
+          "condition_video_keep": "first"
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 37.2117,
+          "peak_memory_mb_mean": 30135.6
+        }
+      }
+    ]
+  }
+]

tests/dfx/perf/tests/test_higgs_audio_v3.json

@@ -1,6 +1,18 @@
 [
   {
     "test_name": "test_higgs_audio_v3",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "bosonai/higgs-audio-v3-tts-4b",
       "stage_overrides": {

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -1,139 +1,152 @@
 [
-    {
-        "test_name": "test_hunyuan_image3_ti2i_inline",
-        "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
-        "server_type": "vllm-omni",
-        "benchmark_endpoint": "/v1/images/edits",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config-inline": {
-                    "type": "yaml",
-                    "content": {
-                        "pipeline": "hunyuan_image_3_moe",
-                        "async_chunk": false,
-                        "trust_remote_code": true,
-                        "connectors": {
-                            "rdma_connector": {
-                                "name": "MooncakeTransferEngineConnector",
-                                "extra": {
-                                    "host": "auto",
-                                    "zmq_port": 50051,
-                                    "protocol": "rdma",
-                                    "device_name": "",
-                                    "memory_pool_size": 17179869184,
-                                    "memory_pool_device": "cpu"
-                                }
-                            }
-                        },
-                        "stages": [
-                            {
-                                "stage_id": 0,
-                                "is_comprehension": true,
-                                "final_output": true,
-                                "final_output_type": "text",
-                                "max_num_seqs": 16,
-                                "gpu_memory_utilization": 0.95,
-                                "enforce_eager": false,
-                                "enable_prefix_caching": false,
-                                "max_num_batched_tokens": 32768,
-                                "devices": "0,1",
-                                "tensor_parallel_size": 2,
-                                "hf_overrides": {
-                                    "rope_parameters": {
-                                        "mrope_section": [
-                                            0,
-                                            32,
-                                            32
-                                        ],
-                                        "rope_type": "default"
-                                    }
-                                },
-                                "omni_kv_config": {
-                                    "need_send_cache": true
-                                },
-                                "output_connectors": {
-                                    "to_stage_1": "rdma_connector"
-                                },
-                                "default_sampling_params": {
-                                    "temperature": 0.0,
-                                    "top_p": 1,
-                                    "top_k": -1,
-                                    "max_tokens": 8192,
-                                    "detokenize": true,
-                                    "skip_special_tokens": false,
-                                    "stop_token_ids": [
-                                        128025
-                                    ]
-                                }
-                            },
-                            {
-                                "stage_id": 1,
-                                "max_num_seqs": 1,
-                                "enforce_eager": true,
-                                "devices": "2,3",
-                                "distributed_executor_backend": "mp",
-                                "omni_kv_config": {
-                                    "need_recv_cache": true
-                                },
-                                "parallel_config": {
-                                    "tensor_parallel_size": 2,
-                                    "vae_patch_parallel_size": 2,
-                                    "enable_expert_parallel": true
-                                },
-                                "input_connectors": {
-                                    "from_stage_0": "rdma_connector"
-                                },
-                                "default_sampling_params": {
-                                    "num_inference_steps": 50,
-                                    "guidance_scale": 0
-                                }
-                            }
-                        ],
-                        "edges": [
-                            {
-                                "from": 0,
-                                "to": 1,
-                                "window_size": -1,
-                                "max_inflight": 1
-                            }
-                        ]
-                    }
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "inline_dataset_5_requests",
-                "dataset": "custom",
-                "dataset-path-inline": {
-                    "type": "jsonl",
-                    "content": [
-                        {
-                            "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
-                            "image_urls": [
-                                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
-                            ],
-                            "num_inference_steps": 8,
-                            "seed": 42,
-                            "guidance_scale": 1.0
-                        }
-                    ]
-                },
-                "task": "ti2i",
-                "bot-task": "think_recaption",
-                "max-concurrency": 16,
-                "num-prompts": 96,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.1874,
-                    "latency_p99": 101.852,
-                    "latency_mean": 79.0385,
-                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 70000
+  {
+    "test_name": "test_hunyuan_image3_ti2i_inline",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/images/edits",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config-inline": {
+          "type": "yaml",
+          "content": {
+            "pipeline": "hunyuan_image_3_moe",
+            "async_chunk": false,
+            "trust_remote_code": true,
+            "connectors": {
+              "rdma_connector": {
+                "name": "MooncakeTransferEngineConnector",
+                "extra": {
+                  "host": "auto",
+                  "zmq_port": 50051,
+                  "protocol": "rdma",
+                  "device_name": "",
+                  "memory_pool_size": 17179869184,
+                  "memory_pool_device": "cpu"
                 }
+              }
+            },
+            "stages": [
+              {
+                "stage_id": 0,
+                "is_comprehension": true,
+                "final_output": true,
+                "final_output_type": "text",
+                "max_num_seqs": 16,
+                "gpu_memory_utilization": 0.95,
+                "enforce_eager": false,
+                "enable_prefix_caching": false,
+                "max_num_batched_tokens": 32768,
+                "devices": "0,1",
+                "tensor_parallel_size": 2,
+                "hf_overrides": {
+                  "rope_parameters": {
+                    "mrope_section": [
+                      0,
+                      32,
+                      32
+                    ],
+                    "rope_type": "default"
+                  }
+                },
+                "omni_kv_config": {
+                  "need_send_cache": true
+                },
+                "output_connectors": {
+                  "to_stage_1": "rdma_connector"
+                },
+                "default_sampling_params": {
+                  "temperature": 0.0,
+                  "top_p": 1,
+                  "top_k": -1,
+                  "max_tokens": 8192,
+                  "detokenize": true,
+                  "skip_special_tokens": false,
+                  "stop_token_ids": [
+                    128025
+                  ]
+                }
+              },
+              {
+                "stage_id": 1,
+                "max_num_seqs": 1,
+                "enforce_eager": true,
+                "devices": "2,3",
+                "distributed_executor_backend": "mp",
+                "omni_kv_config": {
+                  "need_recv_cache": true
+                },
+                "parallel_config": {
+                  "tensor_parallel_size": 2,
+                  "vae_patch_parallel_size": 2,
+                  "enable_expert_parallel": true
+                },
+                "input_connectors": {
+                  "from_stage_0": "rdma_connector"
+                },
+                "default_sampling_params": {
+                  "num_inference_steps": 50,
+                  "guidance_scale": 0
+                }
+              }
+            ],
+            "edges": [
+              {
+                "from": 0,
+                "to": 1,
+                "window_size": -1,
+                "max_inflight": 1
+              }
+            ]
+          }
+        },
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "inline_dataset_5_requests",
+        "dataset": "custom",
+        "dataset-path-inline": {
+          "type": "jsonl",
+          "content": [
+            {
+              "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
+              "image_urls": [
+                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
+              ],
+              "num_inference_steps": 8,
+              "seed": 42,
+              "guidance_scale": 1.0
             }
-        ]
-    }
+          ]
+        },
+        "task": "ti2i",
+        "bot-task": "think_recaption",
+        "max-concurrency": 16,
+        "num-prompts": 96,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.1874,
+          "latency_p99": 101.852,
+          "latency_mean": 79.0385,
+          "peak_memory_mb_max": 70000,
+          "peak_memory_mb_mean": 70000
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp2_cfgp2",
-        "description": "TP=2 CfgP=2 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 2,
-                "cfg-parallel-size": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_p99": 4.7469,
-                    "peak_memory_mb_max": 101100,
-                    "peak_memory_mb_mean": 101100
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp2_cfgp2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=2 CfgP=2 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 2,
+        "cfg-parallel-size": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.21,
+          "latency_p99": 4.7469,
+          "peak_memory_mb_max": 101100,
+          "peak_memory_mb_mean": 101100
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp2_sp2",
-        "description": "TP=2 SP=2 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 2,
-                "ulysses-degree": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.20,
-                    "latency_p99": 5.1025,
-                    "peak_memory_mb_max": 97402,
-                    "peak_memory_mb_mean": 97402
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp2_sp2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=2 SP=2 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 2,
+        "ulysses-degree": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.20,
+          "latency_p99": 5.1025,
+          "peak_memory_mb_max": 97402,
+          "peak_memory_mb_mean": 97402
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp4",
-        "description": "TP=4 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 4,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.213,
-                    "latency_mean": 4.6953,
-                    "latency_p99": 4.8601,
-                    "peak_memory_mb_max": 56912.0,
-                    "peak_memory_mb_mean": 56912.0
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=4 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 4,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.213,
+          "latency_mean": 4.6953,
+          "latency_p99": 4.8601,
+          "peak_memory_mb_max": 56912.0,
+          "peak_memory_mb_mean": 56912.0
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -1,217 +1,273 @@
 [
-    {
-        "test_name": "test_ltx2_baseline_eager",
-        "description": "Single-device baseline with enforce-eager (no torch.compile)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+  {
+    "test_name": "test_ltx2_baseline_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline with enforce-eager (no torch.compile)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_torch_compile",
-        "description": "Single-device with torch.compile (default, no enforce-eager)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_torch_compile",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device with torch.compile (default, no enforce-eager)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cfg2_eager",
-        "description": "CFG-parallel=2 with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cfg2_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CFG-parallel=2 with enforce-eager",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cfg2_compile",
-        "description": "CFG-parallel=2 with torch.compile",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cfg2_compile",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CFG-parallel=2 with torch.compile",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cache_dit_eager",
-        "description": "CacheDiT with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cache-backend": "cache_dit",
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    }
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cache_dit_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CacheDiT with enforce-eager",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cache-backend": "cache_dit",
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -1,331 +1,343 @@
 [
-    {
-        "test_name": "test_qwen3_omni_chunk",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
-                "--async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni_chunk",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "extra_cli_args": [
+        "--async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            135.4379,
+            187.0896,
+            400,
+            435.675,
+            724.6836
+          ],
+          "mean_e2el_ms": [
+            17619.9851,
+            34811.8335,
+            43656.9501,
+            56657.1548,
+            85097.6016
+          ],
+          "mean_audio_rtf": [
+            0.1279,
+            0.2238,
+            0.2735,
+            0.3712,
+            0.5634
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        135.4379,
-                        187.0896,
-                        400,
-                        435.675,
-                        724.6836
-                    ],
-                    "mean_e2el_ms": [
-                        17619.9851,
-                        34811.8335,
-                        43656.9501,
-                        56657.1548,
-                        85097.6016
-                    ],
-                    "mean_audio_rtf": [
-                        0.1279,
-                        0.2238,
-                        0.2735,
-                        0.3712,
-                        0.5634
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        340.329
-                    ],
-                    "mean_e2el_ms": [
-                        4305.477
-                    ],
-                    "mean_audio_rtf": [
-                        0.1305
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        220.0909
-                    ],
-                    "mean_e2el_ms": [
-                        5475.9565
-                    ],
-                    "mean_audio_rtf": [
-                        0.1818
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        450.0155
-                    ],
-                    "mean_e2el_ms": [
-                        8516.8316
-                    ],
-                    "mean_audio_rtf": [
-                        0.2717
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        78.4278,
-                        195.8037,
-                        333.7318,
-                        561.805,
-                        1038.043
-                    ],
-                    "mean_e2el_ms": [
-                        4288.3029,
-                        7429.9429,
-                        9451.6184,
-                        12607.5292,
-                        18153.2547
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        275.3875
-                    ],
-                    "mean_e2el_ms": [
-                        730.6995
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        166.2201
-                    ],
-                    "mean_e2el_ms": [
-                        673.5789
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        357.7781
-                    ],
-                    "mean_e2el_ms": [
-                        953.3493
-                    ]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            340.329
+          ],
+          "mean_e2el_ms": [
+            4305.477
+          ],
+          "mean_audio_rtf": [
+            0.1305
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            220.0909
+          ],
+          "mean_e2el_ms": [
+            5475.9565
+          ],
+          "mean_audio_rtf": [
+            0.1818
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            450.0155
+          ],
+          "mean_e2el_ms": [
+            8516.8316
+          ],
+          "mean_audio_rtf": [
+            0.2717
+          ]
+        }
+      },
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            78.4278,
+            195.8037,
+            333.7318,
+            561.805,
+            1038.043
+          ],
+          "mean_e2el_ms": [
+            4288.3029,
+            7429.9429,
+            9451.6184,
+            12607.5292,
+            18153.2547
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            275.3875
+          ],
+          "mean_e2el_ms": [
+            730.6995
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            166.2201
+          ],
+          "mean_e2el_ms": [
+            673.5789
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            357.7781
+          ],
+          "mean_e2el_ms": [
+            953.3493
+          ]
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

@@ -1,115 +1,127 @@
 [
-    {
-        "test_name": "test_qwen3_omni_3gpu_replica2",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_overrides": {
-                "1": {"num_replicas": 2, "devices": "1,2"},
-                "2": {"num_replicas": 2, "devices": "1,2"}
-            },
-            "extra_cli_args": [
-                "--async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni_3gpu_replica2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 3
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "stage_overrides": {
+          "1": {"num_replicas": 2, "devices": "1,2"},
+          "2": {"num_replicas": 2, "devices": "1,2"}
+      },
+      "extra_cli_args": [
+        "--async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [128, 128, 128, 128],
+        "max_concurrency": [8, 16, 24, 32],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [203, 283, 470, 869],
+            "mean_tpot_ms": [28, 36, 41, 45],
+            "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
+            "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [10],
+        "request_rate": [0.1],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [128, 128, 128, 128],
-                "max_concurrency": [8, 16, 24, 32],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [203, 283, 470, 869],
-                    "mean_tpot_ms": [28, 36, 41, 45],
-                    "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
-                    "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [377],
-                    "mean_audio_ttfp_ms": [684],
-                    "mean_audio_rtf": [0.14]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [228],
-                    "mean_audio_ttfp_ms": [686],
-                    "mean_audio_rtf": [0.16]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [522],
-                    "mean_audio_ttfp_ms": [1119],
-                    "mean_audio_rtf": [0.21]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [377],
+            "mean_audio_ttfp_ms": [684],
+            "mean_audio_rtf": [0.14]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [40],
+        "request_rate": [0.5],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [228],
+            "mean_audio_ttfp_ms": [686],
+            "mean_audio_rtf": [0.16]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [100],
+        "request_rate": [1.0],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [522],
+            "mean_audio_ttfp_ms": [1119],
+            "mean_audio_rtf": [0.21]
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -1,170 +1,182 @@
 [
-    {
-        "test_name": "test_qwen3_omni",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
-                "--no-async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "extra_cli_args": [
+        "--no-async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            109.1218,
+            247.4411,
+            354.9384,
+            447.7885,
+            694.7487
+          ],
+          "mean_e2el_ms": [
+            22413.667,
+            30808.9715,
+            35677.844,
+            46345.8978,
+            68350.4853
+          ],
+          "mean_audio_rtf": [
+            0.1695,
+            0.2033,
+            0.2467,
+            0.3234,
+            0.4438
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        109.1218,
-                        247.4411,
-                        354.9384,
-                        447.7885,
-                        694.7487
-                    ],
-                    "mean_e2el_ms": [
-                        22413.667,
-                        30808.9715,
-                        35677.844,
-                        46345.8978,
-                        68350.4853
-                    ],
-                    "mean_audio_rtf": [
-                        0.1695,
-                        0.2033,
-                        0.2467,
-                        0.3234,
-                        0.4438
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        323.3573
-                    ],
-                    "mean_e2el_ms": [
-                        4623.4173
-                    ],
-                    "mean_audio_rtf": [
-                        0.1406
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        215.4502
-                    ],
-                    "mean_e2el_ms": [
-                        4924.149
-                    ],
-                    "mean_audio_rtf": [
-                        0.1637
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        442.5815
-                    ],
-                    "mean_e2el_ms": [
-                        6441.0986
-                    ],
-                    "mean_audio_rtf": [
-                        0.2091
-                    ]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            323.3573
+          ],
+          "mean_e2el_ms": [
+            4623.4173
+          ],
+          "mean_audio_rtf": [
+            0.1406
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            215.4502
+          ],
+          "mean_e2el_ms": [
+            4924.149
+          ],
+          "mean_audio_rtf": [
+            0.1637
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            442.5815
+          ],
+          "mean_e2el_ms": [
+            6441.0986
+          ],
+          "mean_audio_rtf": [
+            0.2091
+          ]
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -1,180 +1,192 @@
 [
-    {
-        "test_name": "test_qwen3_omni_vllm_text",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "use_omni": false,
-            "extra_cli_args": [
-                "--no-enable-prefix-caching",
-                "--skip-mm-profiling"
-            ]
+  {
+    "test_name": "test_qwen3_omni_vllm_text",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "use_omni": false,
+      "extra_cli_args": [
+        "--no-enable-prefix-caching",
+        "--skip-mm-profiling"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        69.1276,
-                        255.1882,
-                        310.639,
-                        422.0337,
-                        532.0582
-                    ],
-                    "mean_e2el_ms": [
-                        4382.7211,
-                        7510.3815,
-                        9440.4854,
-                        12553.0186,
-                        16932.6065
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        26.3381
-                    ],
-                    "mean_e2el_ms": [
-                        485.2436
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        26.4517
-                    ],
-                    "mean_e2el_ms": [
-                        517.2017
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        27.6972
-                    ],
-                    "mean_e2el_ms": [
-                        561.3684
-                    ]
-                }
-            }
-        ]
-    }
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            69.1276,
+            255.1882,
+            310.639,
+            422.0337,
+            532.0582
+          ],
+          "mean_e2el_ms": [
+            4382.7211,
+            7510.3815,
+            9440.4854,
+            12553.0186,
+            16932.6065
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            26.3381
+          ],
+          "mean_e2el_ms": [
+            485.2436
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            26.4517
+          ],
+          "mean_e2el_ms": [
+            517.2017
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            27.6972
+          ],
+          "mean_e2el_ms": [
+            561.3684
+          ]
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -1,149 +1,185 @@
 [
-    {
-        "test_name": "test_qwen_image_edit_2511_single_device",
-        "description": "Single-device baseline (two input images)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0721,
-                    "latency_mean": 13.8398,
-                    "peak_memory_mb_max": 54731.4,
-                    "peak_memory_mb_mean": 54731.4
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0186,
-                    "latency_mean": 53.6942,
-                    "peak_memory_mb_max": 64402.4,
-                    "peak_memory_mb_mean": 64402.4
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_edit_2511_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (two input images)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0538,
-                    "latency_mean": 18.5591,
-                    "peak_memory_mb_max": 53485.0,
-                    "peak_memory_mb_mean": 53485.0
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0721,
+          "latency_mean": 13.8398,
+          "peak_memory_mb_max": 54731.4,
+          "peak_memory_mb_mean": 54731.4
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0186,
+          "latency_mean": 53.6942,
+          "peak_memory_mb_max": 64402.4,
+          "peak_memory_mb_mean": 64402.4
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0538,
+          "latency_mean": 18.5591,
+          "peak_memory_mb_max": 53485.0,
+          "peak_memory_mb_mean": 53485.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.2776,
-                    "latency_mean": 3.5938,
-                    "peak_memory_mb_max": 55214.0,
-                    "peak_memory_mb_mean": 55214.0
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1051,
-                    "latency_mean": 9.4918,
-                    "peak_memory_mb_max": 64265.6,
-                    "peak_memory_mb_mean": 64265.6
-                }
-            }
-        ]
-    }
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.2776,
+          "latency_mean": 3.5938,
+          "peak_memory_mb_max": 55214.0,
+          "peak_memory_mb_mean": 55214.0
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1051,
+          "latency_mean": 9.4918,
+          "peak_memory_mb_max": 64265.6,
+          "peak_memory_mb_mean": 64265.6
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -1,144 +1,180 @@
 [
-    {
-        "test_name": "test_qwen_image_edit_single_device",
-        "description": "Single-device baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1129,
-                    "latency_mean": 8.8352,
-                    "peak_memory_mb_max": 55356.5,
-                    "peak_memory_mb_mean": 55356.5
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0240,
-                    "latency_mean": 41.5686,
-                    "peak_memory_mb_max": 62762.7,
-                    "peak_memory_mb_mean": 62762.7
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_edit_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0672,
-                    "latency_mean": 14.8517,
-                    "peak_memory_mb_max": 53931.5,
-                    "peak_memory_mb_mean": 53931.5
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1129,
+          "latency_mean": 8.8352,
+          "peak_memory_mb_max": 55356.5,
+          "peak_memory_mb_mean": 55356.5
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0240,
+          "latency_mean": 41.5686,
+          "peak_memory_mb_max": 62762.7,
+          "peak_memory_mb_mean": 62762.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0672,
+          "latency_mean": 14.8517,
+          "peak_memory_mb_max": 53931.5,
+          "peak_memory_mb_mean": 53931.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.2879,
-                    "latency_mean": 3.4701,
-                    "peak_memory_mb_max": 55525.6,
-                    "peak_memory_mb_mean": 55525.6
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1113,
-                    "latency_mean": 8.9656,
-                    "peak_memory_mb_max": 63758.3,
-                    "peak_memory_mb_mean": 63758.3
-                }
-            }
-        ]
-    }
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.2879,
+          "latency_mean": 3.4701,
+          "peak_memory_mb_max": 55525.6,
+          "peak_memory_mb_mean": 55525.6
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1113,
+          "latency_mean": 8.9656,
+          "peak_memory_mb_max": 63758.3,
+          "peak_memory_mb_mean": 63758.3
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -1,49 +1,61 @@
 [
-    {
-        "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.0703,
-                    "latency_mean": 14.208,
-                    "peak_memory_mb_max": 60543.5,
-                    "peak_memory_mb_mean": 60543.5
-                }
-            },
-            {
-                "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.0417,
-                    "latency_mean": 23.914,
-                    "peak_memory_mb_max": 60543.5,
-                    "peak_memory_mb_mean": 60543.5
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_qwen_image_layered_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "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.0703,
+          "latency_mean": 14.208,
+          "peak_memory_mb_max": 60543.5,
+          "peak_memory_mb_mean": 60543.5
+        }
+      },
+      {
+        "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.0417,
+          "latency_mean": 23.914,
+          "peak_memory_mb_max": 60543.5,
+          "peak_memory_mb_mean": 60543.5
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -1,219 +1,279 @@
 [
-    {
-        "test_name": "test_qwen_image_single_device",
-        "description": "Single-device baseline (no parallelism)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": 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.4153,
-                    "latency_mean": 2.4065,
-                    "peak_memory_mb_mean": 52289.9
-                }
-            },
-            {
-                "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.0423,
-                    "latency_mean": 23.5991,
-                    "peak_memory_mb_mean": 61360.5
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (no parallelism)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "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.4175,
-                    "latency_mean": 2.3953,
-                    "peak_memory_mb_mean": 56000
-                }
-            },
-            {
-                "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.0423,
-                    "latency_mean": 23.5768,
-                    "peak_memory_mb_mean": 65000
-                }
-            }
-        ]
+    "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.4153,
+          "latency_mean": 2.4065,
+          "peak_memory_mb_mean": 52289.9
+        }
+      },
+      {
+        "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.0423,
+          "latency_mean": 23.5991,
+          "peak_memory_mb_mean": 61360.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_single_device_step_execution",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "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
+      }
     },
-    {
-        "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "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.1214,
-                    "latency_mean": 8.2211,
-                    "peak_memory_mb_mean": 51704.7
-                }
-            }
-        ]
+    "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.4175,
+          "latency_mean": 2.3953,
+          "peak_memory_mb_mean": 56000
+        }
+      },
+      {
+        "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.0423,
+          "latency_mean": 23.5768,
+          "peak_memory_mb_mean": 65000
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "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.1214,
+          "latency_mean": 8.2211,
+          "peak_memory_mb_mean": 51704.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "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.559,
-                    "latency_mean": 1.794,
-                    "peak_memory_mb_mean": 52397.6775
-                }
-            },
-            {
-                "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.1918,
-                    "latency_mean": 5.2022,
-                    "peak_memory_mb_mean": 61468.8
-                }
-            }
-        ]
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
-        "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true,
-                "step-execution": true,
-                "max-num-seqs": 8
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_high_concurrency",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": [20, 40, 80, 160],
-                "max-concurrency": [1, 2, 4, 8],
-                "warmup-requests": 8,
-                "warmup-concurrency": 8,
-                "warmup-num-inference-steps": 20,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
-                    "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
-                    "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
-                }
-            }
-        ]
-    }
+    "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.559,
+          "latency_mean": 1.794,
+          "peak_memory_mb_mean": 52397.6775
+        }
+      },
+      {
+        "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.1918,
+          "latency_mean": 5.2022,
+          "peak_memory_mb_mean": 61468.8
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "step-execution": true,
+        "max-num-seqs": 8
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_high_concurrency",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": [20, 40, 80, 160],
+        "max-concurrency": [1, 2, 4, 8],
+        "warmup-requests": 8,
+        "warmup-concurrency": 8,
+        "warmup-num-inference-steps": 20,
+        "enable-negative-prompt": true,
+        "baseline": {
+            "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
+            "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
+            "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
+        }
+      }
+    ]
+  }
 ]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -86,3 +86,202 @@ def test_enabled_false_entry_is_skipped():
     # Only the enabled=True entry should appear
     assert len(params) == 1
     assert params[0].get("task") == "voice_clone"
+
+
+def test_resolve_pytest_marks_hardware_dict_with_extra():
+    from tests.dfx.conftest import resolve_pytest_marks
+
+    marks = resolve_pytest_marks(
+        [
+            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+            "full_model",
+            "diffusion",
+            "local_model",
+        ]
+    )
+    names = {m.name for m in marks}
+    assert "H100" in names
+    assert "cuda" in names
+    assert "gpu" in names
+    assert "distributed_cuda" in names
+    assert "full_model" in names
+    assert "diffusion" in names
+    assert "local_model" in names
+
+
+def test_resolve_pytest_marks_rejects_legacy_object_format():
+    from tests.dfx.conftest import resolve_pytest_marks
+
+    with pytest.raises(ValueError, match="mark must be a list"):
+        resolve_pytest_marks(
+            {
+                "hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1},
+                "marks": ["full_model"],
+            }
+        )
+
+
+def test_extract_mark_resource_label():
+    from tests.dfx.conftest import extract_mark_resource_label
+
+    assert extract_mark_resource_label(None) == "na"
+    assert (
+        extract_mark_resource_label(
+            [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "full_model",
+            ]
+        )
+        == "H100"
+    )
+    assert (
+        extract_mark_resource_label(
+            [
+                {
+                    "hardware_marks": {
+                        "res": {"cuda": "H100", "rocm": "MI325"},
+                        "num_cards": 2,
+                    }
+                },
+                "diffusion",
+            ]
+        )
+        == "H100-MI325"
+    )
+
+
+def test_resource_label_for_filename():
+    from tests.dfx.conftest import resource_label_for_filename
+
+    assert resource_label_for_filename("H100") == ""
+    assert resource_label_for_filename("L4") == "L4"
+    assert resource_label_for_filename("910B") == "910B"
+    assert resource_label_for_filename("na") == "na"
+
+
+def test_hardware_json_value():
+    from tests.dfx.conftest import hardware_json_value
+
+    assert hardware_json_value("H100") == "H100"
+    assert hardware_json_value("na") == ""
+    assert hardware_json_value(None) == ""
+
+
+def test_extract_configs_resource_label(monkeypatch):
+    from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label
+
+    monkeypatch.setattr(
+        "tests.dfx.conftest._read_runtime_device_name",
+        lambda *, device_id=0: "NVIDIA H100 80GB HBM3",
+    )
+    get_runtime_resource_label(refresh=True)
+    assert extract_configs_resource_label([]) == "H100"
+    get_runtime_resource_label(refresh=True)
+    monkeypatch.setattr(
+        "tests.dfx.conftest._read_runtime_device_name",
+        lambda *, device_id=0: "Ascend910B2",
+    )
+    assert get_runtime_resource_label(refresh=True) == "910B"
+
+
+def test_load_benchmark_configs_from_dir(tmp_path):
+    from tests.dfx.conftest import load_benchmark_configs
+
+    (tmp_path / "a.json").write_text(
+        json.dumps([{"test_name": "test_a", "server_params": {"model": "m/a"}, "benchmark_params": []}]),
+        encoding="utf-8",
+    )
+    (tmp_path / "b.json").write_text(
+        json.dumps([{"test_name": "test_b", "server_params": {"model": "m/b"}, "benchmark_params": []}]),
+        encoding="utf-8",
+    )
+    configs = load_benchmark_configs(config_dir=tmp_path)
+    assert [c["test_name"] for c in configs] == ["test_a", "test_b"]
+
+
+def test_create_unique_server_pytest_params_applies_marks(tmp_path):
+    from tests.dfx.conftest import create_unique_server_pytest_params
+
+    configs = [
+        {
+            "test_name": "test_with_mark",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "full_model",
+            ],
+            "server_params": {"model": "some/model"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+        {
+            "test_name": "test_without_mark",
+            "server_params": {"model": "other/model"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+    ]
+    params = create_unique_server_pytest_params(configs, tmp_path)
+    by_id = {p.id: p for p in params}
+    assert len(by_id["test_with_mark"].values) == 1
+    assert isinstance(by_id["test_with_mark"].values[0], tuple)
+    assert any(m.name == "H100" for m in by_id["test_with_mark"].marks)
+    assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)
+
+
+def test_is_diffusion_perf_config():
+    from tests.dfx.conftest import is_diffusion_perf_config
+
+    assert not is_diffusion_perf_config(
+        {"test_name": "omni_a", "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "omni"]}
+    )
+    assert is_diffusion_perf_config(
+        {
+            "test_name": "diff_a",
+            "server_type": "vllm-omni",
+            "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "diffusion"],
+        }
+    )
+
+
+def test_benchmark_param_id_suffix_from_task_eval_phase():
+    from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
+
+    params = [
+        {"task": "default_voice", "eval_phase": "latency"},
+        {"task": "default_voice", "eval_phase": "throughput"},
+    ]
+    assert _unique_benchmark_param_id_suffixes(params) == [
+        "default_voice_latency",
+        "default_voice_throughput",
+    ]
+
+
+def test_create_paired_omni_benchmark_pytest_params(tmp_path):
+    from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
+
+    configs = [
+        {
+            "test_name": "test_omni",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "omni",
+            ],
+            "server_params": {"model": "m/omni"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}, {"name": "p1", "num_prompts": 2}],
+        },
+        {
+            "test_name": "test_tts",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "tts",
+            ],
+            "server_params": {"model": "m/tts"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+    ]
+    params = create_paired_omni_benchmark_pytest_params(configs, tmp_path)
+    by_id = {p.id: p for p in params}
+    assert set(by_id) == {"test_omni-p0", "test_omni-p1", "test_tts-p0"}
+    omni_row, bench_row = by_id["test_tts-p0"].values
+    assert omni_row[0] == "test_tts"
+    assert bench_row == ("test_tts", 0)
+    assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
+    assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)

tests/dfx/perf/tests/test_tts.json

@@ -1,6 +1,19 @@
 [
   {
     "test_name": "test_qwen3_tts_base",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100",
+            "npu": "A3"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
       "stage_overrides": {
@@ -98,6 +111,18 @@
   },
   {
     "test_name": "test_qwen3_tts_customvoice",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
       "stage_overrides": {
@@ -342,6 +367,18 @@
   },
   {
     "test_name": "test_voxtral_tts",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "mistralai/Voxtral-4B-TTS-2603"
     },

tests/dfx/perf/tests/test_voxcpm2.json

@@ -1,9 +1,24 @@
 [
   {
     "test_name": "test_voxcpm2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true
+      "trust_remote_code": true,
+      "extra_cli_args": [
+        "--trust-remote-code"
+      ]
     },
     "benchmark_params": [
       {

tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json

@@ -1,107 +1,131 @@
 [
-    {
-        "test_name": "test_wan22_i2v_single_device",
-        "description": "Single-device baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "832x480_frames81_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 832,
-                        "height": 480,
-                        "num_inference_steps": 4,
-                        "num_frames": 81,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0465,
-                    "latency_mean": 21.4354,
-                    "peak_memory_mb_mean": 72557.58
-                }
-            }
-        ]
+  {
+    "test_name": "test_wan22_i2v_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
-        "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-            "serve_args": {
-                "usp": 2,
-                "vae-patch-parallel-size": 2,
-                "use-hsdp": true,
-                "vae-use-slicing": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "832x480_frames81_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 832,
-                        "height": 480,
-                        "num_inference_steps": 4,
-                        "num_frames": 81,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0642,
-                    "latency_mean": 15.5589,
-                    "peak_memory_mb_mean": 45742.88
-                }
-            },
-            {
-                "name": "1280x720_frames121_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 1280,
-                        "height": 720,
-                        "num_inference_steps": 4,
-                        "num_frames": 121,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0128,
-                    "latency_mean": 77.8565,
-                    "peak_memory_mb_mean": 53947.84
-                }
-            }
-        ]
-    }
+    "benchmark_params": [
+      {
+        "name": "832x480_frames81_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 81,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0465,
+          "latency_mean": 21.4354,
+          "peak_memory_mb_mean": 72557.58
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+      "serve_args": {
+        "usp": 2,
+        "vae-patch-parallel-size": 2,
+        "use-hsdp": true,
+        "vae-use-slicing": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames81_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 81,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0642,
+          "latency_mean": 15.5589,
+          "peak_memory_mb_mean": 45742.88
+        }
+      },
+      {
+        "name": "1280x720_frames121_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "num_inference_steps": 4,
+            "num_frames": 121,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0128,
+          "latency_mean": 77.8565,
+          "peak_memory_mb_mean": 53947.84
+        }
+      }
+    ]
+  }
 ]

tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py

@@ -62,8 +62,8 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
         ),
         pytest.param(
             True,
-            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
-            ("Invalid layers", "3", "1"),
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "1"},
+            ("Invalid layers", "2", "10"),
             400,
             id="invalid_layers",
         ),

tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py

@@ -87,8 +87,8 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
         pytest.param(
             {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
         ),
-        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
-        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
+        pytest.param({"layers": 1}, ("layers", "value_error", "2", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "2", "10"), id="layers_above_max"),
         pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
         pytest.param(
             {"seed": 2**32},

tests/diffusion/attention/test_kernels_hub.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+import types
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
+    """Test that when kernels package is not available, the platform fallback logic
+
+    routes to local FLASH_ATTN.
+    """
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    # Temporarily hide/remove kernels module if present
+    monkeypatch.setitem(sys.modules, "kernels", None)
+
+    # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+
+    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
+    kernels_module = types.ModuleType("kernels")
+    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_kernels_hub_execution():
+    """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+
+    num_heads = 8
+    head_dim = 64
+    seq_len = 32
+    batch_size = 1
+
+    torch.manual_seed(42)
+    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
+    k = q.clone()
+    v = q.clone()
+
+    # 1. Test PyTorch SDPA reference
+    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+    attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
+    output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+
+    # 2. Test FlashAttentionHubBackend (FlashAttention 2)
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
+
+    fa_hub_impl = FlashAttentionHubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa_hub.shape == q.shape
+    assert not torch.isnan(output_fa_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
+
+    # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
+    major, _minor = torch.cuda.get_device_capability()
+    if major < 9:
+        pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
+
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
+
+    fa3_hub_impl = FlashAttention3HubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa3_hub.shape == q.shape
+    assert not torch.isnan(output_fa3_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -67,6 +67,50 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(**defaults)
 
 
+def _patch_fake_diffusers_quantization_backends(mocker):
+    class FakePipelineQuantizationConfig:
+        def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+            self.quant_mapping = quant_mapping
+            self.quant_backend = quant_backend
+            self.quant_kwargs = quant_kwargs
+            self.components_to_quantize = components_to_quantize
+
+    class FakeTorchAoConfig:
+        def __init__(self, quant_type, modules_to_not_convert=None):
+            self.quant_type = quant_type
+            self.modules_to_not_convert = modules_to_not_convert
+
+        @classmethod
+        def from_dict(cls, config):
+            return cls(quant_type=config["quant_type"])
+
+    class FakeInt8DynamicActivationInt8WeightConfig:
+        pass
+
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.PipelineQuantizationConfig",
+        FakePipelineQuantizationConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.TorchAoConfig",
+        FakeTorchAoConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_torchao_quant_type_cls",
+        return_value=FakeInt8DynamicActivationInt8WeightConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_diffusers_quantization_config_cls",
+        return_value=FakeTorchAoConfig,
+    )
+    return FakePipelineQuantizationConfig, FakeTorchAoConfig, FakeInt8DynamicActivationInt8WeightConfig
+
+
+class _TransformerBackedPipeline(DiffusionPipeline):
+    def __init__(self, transformer):
+        pass
+
+
 def _make_batch(**overrides) -> DiffusionRequestBatch:
     """Wrap a single request in a DiffusionRequestBatch, matching the forward contract."""
     return DiffusionRequestBatch(requests=[_make_request(**overrides)])
@@ -98,7 +142,7 @@ class TestPipelineArgumentsHandling:
 
     @pytest.mark.parametrize(
         "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"],
+        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
         if feature_id == "cfg_parallel":
@@ -128,14 +172,188 @@ class TestPipelineArgumentsHandling:
             )
         elif feature_id == "enforce_eager":
             od_config = _make_od_config(enforce_eager=True)
-        elif feature_id == "quantization":
-            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8"))
+        elif feature_id == "unsupported_quantization":
+            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="gguf"))
         else:
             raise ValueError(f"Unknown feature ID: {feature_id}")
 
         with pytest.raises(NotImplementedError):
             DiffusersAdapterPipeline(od_config=od_config)
 
+    def test_adapter_load_weights_injects_supported_quantization_config(self, mocker):
+        (
+            FakePipelineQuantizationConfig,
+            FakeTorchAoConfig,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        ) = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value={
+                "_class_name": "TransformerBackedPipeline",
+                "scheduler": ["diffusers", "SchedulerMixin"],
+                "transformer": ["diffusers", "Transformer2DModel"],
+                "transformer_2": ["diffusers", "Transformer2DModel"],
+                "vae": ["diffusers", "AutoencoderKL"],
+            },
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_pipeline_cls=_TransformerBackedPipeline,
+            )
+        )
+        adapter.load_weights()
+
+        pipeline_quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        torchao_config = pipeline_quant_config.quant_mapping["transformer"]
+        assert isinstance(pipeline_quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(torchao_config, FakeTorchAoConfig)
+        assert isinstance(torchao_config.quant_type, FakeInt8DynamicActivationInt8WeightConfig)
+        assert isinstance(
+            pipeline_quant_config.quant_mapping["transformer_2"].quant_type,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        )
+
+    @pytest.mark.parametrize(
+        ("pipeline_config", "diffusers_load_kwargs", "diffusers_pipeline_cls"),
+        [
+            (
+                {
+                    "_class_name": "NoTransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "unet": ["diffusers", "UNet2DConditionModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {},
+                _TransformerBackedPipeline,
+            ),
+            (
+                {
+                    "_class_name": "TransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "transformer": ["diffusers", "Transformer2DModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {"transformer": None},
+                None,
+            ),
+        ],
+    )
+    def test_adapter_load_weights_rejects_converted_quantization_without_transformer_component(
+        self,
+        pipeline_config,
+        diffusers_load_kwargs,
+        diffusers_pipeline_cls,
+        mocker,
+    ):
+        _patch_fake_diffusers_quantization_backends(mocker)
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value=pipeline_config,
+        )
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained"
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_load_kwargs=diffusers_load_kwargs,
+                diffusers_pipeline_cls=diffusers_pipeline_cls,
+            )
+        )
+
+        with pytest.raises(NotImplementedError, match="transformer"):
+            adapter.load_weights()
+        mock_from_pretrained.assert_not_called()
+
+    def test_adapter_rejects_dduf_file_with_converted_quantization(self):
+        with pytest.raises(NotImplementedError, match="dduf_file"):
+            DiffusersAdapterPipeline(
+                od_config=_make_od_config(
+                    quantization_config=SimpleNamespace(quant_method="int8"),
+                    diffusers_load_kwargs={"dduf_file": "model.dduf"},
+                )
+            )
+
+    def test_adapter_load_weights_preserves_diffusers_native_quantization_config(self, mocker):
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        diffusers_native_quantization_config = object()
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="gguf"),
+                diffusers_load_kwargs={
+                    "dduf_file": "model.dduf",
+                    "quantization_config": diffusers_native_quantization_config,
+                },
+            )
+        )
+        adapter.load_weights()
+
+        kwargs = mock_from_pretrained.call_args.kwargs
+        assert kwargs["quantization_config"] is diffusers_native_quantization_config
+        assert kwargs["dduf_file"] == "model.dduf"
+
+    def test_adapter_load_weights_builds_diffusers_native_quantization_config_from_dict(self, mocker):
+        FakePipelineQuantizationConfig, FakeTorchAoConfig, _ = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                diffusers_load_kwargs={
+                    "quantization_config": {
+                        "quant_mapping": {
+                            "transformer": {
+                                "quant_method": "torchao",
+                                "quant_type": "fake-int8-config",
+                            },
+                        },
+                    },
+                },
+            )
+        )
+        adapter.load_weights()
+
+        quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        assert isinstance(quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(quant_config.quant_mapping["transformer"], FakeTorchAoConfig)
+        assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"
+
     def test_adapter_guard_unknown_output_type(self, mocker):
         """Test that the adapter wraps an unknown output type as-is.
         This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass."""

tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional real-backend checks for Diffusers quantization config conversion."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _load_quantization_utils():
+    module_path = (
+        Path(__file__).resolve().parents[3] / "vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py"
+    )
+    spec = importlib.util.spec_from_file_location("diffusers_quantization_utils_real_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_convert_diffusers_quantization_config_from_real_torchao_dict():
+    diffusers = pytest.importorskip("diffusers")
+    torchao_quantization = pytest.importorskip("torchao.quantization")
+    quant_type_cls = getattr(torchao_quantization, "Int8DynamicActivationInt8WeightConfig", None)
+    if quant_type_cls is None:
+        pytest.skip("torchao does not expose Int8DynamicActivationInt8WeightConfig")
+
+    quantization_utils = _load_quantization_utils()
+
+    real_torchao_config = diffusers.TorchAoConfig(quant_type=quant_type_cls())
+    serialized_torchao_config = real_torchao_config.to_dict()
+    assert serialized_torchao_config["quant_type"].keys() == {"default"}
+
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": serialized_torchao_config,
+            },
+        },
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    component_config = quant_config.quant_mapping["transformer"]
+    assert isinstance(quant_config, diffusers.PipelineQuantizationConfig)
+    assert isinstance(component_config, diffusers.TorchAoConfig)
+    assert isinstance(
+        component_config.quant_type,
+        quant_type_cls,
+    )

tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused tests for Diffusers backend quantization conversion helpers."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _FakePipelineQuantizationConfig:
+    def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+        self.quant_mapping = quant_mapping
+        self.quant_backend = quant_backend
+        self.quant_kwargs = quant_kwargs
+        self.components_to_quantize = components_to_quantize
+
+
+class _FakeTorchAoConfig:
+    def __init__(self, quant_type):
+        self.quant_type = quant_type
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(quant_type=config["quant_type"])
+
+
+class _FakeFloat8DynamicActivationFloat8WeightConfig:
+    pass
+
+
+class _FakeInt8DynamicActivationInt8WeightConfig:
+    pass
+
+
+@pytest.fixture(autouse=True)
+def patch_quantization_backends(monkeypatch):
+    monkeypatch.setattr(
+        quantization_utils,
+        "PipelineQuantizationConfig",
+        _FakePipelineQuantizationConfig,
+    )
+    monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig)
+
+    def fake_get_torchao_quant_type_cls(class_name: str):
+        return {
+            "Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig,
+            "Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig,
+        }[class_name]
+
+    monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls)
+    monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig)
+
+
+def _quant_config(method: str, **kwargs):
+    return SimpleNamespace(get_name=lambda: method, **kwargs)
+
+
+@pytest.mark.parametrize(
+    ("method", "expected_type"),
+    [
+        ("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig),
+        ("int8", _FakeInt8DynamicActivationInt8WeightConfig),
+    ],
+)
+def test_apply_injects_converted_quantization_config(method, expected_type):
+    od_config = SimpleNamespace(quantization_config=_quant_config(method))
+    load_kwargs = {}
+
+    injected = quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {"transformer": ["diffusers", "Transformer2DModel"]},
+    )
+
+    assert injected is True
+    torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"]
+    assert isinstance(torchao_config.quant_type, expected_type)
+
+
+def test_apply_injects_converted_quantization_config_for_transformer_2():
+    od_config = SimpleNamespace(quantization_config=_quant_config("int8"))
+    load_kwargs = {}
+
+    quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {
+            "transformer": ["diffusers", "Transformer2DModel"],
+            "transformer_2": ["diffusers", "Transformer2DModel"],
+        },
+    )
+
+    assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"]
+
+
+@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"])
+def test_unsupported_methods_fail_explicitly(method):
+    with pytest.raises(NotImplementedError, match=method):
+        quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method))
+
+
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        _quant_config("fp8", activation_scheme="static"),
+        _quant_config("int8", is_checkpoint_int8_serialized=True),
+        _quant_config("fp8", weight_block_size=[128, 128]),
+        _quant_config("int8", ignored_layers=["transformer.proj_out"]),
+    ],
+)
+def test_ambiguous_mappings_fail_explicitly(quant_config):
+    with pytest.raises(NotImplementedError):
+        quantization_utils.ensure_supported_diffusers_quantization(quant_config)
+
+
+def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker):
+    od_config = SimpleNamespace(quantization_config=_quant_config("fp8"))
+    existing = object()
+    load_kwargs = {"quantization_config": existing}
+    mock_warning = mocker.patch.object(quantization_utils.logger, "warning")
+
+    injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {})
+
+    assert injected is False
+    assert load_kwargs["quantization_config"] is existing
+    assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0]
+
+
+def test_convert_diffusers_quantization_config_from_dict():
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": {
+                    "quant_method": "torchao",
+                    "quant_type": "fake-int8-config",
+                }
+            }
+        }
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    assert isinstance(quant_config, _FakePipelineQuantizationConfig)
+    assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig)
+    assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -19,6 +19,10 @@ class _FakeGroup:
         self.group_ranks = group_ranks
         self.parallel_mode = parallel_mode
         self.device_group = object()
+        self.device_communicator = kwargs.get("device_communicator")
+        reduce_scatter = kwargs.get("reduce_scatter")
+        if reduce_scatter is not None:
+            self.reduce_scatter = reduce_scatter
         self.ulysses_group = kwargs.get("ulysses_group")
         self.ring_group = kwargs.get("ring_group")
         self.local_group = next(group for group in group_ranks if local_rank in group)
@@ -55,6 +59,23 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         created_groups.append(group)
         return group
 
+    def fake_init_vllm_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        group_name,
+    ):
+        del backend
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            f"vllm_{group_name}",
+            device_communicator=object(),
+            reduce_scatter=lambda tensor, **kwargs: tensor,
+        )
+        created_groups.append(group)
+        return group
+
     fake_world_group = SimpleNamespace(
         rank_in_group=local_rank,
         local_rank=local_rank,
@@ -68,6 +89,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
     for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
@@ -87,7 +109,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
     assert parallel_state.vllm_parallel_state._PCP.world_size == 2
     assert parallel_state._DP.world_size == 2
     assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
@@ -96,6 +118,13 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._TP.world_size == 2
     assert parallel_state._PP.world_size == 2
 
+    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
+    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
+
     assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
     assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
     assert parallel_state.vllm_parallel_state._EP.local_group == [
@@ -155,7 +184,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
+    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 

tests/diffusion/distributed/test_hsdp.py

@@ -1,19 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
-These tests verify HSDP configuration logic without requiring a distributed environment.
-"""
+import gc
+import os
+import socket
 
 import pytest
+import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.tensor import DeviceMesh, DTensor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    _unshardable_parameters,
+    shard_model,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
+class _PackedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
+        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
+
+
+class _PackedModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block = _PackedBlock()
+        self.root_weight = nn.Parameter(torch.ones(2))
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield
+        return
+
+    master_port = _find_free_port()
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+            os.environ.pop(key, None)
+        gc.collect()
+
+
+def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
+    model = _PackedModel()
+    ignored_params = _unshardable_parameters(model)
+    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
+    assert ignored_params == expected_ignored
+
+    packed_weight = model.block.packed_weight
+    input_global_scale = model.block.input_global_scale
+    shard_model(
+        model,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "block"],
+        ignored_params=ignored_params,
+    )
+
+    assert isinstance(model.block.weight, DTensor)
+    assert isinstance(model.root_weight, DTensor)
+    assert model.block.packed_weight is packed_weight
+    assert model.block.input_global_scale is input_global_scale
+    assert not isinstance(model.block.packed_weight, DTensor)
+    assert not isinstance(model.block.input_global_scale, DTensor)
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 

tests/diffusion/layers/test_rope_broadcast.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shape broadcasting regression tests for RotaryEmbedding.forward_native.
+
+Covers the batch-dimension shape mismatch fix: forward_native now strips the
+batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
+forward_hip, and apply_rotary_emb_mindiesd paths.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_inputs(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+    cos_batch: int | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Generate x, cos, sin tensors for RoPE testing.
+
+    Args:
+        batch: Batch dimension of x.
+        seq_len: Sequence length.
+        num_heads: Number of attention heads.
+        head_dim: Head dimension (must be even).
+        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
+        dtype: Data type for the tensors.
+    """
+    if cos_batch is None:
+        cos_batch = batch
+
+    torch.manual_seed(42)
+    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
+    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    return x, cos, sin
+
+
+class TestRotaryEmbeddingNativeShapeRegression:
+    """Regression tests for forward_native shape handling (#8e297d81)."""
+
+    @pytest.fixture
+    def rope_neox(self) -> RotaryEmbedding:
+        """NeoX-style (non-interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=True)
+
+    @pytest.fixture
+    def rope_gptj(self) -> RotaryEmbedding:
+        """GPT-J style (interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=False)
+
+    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────
+
+    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
+        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
+        # Manually drop batch dim to create 2D cos/sin
+        cos_2d = cos[0]
+        sin_2d = sin[0]
+        assert cos_2d.dim() == 2
+        output = rope_neox.forward_native(x, cos_2d, sin_2d)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()
+        assert not torch.isinf(output).any()
+
+    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────
+
+    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
+    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
+        """3D cos/sin with batch dim should not raise errors."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+
+    @pytest.mark.parametrize("interleaved", [False, True])
+    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
+        """3D cos/sin with identical batch slices == 2D cos/sin result."""
+        rope = RotaryEmbedding(is_neox_style=not interleaved)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+
+        # 3D variant: expand the same 2D data across batch
+        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
+        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)
+
+        out_2d = rope.forward_native(x, cos_2d, sin_2d)
+        out_3d = rope.forward_native(x, cos_3d, sin_3d)
+
+        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)
+
+    # ── Batch size mismatch — the regression scenario ─────────────────
+
+    @pytest.mark.parametrize(
+        "x_batch,cos_batch",
+        [
+            (2, 1),  # cos has fewer batch elements
+            (1, 2),  # cos has more batch elements (takes [0])
+            (4, 1),  # large x batch, single cos
+        ],
+    )
+    def test_batch_size_mismatch_regression(
+        self,
+        rope_neox: RotaryEmbedding,
+        x_batch: int,
+        cos_batch: int,
+    ) -> None:
+        """forward_native works when cos batch dim differs from x batch dim.
+
+        This is the exact regression scenario — before the fix, a 3D cos/sin
+        with a different batch size than x would cause a shape mismatch error
+        inside apply_rotary_emb_torch during broadcast.
+        """
+        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert output.dtype == x.dtype
+
+    # ── Batch consistency — all elements get same rotation ────────────
+
+    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 2D cos, identical x slices → identical output slices."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
+        cos = torch.randn(S, D // 2)
+        sin = torch.randn(S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        # All output batch slices must be identical
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 3D cos (identical slices), identical x → identical output."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)
+        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
+        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    # ── Output shape / dtype invariants ───────────────────────────────
+
+    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output shape must match input shape for all cos/sin dimensionalities."""
+        B, S, H, D = 2, 32, 4, 128
+        x = torch.randn(B, S, H, D)
+
+        # 2D cos/sin
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)
+
+        # 3D cos/sin, matching batch
+        cos_3d = torch.randn(B, S, D // 2)
+        sin_3d = torch.randn(B, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)
+
+        # 3D cos/sin, different batch
+        cos_3d_1 = torch.randn(1, S, D // 2)
+        sin_3d_1 = torch.randn(1, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)
+
+    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output dtype must match input dtype."""
+        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
+            output = rope_neox.forward_native(x, cos, sin)
+            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"
+
+    def test_output_no_nan_inf(
+        self,
+        rope_neox: RotaryEmbedding,
+        rope_gptj: RotaryEmbedding,
+    ) -> None:
+        """Output must be finite for both RoPE styles."""
+        for rope in (rope_neox, rope_gptj):
+            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
+            output = rope.forward_native(x, cos, sin)
+            assert not torch.isnan(output).any(), f"NaN in {rope}"
+            assert not torch.isinf(output).any(), f"Inf in {rope}"
+
+    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
+        """Output should be contiguous."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.is_contiguous()
+
+    # ── CFG-style doubled batch ───────────────────────────────────────
+
+    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """CFG (classifier-free guidance) doubles batch — cos may have
+        original batch size while x has 2× batch."""
+        B, S, H, D = 2, 16, 8, 64
+        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
+        x = torch.randn(B * 2, S, H, D)
+        # cos/sin only have the original batch
+        cos = torch.randn(B, S, D // 2)
+        sin = torch.randn(B, S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == (B * 2, S, H, D)
+
+    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────
+
+    def test_wan_forward_native_unchanged(self) -> None:
+        """RotaryEmbeddingWan.forward_native still uses its own implementation
+        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
+        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+
+        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos = torch.randn(B, S, 1, D // 2)
+        sin = torch.randn(B, S, 1, D // 2)
+
+        output = rope.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()

tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py

@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Flux2-Klein GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
-    Flux2KleinGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "Flux2Pipeline",
-    model_type: str = "flux",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_flux2_klein_adapter_selected_for_flux_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, Flux2KleinGGUFAdapter)
-
-
-def test_flux2_klein_adapter_matches_flux_model_type():
-    assert Flux2KleinGGUFAdapter.is_compatible(
-        _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
-        object(),
-        _make_source(),
-    )
-
-
-def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight", torch.ones((1, 1))),
-                ("time_in.in_layer.qweight_type", torch.tensor(1)),
-                ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
-                ("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "x_embedder.qweight" in names
-    assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
-    assert "transformer_blocks.0.attn.norm_q.weight" in names
-    assert "proj_out.qweight" in names
-
-
-def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                (
-                    "final_layer.adaLN_modulation.1.scale",
-                    torch.tensor([1.0, 2.0, 3.0, 4.0]),
-                )
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert len(weights) == 1
-    assert weights[0][0] == "norm_out.linear.weight"
-    assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))

tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py

@@ -1,184 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
-    QwenImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "QwenImagePipeline",
-    model_type: str = "qwen_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-class _FakeTransformer:
-    def __init__(self) -> None:
-        self._params = {
-            "img_in.qweight": 1,
-            "img_in.qweight_type": 1,
-            "txt_in.qweight": 1,
-            "txt_in.qweight_type": 1,
-            "proj_out.qweight": 1,
-            "proj_out.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
-            "transformer_blocks.0.img_mod.1.qweight": 1,
-            "transformer_blocks.0.img_mod.1.qweight_type": 1,
-            "transformer_blocks.0.txt_mod.1.qweight": 1,
-            "transformer_blocks.0.txt_mod.1.qweight_type": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight_type": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
-            "transformer_blocks.0.attn.to_out.weight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
-        }
-
-    def named_parameters(self):
-        return list(self._params.items())
-
-    def named_buffers(self):
-        return []
-
-    def get_submodule(self, _prefix: str):
-        return self
-
-
-def test_qwen_adapter_selected_for_qwen_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-    assert isinstance(adapter, QwenImageGGUFAdapter)
-
-
-def test_qwen_adapter_matches_multiple_pipeline_variants():
-    for model_class_name in (
-        "QwenImagePipeline",
-        "QwenImageEditPipeline",
-        "QwenImageEditPlusPipeline",
-        "QwenImageLayeredPipeline",
-    ):
-        assert QwenImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name=model_class_name),
-            _FakeTransformer(),
-            _make_source(),
-        )
-
-
-def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.attn.to_q.qweight_type", 1),
-                ("transformer_blocks.0.attn.to_q.qweight", 2),
-                ("transformer_blocks.0.attn.to_k.qweight_type", 3),
-                ("transformer_blocks.0.attn.to_k.qweight", 4),
-                ("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
-                ("transformer_blocks.0.attn.to_out.0.qweight", 6),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
-
-
-def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight_type", 1),
-                ("img_in.qweight", 2),
-                ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
-                ("time_text_embed.timestep_embedder.linear_1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("img_in.qweight_type", 1) in weights
-    assert ("img_in.qweight", 2) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
-
-
-def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.img_mod.1.qweight_type", 1),
-                ("transformer_blocks.0.img_mod.1.qweight", 2),
-                ("transformer_blocks.0.txt_mod.1.qweight_type", 3),
-                ("transformer_blocks.0.txt_mod.1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights

tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py

@@ -1,83 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Z-Image GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
-    ZImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "ZImagePipeline",
-    model_type: str = "z_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_z_image_adapter_selected_for_z_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, ZImageGGUFAdapter)
-
-
-def test_z_image_adapter_matches_model_type_variants():
-    for model_type in ("z_image", "zimage", "z-image"):
-        assert ZImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name="OtherPipeline", model_type=model_type),
-            object(),
-            _make_source(),
-        )
-
-
-def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
-
-    monkeypatch.setattr(
-        z_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
-                ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
-                ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
-                ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = ZImageGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "all_final_layer.2-1.qweight" in names
-    assert "all_x_embedder.2-1.qweight_type" in names
-    assert "transformer_blocks.0.attention.to_out.0.weight" in names
-    assert "transformer_blocks.0.attention.to_qkv.qweight" in names

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -16,7 +16,6 @@ from vllm.config.load import LoadConfig
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
         loader.load_weights(model)
 
 
-def test_qwen_model_class_selects_qwen_gguf_adapter():
-    od_config = type(
-        "Config",
-        (),
-        {
-            "model_class_name": "QwenImagePipeline",
-            "tf_model_config": {"model_type": "qwen_image"},
-        },
-    )()
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
-
-    assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
-
-
 class _ConfigAwareModel(nn.Module):
     def __init__(self, *, od_config):
         super().__init__()
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
@@ -202,6 +179,59 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_hsdp_processes_quantized_weights_before_sharding(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: events.append("prepare"),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "process", "prepare", "shard"]
+
+
 def test_get_all_weights(prefetch_helios_model, mock_tp_group):
     """Ensure that get all weights on a tiny model resolves to nonempty weights."""
     od_config = OmniDiffusionConfig(

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -1,259 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-class _DummyModel(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.transformer = nn.Linear(2, 2, bias=True)
-        self.vae = nn.Linear(2, 2, bias=False)
-        self.register_buffer("transformer_buffer", torch.ones(1))
-        self.calls: list[list[str]] = []
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="vae",
-                revision=None,
-                prefix="vae.",
-                fall_back_to_pt=True,
-            ),
-        ]
-
-    def load_weights(self, weights):
-        loadable = dict(self.named_parameters())
-        loadable.update(dict(self.named_buffers()))
-        seen: list[str] = []
-        loaded: set[str] = set()
-        for name, tensor in weights:
-            seen.append(name)
-            if name in loadable:
-                target = loadable[name]
-                target.data.copy_(tensor.to(dtype=target.dtype))
-                loaded.add(name)
-        self.calls.append(seen)
-        return loaded
-
-
-def _make_loader() -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
-    loader.load_config = SimpleNamespace(
-        download_dir="cache-dir",
-        ignore_patterns=["*.tmp"],
-    )
-    loader.od_config = None
-    loader.quant_config = None
-    loader.counter_before_loading_weights = 0.0
-    loader.counter_after_loading_weights = 0.0
-    return loader
-
-
-def test_is_gguf_quantization_accepts_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "gguf",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_rejects_non_gguf_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "fp8",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is False
-
-
-def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {"method": "gguf"}
-
-    with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization()
-
-
-def test_is_gguf_quantization_accepts_object_config():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(
-        get_name=lambda: "gguf",
-        gguf_model="weights.gguf",
-    )
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_uses_fallback_object_without_get_name():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_get_model_loadable_names_collects_parameters_and_buffers():
-    loader = _make_loader()
-    model = _DummyModel()
-
-    names = loader._get_model_loadable_names(model)
-
-    assert "transformer.weight" in names
-    assert "transformer.bias" in names
-    assert "transformer_buffer" in names
-
-
-def test_resolve_gguf_model_path_returns_local_file(tmp_path):
-    loader = _make_loader()
-    gguf_file = tmp_path / "model.gguf"
-    gguf_file.write_bytes(b"gguf")
-
-    resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
-
-    assert resolved == str(gguf_file)
-
-
-def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "hf_hub_download",
-        lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
-
-
-def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "download_gguf",
-        lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
-
-
-def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-    loader.od_config = SimpleNamespace(revision="main")
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    class _Adapter:
-        def weights_iterator(self):
-            yield "weight", torch.ones((2, 2))
-            yield "bias", torch.zeros(2)
-
-    monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
-    monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
-
-    weights = list(loader._get_gguf_weights_iterator(source, object()))
-
-    assert weights[0][0] == "transformer.weight"
-    assert torch.equal(weights[0][1], torch.ones((2, 2)))
-    assert weights[1][0] == "transformer.bias"
-    assert torch.equal(weights[1][1], torch.zeros(2))
-
-
-def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
-    )
-
-    def _hf_iter(source):
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert model.calls[0] == ["transformer.weight"]
-    assert model.calls[1] == ["transformer.bias"]
-    assert model.calls[2] == ["vae.weight"]
-
-
-def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter(
-            [
-                ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
-                ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
-            ]
-        ),
-    )
-
-    hf_calls: list[str] = []
-
-    def _hf_iter(source):
-        hf_calls.append(source.subfolder or "")
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert hf_calls == ["vae"]

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -9,6 +9,7 @@ import torch.nn as nn
 
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     ModelOptFp8CheckpointAdapter,
+    ModelOptNvFp4CheckpointAdapter,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -42,6 +43,70 @@ class _QuantizedPackedModelOptModel(nn.Module):
         )
 
 
+class _RemappedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Linear(2, 2, bias=False)
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return {
+            "transformer.orig.proj.weight": "runtime.proj.weight",
+        }.get(name, name)
+
+
+class _QuantizedRemappedModelOptModel(nn.Module):
+    """FP8 target whose params live under a namespace the raw checkpoint keys
+    never use, reachable only via ``remap_checkpoint_key``.
+
+    Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot
+    resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook
+    the adapter treats every tensor as unresolved and silently drops the scales.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        prefix = "transformer.orig.proj."
+        if name.startswith(prefix):
+            return "runtime.proj." + name[len(prefix) :]
+        return name
+
+
+class _RemappedNvFp4Model(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(1, 1, dtype=torch.uint8), requires_grad=False),
+        )
+        for name in ("input_scale", "weight_scale", "weight_scale_2"):
+            self.runtime.proj.register_parameter(name, nn.Parameter(torch.empty(1), requires_grad=False))
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return name.replace("transformer.orig.proj.", "runtime.proj.")
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -92,3 +157,85 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
         "transformer.block.to_q.weight_scale",
         "transformer.block.to_q.input_scale",
     ]
+
+
+def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype():
+    model = _RemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight", fp8_weight),
+                    ("transformer.orig.proj.weight_scale", scale),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["runtime.proj.weight"]
+    assert adapted[0][1].dtype == model.runtime.proj.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
+    """Regression: scales were dropped for a quantized (FP8) target that only the
+    model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are
+    emitted and the FP8 weight passes through unchanged (no dequantization)."""
+    model = _QuantizedRemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    input_scale = torch.tensor([1.0], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight_scale", weight_scale),
+                    ("transformer.orig.proj.input_scale", input_scale),
+                    ("transformer.orig.proj.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    emitted = dict(adapted)
+    assert len(adapted) == len(emitted) == 3
+
+    # Both scales are retained and emitted under the remapped runtime names.
+    assert torch.equal(emitted["runtime.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["runtime.proj.input_scale"], input_scale)
+
+    # The FP8 weight is passed through unchanged under the remapped target name
+    # (the quantized param is FP8, so it must NOT be dequantized). Compare via
+    # float32 since torch.equal has no FP8 CPU kernel.
+    assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
+    assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))
+
+
+def test_modelopt_nvfp4_adapter_remaps_quantized_weights_and_scales():
+    adapter = ModelOptNvFp4CheckpointAdapter(_RemappedNvFp4Model(), _make_source())
+    prefix = "transformer.orig.proj"
+    checkpoint_tensors = [
+        (f"{prefix}.input_scale", torch.tensor([1.0])),
+        (f"{prefix}.weight_scale", torch.tensor([0.5])),
+        (f"{prefix}.weight_scale_2", torch.tensor([0.25])),
+        (f"{prefix}.weight", torch.tensor([[1]], dtype=torch.uint8)),
+    ]
+
+    adapted = list(adapter.adapt(iter(checkpoint_tensors)))
+
+    assert [name for name, _ in adapted] == [
+        name.replace(f"{prefix}.", "runtime.proj.") for name, _ in checkpoint_tensors
+    ]
+
+
+def test_modelopt_nvfp4_adapter_rejects_unconsumed_pre_quant_scale():
+    adapter = ModelOptNvFp4CheckpointAdapter(nn.Module(), _make_source())
+
+    with pytest.raises(ValueError, match="does not consume pre_quant_scale"):
+        list(adapter.adapt(iter([("transformer.proj.pre_quant_scale", torch.tensor([1.0]))])))

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -10,9 +10,38 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture(autouse=True)
+def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Provide the TP metadata required by vLLM parallel linear layers."""
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+
+    from vllm_omni.diffusion.models.cosmos3 import (
+        transformer_cosmos3,
+        transformer_cosmos3_edge,
+    )
+
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1)
+
+
+@pytest.fixture
+def accelerator_device() -> torch.device:
+    """Provide an accelerator device, skipping when none is available."""
+    if current_omni_platform.get_device_count() == 0:
+        pytest.skip("Accelerator required for this test")
+    return current_omni_platform.get_torch_device(0)
+
+
 def _tiny_cosmos3_config(**overrides):
     config = {
         "hidden_size": 8,
@@ -30,6 +59,22 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
+def _tiny_cosmos3_edge_config(**overrides):
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        COSMOS3_EDGE_BACKBONE_TYPE,
+    )
+
+    config = _tiny_cosmos3_config(
+        latent_patch_size=2,
+        latent_channel=48,
+        temporal_compression_factor=4,
+        backbone_type=COSMOS3_EDGE_BACKBONE_TYPE,
+        qk_norm_for_text=False,
+    )
+    config.update(overrides)
+    return config
+
+
 def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
         compute_mrope_position_ids_action,
@@ -68,6 +113,16 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     assert action_offset == 14
 
 
+def test_timestep_embedder_stores_frequencies_in_fp32() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        TimestepEmbedder,
+    )
+
+    embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16)
+
+    assert embedder.freqs.dtype == torch.float32
+
+
 @pytest.mark.parametrize(
     ("key", "value"),
     [
@@ -87,29 +142,245 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) ->
     Cosmos3VFMTransformer._validate_supported_config(None)
 
 
-def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+def test_edge_config_resolves_nemotron_defaults() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
 
-    model = object.__new__(Cosmos3VFMTransformer)
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                layer_norm_epsilon=1e-5,
+                rope_scaling={},
+                rope_parameters={"mrope_section": [24, 20, 20]},
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    assert model.rms_norm_eps == 1e-5
+    assert model.rope_theta == 100_000_000
+    assert model.mrope_section == [24, 20, 20]
+    assert model.latent_channel_size == 48
+    assert model.latent_patch_size == 2
+    assert model.temporal_compression_factor == 4
+
+
+def test_edge_config_requires_backbone_type() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    config = _tiny_cosmos3_edge_config()
+    config.pop("backbone_type")
+
+    with pytest.raises(ValueError, match="must declare backbone_type"):
+        Cosmos3EdgeVFMTransformer(
+            SimpleNamespace(
+                tf_model_config=config,
+                dtype=torch.float32,
+            )
+        )
+
+
+def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+        Cosmos3Relu2MLP,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                use_k_norm_und_for_gen=False,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    layer = model.language_model.layers[0]
+    gen_layer = model.gen_layers[0]
+    assert isinstance(layer.mlp, Cosmos3Relu2MLP)
+    assert not hasattr(layer.mlp, "gate_proj")
+    assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP)
+    assert not hasattr(gen_layer.mlp, "gate_proj")
+    assert not hasattr(layer.self_attn, "norm_q")
+    assert not hasattr(layer.self_attn, "norm_k")
+    assert layer.self_attn.k_norm_und_for_gen is None
+
+
+def test_edge_creates_gen_k_norm_when_enabled() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                qk_norm_for_diffusion=True,
+                use_k_norm_und_for_gen=True,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen
+    assert model.use_k_norm_und_for_gen is True
+    assert isinstance(k_norm, RMSNorm)
+    assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm)
+    assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm)
+
+
+def test_edge_supports_no_gen_qk_norm_variant() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                qk_norm_for_diffusion=False,
+                use_k_norm_und_for_gen=False,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    assert model.qk_norm_for_diffusion is False
+    assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None
+    assert model.gen_layers[0].cross_attention.qk_norm is False
+    assert not hasattr(model.gen_layers[0].cross_attention, "norm_q")
+    assert not hasattr(model.gen_layers[0].cross_attention, "norm_k")
+
+
+def test_edge_validates_required_relu2_weights() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = object.__new__(Cosmos3EdgeVFMTransformer)
     nn.Module.__init__(model)
-    model.language_model = nn.Module()
-    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
-    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
-    model.norm_moe_gen = nn.LayerNorm(2)
+    model.num_hidden_layers = 1
+    model.qk_norm_for_diffusion = True
+    model.use_k_norm_und_for_gen = True
 
-    matched = [
-        name
-        for name, module in model.named_modules()
-        if any(condition(name, module) for condition in model._hsdp_shard_conditions)
-    ]
-    assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
-    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
-    assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+    complete = {
+        "transformer.language_model.layers.0.mlp.up_proj.weight",
+        "transformer.language_model.layers.0.mlp.down_proj.weight",
+        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
+        "transformer.gen_layers.0.mlp.up_proj.weight",
+        "transformer.gen_layers.0.mlp.down_proj.weight",
+    }
+    model.validate_loaded_weights(complete)
 
-    model.latent_patch_size = 2
-    model.latent_channel_size = 3
-    latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
-    torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents)
+    mlp_weights = complete - {
+        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
+    }
+    for missing_weight in mlp_weights:
+        with pytest.raises(ValueError, match="missing required weights"):
+            model.validate_loaded_weights(complete - {missing_weight})
+
+    missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"}
+    with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"):
+        model.validate_loaded_weights(missing_k_norm)
+
+    model.use_k_norm_und_for_gen = False
+    model.validate_loaded_weights(missing_k_norm)
+
+
+def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeCausalAttention,
+    )
+
+    class ScaleProjection(nn.Module):
+        def __init__(self, scale: float) -> None:
+            super().__init__()
+            self.scale = scale
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x * self.scale
+
+    class IdentityOutput(nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x
+
+    class RawKAttention(nn.Module):
+        def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+            del v
+            return q + k
+
+    attn = object.__new__(Cosmos3EdgeCausalAttention)
+    nn.Module.__init__(attn)
+    attn.hidden_size = 4
+    attn.num_heads = 2
+    attn.num_kv_heads = 2
+    attn.head_dim = 2
+    attn.num_heads_local = 2
+    attn.num_kv_heads_local = 2
+    attn.to_q = ScaleProjection(1.0)
+    attn.to_k = ScaleProjection(1.0)
+    attn.to_v = ScaleProjection(1.0)
+    attn.to_out = IdentityOutput()
+    attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5)
+    attn.attn = RawKAttention()
+
+    hidden_states = torch.tensor(
+        [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]],
+        dtype=torch.float32,
+    )
+    freqs_cos = torch.ones(1, 2, 1, 2)
+    freqs_sin = torch.zeros(1, 2, 1, 2)
+
+    reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin)
+    attn.to_k.scale = 3.0
+    scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin)
+
+    torch.testing.assert_close(gen_k, scaled_gen_k)
+    torch.testing.assert_close(gen_v, scaled_gen_v)
+    assert not torch.allclose(reasoner_out, scaled_reasoner_out)
+
+
+def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeLanguageModel,
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    for transformer_cls, language_model_cls in (
+        (Cosmos3VFMTransformer, Cosmos3LanguageModel),
+        (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel),
+    ):
+        model = object.__new__(transformer_cls)
+        nn.Module.__init__(model)
+        model.language_model = nn.Module()
+        model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
+        model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+        model.norm_moe_gen = nn.LayerNorm(2)
+
+        matched = [
+            name
+            for name, module in model.named_modules()
+            if any(condition(name, module) for condition in model._hsdp_shard_conditions)
+        ]
+        assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
+        assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"]
+        assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"]
+        assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+
+        model.latent_patch_size = 2
+        model.latent_channel_size = 3
+        latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
+        torch.testing.assert_close(
+            model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5),
+            latents,
+        )
 
 
 def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -131,6 +402,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_qwen_gen_mlp_remains_gated() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
+
+    model = Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32)
+    )
+
+    assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP)
+    assert hasattr(model.gen_layers[0].mlp, "gate_proj")
+
+
+def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    class ToyReasonerLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, freqs):
+            del freqs
+            kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1)
+            return hidden + self.weight.to(hidden.dtype), kv, kv
+
+    class ToyGeneratorLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, **kwargs):
+            del kwargs
+            return hidden + self.weight.to(hidden.dtype)
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([ToyReasonerLayer()])
+    model.gen_layers = nn.ModuleList([ToyGeneratorLayer()])
+    model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"}
+    assert model.device == torch.device("cpu")
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+    # The reasoner runs once then the generator component stays resident for GEN.
+    assert model._active_model_cpu_offload_component == "generator"
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+
+
+def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
+    monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.to(accelerator_device)
+
+    reasoner_param = model.language_model.layers[0].weight
+    generator_param = model.gen_layers[0].weight
+
+    model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    # On enable, every group is parked on CPU until a phase activates it.
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("reasoner")
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("generator")
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device == accelerator_device
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+    # Disable restores both components onto the device.
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device == accelerator_device
+
+
 def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/flux/test_flux_pipeline.py

@@ -21,6 +21,8 @@ def _make_flux_sampling(**overrides):
         "true_cfg_scale": 4.0,
         "num_outputs_per_prompt": 0,
         "latents": None,
+        "output_type": None,
+        "max_sequence_length": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -131,7 +133,7 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                     "negative_prompt_embeds": negative_prompt_embeds_a,
                     "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a,
                 },
-                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a),
+                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"),
             ),
             SimpleNamespace(
                 request_id="flux-prompt-b",
@@ -145,12 +147,12 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                         "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b],
                     },
                 },
-                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b),
+                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"),
             ),
         ]
     )
 
-    outputs = pipeline.forward(batch, output_type="latent")
+    outputs = pipeline.forward(batch)
 
     assert encode_calls[0]["prompt"] is None
     assert encode_calls[0]["prompt_2"] is None

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,6 +50,48 @@ class TestSetForwardContextNumTokens:
         mock_get.assert_not_called()
 
 
+class TestSetForwardContextDPMetadata:
+    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
+
+    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = None
+        fake_group = mocker.MagicMock()
+        fake_group.world_size = 2
+        fake_group.cpu_group = object()
+
+        def _all_gather_object(output, obj, group):
+            assert group is fake_group.cpu_group
+            output[:] = [obj, obj + 1]
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
+        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is not None
+        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
+
+    def test_keeps_existing_dp_metadata(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        existing_metadata = object()
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = existing_metadata
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is existing_metadata
+        mock_all_gather.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

tests/diffusion/models/hunyuan_image3/test_multi_resolution.py

@@ -524,11 +524,11 @@ class TestAr2DiffusionIntegration:
         self,
         token_ids: list[int],
         prompt_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, Any]:
+    ) -> dict[str, Any] | None:
         """Invoke ``ar2diffusion`` with a single fake AR output.
 
-        Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
-        assert on the resolved height/width.
+        Returns the diffusion request payload so callers can assert on the
+        resolved height/width.
         """
         from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
             ar2diffusion,
@@ -601,8 +601,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001, token_id],
             prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == expected_h, (
             f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
         )
@@ -617,8 +617,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001],  # no ratio tokens
             prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == 768
         assert diffusion_input["width"] == 1024
 
@@ -633,9 +633,9 @@ class TestAr2DiffusionIntegration:
                 token_ids=[150000],
                 prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
             )
-        assert len(result) == 1
-        assert result[0]["height"] == 512
-        assert result[0]["width"] == 512
+        assert result is not None
+        assert result["height"] == 512
+        assert result["width"] == 512
 
     # -- cache behaviour -----------------------------------------------------
 

tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py

@@ -0,0 +1,436 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 forward-parallel behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCFGParallelIntegration:
+    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
+        """LTX23Pipeline must use the shared CFG parallel implementation."""
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        assert issubclass(LTX23Pipeline, CFGParallelMixin)
+
+    def test_ltx2_transformer_has_dit_cache_config(self):
+        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
+        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+        assert isinstance(adapter_config, CacheDiTAdapterConfig)
+        assert adapter_config.has_separate_cfg
+
+
+class TestCFGParallelHelpers:
+    """Test LTX-2.3 CFG helper math without loading model weights."""
+
+    def test_combine_cfg_noise_matches_x0_space_formula(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = object.__new__(LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        video_combined, audio_combined = pipe.combine_cfg_noise(
+            (video_pos, audio_pos),
+            (video_neg, audio_neg),
+            scale,
+            video_latents=video_sample,
+            audio_latents=audio_sample,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+        x0_video_cond = video_sample - video_pos * video_sigma
+        x0_video_uncond = video_sample - video_neg * video_sigma
+        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
+        expected_video = (video_sample - x0_video_guided) / video_sigma
+
+        x0_audio_cond = audio_sample - audio_pos * audio_sigma
+        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
+        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
+        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
+        assert torch.allclose(video_combined, expected_video)
+        assert torch.allclose(audio_combined, expected_audio)
+
+    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
+        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if tensor.shape == video_pos.shape:
+                    return [video_pos, video_neg]
+                return [audio_pos, audio_neg]
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_sample,
+            video_pos,
+            video_neg,
+            video_sigma,
+            scale,
+        )
+        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_sample,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            scale,
+        )
+
+        for rank, expected_branch in ((0, "positive"), (1, "negative")):
+            calls = []
+            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+
+            def fake_predict_noise(**kwargs):
+                calls.append(kwargs["branch"])
+                if kwargs["branch"] == "positive":
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
+            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
+                true_cfg_scale=scale,
+                positive_kwargs={"branch": "positive"},
+                negative_kwargs={"branch": "negative"},
+                cfg_normalize=False,
+                video_latents=video_sample,
+                audio_latents=audio_sample,
+                video_sigma=video_sigma,
+                audio_sigma=audio_sigma,
+            )
+
+            assert calls == [expected_branch]
+            torch.testing.assert_close(video_combined, expected_video)
+            torch.testing.assert_close(audio_combined, expected_audio)
+
+        assert "_cfg_video_latents" not in pipe.__dict__
+        assert "_cfg_audio_latents" not in pipe.__dict__
+
+
+class TestCFGParallelForwardPath:
+    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
+
+    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 4
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
+
+        class StopAtEncodePromptError(Exception):
+            pass
+
+        captured = {}
+
+        def fake_encode_prompt(**kwargs):
+            captured.update(kwargs)
+            raise StopAtEncodePromptError
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        prompt_embeds_a = torch.zeros(2, 3)
+        prompt_embeds_b = torch.ones(2, 3)
+        negative_prompt_embeds_a = torch.full((2, 3), 2.0)
+        negative_prompt_embeds_b = torch.full((2, 3), 3.0)
+        prompt_attention_mask_a = torch.tensor([True, True])
+        prompt_attention_mask_b = torch.tensor([True, False])
+        negative_attention_mask_a = torch.tensor([False, True])
+        negative_attention_mask_b = torch.tensor([False, False])
+
+        requests = [
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-a",
+                    "negative_prompt": "negative-a",
+                    "prompt_embeds": prompt_embeds_a,
+                    "negative_prompt_embeds": negative_prompt_embeds_a,
+                    "prompt_attention_mask": prompt_attention_mask_a,
+                    "negative_prompt_attention_mask": negative_attention_mask_a,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-a",
+            ),
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-b",
+                    "negative_prompt": "negative-b",
+                    "prompt_embeds": prompt_embeds_b,
+                    "negative_prompt_embeds": negative_prompt_embeds_b,
+                    "attention_mask": prompt_attention_mask_b,
+                    "negative_attention_mask": negative_attention_mask_b,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-b",
+            ),
+        ]
+
+        with pytest.raises(StopAtEncodePromptError):
+            pipe.forward(DiffusionRequestBatch(requests=requests))
+
+        assert captured["prompt"] is None
+        assert captured["negative_prompt"] is None
+        torch.testing.assert_close(
+            captured["prompt_embeds"],
+            torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_embeds"],
+            torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["prompt_attention_mask"],
+            torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_attention_mask"],
+            torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0),
+        )
+
+    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    @pytest.mark.parametrize(
+        ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
+        [(1.0, 1, 1.0), (None, 24, 24.0)],
+    )
+    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+        self,
+        monkeypatch,
+        cfg_rank,
+        expected_prompt_value,
+        frame_rate_input,
+        audio_sampling_rate,
+        expected_frame_rate,
+    ):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 1
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.audio_sampling_rate = audio_sampling_rate
+        pipe.audio_hop_length = 1
+        pipe.audio_vae_temporal_compression_ratio = 1
+        pipe.audio_vae_mel_compression_ratio = 1
+        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+        pipe.tokenizer = SimpleNamespace(padding_side="left")
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(mel_bins=2, latent_channels=1),
+        )
+
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
+                    return [video_pos, video_neg]
+                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
+                    return [audio_pos, audio_neg]
+                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
+            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
+            return torch.tensor([1.0, 0.5], device=device), 2
+
+        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+
+        class FakeScheduler:
+            def __init__(self, name="video", calls=None):
+                self.name = name
+                self.calls = [] if calls is None else calls
+                self.config = {
+                    "max_image_seq_len": 4096,
+                    "base_image_seq_len": 1024,
+                    "base_shift": 0.95,
+                    "max_shift": 2.05,
+                }
+                self.sigmas = torch.tensor([0.25, 0.25])
+
+            def __deepcopy__(self, memo):
+                return FakeScheduler("audio", self.calls)
+
+            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
+                return (latents - noise_pred,)
+
+        class FakeConnectors:
+            def to(self, device):
+                return self
+
+            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
+                assert padding_side == "left"
+                assert prompt_embeds.shape[0] == 2
+                return prompt_embeds, prompt_embeds, prompt_attention_mask
+
+        rope_video_fps: list[float] = []
+
+        class FakeRope:
+            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                rope_video_fps.append(fps)
+                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
+
+            def prepare_audio_coords(self, batch_size, num_frames, device):
+                return torch.zeros(batch_size, num_frames, 1, device=device)
+
+        class FakeTransformer:
+            def __init__(self):
+                self.config = SimpleNamespace(in_channels=2)
+                self.rope = FakeRope()
+                self.audio_rope = FakeRope()
+                self.calls = []
+
+            def __call__(self, **kwargs):
+                self.calls.append(kwargs)
+                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
+                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
+                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
+                assert kwargs["hidden_states"].shape == (1, 1, 2)
+                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
+                if cfg_rank == 0:
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+        class DummyProgress:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+            def update(self):
+                pass
+
+        pipe.scheduler = FakeScheduler()
+        pipe.connectors = FakeConnectors()
+        pipe.transformer = FakeTransformer()
+        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
+
+        def fake_encode_prompt(**kwargs):
+            return (
+                torch.ones(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+                torch.zeros(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+            )
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        video_latents = torch.tensor([[[1.0, -2.0]]])
+        audio_latents = torch.tensor([[[0.5, 3.0]]])
+        req = OmniDiffusionRequest(
+            prompt={"prompt": "prompt", "negative_prompt": "negative"},
+            sampling_params=OmniDiffusionSamplingParams(
+                height=32,
+                width=32,
+                num_frames=1,
+                frame_rate=frame_rate_input,
+                num_inference_steps=2,
+                guidance_scale=4.0,
+                latents=video_latents,
+                audio_latents=audio_latents,
+                output_type="latent",
+            ),
+            request_id="ltx23-cfg-parallel-forward-test",
+        )
+
+        output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
+
+        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
+        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert len(pipe.transformer.calls) == 2
+        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
+
+        video_out, audio_out = output.output
+        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+
+        # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
+        # to the model's own 24.0 default, not crash on None; a provided rate is passed through.
+        assert rope_video_fps
+        assert all(fps == expected_frame_rate for fps in rope_video_fps)

tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the LTX-2.3 image-to-video pipeline."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_ltx23_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    return pipe
+
+
+class TestLTX23ImageToVideoForwardStages:
+    def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
+        image = torch.zeros(3, 8, 8)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "make the image move",
+                        "negative_prompt": "jitter",
+                        "multi_modal_data": {"image": image},
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-i2v-forward-stage-delegation",
+                )
+            ]
+        )
+        seen = {}
+
+        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+            seen["req"] = req_arg
+            seen["request_inputs"] = request_inputs
+            seen["kwargs"] = kwargs
+            return ["i2v-delegated"]
+
+        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+
+        output = pipe.forward(req, noise_scale=0.5)
+
+        assert output == ["i2v-delegated"]
+        assert seen["req"] is req
+        assert seen["request_inputs"].prompt == ["make the image move"]
+        assert seen["request_inputs"].negative_prompt == ["jitter"]
+        assert seen["kwargs"]["image"] is image
+        assert seen["kwargs"]["noise_scale"] == 0.5
+
+    def test_denoise_timestep_kwargs_masks_video_only(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        ts = torch.tensor([2.0, 4.0])
+        denoise_ctx = SimpleNamespace(
+            conditioning_mask=torch.tensor([[1.0, 0.0]]),
+            conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        )
+
+        kwargs = pipe._denoise_timestep_kwargs(
+            ts,
+            SimpleNamespace(cfg_parallel_ready=False),
+            denoise_ctx,
+        )
+
+        torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
+        torch.testing.assert_close(kwargs["audio_timestep"], ts)
+        torch.testing.assert_close(kwargs["sigma"], ts)
+
+
+class TestLTX23ImageToVideoPipeline:
+    def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline)
+        assert LTX23ImageToVideoPipeline.support_image_input is True
+
+    def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = object()
+
+        assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
+        with pytest.raises(ValueError, match="exactly one image per prompt"):
+            LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
+
+    def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = torch.zeros(1, 3, 4, 4)
+        additional = {
+            "preprocessed_image": None,
+            "pixel_values": image,
+            "image": torch.ones_like(image),
+        }
+
+        assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
+
+    def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            raise AssertionError("packed I2V latents should not be noised")
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, latents)
+
+    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            return torch.ones(shape, device=device, dtype=dtype)
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
+
+    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        class FakeScheduler:
+            def step(self, noise_pred, t, latents, return_dict=False):
+                return (latents + noise_pred + t,)
+
+        pipe.scheduler = FakeScheduler()
+        latents = torch.tensor([[[1.0], [2.0], [3.0]]])
+        noise_pred = torch.full_like(latents, 10.0)
+
+        out = pipe._step_video_latents_i2v(
+            noise_pred,
+            latents,
+            torch.tensor(0.5),
+            latent_num_frames=3,
+            latent_height=1,
+            latent_width=1,
+        )
+
+        torch.testing.assert_close(out[:, :1], latents[:, :1])
+        torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
+
+
+class TestLTX23ImageToVideoModule:
+    """Test that pipeline_ltx2_3_image2video.py exposes I2V entry points."""
+
+    def test_i2v_classes_importable(self):
+        """I2V classes must be importable from the implementation module."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert LTX23ImageToVideoPipeline is not None
+
+    def test_post_process_func_importable(self):
+        """get_ltx2_post_process_func must be importable for registry lookup."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
+
+        assert callable(get_ltx2_post_process_func)
+
+    def test_i2v_class_matches_package_export(self):
+        """Package export must point at the I2V implementation module."""
+        from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
+            LTX23ImageToVideoPipeline as ModuleExported,
+        )
+
+        assert PackageExported is ModuleExported

tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestLTX23OutputRank:
+    def test_single_process_rank_is_output_rank(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False)
+
+        assert ltx23._is_output_rank() is True
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True
+
+    def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1)
+
+        assert ltx23._is_output_rank() is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True
+
+
+class TestLTX23VaeDistributedDecode:
+    """Test LTX-2.3 distributed VAE helpers without loading weights."""
+
+    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+
+        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
+        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+
+    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            DistributedOperator,
+            GridSpec,
+            TileTask,
+        )
+
+        z = torch.zeros(1, 1, 1, 1, 1)
+        tile_output_shapes = {
+            0: (1, 1, 1, 2, 2),
+            1: (1, 1, 1, 2, 1),
+            2: (1, 1, 1, 1, 2),
+            3: (1, 1, 1, 1, 1),
+        }
+        tasks = [
+            TileTask(0, (0, 0), z, workload=4),
+            TileTask(1, (0, 1), z, workload=2),
+            TileTask(2, (1, 0), z, workload=2),
+            TileTask(3, (1, 1), z, workload=1),
+        ]
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "max_tile_output_shape": (1, 1, 1, 2, 2),
+                "tile_output_shapes": tile_output_shapes,
+            },
+            output_dtype=torch.float32,
+        )
+        seen = {}
+
+        def exec_tile(task):
+            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
+
+        def merge_tiles(coord_tensor_map, passed_grid_spec):
+            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
+            assert passed_grid_spec is grid_spec
+            return torch.stack(
+                [
+                    coord_tensor_map[(0, 0)].flatten()[0],
+                    coord_tensor_map[(0, 1)].flatten()[0],
+                    coord_tensor_map[(1, 0)].flatten()[0],
+                    coord_tensor_map[(1, 1)].flatten()[0],
+                ]
+            )
+
+        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
+
+        rank0_executor = object.__new__(LTX2VaeExecutor)
+        rank0_executor.parallel_size = 2
+        rank0_executor.world_size = 2
+        rank0_executor.rank = 0
+
+        def gather_rank0(local_tile_tensor):
+            assigned = rank0_executor._balance_tasks(tasks, 2)
+            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
+            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
+                rank1_results,
+                list(local_tile_tensor.shape),
+                z.device,
+                torch.float32,
+            )
+            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
+            return [local_tile_tensor, rank1_tile_tensor]
+
+        def fail_final_sync(*_args, **_kwargs):
+            raise AssertionError("broadcast_result=False should not sync the final result")
+
+        rank0_executor.gather_tensors = gather_rank0
+        rank0_executor._sync_final_result = fail_final_sync
+
+        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
+
+        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
+        assert seen["merged_shapes"] == {
+            (0, 0): (1, 1, 1, 2, 2),
+            (0, 1): (1, 1, 1, 2, 1),
+            (1, 0): (1, 1, 1, 1, 2),
+            (1, 1): (1, 1, 1, 1, 1),
+        }
+
+        non_rank0_executor = object.__new__(LTX2VaeExecutor)
+        non_rank0_executor.parallel_size = 2
+        non_rank0_executor.world_size = 2
+        non_rank0_executor.rank = 1
+
+        def gather_rank1(local_tile_tensor):
+            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
+            return None
+
+        def fail_non_rank0_merge(*_args, **_kwargs):
+            raise AssertionError("non-rank0 should not merge gathered tiles")
+
+        non_rank0_executor.gather_tensors = gather_rank1
+        non_rank0_executor._sync_final_result = fail_final_sync
+
+        empty_result = non_rank0_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda _z: (tasks, grid_spec),
+                exec=exec_tile,
+                merge=fail_non_rank0_merge,
+            ),
+            broadcast_result=False,
+        )
+
+        assert tuple(empty_result.shape) == (0,)
+        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
+
+
+class TestLTX23VaeTiling:
+    """Test LTX-2.3 video VAE tile helpers without loading weights."""
+
+    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+
+        vae = SimpleNamespace(
+            spatial_compression_ratio=32,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=448,
+            tile_sample_stride_width=448,
+            temporal_compression_ratio=8,
+            dtype=torch.float32,
+        )
+
+        z = torch.zeros(1, 2, 5, 16, 24)
+        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
+
+        assert grid_spec.grid_shape == (2, 2)
+        assert grid_spec.split_dims == (3, 4)
+        assert grid_spec.tile_spec["sample_height"] == 512
+        assert grid_spec.tile_spec["sample_width"] == 768
+        assert grid_spec.tile_spec["blend_height"] == 64
+        assert grid_spec.tile_spec["blend_width"] == 64
+        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
+        assert grid_spec.tile_spec["tile_output_shapes"] == {
+            0: (1, 3, 33, 512, 512),
+            1: (1, 3, 33, 512, 320),
+            2: (1, 3, 33, 64, 512),
+            3: (1, 3, 33, 64, 320),
+        }
+        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+        assert [tuple(task.tensor.shape) for task in tasks] == [
+            (1, 2, 5, 16, 16),
+            (1, 2, 5, 16, 10),
+            (1, 2, 5, 2, 16),
+            (1, 2, 5, 2, 10),
+        ]
+        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
+
+    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
+
+        class FakeVae:
+            def __init__(self):
+                self.blend_calls = []
+
+            def clear_cache(self):
+                pass
+
+            def blend_v(self, _previous, current, blend_height):
+                self.blend_calls.append(("v", blend_height))
+                return current
+
+            def blend_h(self, _previous, current, blend_width):
+                self.blend_calls.append(("h", blend_width))
+                return current
+
+        fake_vae = FakeVae()
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "sample_height": 10,
+                "sample_width": 10,
+                "blend_height": 1,
+                "blend_width": 2,
+                "tile_sample_stride_height": 5,
+                "tile_sample_stride_width": 5,
+            },
+        )
+        tiles = {
+            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
+            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
+            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
+            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
+        }
+
+        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
+
+        assert merged.shape == (1, 3, 2, 10, 10)
+        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
+        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
+        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
+
+    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
+        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
+
+        z = torch.zeros(1, 2, 1, 16, 24)
+        expected = torch.ones(1, 3, 1, 512, 768)
+        seen = {}
+
+        class FakeExecutor:
+            def execute(self, tensor, operator, broadcast_result=True):
+                seen["tensor"] = tensor
+                seen["operator"] = operator
+                seen["broadcast_result"] = broadcast_result
+                return expected
+
+        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
+        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
+        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
+        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
+
+        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
+            vae,
+            z,
+            temb=torch.tensor(0.5),
+            return_dict=False,
+        )
+
+        assert len(output) == 1
+        assert output[0] is expected
+        assert seen["tensor"] is z
+        assert seen["broadcast_result"] is False
+        assert seen["operator"].split.__name__ == "tile_split"
+        assert seen["operator"].merge.__name__ == "tile_merge"

tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py

@@ -1,4 +1,3 @@
-import inspect
 from types import SimpleNamespace
 
 import pytest
@@ -193,6 +192,7 @@ def _make_request_batch_prompt_sampling(**overrides):
         "true_cfg_scale": None,
         "guidance_scale_provided": False,
         "guidance_scale": 1.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -344,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl
 
     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/qwen_image/test_qwen_image_postprocess.py

@@ -0,0 +1,50 @@
+# 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.qwen_image.pipeline_qwen_image as qwen_image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_postprocess_preserves_output_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
+) -> None:
+    model_path = tmp_path / "model"
+    vae_path = model_path / "vae"
+    vae_path.mkdir(parents=True)
+    (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8")
+
+    processed = object()
+    captured: dict[str, object] = {}
+
+    class FakeImageProcessor:
+        def __init__(self, *, vae_scale_factor: int) -> None:
+            captured["vae_scale_factor"] = vae_scale_factor
+
+        def postprocess(self, images: torch.Tensor) -> object:
+            captured["images"] = images
+            return processed
+
+    monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor)
+    postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path)))
+
+    image = torch.zeros(1, 3, 4, 4)
+    output = postprocess(
+        {
+            "payload": {"image": image},
+            "metadata": {"prompt_embeddings": {"present": True}},
+        }
+    )
+
+    assert captured["vae_scale_factor"] == 16
+    assert captured["images"] is image
+    assert output == {
+        "payload": {"image": processed},
+        "metadata": {"prompt_embeddings": {"present": True}},
+    }

tests/diffusion/models/sd3/test_sd3_pipeline.py

@@ -21,6 +21,7 @@ def _make_sd3_sampling(**overrides):
         "generator": None,
         "latents": None,
         "guidance_scale": 4.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -20,6 +20,12 @@ class _StubTransformer(nn.Module):
         return torch.float32
 
 
+class _StubTextEncoder(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)
@@ -41,12 +47,30 @@ def _noop_progress_bar(*args, **kwargs):
     yield _Bar()
 
 
+def _stub_encode_prompt(
+    prompt,
+    negative_prompt=None,
+    do_classifier_free_guidance=True,
+    num_videos_per_prompt=1,
+    max_sequence_length=512,
+    device=None,
+    dtype=None,
+):
+    del negative_prompt, do_classifier_free_guidance, device, dtype
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    n = batch_size * num_videos_per_prompt
+    hidden_size = 8
+    prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size)
+    return prompt_embeds, None
+
+
 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.text_encoder = _StubTextEncoder()
     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)
@@ -61,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline._num_timesteps = None
     pipeline._current_timestep = None
     pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = _stub_encode_prompt  # type: ignore[method-assign]
     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:
+def test_forward_delegates_denoising_to_diffuse() -> None:
     pipeline = _make_pipeline()
-
-    prompt_embeds = torch.randn(1, 8)
     captured: dict[str, object] = {}
 
     def _fake_diffuse(**kwargs):
@@ -86,8 +109,8 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             width=None,
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=False,
-            guidance_scale=None,
+            guidance_scale_provided=True,
+            guidance_scale=1.0,
             guidance_scale_2=None,
             boundary_ratio=None,
             generator=None,
@@ -95,14 +118,16 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             num_outputs_per_prompt=1,
             max_sequence_length=32,
             latents=None,
+            output_type="latent",
             extra_args={},
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
 
-    output = pipeline.forward(batch, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0)
+    output = pipeline.forward(batch)
 
     assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
     assert captured["guidance_low"] == 1.0
     assert captured["guidance_high"] == 1.0

tests/diffusion/offloader/test_module_collector.py

@@ -149,6 +149,40 @@ class MultiVaePipeline(nn.Module, SupportsComponentDiscovery):
         self.audio_vae = nn.Linear(10, 10)
 
 
+class SiblingDiTsPipeline(nn.Module, SupportsComponentDiscovery):
+    """Two independent DiTs (e.g. a MoE transformer pair)."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.transformer_2 = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class NestedDiTPipeline(nn.Module, SupportsComponentDiscovery):
+    """Cosmos3-like: a DiT plus one of its own submodules listed as a DiT.
+
+    The nested entry exists so offloading can treat the inner stack as an
+    independent ring; the outer DiT's _hsdp_shard_conditions still own all
+    nested blocks for sharding. Order is nested-first, matching Cosmos3.
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Module()
+        self.transformer.language_model = nn.Linear(10, 10)
+        self.transformer.gen_layers = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -238,3 +272,35 @@ class TestProtocolDiscovery:
         assert len(result.vaes) == 2
         assert result.vaes[0] is pipeline.vae
         assert result.vaes[1] is pipeline.audio_vae
+
+
+class TestOutermostDits:
+    """PipelineModules.outermost_dits() drops DiTs nested inside another DiT.
+
+    HSDP shards each returned DiT via its own _hsdp_shard_conditions, so a DiT
+    nested inside another discovered DiT must be excluded: it is already covered
+    by its ancestor's conditions and may not declare conditions of its own.
+    """
+
+    def test_keeps_independent_sibling_dits(self):
+        result = ModuleDiscovery.discover(SiblingDiTsPipeline())
+
+        names, modules = result.outermost_dits()
+
+        assert names == ["transformer", "transformer_2"]
+        assert [id(m) for m in modules] == [id(m) for m in result.dits]
+
+    def test_drops_nested_dit_keeping_only_ancestor(self):
+        pipeline = NestedDiTPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        # Discovery sees both, nested-first (matches Cosmos3's _dit_modules order).
+        assert result.dit_names == ["transformer.language_model", "transformer"]
+
+        names, modules = result.outermost_dits()
+
+        # Only the ancestor survives; sharding the nested condition-less stack
+        # again would double-wrap its blocks and raise in apply_hsdp_to_model.
+        assert names == ["transformer"]
+        assert len(modules) == 1
+        assert modules[0] is pipeline.transformer

tests/diffusion/offloader/test_sequential_backend.py

@@ -7,7 +7,8 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -41,6 +42,40 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
+def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.enable_args = None
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            self.enable_args = kwargs
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert backend.enabled is True
+    assert pipeline.enable_args == {
+        "device": torch.device("cpu"),
+        "pin_memory": False,
+        "use_hsdp": False,
+    }
+
+    backend.disable()
+
+    assert backend.enabled is False
+    assert pipeline.disable_called is True
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""

tests/diffusion/quantization/test_bitsandbytes_config.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BitsAndBytes quantization config."""
+
+import importlib.util
+import sys
+import types
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from torch.nn import Module, Parameter
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.")
+
+bitsandbytes_available = pytest.mark.skipif(
+    importlib.util.find_spec("bitsandbytes") is None,
+    reason="bitsandbytes package not installed",
+)
+
+
+def _ensure_bitsandbytes_importable(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Inject a stub bitsandbytes package when the optional dep is absent.
+
+    ``mocker.patch("bitsandbytes...")`` imports the target module; without a
+    stub, core-model CI without bitsandbytes raises ModuleNotFoundError.
+    """
+    if importlib.util.find_spec("bitsandbytes") is not None:
+        return
+    bnb = types.ModuleType("bitsandbytes")
+    bnb_functional = types.ModuleType("bitsandbytes.functional")
+    bnb.functional = bnb_functional
+    monkeypatch.setitem(sys.modules, "bitsandbytes", bnb)
+    monkeypatch.setitem(sys.modules, "bitsandbytes.functional", bnb_functional)
+
+
+def test_bitsandbytes_config_creation():
+    """Test that BitsAndBytes config can be created."""
+    config = build_quant_config("bitsandbytes")
+    assert config is not None
+    assert config.get_name() == "bitsandbytes"
+
+
+def test_bitsandbytes_config_with_custom_params():
+    """Test BitsAndBytes config with custom parameters."""
+    config = build_quant_config(
+        "bitsandbytes",
+        quant_type="fp4",
+        compress_statistics=False,
+        ignored_layers=["to_out"],
+    )
+    assert config is not None
+    assert config.quant_type == "fp4"
+    assert config.compress_statistics is False
+    assert "to_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list includes bitsandbytes."""
+    assert "bitsandbytes" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_quantization_integration():
+    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(model="test", quantization_config="bitsandbytes")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "bitsandbytes"
+
+    config2 = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "method": "bitsandbytes",
+            "quant_type": "nf4",
+            "compress_statistics": True,
+        },
+    )
+    assert config2.quantization_config is not None
+    assert config2.quantization_config.get_name() == "bitsandbytes"
+    assert config2.quantization_config.quant_type == "nf4"
+    assert config2.quantization_config.compress_statistics is True
+
+
+def test_quantization_dict_not_mutated():
+    """Test that passing a dict to quantization_config doesn't mutate it."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    original_dict = {"method": "bitsandbytes", "quant_type": "nf4"}
+    dict_copy = original_dict.copy()
+
+    OmniDiffusionConfig(model="test", quantization_config=original_dict)
+
+    assert original_dict == dict_copy
+
+
+def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Test get_quant_method routing for CUDA."""
+    from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+    config = build_quant_config("bitsandbytes")
+
+    def _fake_init(self, quant_config):
+        pass
+
+    layer = mocker.Mock(spec=LinearBase)
+    mocker.patch.object(BnBOnlineLinearMethod, "__init__", _fake_init)
+
+    prefix = "test_layer"
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, BnBOnlineLinearMethod)
+
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+class TestBnBOnlineLinearMethod:
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        config = mocker.Mock()
+        config.quant_type = "nf4"
+        config.compress_statistics = True
+        return config
+
+    @pytest.fixture
+    def mock_deps(self, mocker, monkeypatch: pytest.MonkeyPatch):
+        _ensure_bitsandbytes_importable(monkeypatch)
+        mock_qweight = torch.ones((64, 32), dtype=torch.uint8)
+        mock_quant_state = mocker.Mock()
+        mock_quant = mocker.patch(
+            "bitsandbytes.functional.quantize_4bit",
+            return_value=(mock_qweight, mock_quant_state),
+            create=True,
+        )
+        # x is (2, 16, 32) -> flattened (32, 32); matmul must return (32, 64)
+        # so apply can reshape back to (2, 16, 64).
+        mock_matmul = mocker.patch(
+            "bitsandbytes.matmul_4bit",
+            return_value=torch.randn(32, 64),
+            create=True,
+        )
+        return {
+            "quant": mock_quant,
+            "matmul": mock_matmul,
+            "mock_qweight": mock_qweight,
+            "mock_quant_state": mock_quant_state,
+        }
+
+    def test_process_weights_after_loading(self, mock_deps, mock_quant_config, mocker):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.randn(64, 32))
+        mocker.patch.object(torch.Tensor, "cuda", return_value=layer.weight.data)
+
+        method.process_weights_after_loading(layer)
+
+        mock_deps["quant"].assert_called_once()
+        assert torch.equal(layer.weight.data, mock_deps["mock_qweight"])
+        assert layer.quant_state is mock_deps["mock_quant_state"]
+        assert layer.bnb_shape == (64, 32)
+
+    def test_apply(self, mock_deps, mock_quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.ones((64, 32), dtype=torch.uint8), requires_grad=False)
+        layer.quant_state = mock_deps["mock_quant_state"]
+
+        x = torch.randn(2, 16, 32, dtype=torch.float16)
+        bias = torch.randn(64)
+        output = method.apply(layer, x, bias)
+
+        mock_deps["matmul"].assert_called_once()
+        assert output.shape == (2, 16, 64)
+        assert output.dtype == torch.float16
+
+
+@pytest.fixture
+def quant_config():
+    from vllm_omni.quantization.bitsandbytes_config import DiffusionBitsAndBytesConfig
+
+    return DiffusionBitsAndBytesConfig(
+        quant_type="nf4",
+        compress_statistics=True,
+    )
+
+
+@cuda_available
+@bitsandbytes_available
+class TestCudaBnBSmoke:
+    """Smoke tests using real bitsandbytes CUDA kernels."""
+
+    @pytest.fixture
+    def real_layer(self):
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        layer.logical_widths = [128]
+        layer.input_size_per_partition = 64
+        layer.output_size_per_partition = 128
+        layer.orig_dtype = torch.float16
+        return layer
+
+    def test_real_cuda_quantize_4bit_shape_contract(self, quant_config):
+        import bitsandbytes.functional as bnb_F
+
+        weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
+        qweight, quant_state = bnb_F.quantize_4bit(
+            weight,
+            quant_type=quant_config.quant_type,
+            compress_statistics=quant_config.compress_statistics,
+        )
+
+        assert qweight.dtype == torch.uint8
+        assert quant_state is not None
+
+    def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+        method.process_weights_after_loading(real_layer)
+
+        assert real_layer.weight.dtype == torch.uint8
+        assert hasattr(real_layer, "quant_state")
+        assert real_layer.bnb_shape == (128, 64)
+
+    def test_real_cuda_bnb_apply_forward(self, quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        method.process_weights_after_loading(layer)
+
+        x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda")
+        output = method.apply(layer, x)
+
+        assert output.shape == (2, 16, 128)
+        assert output.dtype == torch.float16

tests/diffusion/quantization/test_fp8_config.py

@@ -154,21 +154,6 @@ def test_component_config_with_default():
     assert config.resolve("vae.encoder") is None
 
 
-def test_gguf_config():
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
-
-    config = build_quant_config(
-        {
-            "method": "gguf",
-            "gguf_model": "path/to/model.gguf",
-        }
-    )
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "path/to/model.gguf"
-    assert config.get_name() == "gguf"
-
-
 def test_integration_string():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
 def test_supported_methods_includes_vllm():
     from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
 
-    for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
+    for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
         assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
 
 

tests/diffusion/quantization/test_gguf_config.py

@@ -1,128 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for diffusion GGUF quantization config and linear method."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-
-from vllm_omni.quantization import build_quant_config
-from vllm_omni.quantization.gguf_config import (
-    UNQUANTIZED_TYPES,
-    DiffusionGGUFConfig,
-    DiffusionGGUFLinearMethod,
-    dequant_gemm_gguf,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def test_gguf_config_creation_and_delegation():
-    config = DiffusionGGUFConfig(
-        gguf_model="weights.gguf",
-        unquantized_modules=["proj_out"],
-    )
-
-    assert config.gguf_model == "weights.gguf"
-    assert config.unquantized_modules == ["proj_out"]
-    assert config.get_name() == "gguf"
-
-
-def test_build_quant_config_builds_gguf_config():
-    config = build_quant_config(
-        "gguf",
-        gguf_model="repo/model-Q4_0.gguf",
-        unquantized_modules=["norm_out"],
-    )
-
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "repo/model-Q4_0.gguf"
-    assert config.unquantized_modules == ["norm_out"]
-
-
-def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
-
-    assert isinstance(method, DiffusionGGUFLinearMethod)
-
-
-def test_gguf_vllm_config_respects_unquantized_modules():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
-
-    assert isinstance(method, UnquantizedLinearMethod)
-
-
-def test_gguf_vllm_config_returns_none_for_non_linear_layers():
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
-
-    assert method is None
-
-
-def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
-    qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
-
-    out = dequant_gemm_gguf(x, qweight, qweight_type)
-
-    assert torch.allclose(out, x @ qweight.T)
-
-
-def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    layer = SimpleNamespace(
-        qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
-        qweight_type=SimpleNamespace(weight_type=qweight_type),
-    )
-    x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
-    bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
-
-    out = method.apply(layer, x, bias)
-
-    expected = x @ layer.qweight.T + bias
-    assert torch.allclose(out, expected)
-
-
-def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    qweight = torch.nn.Parameter(
-        torch.tensor(
-            [
-                [1.0, 0.0],
-                [0.0, 1.0],
-                [2.0, 0.0],
-                [0.0, 2.0],
-            ],
-            dtype=torch.float32,
-        ),
-        requires_grad=False,
-    )
-    qweight.shard_id = ["left", "right"]
-    qweight.shard_offset_map = {
-        "left": (0, 2, 2),
-        "right": (2, 4, 2),
-    }
-    layer = SimpleNamespace(
-        qweight=qweight,
-        qweight_type=SimpleNamespace(
-            shard_weight_type={
-                "left": qweight_type,
-                "right": qweight_type,
-            }
-        ),
-    )
-    x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
-
-    out = method.apply(layer, x)
-
-    expected_q = x @ qweight[:2].T
-    expected_k = x @ qweight[2:4].T
-    assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))

tests/diffusion/quantization/test_mxfp4_config.py

@@ -537,3 +537,118 @@ def test_single_scale_row_parallel_tp2_loader_simulation():
 
     for rank in range(_TP2):
         assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 (gfx950) — get_quant_method dispatch
+#
+# These run on CPU: the platform check, gcnArchName probe and aiter op
+# registration are all mocked (stdlib unittest.mock + monkeypatch, so no
+# pytest-mock dependency).  They cover only the dispatch + weight-allocation
+# logic added by the ROCm PR — the AITER GEMM / quant kernels require real
+# gfx950 hardware and are intentionally NOT exercised here.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def _rocm_platform(monkeypatch: pytest.MonkeyPatch):
+    """Make current_omni_platform report ROCm, and stub the aiter
+    custom-op registration so ROCmMxfp4*Method can be constructed without aiter."""
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization import mxfp4_config
+
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr(mxfp4_config, "_register_rocm_mxfp4_op", lambda: None)
+
+
+def _patch_gcn_arch(monkeypatch: pytest.MonkeyPatch, arch: str) -> None:
+    """Patch torch.cuda.get_device_properties(...).gcnArchName to return `arch`."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *a, **k: SimpleNamespace(gcnArchName=arch),
+    )
+
+
+def _fake_linear_layer():
+    """A stand-in that passes isinstance(layer, LinearBase) without a real layer."""
+    from unittest.mock import MagicMock
+
+    from vllm.model_executor.layers.linear import LinearBase
+
+    return MagicMock(spec=LinearBase)
+
+
+def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
+    """ROCm + gfx950 + online checkpoint must return ROCmMxfp4OnlineLinearMethod."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    method = cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+    assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
+
+
+def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
+    """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
+
+    assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
+
+
+def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
+    """MXFP4 on ROCm requires gfx950; any other arch must raise."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx942:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    with pytest.raises(NotImplementedError, match="gfx950"):
+        cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 — create_weights (online lazy meta-device placeholder)
+#
+# The base ROCmMxfp4LinearMethod is abstract (no create_weights); the online
+# subclass gets create_weights from _LazyWeightMixin, which registers a BF16
+# weight on the meta device to be materialised at load time.
+# ---------------------------------------------------------------------------
+
+
+def test_rocm_create_weights_column_parallel_tp2(_rocm_platform):
+    """Column-parallel TP=2: meta BF16 weight has the output halved, input full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
+    assert layer.weight.dtype == torch.bfloat16
+    assert layer.weight.device.type == "meta"
+    assert layer.weight.input_dim == 1
+    assert layer.weight.output_dim == 0
+    assert layer.logical_widths == [_TP2_N // _TP2]
+    assert layer.weight_block_size is None
+
+
+def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
+    """Row-parallel TP=2: meta BF16 weight has the input halved, output full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.input_size_per_partition == _TP2_K // _TP2
+    assert layer.output_size_per_partition == _TP2_N

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
         candidate_quantization_config_json=None,
         candidate_ignored_layers=None,
         ignored_layers=None,
-        candidate_gguf_model=None,
         candidate_load_format="default",
     )
 

tests/diffusion/test_diffusion_model_runner.py

@@ -509,6 +509,25 @@ def test_split_diffusion_output_by_request_slices_tuple_outputs():
     assert torch.equal(outputs[1].output[1], torch.tensor([20]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
+    """execute_model runs the pipeline forward after the KV receive step."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert runner.pipeline.forward_calls == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):

tests/diffusion/test_diffusion_output_formatter.py

@@ -51,21 +51,23 @@ def _timings() -> DiffusionStepTimings:
     )
 
 
-def test_formatter_preserves_single_video_audio_actions_and_custom_output(
+def test_formatter_preserves_single_video_audio_actions_and_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
-    custom_output = {"from_output": "kept"}
     postprocess_output = normalize_diffusion_postprocess_output(
         {
-            "video": ["frame-0"],
-            "audio": "audio-0",
-            "actions": "action-0",
-            "custom_output": {"from_postprocess": "merged"},
-            "audio_sample_rate": 48000,
-            "fps": 24.0,
-        },
-        custom_output,
+            "payload": {
+                "video": ["frame-0"],
+                "audio": "audio-0",
+                "actions": "action-0",
+            },
+            "metadata": {
+                "audio": {"sample_rate": 48000},
+                "video": {"fps": 24.0},
+                "actions": {"action_mode": "policy"},
+            },
+        }
     )
 
     results = format_diffusion_outputs(
@@ -73,7 +75,6 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
         od_config=_config(),
         diffusion_output=DiffusionOutput(
             output=None,
-            custom_output=custom_output,
             stage_durations={"execute": 1.25},
             peak_memory_mb=321.0,
         ),
@@ -87,12 +88,12 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     assert result.images == ["frame-0"]
     assert result.prompt == "prompt-0"
     assert result.final_output_type == "image"
-    assert result.custom_output == {
-        "from_output": "kept",
-        "from_postprocess": "merged",
-    }
-    assert custom_output == {"from_output": "kept"}
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 48000},
+            "video": {"fps": 24.0},
+            "actions": {"action_mode": "policy"},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 48000,
         "fps": 24.0,
@@ -110,18 +111,123 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     }
 
 
-def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_formatter_normalizes_payload_metadata_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
     postprocess_output = normalize_diffusion_postprocess_output(
-        "caption",
-        {"text_output": "caption"},
+        {
+            "payload": {
+                "video": ["frame-0"],
+                "actions": "action-0",
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": object(),
+                },
+            },
+        }
+    )
+
+    assert postprocess_output.outputs == {
+        "video": ["frame-0"],
+        "actions": "action-0",
+    }
+    assert postprocess_output.primary_key == "video"
+    assert postprocess_output.metadata == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
+    assert "internal" not in postprocess_output.metadata
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.multimodal_output == {
+        "metadata": postprocess_output.metadata,
+        "actions": "action-0",
+    }
+
+
+def test_formatter_maps_trajectory_payload_to_request_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    latents = torch.zeros(2, 1, 4)
+    timesteps = torch.tensor([1.0, 0.5])
+    log_probs = torch.tensor([-0.1, -0.2])
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {
+                "image": "image-0",
+                "trajectory": {
+                    "latents": latents,
+                    "timesteps": timesteps,
+                    "log_probs": log_probs,
+                },
+            },
+            "metadata": {"trajectory": {"type": "denoising"}},
+        }
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == ["image-0"]
+    assert result.latents is latents
+    assert result.trajectory_latents is latents
+    assert result.trajectory_timesteps is timesteps
+    assert result.trajectory_log_probs is log_probs
+    assert result.multimodal_output == {
+        "metadata": {"trajectory": {"type": "denoising"}},
+        "trajectory": {
+            "latents": latents,
+            "timesteps": timesteps,
+            "log_probs": log_probs,
+        },
+    }
+
+
+def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {"text": "caption"},
+            "metadata": {"text": {"text_output": "caption"}},
+        }
     )
 
     [result] = format_diffusion_outputs(
         request=_request("describe this"),
         od_config=_config(),
-        diffusion_output=DiffusionOutput(output="caption"),
-        output_data="caption",
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"payload": {"text": "caption"}},
         postprocess_output=postprocess_output,
         timings=_timings(),
     )
@@ -129,7 +235,10 @@ def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch)
     assert result.images == []
     assert result.prompt == "describe this"
     assert result.final_output_type == "text"
-    assert result.custom_output == {"text_output": "caption"}
+    assert result.multimodal_output == {
+        "metadata": {"text": {"text_output": "caption"}},
+        "text": "caption",
+    }
 
 
 def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
@@ -144,7 +253,7 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
         "_try_load_model_cls",
         lambda _: AudioModel,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak"),
@@ -175,8 +284,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
             "actions": "action-0",
             "audio_sample_rate": 16000,
             "fps": 30.0,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -192,6 +300,10 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
     assert result.prompt == "watch and listen"
     assert result.final_output_type == "image"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 16000},
+            "video": {"fps": 30.0},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 16000,
         "fps": 30.0,
@@ -207,8 +319,7 @@ def test_formatter_preserves_audio_only_postprocess_dict(
         {
             "audio": "waveform",
             "audio_sample_rate": 24000,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -224,6 +335,9 @@ def test_formatter_preserves_audio_only_postprocess_dict(
     assert result.prompt == "speak"
     assert result.final_output_type == "audio"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 24000},
+        },
         "audio": "waveform",
         "audio_sample_rate": 24000,
     }
@@ -238,7 +352,7 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs(
         "_try_load_model_cls",
         lambda _: None,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak", num_outputs_per_prompt=2),
@@ -264,10 +378,8 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
             "video": ["frame-0", "frame-1"],
             "audio": ["audio-0", "audio-1"],
             "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
-            "custom_output": {"shared": True},
             "fps": 12.5,
-        },
-        {},
+        }
     )
 
     results = format_diffusion_outputs(
@@ -282,7 +394,7 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
     assert len(results) == 1
     assert results[0].images == ["frame-0", "frame-1"]
     assert results[0].prompt == "prompt-0"
-    assert results[0].custom_output == {"shared": True}
+    assert results[0].multimodal_output["metadata"] == {"video": {"fps": 12.5}}
     assert results[0].multimodal_output["audio"] == ["audio-0", "audio-1"]
     assert results[0].multimodal_output["fps"] == 12.5
     torch.testing.assert_close(

tests/diffusion/test_diffusion_output_metadata.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.outputs.output_metadata import (
+    strip_internal_metadata,
+    validate_diffusion_metadata,
+    validate_public_diffusion_metadata,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_validate_accepts_canonical_metadata_fields() -> None:
+    validate_diffusion_metadata(
+        {
+            "video": {
+                "fps": 24.0,
+                "video_fps_multiplier": 2,
+                "shape": [1, 3, 16, 64, 64],
+            },
+            "audio": {
+                "sample_rate": 48000,
+                "audiox_task": "tts",
+            },
+            "actions": {
+                "raw_action_dim": 7,
+                "action_mode": "policy",
+                "domain_id": 3,
+            },
+            "transfer": {
+                "controls": {"edge": torch.zeros(1, 3, 4, 8, 8)},
+                "hints": ["edge"],
+            },
+            "text": {
+                "text_output": "caption",
+                "think_text": "reasoning",
+            },
+            "common": {
+                "action_only_output": True,
+            },
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    ("metadata", "error_type"),
+    [
+        ({"video": {"fps": 0}}, ValueError),
+        ({"video": {"fps": "24"}}, TypeError),
+        ({"video": {"video_fps_multiplier": 0}}, ValueError),
+        ({"video": {"shape": [1, 0, 64]}}, ValueError),
+        ({"audio": {"sample_rate": 0}}, ValueError),
+        ({"audio": {"audiox_task": ""}}, ValueError),
+        ({"actions": {"raw_action_dim": 0}}, ValueError),
+        ({"actions": {"action_mode": ""}}, ValueError),
+        ({"actions": {"domain_id": "7"}}, TypeError),
+        ({"transfer": {"controls": {"edge": "not-a-tensor"}}}, TypeError),
+        ({"transfer": {"hints": ["edge", 1]}}, TypeError),
+        ({"text": {"text_output": 1}}, TypeError),
+        ({"text": {"think_text": 1}}, TypeError),
+        ({"common": {"action_only_output": "true"}}, TypeError),
+    ],
+)
+def test_validate_rejects_invalid_canonical_metadata(metadata: dict[str, object], error_type: type[Exception]) -> None:
+    with pytest.raises(error_type):
+        validate_diffusion_metadata(metadata)
+
+
+def test_validate_allows_unknown_model_metadata() -> None:
+    validate_diffusion_metadata(
+        {
+            "model_specific": {"opaque": object()},
+            "video": {"model_key": object()},
+        }
+    )
+
+
+def test_public_metadata_rejects_internal_group() -> None:
+    with pytest.raises(ValueError, match="internal"):
+        validate_public_diffusion_metadata({"internal": {"secret": object()}})
+
+
+def test_strip_internal_metadata_removes_private_group() -> None:
+    metadata = {
+        "actions": {"action_mode": "policy"},
+        "internal": {"robolab_action_postprocess": object()},
+    }
+
+    public_metadata = strip_internal_metadata(metadata)
+
+    assert public_metadata == {"actions": {"action_mode": "policy"}}
+    assert "internal" not in public_metadata

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -16,7 +16,6 @@ from unittest.mock import Mock, patch
 import pytest
 
 from vllm_omni.diffusion.registry import (
-    _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
     _DIFFUSION_IR_OP_PRIORITY_FUNCS,
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
@@ -62,7 +61,6 @@ class TestRegisterDiffusionModel:
         original_models = _DIFFUSION_MODELS.copy()
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
-        original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
         original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
@@ -71,8 +69,6 @@ class TestRegisterDiffusionModel:
         _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
         _DIFFUSION_POST_PROCESS_FUNCS.clear()
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
 
@@ -84,7 +80,6 @@ class TestRegisterDiffusionModel:
             class_name="TestPipeline",
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
-            action_post_process_func_name="test_action_post_process",
             ir_op_priority_func_name="test_ir_op_priority",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
@@ -95,9 +90,21 @@ class TestRegisterDiffusionModel:
         )
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
-        assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
         assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
 
+    def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
+        """Deprecated action hook keyword is accepted but not registered."""
+        register_diffusion_model(
+            model_arch="LegacyActionPipeline",
+            module_name="test_plugin.diffusion.pipeline",
+            class_name="LegacyActionPipeline",
+            post_process_func_name="test_post_process",
+            action_post_process_func_name="test_action_post_process",
+        )
+
+        assert "LegacyActionPipeline" in _DIFFUSION_MODELS
+        assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"
+
 
 class TestWorkerUsesHook:
     """Test that DiffusionWorker resolves model runner via platform hook."""

tests/diffusion/test_diffusion_step_pipeline.py

@@ -30,6 +30,9 @@ from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
@@ -97,6 +100,51 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _ProfilingStepPipeline(_StepPipeline):
+    enable_diffusion_pipeline_profiler = True
+
+    def __init__(self):
+        super().__init__()
+        self._stage_durations: dict[str, float] = {}
+
+    @property
+    def stage_durations(self) -> dict[str, float]:
+        return dict(self._stage_durations)
+
+    def clear_profiler_records(self) -> None:
+        self._stage_durations.clear()
+
+    def prepare_encode(self, state, **kwargs):
+        result = super().prepare_encode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0
+        return result
+
+    def denoise_step(self, input_batch, **kwargs):
+        result = super().denoise_step(input_batch, **kwargs)
+        self._stage_durations["QwenImagePipeline.diffuse"] = 2.0
+        return result
+
+    def post_decode(self, state, **kwargs):
+        result = super().post_decode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0
+        return result
+
+
+class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin):
+    _PROFILER_TARGETS: list[str] = []
+
+    def __init__(self):
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=True,
+        )
+
+    def forward(self):
+        return None
+
+    def denoise_step(self):
+        return "ok"
+
+
 class _InterruptingStepPipeline(_StepPipeline):
     interrupt = True
 
@@ -130,6 +178,9 @@ class _FakePeakMemoryPlatform:
         index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
         return int((self._reserved_mb[index] - 100) * 1024**2)
 
+    def is_available(self) -> bool:
+        return True
+
 
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
@@ -437,6 +488,16 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
     torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
 
 
+@pytest.mark.cpu
+def test_step_profiler_reports_denoise_step_as_diffuse():
+    pipeline = _AutoDenoiseProfilerPipeline()
+
+    assert pipeline.denoise_step() == "ok"
+
+    assert any(key.endswith(".diffuse") for key in pipeline.stage_durations)
+    assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations)
+
+
 @pytest.mark.cpu
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
@@ -469,6 +530,53 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _ProfilingStepPipeline()
+        req = _make_step_request()
+        reset_calls = []
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "is_available",
+            lambda: True,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "reset_peak_memory_stats",
+            lambda: reset_calls.append(True),
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_reserved",
+            lambda: 2 * 1024**2,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_allocated",
+            lambda: 1024**2,
+        )
+
+        DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(req, step_id=0),
+        )
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_cached_scheduler_output(step_id=1),
+        )
+
+        output = result.get_request_output("req-1")
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.peak_memory_mb == 2
+        assert output.result.stage_durations == {
+            "QwenImagePipeline.text_encoder.forward": 1.0,
+            "QwenImagePipeline.diffuse": 4.0,
+            "QwenImagePipeline.vae.decode": 3.0,
+        }
+        assert reset_calls == [True]
+
     def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch):
         runner = _make_runner()
         req = _make_step_request()

tests/diffusion/test_diffusion_streaming_output.py

@@ -85,19 +85,20 @@ class _PipelineBackedEngine:
             outputs = self.pipeline.step_outputs(request)
         except Exception as exc:
             outputs = [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
+                _streaming_diffusion_output(chunk=0, finished=False),
                 DiffusionOutput(error=str(exc), finished=True),
             ]
         for output in outputs:
             if output.error is not None:
                 yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
                 continue
-            custom_output = output.custom_output or {}
+            payload, metadata = _streaming_payload_and_metadata(output)
+            stream_metadata = metadata.get("stream", {})
             yield [
                 OmniRequestOutput.from_diffusion(
                     request_id=request.request_id,
-                    images=list(custom_output.get("images") or []),
-                    custom_output=custom_output,
+                    images=list(payload.get("image") or []),
+                    multimodal_output={"metadata": {"stream": stream_metadata}},
                     finished=output.finished,
                 )
             ]
@@ -106,6 +107,30 @@ class _PipelineBackedEngine:
         del request_id
 
 
+def _streaming_diffusion_output(
+    *,
+    chunk: int,
+    finished: bool,
+    images: list | None = None,
+) -> DiffusionOutput:
+    payload = {"image": images or []}
+    metadata = {"stream": {"chunk": chunk}}
+    return DiffusionOutput(
+        output={
+            "payload": payload,
+            "metadata": metadata,
+        },
+        finished=finished,
+    )
+
+
+def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
+    envelope = output.output if isinstance(output.output, dict) else {}
+    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
+    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
+    return payload, metadata
+
+
 class TestPipelineStreamingOutputToStageDiffusionClient:
     """Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
 
@@ -114,8 +139,8 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
 
@@ -123,7 +148,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
 
         assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
         assert [output.finished for output in outputs] == [False, True]
-        assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+        assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
         assert pipeline.requests[0].request_id == "req-stream"
 
     @pytest.mark.asyncio
@@ -136,7 +161,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         assert len(outputs) == 2
         assert outputs[0].request_id == "req-error"
         assert outputs[0].finished is False
-        assert outputs[0].custom_output["chunk"] == 0
+        assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0
         assert outputs[1].request_id == "req-error"
         assert outputs[1].finished is True
         assert outputs[1].error is not None
@@ -227,8 +252,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)
@@ -247,7 +272,7 @@ class TestPipelineStreamingOutputToEntrypoint:
 
             await _wait_for(lambda: len(pipeline.requests) == 1)
             assert pipeline.requests[0].request_id.startswith("req-omni-")
-            assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+            assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
             assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
         finally:
             await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@@ -258,8 +283,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True),
+                _streaming_diffusion_output(chunk=0, images=[frames], finished=False),
+                _streaming_diffusion_output(chunk=1, images=[frames], finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -937,6 +937,76 @@ class TestDistributedReceive:
 # ── TP auto-detect ───────────────────────────────────────────────────
 
 
+class TestBuildCfgRankLocalPayloads:
+    """Guard _build_cfg_rank_local_payloads padding logic for the
+    cfg_size==2 / no branch_roles scenario (hunyuan-image shared-prefix reuse)."""
+
+    def test_cfg_size2_no_branch_roles_copies_main_payload_to_follower(self):
+        """When AR doesn't split KV by branch, cfg follower (rank 1) should
+        receive a copy of the positive KV, not None."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # Rank 0: main payload (always present)
+        assert payloads[0] is not None
+        assert "past_key_values" in payloads[0]
+        # Rank 1: must also be non-None — same positive KV for shared-prefix reuse
+        assert payloads[1] is not None
+        assert payloads[1]["past_key_values"] is req.past_key_values
+        assert payloads[1]["kv_metadata"] == {"seq_len": 10}
+
+    def test_cfg_size3_no_branch_roles_keeps_none_padding(self):
+        """When cfg_size > 2 and no branch_roles, padding stays None (safety guard)."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=3)
+
+        assert len(payloads) == 3
+        assert payloads[0] is not None
+        assert payloads[1] is None
+        assert payloads[2] is None
+
+    def test_empty_main_payload_keeps_none_padding(self):
+        """When main_payload is empty, padding stays None even with cfg_size==2."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # No past_key_values → main_payload is empty dict → falsy → padding stays None
+        assert payloads[1] is None
+
+    def test_with_branch_roles_uses_original_behavior(self):
+        """When branch_roles exist, follower payload comes from branch KV, not padding."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(
+                cfg_text_past_key_values=SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+                cfg_text_kv_metadata={"source": "cfg_text"},
+            ),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        assert payloads[0] is not None
+        assert payloads[1] is not None
+        assert payloads[1]["sp.cfg_active_branch"] == "cfg_text"
+
+
 class TestAutoDetectTP:
     def test_auto_detect_when_config_defaults(self):
         """When config from_tp/to_tp == 1 (default), manager should auto-detect."""

tests/docs/test_generate_examples.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+ROOT_DIR = Path(__file__).parents[2]
+HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py"
+SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH)
+assert SPEC and SPEC.loader
+generate_examples = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(generate_examples)
+
+
+def test_model_example_titles_match_shared_display_names():
+    for slug, display_name in generate_examples.MODEL_DISPLAY_NAMES.items():
+        for category, mode_title in generate_examples.SERVING_MODE_TITLES.items():
+            example_path = ROOT_DIR / "examples" / category / slug
+            example = generate_examples.Example(example_path, category)
+
+            assert example.title == f"{display_name}: {mode_title}"
+            assert example.nav_title == display_name
+
+
+def test_model_example_title_mismatch_is_rejected(tmp_path):
+    example_path = tmp_path / "mimo_audio"
+    example_path.mkdir()
+    (example_path / "README.md").write_text("# A different title\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="Model example title mismatch"):
+        generate_examples.Example(example_path, "offline_inference")
+
+
+def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
+    display_names_path = tmp_path / "model_display_names.yml"
+    display_names_path.write_text("- not a mapping\n", encoding="utf-8")
+    monkeypatch.setattr(generate_examples, "MODEL_DISPLAY_NAMES_FILE", display_names_path)
+
+    with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
+        generate_examples.load_model_display_names()

tests/dreamzero/test_pipeline_state.py

@@ -4,7 +4,6 @@
 from collections import OrderedDict
 
 import pytest
-import torch
 
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
@@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
     assert "session-b" not in pipeline._states
 
 
-def test_dreamzero_state_cache_access_requires_initialization() -> None:
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
     state = DreamZeroState()
 
-    with pytest.raises(RuntimeError, match="KV caches not initialized"):
-        state.get_kv_caches()
-
-    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
-        state.get_crossattn_caches()
-
-    with pytest.raises(RuntimeError, match="create_kv_caches first"):
-        state.update_kv_cache(0, torch.empty(0))
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)

tests/e2e/accuracy/test_cosmos3_similarity.py

@@ -42,8 +42,6 @@ def _server_args() -> list[str]:
     return [
         "--num-gpus",
         "1",
-        "--model-class-name",
-        "Cosmos3OmniDiffusersPipeline",
         "--stage-init-timeout",
         "900",
         "--init-timeout",

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -331,7 +331,10 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
         if ro and getattr(ro, "outputs", None):
             cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
         if not cot_text:
-            ar_text = getattr(out, "custom_output", {}).get("ar_generated_text")
+            multimodal_output = getattr(out, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            ar_text = text_metadata.get("ar_generated_text") if isinstance(text_metadata, dict) else None
             if isinstance(ar_text, list):
                 cot_text = "\n".join(text for text in ar_text if text)
             else:

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -48,6 +48,7 @@ _DEPLOY_CONFIG = {
             "trust_remote_code": True,
             "devices": "0,1,2,3",  # set dynamically by _write_deploy_config
             "vae_use_slicing": False,
+            "moe_backend": "flashinfer_cutlass",
             "vae_use_tiling": False,
             "parallel_config": {
                 "pipeline_parallel_size": 1,

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -288,7 +288,7 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
         else:
             # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
             # Return a minimal dummy output to avoid crashing.
-            return DiffusionOutput(output=None, custom_output={})
+            return DiffusionOutput(output=None)
 
         if isinstance(negative_prompt_ids, list):
             negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
@@ -399,14 +399,20 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
             image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
 
         return DiffusionOutput(
-            output=_maybe_to_cpu(image),
+            output={
+                "payload": {
+                    "image": _maybe_to_cpu(image),
+                },
+                "metadata": {
+                    "prompt_embeddings": {
+                        "prompt_embeds": _maybe_to_cpu(prompt_embeds),
+                        "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
+                        "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
+                        "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
+                    },
+                },
+            },
             trajectory_latents=_maybe_to_cpu(all_latents),
             trajectory_log_probs=_maybe_to_cpu(all_log_probs),
             trajectory_timesteps=_maybe_to_cpu(all_timesteps),
-            custom_output={
-                "prompt_embeds": _maybe_to_cpu(prompt_embeds),
-                "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
-                "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
-                "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
-            },
         )

tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py

@@ -390,20 +390,23 @@ class vLLMOmniHttpServerLocal:
 
         diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0
 
-        mm_output = final_res.custom_output or {}
+        metadata = (
+            final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {}
+        )
+        prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {}
 
         if sampling_params.get("logprobs", False):
-            all_log_probs = mm_output.get("all_log_probs")
+            all_log_probs = final_res.trajectory_log_probs
             log_probs = all_log_probs[0] if all_log_probs is not None else None
         else:
             log_probs = None
 
-        all_latents = mm_output.get("all_latents")
-        all_timesteps = mm_output.get("all_timesteps")
-        prompt_embeds = mm_output.get("prompt_embeds")
-        prompt_embeds_mask = mm_output.get("prompt_embeds_mask")
-        negative_prompt_embeds = mm_output.get("negative_prompt_embeds")
-        negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask")
+        all_latents = final_res.trajectory_latents
+        all_timesteps = final_res.trajectory_timesteps
+        prompt_embeds = prompt_embeddings.get("prompt_embeds")
+        prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask")
+        negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds")
+        negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask")
 
         extra_fields = {
             "all_latents": all_latents[0] if all_latents is not None else None,

tests/e2e/offline_inference/test_flux2_klein.py

@@ -18,6 +18,8 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
+# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
+# This model will also use the tiny model builder when run at the core model level.
 MODEL = "black-forest-labs/FLUX.2-klein-4B"
 
 _OMNI_RUNNER_PARAM = (MODEL, None)
@@ -70,41 +72,6 @@ def _send_inpaint_with_generator(
     )
 
 
-# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "a cup of coffee on the table",
-            "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
-        }
-    )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
-    input_image, mask_image = _create_test_inputs()
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "Fill in the masked area with a beautiful garden",
-            "multi_modal_data": {"image": input_image, "mask_image": mask_image},
-            "sampling_params": 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,
-            ),
-        }
-    )
-
-
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"}, num_cards=1)

tests/e2e/offline_inference/test_krea2_expansion.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``).
+
+Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default (override with the
+``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the Raw checkpoint or a local diffusers
+directory). They cover a basic functional smoke plus the layerwise-CPU-offload path that the pipeline declares
+via ``SupportsComponentDiscovery`` / ``_layerwise_offload_blocks_attrs``.
+"""
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunnerHandler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
+# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
+LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
+PROMPT = "a fox in the snow, photorealistic"
+# darkbrush is a trigger-word LoRA: the trained monochrome ink-wash style only manifests when
+# its trigger phrase is present, so the LoRA test prompt appends it for a strong, stable signal.
+LORA_TRIGGER = "monochrome ink wash style"
+LORA_PROMPT = f"a fox in the snow, {LORA_TRIGGER}"
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.full_model,
+]
+
+
+def _sampling() -> OmniDiffusionSamplingParams:
+    # Small resolution + few steps to keep the L4 case light. guidance_scale resolution is checkpoint-aware inside
+    # the pipeline (distilled -> no-CFG, Raw -> CFG), so this stays agnostic to which checkpoint KREA2_MODEL points at.
+    return OmniDiffusionSamplingParams(
+        height=512,
+        width=512,
+        num_inference_steps=8,
+        guidance_scale=0.0,
+        seed=42,
+    )
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Basic functional smoke: a single prompt produces a decoded image."""
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize(
+    "omni_runner",
+    [(MODEL, None, {"enable_layerwise_offload": True})],
+    indirect=True,
+)
+def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs)."""
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.cache
+@pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True)
+def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``).
+
+    Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware
+    (False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint
+    exercises the no-CFG branch.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize(
+    "omni_runner",
+    # Standalone HSDP: use_hsdp=True with an explicit hsdp_shard_size shards the DiT weights
+    # across 2 ranks (replicate_size defaults to 1 -> world_size = 1 * 2 = 2). HSDP must keep
+    # tensor_parallel_size/data_parallel_size at 1 (see vllm_omni/diffusion/data.py __post_init__).
+    # Also validated locally at 8 GPUs (hsdp_shard_size=8); CI shard provisions 2.
+    [(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})],
+    indirect=True,
+)
+def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs.
+
+    Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda)
+    shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize(
+    "omni_runner",
+    # VAE patch-parallel decode: partition the VAE decode across the TP world. Mirrors the online
+    # "tp_vae_patch" case (tensor_parallel_size=2 + vae_patch_parallel_size=2) from the wan2.2/z-image
+    # expansion suites, so vae_patch_parallel_size must equal the world size established by TP.
+    # Also validated locally at 8 GPUs (tp=8, vae_patch=8); CI shard provisions 2.
+    [(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})],
+    indirect=True,
+)
+def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise VAE patch-parallel (decode) across 2 GPUs.
+
+    Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the
+    multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with
+    fewer than 2 CUDA devices.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray:
+    sp = _sampling()
+    sp.lora_request = lora_request
+    sp.lora_scale = lora_scale
+    resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp})
+    return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16)
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation.
+
+    Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``,
+    264 modules matched via the ReplicatedLinear projections) passed per-request through
+    ``OmniDiffusionSamplingParams.lora_request`` / ``lora_scale`` — the same fields the
+    diffusion LoRA manager reads to activate/deactivate the adapter.
+    """
+    lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA)
+
+    baseline = _generate(omni_runner_handler, None)
+    img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0)
+    img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0)
+    restored = _generate(omni_runner_handler, None)
+
+    diff_1x = np.abs(baseline - img_1x).mean()
+    diff_2x = np.abs(baseline - img_2x).mean()
+    diff_restored = np.abs(baseline - restored).mean()
+
+    # (a) Adapter has a visible effect at both scales.
+    assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
+    assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
+    # (b) Scale changes the output (stronger adapter -> larger perturbation).
+    assert not np.isclose(diff_1x, diff_2x, atol=1.0), f"LoRA scale had no effect: 1x={diff_1x:.2f}, 2x={diff_2x:.2f}"
+    # (c) Passing lora_request=None cleanly deactivates: same seed -> back to baseline (modulo fp drift).
+    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -22,6 +22,7 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
+        "vae_use_tiling": True,
     },
 )
 
@@ -32,6 +33,7 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
+        "vae_use_tiling": True,
     },
 )
 
@@ -45,10 +47,6 @@ pytestmark = [
     ),
 ]
 
-_SKIP_ISSUE_3236 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
-)
-
 
 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -74,7 +72,6 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     container.close()
 
 
-@_SKIP_ISSUE_3236
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
     """End-to-end test for MagiHuman generating video and audio."""

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -5,7 +5,7 @@ Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
 match a golden reference.
 
 Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
-Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml``
+Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
 
 Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
   Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
@@ -36,9 +36,9 @@ _VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
 _AR_PATCH_SIZE = 16
 
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
-T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
 
-_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG)
+_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
 
 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"

tests/e2e/offline_inference/test_omnivoice_expansion.py

@@ -18,13 +18,15 @@ from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "k2-fsa/OmniVoice"
-STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")
+DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
 
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
+# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
+# The migrated test passes deploy_config through extra Omni kwargs.
 _OMNI_RUNNER_PARAM = (
     MODEL,
-    STAGE_CONFIG,
+    None,
     {
+        "deploy_config": DEPLOY_CONFIG,
         "trust_remote_code": True,
         "log_stats": True,
     },

tests/e2e/offline_inference/test_qwen3_tts_customvoice.py

@@ -34,12 +34,10 @@ def get_cuda_graph_config():
             "stages": {
                 0: {
                     "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,
                 },

tests/e2e/online_serving/test_flux2_klein.py

@@ -1,62 +0,0 @@
-"""
-Tests for Flux2 Klein; currently Dev is implemented separately,
-but ideally these models will fold together in the future.
-"""
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-
-FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
-POSITIVE_PROMPT = "A cat sitting on a windowsill"
-NEGATIVE_PROMPT = "blurry, low quality"
-
-
-# For now, we only test a good configuration for this model to keep the CI
-# light, but in the near future may explore using random models to check
-# cross-feature compatibility more generally.
-def _get_diffusion_feature_cases(model: str):
-    return [
-        # FP8 / Hybrid sequence parallelism
-        pytest.param(
-            OmniServerParams(
-                model=model,
-            ),
-            id="default",
-            marks=FOUR_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases(
-        model="black-forest-labs/FLUX.2-klein-4B",
-    ),
-    indirect=True,
-)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 2,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "true_cfg_scale": 4.0,
-            "seed": 42,
-        },
-    }
-
-    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_flux2_klein_expansion.py

@@ -1,163 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end tests for Flux2 Klein inpainting in online serving mode.
-
-Uses /v1/images/edits endpoint which is the correct API for image inpainting.
-"""
-
-import base64
-from io import BytesIO
-
-import httpx
-import pytest
-from PIL import Image, ImageDraw
-
-from tests.helpers.runtime import OmniServer, OmniServerParams
-
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
-
-MODEL = "black-forest-labs/FLUX.2-klein-4B"
-
-_HEIGHT = 512
-_WIDTH = 512
-_NUM_INFERENCE_STEPS = 4
-
-
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
-            ),
-            id="tp2_cpu_offload",
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2"],
-            ),
-            id="tp2_basic",
-        ),
-    ]
-
-
-def _image_to_base64_jpeg(image: Image.Image) -> str:
-    buffer = BytesIO()
-    image.save(buffer, format="JPEG")
-    buffer.seek(0)
-    return base64.b64encode(buffer.read()).decode("utf-8")
-
-
-def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
-    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 _image_to_base64_jpeg(mask)
-
-
-def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
-    return list(img1.getdata()) == list(img2.getdata())
-
-
-def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
-    url = f"http://{host}:{port}/v1/images/edits"
-    files = {
-        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
-        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
-    }
-    data = {"prompt": prompt, "model": model, **kwargs}
-    with httpx.Client(timeout=60.0) as client:
-        response = client.post(url, files=files, data=data)
-        response.raise_for_status()
-        return response.json()
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="Fill in the masked area with a beautiful garden",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 1
-    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
-    img = Image.open(BytesIO(base64.b64decode(img_data)))
-    assert img.size == (_WIDTH, _HEIGHT)
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-    prompt = "A red flower in a field"
-
-    result1 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    result2 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
-    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
-
-    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
-    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
-
-    assert _compare_images(img1, img2), (
-        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
-    )
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="A beautiful landscape",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=2,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 2

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
         windows stitch into a coherent PCM stream. The byte-count gate is the same as
         the sync paths; per-chunk audio content is verified offline against Whisper.
 
-        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
         check still catches catastrophic codec failure (silence, white noise, sample
         scramble all give HNR << 0) while allowing for the sliding-window codec's
         slightly-noisier-than-sync output. The default 1.0 dB threshold has only

tests/e2e/online_serving/test_krea2_expansion.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline.
+
+Exercises the OpenAI-compatible ``/v1/images/generations`` route (served via
+``vllm serve <model> --omni --enforce-eager``) rather than the offline engine path
+covered by ``tests/e2e/offline_inference/test_krea2_expansion.py``.
+
+Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default
+(override with the ``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the
+Raw checkpoint or a local diffusers directory). Krea 2 shares the Qwen-Image VAE and a
+similar single-stream MMDiT, so this test mirrors the structure/fixtures/markers of
+``test_qwen_image_expansion.py``.
+
+Coverage mirrors the offline suite: a basic functional smoke plus the layerwise-CPU-offload
+path that the pipeline declares via ``SupportsComponentDiscovery`` /
+``_layerwise_offload_blocks_attrs``. Sizing is kept light (512x512, few steps) so the case
+fits single-GPU H100 CI rather than the full 2048x2048 distilled resolution.
+"""
+
+import base64
+import os
+from io import BytesIO
+
+import numpy as np
+import pytest
+import requests
+from PIL import Image
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
+# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
+LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
+PROMPT = "a fox in the snow, photorealistic"
+# darkbrush is a trigger-word LoRA; include its trigger phrase so the LoRA case has a strong signal.
+LORA_PROMPT = "a fox in the snow, monochrome ink wash style"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Return single-card L4 online-serving cases for Krea 2.
+
+    Mirrors the offline suite: a basic functional smoke plus layerwise CPU offload.
+    """
+    return [
+        # Basic functional smoke (single-card).
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enforce-eager"],
+            ),
+            id="single_card_basic",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        # Layerwise CPU offload (single-card): SupportsComponentDiscovery +
+        # _layerwise_offload_blocks_attrs.
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enforce-eager", "--enable-layerwise-offload"],
+            ),
+            id="single_card_layerwise_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(MODEL),
+    indirect=True,
+)
+def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Serve Krea 2 and validate the OpenAI-compatible image-generation route.
+
+    guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG,
+    Raw -> CFG), so guidance_scale=0.0 stays agnostic to which checkpoint KREA2_MODEL points at.
+    Validation is delegated to assert_diffusion_response, which checks output dimensions and
+    basic correctness.
+    """
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 8,
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+def _basic_lora_payload(lora: dict | None = None) -> dict:
+    payload = {
+        "prompt": LORA_PROMPT,
+        "n": 1,
+        "size": "512x512",
+        "num_inference_steps": 8,
+        "guidance_scale": 0.0,
+        "seed": 42,
+    }
+    if lora is not None:
+        payload["lora"] = lora
+    return payload
+
+
+def _post_image(server: OmniServer, payload: dict) -> np.ndarray:
+    url = f"http://{server.host}:{server.port}/v1/images/generations"
+    resp = requests.post(url, json=payload, headers={"Authorization": "Bearer EMPTY"}, timeout=900)
+    resp.raise_for_status()
+    b64 = resp.json()["data"][0]["b64_json"]
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    img.load()
+    return np.asarray(img.convert("RGB"), dtype=np.int16)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        pytest.param(
+            OmniServerParams(model=MODEL, server_args=["--enforce-eager"]),
+            id="single_card_lora",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        )
+    ],
+    indirect=True,
+)
+def test_krea2_lora(omni_server: OmniServer):
+    """Serve Krea 2 and validate per-request LoRA over ``/v1/images/generations``.
+
+    Passes ``LORA`` via the top-level ``lora`` object (``{name, path, scale}``) the diffusion
+    image route parses into a ``LoRARequest`` — no server-side LoRA flag needed. Confirms the
+    adapter has a visible effect and that omitting ``lora`` cleanly restores the base output.
+    """
+    baseline = _post_image(omni_server, _basic_lora_payload())
+    with_lora = _post_image(omni_server, _basic_lora_payload({"name": "darkbrush", "path": LORA, "scale": 1.0}))
+    restored = _post_image(omni_server, _basic_lora_payload())
+
+    diff_lora = np.abs(baseline - with_lora).mean()
+    diff_restored = np.abs(baseline - restored).mean()
+
+    assert diff_lora > 0.5, f"LoRA had no visible effect over the serving path: diff={diff_lora}"
+    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"

tests/e2e/online_serving/test_qwen3_omni.py

@@ -226,3 +226,27 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
     }
 
     _run_prefix_cache_check(openai_client, request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
+    """Ensure Thinker-talker models reject /v1/completions; we do this because the
+    thinker-talker handoff implementations currently use ChatML <|im_start|> and
+    <|im_end|> markers to segment the input sequence; when we don't have them,
+    the talker does not get any embeddings, which breaks the server.
+    """
+    responses = openai_client.send_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "prompt": "Hello, how are you?",
+                "max_tokens": 10,
+            },
+        },
+        err_code=400,
+    )
+    assert not responses[0].success

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -7,6 +7,7 @@ E2E Online tests for Qwen3-Omni model.
 import os
 
 import pytest
+from openai import BadRequestError
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -410,6 +411,44 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
+    """
+    Regression: ``use_audio_in_video`` with default video loader sampling (no
+    ``media_io_kwargs.video.fps``) after a prior request that pinned fps.
+
+    Before the fix the second request could partial-hit mm-cache (audio cached,
+    video re-processed without audio) and fail with ``StopIteration`` in
+    ``replace_multimodal_special_tokens``.
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text="Reply with one word: OK.",
+    )
+
+    base_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "use_audio_in_video": True,
+        "modalities": ["text"],
+    }
+
+    # Prime mm-cache with explicit loader fps (historical failure mode).
+    openai_client.send_omni_request(
+        {
+            **base_config,
+            "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
+        }
+    )
+
+    # Default loader sampling: must succeed without explicit media_io fps.
+    openai_client.send_omni_request(dict(base_config))
+
+
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, openai_client) -> None:
@@ -572,3 +611,21 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     text = responses[0].text_content if responses else ""
     word_count = len(text.split())
     assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
+def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
+    """Ensure invalid audio format is rejected with 400 before streaming starts."""
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt(),
+    )
+    with pytest.raises(BadRequestError, match="audio format"):
+        openai_client.client.chat.completions.create(
+            model=omni_server.model,
+            messages=messages,
+            modalities=["text", "audio"],
+            audio={"voice": "alloy", "format": "aac"},
+            stream=True,
+        )

tests/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
 Kept cases (maximum feature coverage, no cpu-offload):
   sp_001              : cache_dit + Ulysses-SP 2          (2×H100)
   cfg_parallel_002    : cache_dit + CFG-Parallel 2        (2×H100)
-  layers_guard_001    : layerwise offload + layers=3      (1×H100, issue #1969 guard)
+  layers_guard_001    : layerwise offload + layers=2      (1×H100, issue #1969 guard)
 
 Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
 """
@@ -116,8 +116,8 @@ LAYERS_GUARD_CASES = [
             model=MODEL,
             server_args=["--enable-layerwise-offload"],
         ),
-        3,
-        id="layers_guard_001_layers3",
+        2,
+        id="layers_guard_001_layers2",
         marks=SINGLE_CARD_FEATURE_MARKS,
     ),
 ]