Merge branch 'main' into fix/tp-local-kv-recv-consensus

docker/Dockerfile.npu

@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
+RUN python3 -m pip install -v --no-deps \
+      "stepaudio2-minicpmo>=0.1.1" \
+      "step-audio2>=1.0.0"
+
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.a3

@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
+RUN python3 -m pip install -v --no-deps \
+      "stepaudio2-minicpmo>=0.1.1" \
+      "step-audio2>=1.0.0"
+
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.ci

@@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "psutil>=7.2.0" \
       "soundfile>=0.13.1" \
       "pyttsx3>=2.99" \
+      "opencc==1.4.1" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
       "jiwer>=3.0.0" \

docker/Dockerfile.npu.ci.a3

@@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "psutil>=7.2.0" \
       "soundfile>=0.13.1" \
       "pyttsx3>=2.99" \
+      "opencc==1.4.1" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
       "jiwer>=3.0.0" \

docs/design/metrics.md

@@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur
 
 `OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
 
-The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
+The orchestrator side does not add its own throttle on top. Scheduler gauges are
+recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's
+scheduler passed its own 1 Hz gate), while output-batch request/token metrics
+are recorded from `IterationStats` whenever a non-empty output batch is
+processed. A previous global `_last_stats_ts` on the orchestrator was removed
+because it starved every replica other than the first to emit within each
+second.
 
 ## Metric Definitions
 

tests/benchmarks/test_daily_omni_pack_modes.py

@@ -0,0 +1,249 @@
+"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
+
+The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
+OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly
+interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM
+system message (not the Qwen system prompt). These tests pin that contract so a refactor
+cannot silently regress the score.
+
+vllm stubs are installed by tests/benchmarks/conftest.py before collection.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Load the data module directly (bypasses vllm_omni.__init__ heavy imports).
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_MODULE_PATH = _REPO_ROOT / "vllm_omni" / "benchmarks" / "data_modules" / "daily_omni_dataset.py"
+_MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
+
+if _MODULE_NAME not in sys.modules:
+    _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    _mod = importlib.util.module_from_spec(_spec)
+    sys.modules[_MODULE_NAME] = _mod
+    _spec.loader.exec_module(_mod)
+
+from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (  # noqa: E402
+    DailyOmniDataset,
+    _uniform_sample_indices,
+)
+
+_VIDEO_ID = "vid001"
+_NUM_FRAMES = 5
+
+
+@pytest.fixture
+def video_dir(tmp_path: Path) -> Path:
+    """Official Daily-Omni ``Videos/`` layout with placeholder media files."""
+    per_video = tmp_path / "Videos" / _VIDEO_ID
+    per_video.mkdir(parents=True)
+    (per_video / f"{_VIDEO_ID}_video.mp4").write_bytes(b"fake-mp4")
+    (per_video / f"{_VIDEO_ID}_audio.wav").write_bytes(b"fake-wav")
+    return tmp_path / "Videos"
+
+
+@pytest.fixture
+def qa_json(tmp_path: Path) -> Path:
+    path = tmp_path / "qa.json"
+    path.write_text(
+        json.dumps(
+            [
+                {
+                    "video_id": _VIDEO_ID,
+                    "question": "What happens first?",
+                    "choice": {"A": "a", "B": "b", "C": "c", "D": "d"},
+                    "answer": "A",
+                    "task_type": "AV Event Alignment",
+                    "video_duration": "30s",
+                }
+            ]
+        ),
+        encoding="utf-8",
+    )
+    return path
+
+
+def _make_dataset(qa_json: Path, video_dir: Path, **kwargs) -> DailyOmniDataset:
+    return DailyOmniDataset(
+        qa_json_path=str(qa_json),
+        video_dir=str(video_dir),
+        disable_shuffle=True,
+        **kwargs,
+    )
+
+
+def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) -> None:
+    """Replace decord/ffmpeg extraction with deterministic in-memory frames + segments."""
+    from PIL import Image
+
+    def fake_extract(video_path, *, audio_path, include_audio, **_kw):
+        counter.append(1)
+        frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
+        segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+        return frames, segments
+
+    monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
+
+
+# ---------------------------------------------------------------------------
+# _uniform_sample_indices
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)])
+def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None:
+    """Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints)."""
+    got = _uniform_sample_indices(n, k)
+    if n <= k:
+        assert got == list(range(n))
+    else:
+        gap = n / k
+        expected = [int(i * gap + gap / 2) for i in range(k)]
+        assert got == expected
+        # Endpoints differ from linspace for typical Daily-Omni downsample sizes.
+        linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
+        if n == 1000 and k == 64:
+            assert got[0] == 7 and got[-1] == 992
+            assert linspace[0] == 0 and linspace[-1] == 999
+            assert got != linspace
+
+
+def test_uniform_sample_indices_non_positive_k() -> None:
+    assert _uniform_sample_indices(10, 0) == []
+    assert _uniform_sample_indices(10, -1) == []
+
+
+# ---------------------------------------------------------------------------
+# minicpm-interleave packing
+# ---------------------------------------------------------------------------
+
+
+def test_minicpm_interleave_alternates_image_and_audio(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, extra, position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert position == "first"
+    assert [p["type"] for p in parts] == ["image_url", "audio_url"] * _NUM_FRAMES
+    assert extra["mm_processor_kwargs"] == {
+        "use_audio_in_video": False,
+        "max_slice_nums": 1,
+        "use_image_id": False,
+    }
+
+
+def test_minicpm_interleave_visual_mode_emits_frames_only(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="visual")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert [p["type"] for p in parts] == ["image_url"] * _NUM_FRAMES
+
+
+def test_minicpm_interleave_writes_file_urls_under_video_dir(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    segment_dir = video_dir / ".minicpm_daily_omni_interleave" / _VIDEO_ID
+    for part in parts:
+        url = part[part["type"]]["url"]
+        assert url.startswith("file://")
+        assert Path(url[len("file://") :]).parent == segment_dir.resolve()
+    assert sorted(p.name for p in segment_dir.glob("frame_*.jpg")) == [f"frame_{i:04d}.jpg" for i in range(_NUM_FRAMES)]
+
+
+def test_minicpm_interleave_inline_mode_emits_data_urls(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(
+        qa_json,
+        video_dir,
+        pack_mode="minicpm-interleave",
+        input_mode="all",
+        inline_local_video=True,
+    )
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert parts[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+    assert parts[1]["audio_url"]["url"].startswith("data:audio/wav;base64,")
+    assert not (video_dir / ".minicpm_daily_omni_interleave").exists()
+
+
+def test_minicpm_interleave_caches_file_parts_but_not_inline_payloads(monkeypatch, qa_json, video_dir) -> None:
+    """Inline base64 must not be retained: ~700 Daily-Omni videos would pin GBs of RSS."""
+    file_calls: list[int] = []
+    ds_file = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds_file, file_calls)
+    ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    assert len(file_calls) == 1
+    assert ds_file._minicpm_interleave_cache
+
+    inline_calls: list[int] = []
+    ds_inline = _make_dataset(
+        qa_json,
+        video_dir,
+        pack_mode="minicpm-interleave",
+        input_mode="all",
+        inline_local_video=True,
+    )
+    _stub_extract(monkeypatch, ds_inline, inline_calls)
+    ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    assert len(inline_calls) == 2
+    assert not ds_inline._minicpm_interleave_cache
+
+
+def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path) -> None:
+    empty = tmp_path / "empty_videos"
+    empty.mkdir()
+    ds = _make_dataset(qa_json, empty, pack_mode="minicpm-interleave", input_mode="all")
+
+    parts, extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert parts is None
+    assert extra is None
+
+
+# ---------------------------------------------------------------------------
+# System prompt selection
+# ---------------------------------------------------------------------------
+
+
+def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None:
+    """Official MiniCPM omni mode still sends role=system with empty content."""
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"})
+
+    assert [m["role"] for m in messages] == ["system", "user"]
+    assert messages[0]["content"][0]["text"] == ""
+
+
+def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="qwen", input_mode="all")
+
+    messages = ds._build_daily_omni_openai_messages([], "q?", {"A": "a"})
+
+    assert [m["role"] for m in messages] == ["system", "user"]
+    assert "Qwen" in messages[0]["content"][0]["text"]
+
+
+def test_invalid_pack_mode_rejected(qa_json, video_dir) -> None:
+    with pytest.raises(ValueError, match="pack_mode must be"):
+        _make_dataset(qa_json, video_dir, pack_mode="bogus")

tests/config/test_config_factory.py

@@ -2214,6 +2214,32 @@ class TestPlatformOverrides:
         assert rocm.stages[0].enforce_eager is None
         assert rocm.stages[1].enforce_eager is True
 
+    def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
+        pipeline = resolve_pipeline_config("minicpmo_4_5")
+        assert isinstance(pipeline, PipelineConfig)
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
+
+        cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
+        cuda_stages = merge_pipeline_deploy(pipeline, cuda)
+        assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
+
+        # The CUDA memory budget must not leak into the existing NPU profile.
+        npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu")
+        npu_stages = merge_pipeline_deploy(pipeline, npu)
+        assert "kv_cache_memory_bytes" not in npu_stages[1].yaml_engine_args
+
+        # Multi-GPU replica profiles place Talker on dedicated devices and
+        # preserve their previous automatic KV-cache sizing.
+        for filename in (
+            "minicpmo_4_5_3gpu_stage1_replicas.yaml",
+            "minicpmo_4_5_4gpu_stage1_replicas.yaml",
+        ):
+            replica = _apply_platform_overrides(
+                load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda"
+            )
+            replica_stages = merge_pipeline_deploy(pipeline, replica)
+            assert replica_stages[1].yaml_engine_args["kv_cache_memory_bytes"] is None
+
     def test_npu_overrides(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
+from vllm_omni.engine.serialization import deserialize_additional_information
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine_request() -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id="req",
+        prompt_token_ids=[1],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=0,
+        force_kv_transfer=force_kv_transfer,
+    )
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    metadata = deserialize_additional_information(tagged.additional_information)
+    return result, metadata
+
+
+def test_stage_zero_request_omits_kv_transfer():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+
+    assert omits_transfer
+    assert "omni_force_kv_transfer" not in metadata
+
+
+def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+
+    assert not omits_transfer
+    assert metadata["omni_final_stage_id"] == 0
+    assert metadata["omni_force_kv_transfer"] is True

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -21,6 +21,7 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -93,6 +94,17 @@ def _make_update(**kwargs) -> StreamingUpdate:
     return StreamingUpdate(**defaults)
 
 
+def test_generation_scheduler_records_prefill_stats_for_metrics() -> None:
+    request = _make_request(prompt_token_ids=[1, 2, 3, 4])
+    request.prefill_stats = PrefillStats()
+
+    OmniGenerationScheduler._record_prefill_stats(request)
+
+    assert request.prefill_stats.num_prompt_tokens == 4
+    assert request.prefill_stats.num_computed_tokens == 4
+    assert request.prefill_stats.num_cached_tokens == 0
+
+
 def test_resumable_generation_stop_marks_segment_boundary() -> None:
     session = _make_request(request_id="req-generation-segment")
     session.status = RequestStatus.RUNNING

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -9,7 +9,9 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -91,7 +93,20 @@ def _make_pipeline() -> Wan22Pipeline:
     return pipeline
 
 
-def test_forward_delegates_denoising_to_diffuse() -> None:
+@pytest.mark.parametrize(
+    ("sampling_params_kwargs", "expected_low", "expected_high"),
+    [
+        ({}, 4.0, 4.0),
+        ({"guidance_scale": 0.0}, 0.0, 0.0),
+        ({"guidance_scale": 1.0}, 1.0, 1.0),
+        ({"guidance_scale": 3.0, "guidance_scale_2": 5.0}, 3.0, 5.0),
+    ],
+)
+def test_forward_delegates_denoising_to_diffuse(
+    sampling_params_kwargs: dict[str, float],
+    expected_low: float,
+    expected_high: float,
+) -> None:
     pipeline = _make_pipeline()
     captured: dict[str, object] = {}
 
@@ -101,25 +116,15 @@ def test_forward_delegates_denoising_to_diffuse() -> None:
 
     pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
 
-    mock_req = SimpleNamespace(
+    mock_req = OmniDiffusionRequest(
         prompt="prompt",
         request_id="test-req",
-        sampling_params=SimpleNamespace(
-            height=None,
-            width=None,
+        sampling_params=OmniDiffusionSamplingParams(
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=True,
-            guidance_scale=1.0,
-            guidance_scale_2=None,
-            boundary_ratio=None,
-            generator=None,
-            seed=None,
-            num_outputs_per_prompt=1,
             max_sequence_length=32,
-            latents=None,
             output_type="latent",
-            extra_args={},
+            **sampling_params_kwargs,
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
@@ -129,8 +134,8 @@ def test_forward_delegates_denoising_to_diffuse() -> None:
     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
+    assert captured["guidance_low"] == expected_low
+    assert captured["guidance_high"] == expected_high
     assert captured["boundary_timestep"] == pytest.approx(875.0)
     assert captured["latent_condition"] is None
     assert captured["first_frame_mask"] is None

tests/diffusion/test_diffusion_request.py

@@ -72,3 +72,40 @@ def test_tp_seed_same_across_ranks_and_varies_across_requests():
 
     # Seeds must vary across requests (non-determinism preserved).
     assert len(set(seeds)) == n_requests, f"Expected {n_requests} unique seeds but got {len(set(seeds))}: {seeds}"
+
+
+def _request_with_guidance(**sampling_kwargs) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompt={"prompt": "a cup of coffee on a table"},
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1, **sampling_kwargs),
+        request_id="request-test",
+    )
+
+
+def test_omitted_guidance_scale_is_unset_and_resolved_to_one():
+    # No guidance_scale passed: the caller did not set it, so it must NOT count as
+    # "provided" (each pipeline applies its own default), and __post_init__ resolves the
+    # None "unset" sentinel to a concrete float so no consumer ever observes None.
+    req = _request_with_guidance()
+
+    assert req.sampling_params.guidance_scale_provided is False
+    assert req.sampling_params.guidance_scale == 1.0
+
+
+def test_explicit_zero_guidance_scale_is_honored():
+    # Regression (#4998): an explicit guidance_scale=0 (disable CFG) must be honored,
+    # not mistaken for "unset" and overridden by the pipeline's model default. This is
+    # the sentinel collision that the None default + ``is not None`` check fixes.
+    req = _request_with_guidance(guidance_scale=0.0)
+
+    assert req.sampling_params.guidance_scale_provided is True
+    assert req.sampling_params.guidance_scale == 0.0
+
+
+@pytest.mark.parametrize("value", [0.999, 1.0, 5.0])
+def test_explicit_nonzero_guidance_scale_is_honored(value):
+    # Any explicitly-set value (truthy or not) is honored as provided and preserved.
+    req = _request_with_guidance(guidance_scale=value)
+
+    assert req.sampling_params.guidance_scale_provided is True
+    assert req.sampling_params.guidance_scale == value

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -10,8 +10,10 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
+from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
 from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -37,8 +39,10 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         external_req_id=external_req_id or req_id,
         status=status,
         prompt_token_ids=[],
+        num_prompt_tokens=0,
         num_computed_tokens=0,
         num_output_placeholders=0,
+        prefill_stats=None,
         additional_information=None,
         is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
@@ -485,6 +489,48 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt
     assert request.additional_information["meta"]["finished"].item() is True
 
 
+def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+    request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk")
+    request.prefill_stats = PrefillStats()
+    adapter.request_ids_mapping[request.request_id] = request.external_req_id
+
+    connector.get.return_value = (
+        {
+            "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        },
+        8,
+    )
+    assert adapter._poll_single_request(request) is True
+    OmniGenerationScheduler._record_prefill_stats(request)
+    first_chunk_stats = request.prefill_stats
+    request.prefill_stats = None
+
+    connector.get.return_value = (
+        {
+            "codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+        },
+        8,
+    )
+    assert adapter._poll_single_request(request) is True
+    assert isinstance(request.prefill_stats, PrefillStats)
+    OmniGenerationScheduler._record_prefill_stats(request)
+    second_chunk_stats = request.prefill_stats
+
+    assert first_chunk_stats is not None
+    assert first_chunk_stats.num_prompt_tokens == 2
+    assert second_chunk_stats is not None
+    assert second_chunk_stats.num_prompt_tokens == 3
+
+    prompt_token_stats = PromptTokenStats()
+    prompt_token_stats.update_from_output(first_chunk_stats)
+    prompt_token_stats.update_from_output(second_chunk_stats)
+    assert prompt_token_stats.total == 5
+    assert prompt_token_stats.computed == 5
+
+
 def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter):
     adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"})
     request = _req("req-1", RequestStatus.WAITING)

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -468,6 +468,10 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
         "messages": messages,
         "stream": True,
         "key_words": {"text": ["london"]},
+        # A one-word English clip is too short for whisper's language detection to
+        # be stable; unpinned it renders "London" as 런던 / ランデ / 梁敦 and the
+        # containment check then fails on correct audio.
+        "transcript_language": "en",
     }
 
     # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py).

tests/engine/test_async_omni_engine_input.py

@@ -7,7 +7,9 @@ from vllm.v1.engine import EngineCoreRequest
 from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
+from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -107,6 +109,34 @@ def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: M
     assert torch.equal(info["hidden_states"]["tts"], hidden)
 
 
+def test_cfg_companion_suppresses_payload_but_forces_kv_transfer(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
+    engine.request_queue = mocker.Mock()
+
+    engine._enqueue_cfg_companions(
+        parent_id="req",
+        original_prompt={"prompt": "positive"},
+        stage0_params=params,
+        sampling_params_list=[params],
+    )
+
+    message = engine.request_queue.sync_q.put.call_args.args[0]
+    metadata = deserialize_additional_information(message.prompt.additional_information)
+    assert metadata["omni_final_stage_id"] == 0
+    assert metadata["omni_force_kv_transfer"] is True
+
+
 def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)

tests/engine/test_async_omni_engine_stage_init.py

@@ -172,6 +172,41 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
+def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.stage_configs = [types.SimpleNamespace()]
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.single_stage_mode = False
+    engine.diffusion_batch_size = 1
+    engine.async_chunk = False
+    engine.tokenizer = None
+    engine._single_stage_id_filter = None
+    engine._omni_master_address = None
+    engine._omni_master_port = None
+    engine._omni_dp_size_local = 1
+    engine._omni_heartbeat_timeout = 30.0
+    engine._omni_lb_policy = "random"
+    engine.request_queue = types.SimpleNamespace()
+    engine._log_stats = True
+
+    captured: dict[str, object] = {}
+    runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None)
+
+    def _capture_create_stage_runtime(**kwargs):
+        captured.update(kwargs)
+        return runtime
+
+    monkeypatch.setattr(engine_mod, "create_stage_runtime", _capture_create_stage_runtime)
+
+    engine._initialize_stages(stage_init_timeout=7)
+
+    assert captured["stage_init_timeout"] == 7
+    assert captured["log_stats"] is True
+
+
 def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -342,6 +377,95 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
     assert runtime.stage_pools[1].output_processor is stage1_output_processor
 
 
+def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg).replicas[0]
+    plan.engine_args_dict = {}
+
+    captured: dict[str, object] = {}
+    addresses = types.SimpleNamespace(
+        inputs=["tcp://127.0.0.1:1"],
+        outputs=["tcp://127.0.0.1:2"],
+        frontend_stats_publish_address=None,
+    )
+    resources = types.SimpleNamespace(manager=object(), coordinator=None, addresses=addresses)
+    stage_client = types.SimpleNamespace()
+
+    @contextlib.contextmanager
+    def _capture_launch_stage_replica(**kwargs):
+        captured["log_stats"] = kwargs["log_stats"]
+        yield resources
+
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
+    monkeypatch.setattr(
+        runtime_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        lambda **kwargs: (captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client),
+    )
+
+    assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
+    assert captured["log_stats"] is True
+    assert captured["client_log_stats"] is True
+
+
+def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg)
+    stage_client = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(),
+    )
+    captured: dict[str, object] = {}
+    output_processor = object()
+
+    def _capture_build_llm_stage_output_processor(plan, stage_vllm_config, *, log_stats=False):
+        captured["plan"] = plan
+        captured["stage_vllm_config"] = stage_vllm_config
+        captured["log_stats"] = log_stats
+        return output_processor
+
+    monkeypatch.setattr(
+        runtime_mod,
+        "build_llm_stage_output_processor",
+        _capture_build_llm_stage_output_processor,
+    )
+
+    pools = runtime._assemble_stage_pools([stage_plan], {0: [stage_client]})
+
+    assert pools[0].output_processor is output_processor
+    assert captured == {
+        "plan": stage_plan,
+        "stage_vllm_config": cfg,
+        "log_stats": True,
+    }
+
+
 def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 
@@ -537,10 +661,11 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
         stage_configs=[],
         model="dummy-model",
         config_path="dummy-config",
-        stage_init_timeout=302,
+        stage_init_timeout=1,
         diffusion_batch_size=1,
         async_chunk=False,
     )
+    stage_init_timeout = 302
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
@@ -587,14 +712,14 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
     )
 
     try:
-        runtime._initialize_local_llm_replica(plan, 302)
+        runtime._initialize_local_llm_replica(plan, stage_init_timeout)
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)
         else:
             os.environ[device_env_var] = prev_device_env
 
-    assert captured_timeout == 302
+    assert captured_timeout == stage_init_timeout
 
 
 def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):

tests/engine/test_orchestrator.py

@@ -15,6 +15,7 @@ import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -145,7 +146,7 @@ class FakeStageClient:
         try:
             return self._engine_core_outputs.get_nowait()
         except queue.Empty:
-            return SimpleNamespace(outputs=[])
+            return SimpleNamespace(outputs=[], scheduler_stats=None, finished_requests=None)
 
     def get_diffusion_output_nowait(self):
         try:
@@ -253,12 +254,40 @@ class FakeOutputProcessor:
         return None
 
 
+class IterationStatsOutputProcessor(FakeOutputProcessor):
+    def process_outputs(self, *_args, **kwargs):
+        iteration_stats = kwargs.get("iteration_stats")
+        if iteration_stats is None and len(_args) >= 3:
+            iteration_stats = _args[2]
+        assert iteration_stats is not None
+        iteration_stats.num_generation_tokens += 3
+        return super().process_outputs(*_args, **kwargs)
+
+
+class RecordingOutputProcessor(FakeOutputProcessor):
+    def __init__(self, *, request_outputs: list[object] | None = None) -> None:
+        super().__init__(request_outputs=request_outputs)
+        self.process_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+
+    def process_outputs(self, *args, **kwargs):
+        self.process_calls.append((args, kwargs))
+        return super().process_outputs(*args, **kwargs)
+
+
+class RecordingStatLogger:
+    def __init__(self) -> None:
+        self.records: list[tuple[Any, Any, int]] = []
+
+    def record(self, scheduler_stats, iteration_stats, *, engine_idx: int = 0) -> None:
+        self.records.append((scheduler_stats, iteration_stats, engine_idx))
+
+
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
 
 def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
-    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
+    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
 
 
 def _build_request_output(
@@ -329,6 +358,7 @@ def _build_harness(
     output_processors: list[object] | None = None,
     stage_vllm_configs: list[object] | None = None,
     async_chunk: bool = False,
+    log_stats: bool = False,
     stage_pools: list[StagePool] | None = None,
 ) -> OrchestratorFixture:
     """Build an Orchestrator test harness.
@@ -364,6 +394,7 @@ def _build_harness(
                 rpc_async_queue=rpc_queue.async_q,
                 stage_pools=stage_pools,
                 async_chunk=async_chunk,
+                log_stats=log_stats,
             )
             ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
             await orchestrator.run()
@@ -2137,38 +2168,143 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
-def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None:
-    """Regression: each (stage, replica) must independently publish its wrapped
-    vllm:* stats when its scheduler emits non-None scheduler_stats.
-
-    A previous version of Orchestrator carried a global self._last_stats_ts /
-    _stats_interval_s gate around _stat_logger.record(). Because
-    OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler
-    (one per (stage, replica)), the extra global gate starved every replica
-    other than the first to emit within each second — their {stage, replica}
-    gauges/counters went stale.
-
-    The fix removed the global gate entirely; the only signal needed is
-    'this replica's scheduler emitted non-None scheduler_stats'. This test
-    fails loudly if someone reintroduces the global throttle.
-    """
-    import inspect
-
-    from vllm_omni.engine.orchestrator import Orchestrator
-
-    source = inspect.getsource(Orchestrator)
-    assert "_last_stats_ts" not in source, (
-        "Orchestrator must not gate stat recording on a global timestamp. "
-        "OmniSchedulerMixin.make_stats() already throttles per scheduler "
-        "(per (stage, replica)); an outer global gate starves all but the "
-        "first replica to emit within each 1s window."
+@pytest.mark.asyncio
+async def test_orchestrator_records_iteration_stats_without_scheduler_stats(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor(
+        request_outputs=[_build_request_output("req-stats", token_ids=[1], finished=True)]
     )
-    assert "_stats_interval_s" not in source
-    assert "raw_outputs.scheduler_stats is not None" in source, (
-        "Orchestrator must gate stat recording solely on "
-        "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz "
-        "throttle in OmniSchedulerMixin.make_stats() is the only gate needed."
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    request = SimpleNamespace(request_id="req-stats", prompt_token_ids=[1, 2])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-stats",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-without-scheduler-stats", 1.0))
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert scheduler_stats is None
+        assert iteration_stats is not None
+        assert iteration_stats.num_generation_tokens == 3
+        assert engine_idx == 0
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_records_scheduler_stats_without_outputs(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    scheduler_stats = SimpleNamespace(num_running_reqs=0, kv_cache_usage=0.5)
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=scheduler_stats,
+                finished_requests=None,
+            )
+        )
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        recorded_scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert recorded_scheduler_stats is scheduler_stats
+        assert iteration_stats is None
+        assert engine_idx == 0
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=None,
+                finished_requests={"req-finished"},
+            )
+        )
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+        assert stat_logger.records == []
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_without_stat_logger(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    orchestrator_fixture.orchestrator._stat_logger = None
+
+    try:
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-output", 1.0))
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_preserves_none_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
     )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=None)
+
+    assert processor.process_calls[0][0][2] is None
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+    iteration_stats = IterationStats()
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=iteration_stats)
+
+    assert iteration_stats.num_generation_tokens == 3
 
 
 @pytest.mark.asyncio

tests/engine/test_output_processor.py

@@ -10,9 +10,11 @@ import pytest
 import torch
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.metrics.stats import IterationStats, PrefillStats
 
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.outputs import output_processor
 from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
@@ -616,3 +618,48 @@ def test_no_detokenizer_final_only():
     result = s.make_request_output([], None, FinishReason.STOP, None)
     assert result is not None
     assert AUDIO in result.outputs[0].multimodal_output
+
+
+def test_mm_only_outputs_update_iteration_stats():
+    processor = MultimodalOutputProcessor(
+        tokenizer=None,
+        log_stats=True,
+        engine_core_output_type=AUDIO,
+    )
+    state = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id].append(state.request_id)
+
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=4,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=0,
+    )
+    output = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[10, 11],
+        finish_reason=FinishReason.STOP,
+        events=[
+            EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
+            EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0),
+        ],
+        prefill_stats=prefill_stats,
+        multimodal_output={AUDIO: torch.ones(1, 8)},
+    )
+    iteration_stats = IterationStats()
+
+    processor.process_outputs(
+        [output],
+        engine_core_timestamp=3.0,
+        iteration_stats=iteration_stats,
+    )
+
+    assert iteration_stats.num_prompt_tokens == 4
+    assert iteration_stats.num_generation_tokens == 2
+    assert len(iteration_stats.time_to_first_tokens_iter) == 1
+    assert len(iteration_stats.finished_requests) == 1
+    finished = iteration_stats.finished_requests[0]
+    assert finished.finish_reason == FinishReason.STOP
+    assert finished.num_prompt_tokens == state.prompt_len
+    assert finished.num_generation_tokens == 2

tests/engine/test_single_stage_mode.py

@@ -778,6 +778,7 @@ class TestSingleStageReplicaInitialization:
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
             omni_master_port=26000,
+            log_stats=True,
         )
         runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"}
@@ -811,10 +812,17 @@ class TestSingleStageReplicaInitialization:
         sentinel_client = SimpleNamespace()
 
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect)
+        client_kwargs: dict[str, Any] = {}
+
+        def _capture_make_async_mp_client(**kwargs):
+            client_kwargs.update(kwargs)
+            events.append("attach")
+            return sentinel_client
+
         mocker.patch.object(
             StageEngineCoreClientBase,
             "make_async_mp_client",
-            side_effect=lambda **_: (events.append("attach"), sentinel_client)[1],
+            side_effect=_capture_make_async_mp_client,
         )
 
         result = runtime._initialize_remote_replica(plan, stage_init_timeout=60)
@@ -824,6 +832,7 @@ class TestSingleStageReplicaInitialization:
         assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0
         assert mock_connect.call_args.kwargs["stage_id"] == 1
         assert mock_connect.call_args.kwargs["replica_id"] == 0
+        assert client_kwargs["log_stats"] is True
         assert events == ["enter", "exit", "attach"]
 
     def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):

tests/entrypoints/openai_api/test_image_server.py

@@ -1301,7 +1301,11 @@ def test_generate_images_rejects_model_mismatch(test_client):
     assert "model mismatch" in response.json()["detail"].lower()
 
 
-def test_image_file_response_format_multiple(test_client):
+@pytest.mark.parametrize(
+    ("output_format", "expected_pil_format"),
+    [("png", "PNG"), ("jpeg", "JPEG"), ("webp", "WEBP")],
+)
+def test_image_file_response_format_multiple(test_client, output_format, expected_pil_format):
     """Test response_format=file with n>1 returns ZIP archive"""
     response = test_client.post(
         "/v1/images/generations",
@@ -1309,6 +1313,7 @@ def test_image_file_response_format_multiple(test_client):
             "prompt": "a dog",
             "n": 3,
             "response_format": "file",
+            "output_format": output_format,
         },
     )
 
@@ -1317,23 +1322,32 @@ def test_image_file_response_format_multiple(test_client):
     assert "attachment" in response.headers.get("content-disposition", "")
     assert ".zip" in response.headers.get("content-disposition", "")
 
-    # Verify it's a valid ZIP with 3 PNG files
+    # Verify it's a valid ZIP with 3 correctly labelled image files
     import zipfile
 
     zip_buffer = io.BytesIO(response.content)
     with zipfile.ZipFile(zip_buffer, "r") as zf:
         files = zf.namelist()
         assert len(files) == 3
-        assert all(f.endswith(".png") for f in files)
+        assert all(f.endswith(f".{output_format}") for f in files)
 
-        # Verify each file is a valid PNG
+        # Verify each file's bytes match its advertised format
         for filename in files:
             img_bytes = zf.read(filename)
             img = Image.open(io.BytesIO(img_bytes))
-            assert img.format == "PNG"
+            assert img.format == expected_pil_format
 
 
-def test_image_file_response_format_single(test_client):
+@pytest.mark.parametrize(
+    ("output_format", "expected_media_type", "expected_pil_format"),
+    [
+        ("png", "image/png", "PNG"),
+        ("jpg", "image/jpeg", "JPEG"),
+        ("jpeg", "image/jpeg", "JPEG"),
+        ("webp", "image/webp", "WEBP"),
+    ],
+)
+def test_image_file_response_format_single(test_client, output_format, expected_media_type, expected_pil_format):
     """Test response_format=file with n=1 returns a single image file."""
     response = test_client.post(
         "/v1/images/generations",
@@ -1341,16 +1355,61 @@ def test_image_file_response_format_single(test_client):
             "prompt": "a dog",
             "n": 1,
             "response_format": "file",
+            "output_format": output_format,
         },
     )
 
     assert response.status_code == 200
-    assert response.headers["content-type"] == "image/png"
+    assert response.headers["content-type"] == expected_media_type
     assert "attachment" in response.headers.get("content-disposition", "")
-    assert ".png" in response.headers.get("content-disposition", "")
+    assert f".{output_format}" in response.headers.get("content-disposition", "")
 
     img = Image.open(io.BytesIO(response.content))
-    assert img.format == "PNG"
+    assert img.format == expected_pil_format
+
+
+@pytest.mark.parametrize(
+    ("output_format", "image_count", "expected_media_type", "expected_pil_format"),
+    [
+        ("jpeg", 1, "image/jpeg", "JPEG"),
+        ("webp", 2, "application/zip", "WEBP"),
+    ],
+)
+def test_multistage_image_file_response_uses_requested_format(
+    async_omni_test_client,
+    output_format,
+    image_count,
+    expected_media_type,
+    expected_pil_format,
+):
+    engine = async_omni_test_client.app.state.engine_client
+    engine._images = [Image.new("RGB", (16, 16), color="green") for _ in range(image_count)]
+
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a dog",
+            "n": image_count,
+            "response_format": "file",
+            "output_format": output_format,
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == expected_media_type
+    if image_count == 1:
+        assert f".{output_format}" in response.headers["content-disposition"]
+        assert Image.open(io.BytesIO(response.content)).format == expected_pil_format
+        return
+
+    import zipfile
+
+    with zipfile.ZipFile(io.BytesIO(response.content), "r") as archive:
+        filenames = archive.namelist()
+        assert len(filenames) == image_count
+        assert all(filename.endswith(f".{output_format}") for filename in filenames)
+        for filename in filenames:
+            assert Image.open(io.BytesIO(archive.read(filename))).format == expected_pil_format
 
 
 def make_test_image_bytes(size=(64, 64)) -> bytes:

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -48,6 +48,25 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
     assert mux_calls[0]["audio"] is None
 
 
+@pytest.mark.parametrize("frame_count", [3, 4])
+def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count):
+    video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames(video)
+
+    assert frames.shape == (frame_count, 2, 5, 3)
+    np.testing.assert_array_equal(frames, video.numpy())
+
+
+def test_channel_first_video_tensor_is_converted_to_channel_last():
+    video = torch.arange(3 * 5 * 2 * 6, dtype=torch.uint8).reshape(3, 5, 2, 6)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames(video)
+
+    assert frames.shape == (5, 2, 6, 3)
+    np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy())
+
+
 def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
     muxers = []
 

tests/entrypoints/test_stream_finish_reason.py

@@ -17,6 +17,7 @@ Key invariants tested:
 """
 
 import json
+import time
 from unittest.mock import MagicMock
 
 import pytest
@@ -80,28 +81,37 @@ def _make_audio_omni_output(
     request_id: str = "test-req",
     index: int = 0,
     num_prompt_tokens: int = 3,
+    stage_id: int | None = None,
+    replica_id: int | None = None,
+    audio_samples: int = 0,
 ) -> OmniRequestOutput:
     """Build an OmniRequestOutput for audio (no torch dependency)."""
+    completion = CompletionOutput(
+        index=index,
+        text="",
+        token_ids=[],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    if audio_samples > 0:
+        completion.multimodal_output = {
+            "audio": [MagicMock(numel=MagicMock(return_value=audio_samples))],
+            "sr": 24000,
+        }
     res = RequestOutput(
         request_id=request_id,
         prompt="test",
         prompt_token_ids=list(range(num_prompt_tokens)),
         prompt_logprobs=None,
-        outputs=[
-            CompletionOutput(
-                index=index,
-                text="",
-                token_ids=[],
-                cumulative_logprob=0.0,
-                logprobs=None,
-                finish_reason="stop",
-                stop_reason=None,
-            )
-        ],
+        outputs=[completion],
         finished=True,
     )
     return OmniRequestOutput(
         request_id=request_id,
+        stage_id=stage_id,
+        replica_id=replica_id,
         final_output_type="audio",
         request_output=res,
         finished=True,
@@ -321,6 +331,74 @@ async def test_single_modality_audio_only_one_stop():
     assert finish_reasons[-1] == "stop"
 
 
+@pytest.mark.parametrize(
+    ("output_replica_id", "pool_replica_id", "expected_replica_id", "pool_lookup"),
+    [
+        (7, None, 7, False),
+        (None, 5, 5, True),
+    ],
+)
+@pytest.mark.asyncio
+async def test_streaming_audio_metrics_resolve_replica_id(
+    monkeypatch,
+    output_replica_id,
+    pool_replica_id,
+    expected_replica_id,
+    pool_lookup,
+):
+    import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod
+    from vllm_omni.entrypoints.client_request_state import ClientRequestState
+
+    serving_chat = _build_serving_chat()
+    request = _make_request(modalities=["audio"])
+    req_state = ClientRequestState(
+        request_id="internal-req",
+        external_request_id="test-req",
+    )
+    req_state.request_arrival_ts = time.time() - 1.0
+    serving_chat.engine_client.request_states = {"internal-req": req_state}
+    serving_chat.engine_client.mod_metrics = object()
+    serving_chat.engine_client.engine = MagicMock()
+    serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()]
+    serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.return_value = pool_replica_id
+
+    first_packet_calls = []
+    finalize_calls = []
+
+    def fake_observe_audio_first_packet(*args, **kwargs):
+        first_packet_calls.append(kwargs)
+
+    def fake_observe_audio_streaming_finalize(*args, **kwargs):
+        finalize_calls.append(kwargs)
+
+    monkeypatch.setattr(serving_chat_mod, "observe_audio_first_packet", fake_observe_audio_first_packet)
+    monkeypatch.setattr(serving_chat_mod, "observe_audio_streaming_finalize", fake_observe_audio_streaming_finalize)
+
+    async def result_generator():
+        yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400)
+
+    await _collect_stream(
+        serving_chat.chat_completion_stream_generator(
+            request=request,
+            result_generator=result_generator(),
+            request_id="test-req",
+            model_name="test-model",
+            conversation=[],
+            tokenizer=MagicMock(),
+            request_metadata=MagicMock(),
+        )
+    )
+
+    if pool_lookup:
+        serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_called_once_with("internal-req")
+    else:
+        serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_not_called()
+    assert first_packet_calls[0]["stage_id"] == 2
+    assert first_packet_calls[0]["replica_id"] == expected_replica_id
+    assert finalize_calls[0]["stage_id"] == 2
+    assert finalize_calls[0]["replica_id"] == expected_replica_id
+
+
 # ---------------------------------------------------------------------------
 # Tests: fallback stop chunk when declared modality is not produced
 # ---------------------------------------------------------------------------

tests/helpers/assertions.py

@@ -527,7 +527,7 @@ def _resolve_audio_transcript(
     audio_bytes = getattr(response, "audio_bytes", None)
     if not audio_bytes:
         return None
-    return convert_audio_bytes_to_text(audio_bytes)
+    return convert_audio_bytes_to_text(audio_bytes, language=request_config.get("transcript_language"))
 
 
 def assert_omni_response(response: Any, request_config: dict[str, Any], run_level):
@@ -632,6 +632,7 @@ def _assert_transcript_matches(
     *,
     threshold: float,
     escalation_model: str | None = None,
+    language: str | None = None,
 ) -> None:
     """Assert spoken audio matches ``expected_text``.
 
@@ -647,7 +648,8 @@ def _assert_transcript_matches(
     clip is re-transcribed with that stronger ASR and the assertion is decided on
     its verdict -- so a weak whisper-``small`` mishear on a short clip does not
     flake the gate, while a genuine model artifact still fails (the strong ASR
-    mismatches too).
+    mismatches too). ``language`` applies to that escalated pass too, so a request
+    that pins a language does not silently fall back to auto-detection on retry.
     """
     expected = str(expected_text).strip().lower()
     similarity = cosine_similarity_text(transcript.strip().lower(), expected)
@@ -664,7 +666,7 @@ def _assert_transcript_matches(
             f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "
             f"escalating to whisper-{escalation_model} to rule out an ASR mishear"
         )
-        strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model)
+        strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model, language=language)
         strong_similarity = cosine_similarity_text(strong_transcript.strip().lower(), expected)
         print(
             f"audio content (whisper-{escalation_model}): {strong_transcript}\n"
@@ -741,6 +743,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
                     expected_text,
                     threshold=0.9,
                     escalation_model=request_config.get("transcript_escalation_model"),
+                    language=request_config.get("transcript_language"),
                 )
         _assert_preset_voice_gender_from_audio(
             response.audio_bytes,

tests/helpers/media.py

@@ -652,7 +652,9 @@ def _serialize_whisper_model_download(model_size: str = "small"):
         f.close()
 
 
-def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "small") -> str:
+def _whisper_transcribe_in_current_process(
+    output_path: str, model_size: str = "small", language: str | None = None
+) -> str:
     import whisper
 
     device_index = None
@@ -684,6 +686,9 @@ def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "
             temperature=0.0,
             word_timestamps=True,
             condition_on_previous_text=False,
+            # None keeps whisper's auto-detection. Do not default this to a
+            # language: callers include non-English audio tests.
+            language=language,
         )["text"]
     finally:
         del model
@@ -694,15 +699,15 @@ def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "
     return text or ""
 
 
-def convert_audio_file_to_text(output_path: str, model_size: str = "small") -> str:
+def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
     """Convert an audio file to text in an isolated subprocess."""
     ctx = multiprocessing.get_context("spawn")
     with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
-        future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size)
+        future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language)
         return future.result()
 
 
-def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") -> str:
+def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str:
     output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
     os.close(output_fd)
     if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
@@ -710,7 +715,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") ->
     data, samplerate = sf.read(io.BytesIO(raw_bytes))
     sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
     print(f"audio data is saved: {output_path}")
-    return convert_audio_file_to_text(output_path, model_size)
+    return convert_audio_file_to_text(output_path, model_size, language)
 
 
 __all__ = [

tests/helpers/minicpmo_4_5_duplex.py

@@ -24,6 +24,19 @@ DEPLOY_CONFIG = modify_stage_config(
             1: {"kv_cache_memory_bytes": 512 * 1024 * 1024},
             2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
         },
+        # Platform overrides are applied after ordinary stage settings. Keep
+        # the constrained duplex test's Talker budget ahead of the base CUDA
+        # profile's single-GPU default.
+        "platforms": {
+            "cuda": {
+                "stages": [
+                    {
+                        "stage_id": 1,
+                        "kv_cache_memory_bytes": 512 * 1024 * 1024,
+                    }
+                ]
+            }
+        },
     },
 )
 CORE_DEPLOY_CONFIG = modify_stage_config(

tests/helpers/tests/test_assertions.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import pytest
 
-from tests.helpers.assertions import _assert_transcript_matches
+from tests.helpers import assertions
+from tests.helpers.assertions import _assert_transcript_matches, _resolve_audio_transcript
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -25,3 +28,63 @@ def test_short_transcript_unrelated_text_still_fails():
             expected_text="how are you",
             threshold=0.9,
         )
+
+
+def _capture_transcribe(monkeypatch) -> dict:
+    captured: dict = {}
+
+    def fake_convert(raw_bytes, model_size="small", language=None):
+        captured["model_size"] = model_size
+        captured["language"] = language
+        return "London"
+
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", fake_convert)
+    return captured
+
+
+def test_resolve_transcript_honors_declared_language(monkeypatch):
+    captured = _capture_transcribe(monkeypatch)
+    response = SimpleNamespace(audio_content=None, audio_bytes=b"fake-wav")
+
+    _resolve_audio_transcript(
+        response,
+        {"modalities": ["text", "audio"], "transcript_language": "en"},
+        "full_model",
+        speech_api=False,
+    )
+
+    assert captured["language"] == "en"
+
+
+def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch):
+    captured = _capture_transcribe(monkeypatch)
+    response = SimpleNamespace(audio_content=None, audio_bytes=b"fake-wav")
+
+    _resolve_audio_transcript(
+        response,
+        {"modalities": ["text", "audio"]},
+        "full_model",
+        speech_api=False,
+    )
+
+    assert captured["language"] is None
+    assert captured["model_size"] == "small"
+
+
+def test_escalated_transcript_keeps_declared_language(monkeypatch):
+    # The escalated pass must honour the same language, otherwise a request that
+    # pins one silently falls back to auto-detection on retry.
+    captured = _capture_transcribe(monkeypatch)
+
+    with pytest.raises(AssertionError, match="after ASR escalation"):
+        _assert_transcript_matches(
+            "totally different words",
+            audio_bytes=b"fake-wav",
+            expected_text="how are you",
+            threshold=0.9,
+            escalation_model="large-v3",
+            language="en",
+        )
+
+    assert captured["model_size"] == "large-v3"
+    assert captured["language"] == "en"

tests/helpers/tests/test_media.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from tests.helpers import media
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _stub_transcribe(monkeypatch) -> dict:
+    """Capture the kwargs the helper hands to whisper, with no model and no device."""
+    captured: dict = {}
+
+    class FakeModel:
+        def transcribe(self, path, **kwargs):
+            captured.update(kwargs)
+            return {"text": "London"}
+
+    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=lambda size, device=None: FakeModel()))
+    # Keep the unit test off the accelerator probe so it stays hermetic.
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms",
+        SimpleNamespace(current_omni_platform=SimpleNamespace(is_available=lambda: False)),
+    )
+    return captured
+
+
+def test_transcribe_forwards_requested_language(monkeypatch):
+    captured = _stub_transcribe(monkeypatch)
+
+    media._whisper_transcribe_in_current_process("/tmp/does-not-matter.wav", "small", language="en")
+
+    assert captured["language"] == "en"
+
+
+def test_transcribe_defaults_to_auto_language(monkeypatch):
+    # Auto-detect must stay the default: forcing a language globally would break
+    # the non-English audio tests (e.g. the Chinese Qwen3-Omni prompts).
+    captured = _stub_transcribe(monkeypatch)
+
+    media._whisper_transcribe_in_current_process("/tmp/does-not-matter.wav", "small")
+
+    assert captured.get("language") is None
+
+
+def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path):
+    """Cover the two hops the tests above skip: bytes -> file -> executor.submit.
+
+    The real call crosses a spawn ProcessPoolExecutor, so capture what gets
+    submitted rather than what the worker eventually does.
+    """
+    submitted: dict = {}
+
+    class FakeExecutor:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+        def submit(self, fn, *args):
+            submitted["fn"] = fn
+            submitted["args"] = args
+            return SimpleNamespace(result=lambda: "London")
+
+    monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", lambda *a, **kw: FakeExecutor())
+
+    wav = tmp_path / "clip.wav"
+    sf.write(wav, np.zeros(2400, dtype=np.float32), 24000)
+
+    assert media.convert_audio_bytes_to_text(wav.read_bytes(), "small", "en") == "London"
+
+    assert submitted["fn"] is media._whisper_transcribe_in_current_process
+    assert submitted["args"][1:] == ("small", "en")

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -170,7 +170,12 @@ class TestDeployTopology:
                 0.15,
             ]
             assert sum(memory_utilizations) <= 0.9 + 1e-6
-            assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == {"video": {"count": 1, "num_frames": 32}}
+            # Daily-Omni minicpm-interleave: up to 64 image/audio items (+ optional video).
+            assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == {
+                "image": 64,
+                "audio": 64,
+                "video": 1,
+            }
         elif filename in {"minicpmo_4_5_2gpu.yaml"}:
             assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [
                 0.9,

tests/worker/test_omni_connector_mixin.py

@@ -17,6 +17,9 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVTransferManager,
+)
 from vllm_omni.outputs import OmniConnectorOutput
 from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
@@ -1198,6 +1201,38 @@ class TestRankAwareKVRouting(unittest.TestCase):
         )
         host.shutdown_omni_connectors()
 
+    def test_single_rank_hooks_use_default_manager_key(self):
+        host = self._make_host(from_tp=1, to_tp=1, local_rank=0)
+        expected = ["omni_0_to_1_kv_cache_req"]
+
+        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
+        host.shutdown_omni_connectors()
+
+    def test_init_uses_effective_manager_tp2_topology(self):
+        model_config = _make_model_config(stage_id=0)
+        model_config.omni_kv_config = {
+            "need_send_cache": True,
+            "omni_from_stage": "0",
+            "omni_to_stage": "1",
+            "rank_mapping": {"from_tp": 2, "to_tp": 2},
+        }
+        with patch(
+            "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank",
+            return_value=1,
+        ):
+            manager = OmniKVTransferManager.from_model_config(model_config)
+
+        host = MixinHost()
+        host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager)
+
+        expected = ["req_0_0_1_1"]
+        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2)
+        self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2)
+        host.shutdown_omni_connectors()
+
     def test_send_keys_route_from_rank_gt_to_rank(self):
         host = self._make_host(from_tp=4, to_tp=2, local_rank=3)
         self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"])

vllm_omni/benchmarks/data_modules/daily_omni_dataset.py

@@ -52,11 +52,14 @@ Usage:
 """
 
 import base64
+import io
 import json
 import logging
+import math
 import os
 import shutil
 import tarfile
+import wave
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Literal
@@ -77,6 +80,10 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+# Official MiniCPM ``get_video_frame_audio_segments`` default (env ``MAX_NUM_FRAMES``).
+_MINICPM_MAX_NUM_FRAMES = int(os.environ.get("MAX_NUM_FRAMES", "64"))
+_MINICPM_AUDIO_SR = 16000
+
 
 def _daily_omni_hf_cache_root() -> Path:
     return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")).expanduser().resolve()
@@ -203,12 +210,62 @@ class _ListDatasetIterator:
 # Aligns with Lliar-liar/Daily-Omni CLI ``--input_mode`` (test_model/*/testmodel.py).
 DailyOmniInputMode = Literal["all", "visual", "audio"]
 
+# How multimodal parts are packed into OpenAI chat messages:
+# - ``qwen``: one ``video_url`` + one ``audio_url`` (upstream Daily-Omni / Qwen protocol).
+# - ``minicpm-interleave``: 1fps ``[image_i, audio_i, ...]`` matching OpenBMB MiniCPM-o
+#   ``get_video_frame_audio_segments`` + interleaved omni contents (paper ~80% recipe).
+DailyOmniPackMode = Literal["qwen", "minicpm-interleave"]
+
 # ``build_conversation()`` in Daily-Omni ``test_model/Qwen2.5-Omni/testmodel.py`` (verbatim).
 DAILY_OMNI_SYSTEM_TEXT = (
     "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
     "capable of perceiving auditory and visual inputs, as well as generating text and speech."
 )
 
+# MiniCPM-o ``get_sys_prompt(mode="omni")`` without ref-audio is an empty system string
+# (still sent as role=system; see ``_build_daily_omni_openai_messages``).
+MINICPM_OMNI_SYSTEM_TEXT = ""
+
+
+def _uniform_sample_indices(n: int, k: int) -> list[int]:
+    """Evenly pick ``k`` indices from ``0..n-1`` (OpenBMB ``uniform_sample`` midpoint bins).
+
+    Matches OmniEvalKit / MiniCPM ``uniform_sample``: midpoint of each equal-width bin,
+    not ``np.linspace`` endpoints. See
+    https://github.com/OpenBMB/OmniEvalKit/blob/1adac0577258d539efc03f16d8a0f4b3f7df6c19/o_e_Kit/utils/utils.py#L50-L54
+    """
+    if k <= 0:
+        return []
+    if n <= k:
+        return list(range(n))
+    gap = n / k
+    return [int(i * gap + gap / 2) for i in range(k)]
+
+
+def _numpy_to_wav_bytes(audio_np: Any, sample_rate: int = _MINICPM_AUDIO_SR) -> bytes:
+    """Encode a mono float/int numpy waveform as 16-bit PCM WAV bytes."""
+    import numpy as np
+
+    arr = np.asarray(audio_np)
+    if arr.ndim > 1:
+        arr = arr.reshape(-1)
+    if arr.dtype != np.int16:
+        arr = np.clip(arr.astype(np.float64), -1.0, 1.0)
+        arr = (arr * 32767.0).astype(np.int16)
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(arr.tobytes())
+    return buf.getvalue()
+
+
+def _pil_to_jpeg_bytes(image: Any, quality: int = 85) -> bytes:
+    buf = io.BytesIO()
+    image.save(buf, format="JPEG", quality=quality)
+    return buf.getvalue()
+
 
 @dataclass
 class DailyOmniSampleRequest(SampleRequest):
@@ -262,6 +319,11 @@ class DailyOmniDataset(BenchmarkDataset):
             ``all`` — video + WAV (default; official audio-visual protocol);
             ``visual`` — video only;
             ``audio`` — extracted WAV only (requires ``{video_id}/{video_id}_audio.wav`` under ``video_dir``).
+        pack_mode: Multimodal packing recipe:
+            ``qwen`` — one ``video_url`` + optional ``audio_url`` (default; Daily-Omni/Qwen);
+            ``minicpm-interleave`` — MiniCPM-o paper recipe: 1fps frames interleaved with
+            matching 1s audio segments as ``image_url``/``audio_url`` pairs (required to approach
+            OpenBMB ~80% Daily-Omni; needs local video+WAV under ``video_dir``).
         max_duration_seconds: Reserved for future ffprobe-based filtering; currently **not applied**
             when building requests (metadata ``video_duration`` is still passed through for eval).
         dataset_subset: Optional HuggingFace subset name (``load_dataset(..., name=...)``); used by bench
@@ -270,7 +332,8 @@ class DailyOmniDataset(BenchmarkDataset):
         inline_local_video: If True, embed local MP4 as ``data:video/mp4;base64,...`` in requests so
             the API server does not need ``--allowed-local-media-path`` (large JSON; use for small runs).
             When ``input_mode`` is ``audio`` or ``all``, local WAV is embedded the same way
-            (``data:audio/wav;base64,...``).
+            (``data:audio/wav;base64,...``). For ``minicpm-interleave``, embeds JPEG/WAV segment
+            data URLs instead of the raw MP4.
         trust_remote_code: Whether to trust remote code when loading HuggingFace dataset
             (online mode only).
     """
@@ -291,6 +354,7 @@ class DailyOmniDataset(BenchmarkDataset):
         random_seed: int = 0,
         video_dir: str | None = None,
         input_mode: DailyOmniInputMode = "all",
+        pack_mode: DailyOmniPackMode = "qwen",
         inline_local_video: bool = False,
         trust_remote_code: bool = False,
         max_duration_seconds: float | None = None,
@@ -300,6 +364,8 @@ class DailyOmniDataset(BenchmarkDataset):
     ) -> None:
         if input_mode not in ("all", "visual", "audio"):
             raise ValueError(f"input_mode must be 'all', 'visual', or 'audio', got {input_mode!r}")
+        if pack_mode not in ("qwen", "minicpm-interleave"):
+            raise ValueError(f"pack_mode must be 'qwen' or 'minicpm-interleave', got {pack_mode!r}")
 
         # Validate arguments: need either local JSON or HF path
         if qa_json_path is None and dataset_path is None:
@@ -318,11 +384,14 @@ class DailyOmniDataset(BenchmarkDataset):
         self.video_dir = Path(video_dir) if video_dir else None
         self.inline_local_video = inline_local_video
         self.input_mode: DailyOmniInputMode = input_mode
+        self.pack_mode: DailyOmniPackMode = pack_mode
         self.max_duration_seconds = max_duration_seconds
         self.trust_remote_code = trust_remote_code
 
         #: In-process cache of ffprobe durations only (no disk persistence).
         self._video_durations: dict[str, float] = {}
+        #: Cache of MiniCPM 1fps interleaved expansions: video_id -> content parts.
+        self._minicpm_interleave_cache: dict[str, list[dict[str, Any]]] = {}
 
         # Initialize parent BenchmarkDataset
         super().__init__(
@@ -336,11 +405,13 @@ class DailyOmniDataset(BenchmarkDataset):
 
         # Verify dataset info
         logger.info(
-            "Loaded Daily-Omni dataset: mode=%s, source=%s, random_seed=%d, input_mode=%s, max_duration=%s",
+            "Loaded Daily-Omni dataset: mode=%s, source=%s, random_seed=%d, "
+            "input_mode=%s, pack_mode=%s, max_duration=%s",
             "local_json" if self.qa_json_path else "huggingface",
             str(self.qa_json_path) if self.qa_json_path else f"{dataset_path}/{dataset_split}",
             random_seed,
             input_mode,
+            pack_mode,
             f"{max_duration_seconds}s" if max_duration_seconds else "unlimited",
         )
 
@@ -575,7 +646,8 @@ class DailyOmniDataset(BenchmarkDataset):
         messages = self._build_daily_omni_openai_messages(mm_payload, question, choice)
         user_text = self._official_daily_omni_user_prompt(question, choice)
         # Text-only length estimate (same as before: no MM token count in bench).
-        prompt_len = len(tokenizer.encode(f"{DAILY_OMNI_SYSTEM_TEXT}\n{user_text}"))
+        system_text = self._system_text_for_pack_mode()
+        prompt_len = len(tokenizer.encode(f"{system_text}\n{user_text}" if system_text else user_text))
 
         return DailyOmniSampleRequest(
             prompt=user_text,
@@ -766,10 +838,11 @@ class DailyOmniDataset(BenchmarkDataset):
         video_id: str,
         video_url: str | None,
     ) -> tuple[dict[str, Any] | list[dict[str, Any]] | None, dict[str, Any] | None, Literal["first", "last"]]:
-        """Build ``multi_modal_data`` and request extras for the active ``input_mode``.
+        """Build multimodal OpenAI content parts and request extras for the active modes."""
+        if self.pack_mode == "minicpm-interleave":
+            return self._compose_minicpm_interleaved_multimodal(video_id, video_url)
 
-        Mirrors upstream Daily-Omni: separate video + WAV with ``use_audio_in_video=False``.
-        """
+        # Qwen / upstream Daily-Omni: separate video + WAV with ``use_audio_in_video=False``.
         extra: dict[str, Any] = {"mm_processor_kwargs": {"use_audio_in_video": False}}
         mode = self.input_mode
 
@@ -787,6 +860,219 @@ class DailyOmniDataset(BenchmarkDataset):
             return None, None, "first"
         return [v, a], extra, "first"
 
+    def _compose_minicpm_interleaved_multimodal(
+        self,
+        video_id: str,
+        video_url: str | None,
+    ) -> tuple[list[dict[str, Any]] | None, dict[str, Any] | None, Literal["first", "last"]]:
+        """Pack media like MiniCPM-o official omni chat: 1fps ``[frame_i, audio_i, ...]``.
+
+        OpenBMB's HF / utils path expands ``video_url`` with ``use_audio=True`` into interleaved
+        PIL frames and 1s audio ndarrays. vLLM-Omni does not implement that expansion, so the
+        bench client builds ordered ``image_url`` / ``audio_url`` parts instead. MiniCPM-o 5.0's
+        OpenAI chat template preserves content-part order, so placeholders stay time-aligned.
+
+        ``max_slice_nums=1`` and ``use_image_id=False`` mirror MiniCPM-o's omni path
+        (``modeling_minicpmo.py`` passes both for interleaved frames): slicing every frame with
+        the config default of 9 would blow past ``max_model_len``, and ``<image_id>N</image_id>``
+        prefixes would label each frame as a separate picture instead of a video timeline.
+        """
+        del video_url  # Hub HTTP URLs are not used; local extract is required.
+        extra: dict[str, Any] = {
+            "mm_processor_kwargs": {
+                "use_audio_in_video": False,
+                "max_slice_nums": 1,
+                "use_image_id": False,
+            }
+        }
+        mode = self.input_mode
+        if mode == "audio":
+            # Fall back to whole-WAV for audio-only ablations.
+            a = self._get_audio_content(video_id)
+            return ([a] if a else None), extra, "first"
+
+        parts = self._get_minicpm_interleaved_parts(video_id, include_audio=(mode == "all"))
+        if not parts:
+            return None, None, "first"
+        return parts, extra, "first"
+
+    def _get_minicpm_interleaved_parts(
+        self,
+        video_id: str,
+        *,
+        include_audio: bool,
+    ) -> list[dict[str, Any]] | None:
+        """Return cached or freshly extracted MiniCPM-style interleaved OpenAI content parts."""
+        cache_key = f"{video_id}|audio={int(include_audio)}|inline={int(self.inline_local_video)}"
+        cached = self._minicpm_interleave_cache.get(cache_key)
+        if cached is not None:
+            return cached
+
+        self._lazy_ensure_hub_media_dir()
+        video_path = self._resolve_local_video_path(video_id)
+        if video_path is None:
+            logger.warning(
+                "MiniCPM interleave pack requires local video under video_dir=%s for video_id=%r",
+                self.video_dir,
+                video_id,
+            )
+            return None
+        audio_path = self._resolve_local_audio_path(video_id) if include_audio else None
+        if include_audio and audio_path is None:
+            logger.warning(
+                "MiniCPM interleave pack (input_mode=all) missing WAV for video_id=%r under %s",
+                video_id,
+                self.video_dir,
+            )
+            return None
+
+        try:
+            frames, audio_segments = self._extract_minicpm_frame_audio_segments(
+                video_path,
+                audio_path=audio_path,
+                include_audio=include_audio,
+            )
+        except Exception:
+            logger.exception(
+                "Failed MiniCPM 1fps extract for video_id=%r path=%s",
+                video_id,
+                video_path,
+            )
+            return None
+
+        if not frames:
+            logger.warning("No frames extracted for video_id=%r", video_id)
+            return None
+        if include_audio and len(audio_segments) != len(frames):
+            logger.warning(
+                "Frame/audio length mismatch for video_id=%r: frames=%d audios=%d",
+                video_id,
+                len(frames),
+                len(audio_segments),
+            )
+            n = min(len(frames), len(audio_segments))
+            frames = frames[:n]
+            audio_segments = audio_segments[:n]
+
+        cache_root: Path | None = None
+        if not self.inline_local_video:
+            assert self.video_dir is not None
+            cache_root = self.video_dir / ".minicpm_daily_omni_interleave" / video_id
+            cache_root.mkdir(parents=True, exist_ok=True)
+
+        parts: list[dict[str, Any]] = []
+        for i, frame in enumerate(frames):
+            jpeg = _pil_to_jpeg_bytes(frame)
+            if self.inline_local_video:
+                b64 = base64.b64encode(jpeg).decode("ascii")
+                parts.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
+                    }
+                )
+            else:
+                assert cache_root is not None
+                frame_path = cache_root / f"frame_{i:04d}.jpg"
+                if not frame_path.is_file() or frame_path.stat().st_size == 0:
+                    frame_path.write_bytes(jpeg)
+                parts.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": frame_path.expanduser().resolve().as_uri()},
+                    }
+                )
+
+            if include_audio:
+                wav = _numpy_to_wav_bytes(audio_segments[i])
+                if self.inline_local_video:
+                    b64 = base64.b64encode(wav).decode("ascii")
+                    parts.append(
+                        {
+                            "type": "audio_url",
+                            "audio_url": {"url": f"data:audio/wav;base64,{b64}"},
+                        }
+                    )
+                else:
+                    assert cache_root is not None
+                    audio_seg_path = cache_root / f"audio_{i:04d}.wav"
+                    if not audio_seg_path.is_file() or audio_seg_path.stat().st_size == 0:
+                        audio_seg_path.write_bytes(wav)
+                    parts.append(
+                        {
+                            "type": "audio_url",
+                            "audio_url": {"url": audio_seg_path.expanduser().resolve().as_uri()},
+                        }
+                    )
+
+        # File URLs are tiny; cache them so repeated QA on the same video skips re-extract.
+        # Inline base64 must not be retained: ~700 Daily-Omni videos would pin GBs of RSS.
+        if not self.inline_local_video:
+            self._minicpm_interleave_cache[cache_key] = parts
+        logger.debug(
+            "MiniCPM interleave packed video_id=%r frames=%d include_audio=%s parts=%d",
+            video_id,
+            len(frames),
+            include_audio,
+            len(parts),
+        )
+        return parts
+
+    @staticmethod
+    def _extract_minicpm_frame_audio_segments(
+        video_path: Path,
+        *,
+        audio_path: Path | None,
+        include_audio: bool,
+        max_num_frames: int = _MINICPM_MAX_NUM_FRAMES,
+    ) -> tuple[list[Any], list[Any]]:
+        """Port of MiniCPM ``get_video_frame_audio_segments`` (stack_frames=1, 1fps)."""
+        import numpy as np
+        from decord import VideoReader, cpu
+        from PIL import Image
+        from vllm.multimodal.media.audio import load_audio
+
+        vr = VideoReader(str(video_path), ctx=cpu(0))
+        avg_fps = float(vr.get_avg_fps()) or 1.0
+        duration = len(vr) / avg_fps
+        num_seconds = max(1, int(math.ceil(duration)))
+        second_timestamps = list(range(num_seconds))
+
+        if duration > max_num_frames:
+            timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))]
+            frame_idx = [min(int(ts * avg_fps), len(vr) - 1) for ts in timestamps]
+            pick = _uniform_sample_indices(len(frame_idx), max_num_frames)
+            frame_idx = [frame_idx[i] for i in pick]
+            timestamps = [timestamps[i] for i in pick]
+        else:
+            frame_idx = [min(int(i * avg_fps), len(vr) - 1) for i in range(num_seconds)]
+            timestamps = second_timestamps
+
+        video = vr.get_batch(frame_idx).asnumpy()
+        frames = [Image.fromarray(v.astype("uint8")).convert("RGB") for v in video]
+
+        audio_segments: list[Any] = []
+        if include_audio:
+            load_path = str(audio_path) if audio_path is not None else str(video_path)
+            audio_np, sr = load_audio(load_path, sr=_MINICPM_AUDIO_SR, mono=True)
+            for i, start_time in enumerate(timestamps):
+                end_time = timestamps[i + 1] if i < len(timestamps) - 1 else duration
+                start_sample = int(start_time * sr)
+                end_sample = int(end_time * sr)
+                segment = audio_np[start_sample:end_sample]
+                if i == len(timestamps) - 1 and len(segment) < 1600:
+                    segment = np.concatenate([segment, np.zeros(1600 - len(segment), dtype=segment.dtype)])
+                if len(segment) == 0:
+                    segment = np.zeros(1600, dtype=np.float32)
+                audio_segments.append(segment)
+
+        return frames, audio_segments
+
+    def _system_text_for_pack_mode(self) -> str:
+        if self.pack_mode == "minicpm-interleave":
+            return MINICPM_OMNI_SYSTEM_TEXT
+        return DAILY_OMNI_SYSTEM_TEXT
+
     @staticmethod
     def _media_desc_for_official_prompt(mode: DailyOmniInputMode) -> str:
         """``media_desc`` in upstream ``build_conversation``."""
@@ -835,14 +1121,19 @@ class DailyOmniDataset(BenchmarkDataset):
         question: str,
         choice: Any,
     ) -> list[dict[str, Any]]:
-        """Map upstream conversation to OpenAI Chat Completions ``messages`` (video_url / audio_url parts)."""
+        """Map upstream conversation to OpenAI Chat Completions ``messages``."""
         user_text = self._official_daily_omni_user_prompt(question, choice)
         mm_list: list[dict[str, Any]] = mm_payload if isinstance(mm_payload, list) else [mm_payload]
         user_content: list[dict[str, Any]] = [*mm_list, {"type": "text", "text": user_text}]
-        return [
-            {"role": "system", "content": [{"type": "text", "text": DAILY_OMNI_SYSTEM_TEXT}]},
-            {"role": "user", "content": user_content},
-        ]
+        system_text = self._system_text_for_pack_mode()
+        messages: list[dict[str, Any]] = []
+        # MiniCPM-o ``get_sys_prompt(mode="omni")`` still emits role=system with empty
+        # content (no ref-audio). Truthiness would drop that role and change chat-template
+        # control tokens; keep the empty system message for minicpm-interleave.
+        if self.pack_mode == "minicpm-interleave" or system_text:
+            messages.append({"role": "system", "content": [{"type": "text", "text": system_text}]})
+        messages.append({"role": "user", "content": user_content})
+        return messages
 
     def sample_by_task_type(
         self,
@@ -900,6 +1191,7 @@ class DailyOmniDataset(BenchmarkDataset):
             f"dataset_split={self.dataset_split!r}, "
             f"video_dir={self.video_dir!r}, "
             f"input_mode={self.input_mode!r}, "
+            f"pack_mode={self.pack_mode!r}, "
             f"inline_local_video={self.inline_local_video!r}, "
             f"max_duration_seconds={self.max_duration_seconds}, "
             f"random_seed={self.random_seed}"

vllm_omni/benchmarks/patch/patch.py

@@ -280,6 +280,7 @@ def get_samples(args, tokenizer):
                 random_seed=args.seed,
                 video_dir=video_dir,
                 input_mode=getattr(args, "daily_omni_input_mode", "all"),
+                pack_mode=getattr(args, "daily_omni_pack_mode", "qwen"),
                 inline_local_video=getattr(args, "daily_omni_inline_local_video", False),
                 trust_remote_code=getattr(args, "trust_remote_code", False),
                 disable_shuffle=getattr(args, "disable_shuffle", False),
@@ -310,6 +311,7 @@ def get_samples(args, tokenizer):
                 random_seed=args.seed,
                 video_dir=video_dir,
                 input_mode=getattr(args, "daily_omni_input_mode", "all"),
+                pack_mode=getattr(args, "daily_omni_pack_mode", "qwen"),
                 inline_local_video=getattr(args, "daily_omni_inline_local_video", False),
                 trust_remote_code=getattr(args, "trust_remote_code", False),
                 no_stream=getattr(args, "no_stream", False),

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -147,7 +147,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return None
 
     def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
-        """True when orchestrator will not run stage 1+ for this request (e.g. text-only).
+        """True when this stage-zero-final request does not need downstream KV.
 
         The result is cached per request to avoid repeated deserialization of
         additional_information on every scheduler tick.
@@ -162,7 +162,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             result = False
         else:
             info = deserialize_additional_information(payload)
-            result = info.get("omni_final_stage_id") == 0
+            result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False))
 
         self._omits_kv_transfer_cache[rid] = result
         return result

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -56,6 +56,24 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             )
         self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
+    @staticmethod
+    def _record_prefill_stats(request: Request) -> None:
+        """Mirror upstream first-prefill prompt accounting for generation stages.
+
+        Async-chunk generation stages can grow ``prompt_token_ids`` after the
+        first schedule; until this path tracks final chunk lengths, this records
+        the first scheduled prompt snapshot. Keep this in sync with vLLM
+        Scheduler.schedule() if this fast path gains prefix-cache or external-KV
+        accounting.
+        """
+        if request.prefill_stats is None:
+            return
+        request.prefill_stats.set(
+            num_prompt_tokens=request.num_prompt_tokens,
+            num_local_cached_tokens=0,
+            num_external_cached_tokens=0,
+        )
+
     def _handle_stopped_request(self, request: Request) -> bool:
         if (
             request.resumable
@@ -147,6 +165,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+            if num_computed_tokens == 0:
+                self._record_prefill_stats(request)
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
             cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
@@ -213,6 +233,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.running.append(request)
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+            if request.num_computed_tokens == 0:
+                self._record_prefill_stats(request)
 
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -24,9 +24,10 @@ stages:
     max_num_batched_tokens: 16384
     devices: "0"
     limit_mm_per_prompt:
-      video:
-        count: 1
-        num_frames: 32
+      # Daily-Omni MiniCPM pack: up to 64 interleaved image/audio pairs (1fps).
+      image: 64
+      audio: 64
+      video: 1
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0
@@ -34,7 +35,8 @@ stages:
       max_tokens: 2048
       seed: 42
       detokenize: true
-      repetition_penalty: 1.1
+      # 1.0 for MCQ accuracy; 1.1 can bias short A–D answers.
+      repetition_penalty: 1.0
 
   - stage_id: 1
     max_num_seqs: 4
@@ -75,6 +77,13 @@ stages:
       detokenize: true
 
 platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Four Talker streams need much less KV than automatic sizing.
+        # Leave single-GPU headroom for concurrent Stage 2 diffusion batches.
+        kv_cache_memory_bytes: 2147483648  # 2 GiB
+
   npu:
     stages:
       - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml

@@ -15,3 +15,10 @@ stages:
     devices: "1,2"
     num_replicas: 2
     async_scheduling: false
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Stage 1 has dedicated GPUs in this profile; retain automatic sizing.
+        kv_cache_memory_bytes: null

vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml

@@ -15,3 +15,10 @@ stages:
     devices: "1,2,3"
     num_replicas: 3
     async_scheduling: false
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Stage 1 has dedicated GPUs in this profile; retain automatic sizing.
+        kv_cache_memory_bytes: null

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -640,7 +640,9 @@ class FluxKontextPipeline(
         width = req.sampling_params.width or width
         num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
         sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = req.sampling_params.guidance_scale or guidance_scale
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale_provided else guidance_scale
+        )
         num_images_per_prompt = (
             req.sampling_params.num_outputs_per_prompt
             if req.sampling_params.num_outputs_per_prompt > 0

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -733,7 +733,7 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsCompon
         )
         width = ar_width or req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor
         num_inference_steps = req.sampling_params.num_inference_steps or 50
-        guidance_scale = req.sampling_params.guidance_scale or 1.5
+        guidance_scale = req.sampling_params.guidance_scale if req.sampling_params.guidance_scale_provided else 1.5
 
         # Ensure dimensions are multiples of vae_scale_factor * patch_size
         multiple_of = self.vae_scale_factor * self._patch_size

vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py

@@ -353,7 +353,7 @@ class FlowMatchingAudioPipeline(
             sampling_params.extra_args = extra_args
 
         num_inference_steps = sampling_params.num_inference_steps or 32
-        guidance_scale = sampling_params.guidance_scale or 3.0
+        guidance_scale = sampling_params.guidance_scale if sampling_params.guidance_scale_provided else 3.0
         generator = self._resolve_diffusion_generator(sampling_params)
 
         with self._stage_timer("consume_payload"):

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -37,7 +37,7 @@ from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.inputs.data import OmniTextPrompt
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -84,6 +84,24 @@ def resolve_wan_flow_shift(req: OmniDiffusionRequest, od_config: OmniDiffusionCo
         raise ValueError(f"Invalid flow_shift={raw_flow_shift!r}. flow_shift must be a float.") from exc
 
 
+def resolve_wan_guidance_scales(
+    sampling_params: OmniDiffusionSamplingParams,
+    default_guidance_scale: float,
+) -> tuple[float, float]:
+    guidance_scale = (
+        sampling_params.guidance_scale if sampling_params.guidance_scale_provided else default_guidance_scale
+    )
+    guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
+    guidance_high = (
+        sampling_params.guidance_scale_2
+        if sampling_params.guidance_scale_2_provided and sampling_params.guidance_scale_2 is not None
+        else (
+            guidance_scale[1] if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1 else guidance_low
+        )
+    )
+    return guidance_low, guidance_high
+
+
 def retrieve_latents(
     encoder_output: torch.Tensor,
     generator: torch.Generator | None = None,
@@ -571,25 +589,10 @@ class Wan22Pipeline(
         width = (width // mod_value) * mod_value
         num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
-        # Respect per-request guidance_scale when explicitly provided.
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-        else:
-            guidance_scale = 4.0
-
         output_type = req.sampling_params.output_type or "np"
         attention_kwargs: dict | None = None
 
-        guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
-        guidance_high = (
-            req.sampling_params.guidance_scale_2
-            if req.sampling_params.guidance_scale_2 is not None
-            else (
-                guidance_scale[1]
-                if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
-                else guidance_low
-            )
-        )
+        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0)
 
         # record guidance for properties
         self._guidance_scale = guidance_low

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     create_transformer_from_config,
     load_transformer_config,
     resolve_wan_flow_shift,
+    resolve_wan_guidance_scales,
     resolve_wan_sample_solver,
     retrieve_latents,
 )
@@ -475,26 +476,10 @@ class Wan22I2VPipeline(
         num_frames = req.sampling_params.num_frames or 81
         num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
-        # Respect per-request guidance_scale when explicitly provided.
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-        else:
-            guidance_scale = 5.0
-
         output_type = req.sampling_params.output_type or "np"
         attention_kwargs: dict | None = None
 
-        # Handle guidance scales
-        guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
-        guidance_high = (
-            req.sampling_params.guidance_scale_2
-            if req.sampling_params.guidance_scale_2 is not None
-            else (
-                guidance_scale[1]
-                if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
-                else guidance_low
-            )
-        )
+        guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0)
 
         self._guidance_scale = guidance_low
         self._guidance_scale_2 = guidance_high

vllm_omni/diffusion/request.py

@@ -41,12 +41,11 @@ class OmniDiffusionRequest:
         if self.sampling_params.generator is None and self.sampling_params.seed is None:
             self.sampling_params.seed = random.randint(0, 2**31 - 1)
 
-        # Detect whether user explicitly provided guidance_scale.
-        # The sentinel default is 0.0 (false-like); any truthy value means
-        # the caller set it intentionally.  We must resolve this BEFORE
-        # auto-filling guidance_scale_2, otherwise the sentinel leaks into
-        # guidance_scale_2.
-        if self.sampling_params.guidance_scale:
+        # Detect whether the caller explicitly provided guidance_scale before
+        # resolving the omitted value.  ``0.0`` is API-valid and must not be
+        # treated as omission because that would unexpectedly enable CFG in
+        # pipelines with a model-specific default above zero.
+        if self.sampling_params.guidance_scale is not None:
             self.sampling_params.guidance_scale_provided = True
         else:
             self.sampling_params.guidance_scale = 1.0

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -409,6 +409,7 @@ class OmniKVTransferManager:
         self._prefetch_executor: ThreadPoolExecutor | None = (
             ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None
         )
+
         self._prefetch_futures: dict[str, Any] = {}
         self._bg_copy_stream: current_omni_platform.Stream | None = None
 
@@ -421,6 +422,11 @@ class OmniKVTransferManager:
             except Exception as e:
                 logger.warning("Failed to eagerly initialize sender connector: %s", e)
 
+    @property
+    def tp_topology(self) -> KVTPTopology:
+        """Return the effective topology after config and runtime detection."""
+        return self._tp_topo
+
     @property
     def topo_config(self) -> _TransferTopoConfig:
         """Lazily build the parallel topology snapshot on first access."""

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping
 from typing import Any
 
 import torch
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.request import Request, RequestStatus
 
 from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
@@ -109,6 +110,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0)
         return max(0, num_computed - num_placeholders)
 
+    @staticmethod
+    def _refresh_generation_chunk_prefill_state(request: Request) -> None:
+        request.num_prompt_tokens = len(request.prompt_token_ids)
+        if getattr(request, "prefill_stats", None) is None:
+            request.prefill_stats = PrefillStats()
+
     @classmethod
     def create_connector(cls, model_config: Any):
         connector_config = getattr(model_config, "stage_connector_config", None)
@@ -249,15 +256,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 new_ids = payload_data.get("codes", {}).get("audio")
                 has_tensor_codes = isinstance(new_ids, torch.Tensor)
                 use_tensor_codes = has_tensor_codes and new_ids.ndim >= 2
+                prompt_token_ids: list[int]
                 if use_tensor_codes:
-                    request.prompt_token_ids = [0] if new_ids.numel() > 0 else []
+                    prompt_token_ids = [0] if new_ids.numel() > 0 else []
                 elif has_tensor_codes:
                     new_ids = new_ids.tolist()
+                    prompt_token_ids = new_ids
                 elif new_ids is None:
                     new_ids = []
-                    request.prompt_token_ids = new_ids
-                if not use_tensor_codes:
-                    request.prompt_token_ids = new_ids
+                    prompt_token_ids = new_ids
+                else:
+                    prompt_token_ids = new_ids
+                request.prompt_token_ids = prompt_token_ids
                 prev_info = getattr(request, "additional_information", None)
                 info = dict(prev_info) if isinstance(prev_info, dict) else {}
                 for key, value in payload_data.items():
@@ -293,6 +303,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # chunk as ready, otherwise Stage1 can consume before the
                     # first DAC frame arrives.
                     return False
+                self._refresh_generation_chunk_prefill_state(request)
 
             # Mark as finished for consumption
             self._finished_load_reqs.add(req_id)

vllm_omni/engine/async_engine_utils.py

@@ -77,12 +77,17 @@ def upgrade_to_omni_request(
 def apply_omni_final_stage_metadata(
     request: EngineCoreRequest,
     final_stage_id: int,
+    *,
+    force_kv_transfer: bool = False,
 ) -> EngineCoreRequest:
-    """Tag a request so the AR scheduler can skip DiT KV for stage zero."""
+    """Tag a request with its final stage and optional KV-transfer override."""
     merged: dict[str, Any] = {}
     if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
         merged = deserialize_additional_information(request.additional_information)
     merged["omni_final_stage_id"] = final_stage_id
+    merged.pop("omni_force_kv_transfer", None)
+    if force_kv_transfer:
+        merged["omni_force_kv_transfer"] = True
     payload = serialize_additional_information(merged)
     return OmniEngineCoreRequest.from_request(
         request,

vllm_omni/engine/async_omni_engine.py

@@ -331,6 +331,7 @@ class AsyncOmniEngine:
             omni_heartbeat_timeout=self._omni_heartbeat_timeout,
             omni_lb_policy=self._omni_lb_policy,
             request_queue=self.request_queue,
+            log_stats=self._log_stats,
         )
         self._runtime.initialize()
 
@@ -797,11 +798,9 @@ class AsyncOmniEngine:
                 supported_tasks=self.supported_tasks,
             )
             request.external_req_id = cid
-            # Companions are stage-0-final: without this worker-visible tag
-            # the AR runner would still emit downstream connector payloads
-            # for them (the orchestrator-side final_stage_id=0 registration
-            # alone does not reach the worker).
-            request = apply_omni_final_stage_metadata(request, 0)
+            # Companions are stage-0-final for ordinary downstream payloads,
+            # but diffusion still needs their CFG KV caches.
+            request = apply_omni_final_stage_metadata(request, 0, force_kv_transfer=True)
 
             # Registration of this companion on stage-0's output processor is
             # deferred to Orchestrator._handle_add_companion, which routes

vllm_omni/engine/orchestrator.py

@@ -362,6 +362,7 @@ class Orchestrator:
         self.async_chunk = bool(async_chunk)
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
+        self.log_stats = log_stats
         self._orch_monitor = create_orch_monitor(
             enabled=enable_orch_monitor,
             replica_sampler=self._sample_replica_metrics,
@@ -446,7 +447,10 @@ class Orchestrator:
         already no-ops on ``scheduler_stats is None`` (which is what
         the upstream scheduler returns when its own log_stats is False),
         so this gate is mainly to keep the ``/metrics`` surface clean
-        when the user did not request stats.
+        when the user did not request stats. When stats are enabled, record()
+        must still receive IterationStats even on ticks where scheduler_stats
+        is None; upstream PrometheusStatLogger records token/request metrics
+        from iteration_stats independently of scheduler gauges.
         """
         self._running_counter = running_counter
         self._transfer_emitter = transfer_emitter
@@ -937,19 +941,17 @@ class Orchestrator:
                                 )
                                 if req_state.streaming.enabled:
                                     await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state)
-                            # OmniSchedulerMixin.make_stats() already throttles
-                            # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats
-                            # being non-None means this replica passed its own gate.
-                            # A second global throttle here would drop stats for
-                            # other (stage, replica) pairs in the same 1s window.
-                            record_stats = self._stat_logger is not None and raw_outputs.scheduler_stats is not None
-                            iteration_stats = IterationStats() if record_stats else None
+                            iteration_stats = (
+                                IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
+                            )
                             raw_output = await pool.process_llm_raw_outputs(
                                 replica_id,
                                 raw_outputs,
                                 iteration_stats=iteration_stats,
                             )
-                            if record_stats:
+                            if self._stat_logger is not None and (
+                                raw_outputs.scheduler_stats is not None or iteration_stats is not None
+                            ):
                                 self._stat_logger.record(
                                     raw_outputs.scheduler_stats,
                                     iteration_stats,

vllm_omni/engine/stage_pool.py

@@ -1081,7 +1081,11 @@ class StagePool:
     async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None:
         """Pull raw EngineCoreOutputs from a stage replica without processing."""
         outputs = await client.get_output_async()
-        if not outputs.outputs:
+        # Keep scheduler-only / finished-only batches. Omni schedulers mirror
+        # upstream by emitting SchedulerStats on throttled ticks even when no
+        # request output is produced, and dropping those batches loses KV/queue
+        # gauges for that interval.
+        if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests:
             return None
         return outputs
 
@@ -1097,7 +1101,6 @@ class StagePool:
             return []
         client = cast(StagePoolLLMClient, raw_client)
         processor = self.output_processor
-        iteration_stats = IterationStats()
         processed = processor.process_outputs(
             raw_outputs.outputs,
             raw_outputs.timestamp,

vllm_omni/engine/stage_runtime.py

@@ -126,6 +126,7 @@ class StageRuntime:
         diffusion_batch_size: int,
         async_chunk: bool,
         tokenizer: str | None = None,
+        log_stats: bool = False,
     ) -> None:
         self._stage_configs = stage_configs
         self._model = model
@@ -134,6 +135,7 @@ class StageRuntime:
         self._diffusion_batch_size = diffusion_batch_size
         self._async_chunk = async_chunk
         self._tokenizer = tokenizer
+        self._log_stats = log_stats
         self._num_stages = len(stage_configs)
 
         # Populated by initialize()
@@ -575,7 +577,7 @@ class StageRuntime:
                 with launch_stage_replica(
                     vllm_config=vllm_config,
                     executor_class=executor_class,
-                    log_stats=False,
+                    log_stats=self._log_stats,
                     stage_id=plan.metadata.stage_id,
                     replica_id=plan.replica_id,
                     stage_config=plan.stage_cfg,
@@ -594,6 +596,7 @@ class StageRuntime:
             stage_client = StageEngineCoreClientBase.make_async_mp_client(
                 vllm_config=vllm_config,
                 executor_class=executor_class,
+                log_stats=self._log_stats,
                 metadata=plan.metadata,
                 client_addresses=self._client_addresses_from_zmq(resources.addresses),
                 engine_manager=resources.manager,
@@ -704,7 +707,11 @@ class StageRuntime:
                 stage_vllm_config = plan.replicas[0].stage_vllm_config
                 if stage_vllm_config is None:
                     raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config")
-                output_processor = build_llm_stage_output_processor(plan, stage_vllm_config)
+                output_processor = build_llm_stage_output_processor(
+                    plan,
+                    stage_vllm_config,
+                    log_stats=self._log_stats,
+                )
 
             stage_pools.append(
                 StagePool(
@@ -742,10 +749,11 @@ class DistStageRuntime(StageRuntime):
         stage_init_timeout: int,
         diffusion_batch_size: int,
         async_chunk: bool,
-        tokenizer: str | None = None,
         single_stage_id_filter: int | None,
         omni_master_address: str,
         omni_master_port: int,
+        tokenizer: str | None = None,
+        log_stats: bool = False,
         omni_dp_size_local: int = 1,
         omni_heartbeat_timeout: float = 30.0,
         omni_lb_policy: str = "random",
@@ -759,6 +767,7 @@ class DistStageRuntime(StageRuntime):
             diffusion_batch_size=diffusion_batch_size,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            log_stats=log_stats,
         )
         self._single_stage_id_filter = single_stage_id_filter
         self._omni_master_address = omni_master_address
@@ -1051,6 +1060,7 @@ class DistStageRuntime(StageRuntime):
             client = StageEngineCoreClientBase.make_async_mp_client(
                 vllm_config=vllm_config,
                 executor_class=ctx.executor_class,
+                log_stats=self._log_stats,
                 metadata=metadata,
                 client_addresses=client_addresses,
                 engine_manager=resources.manager,
@@ -1089,6 +1099,7 @@ def create_stage_runtime(
     omni_heartbeat_timeout: float = 30.0,
     omni_lb_policy: str = "random",
     request_queue: janus.Queue[EngineQueueMessage] | None = None,
+    log_stats: bool = False,
 ) -> StageRuntime:
     """Factory: select StageRuntime or DistStageRuntime."""
     if single_stage_mode:
@@ -1102,6 +1113,7 @@ def create_stage_runtime(
             diffusion_batch_size=diffusion_batch_size,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            log_stats=log_stats,
             single_stage_id_filter=single_stage_id_filter,
             omni_master_address=omni_master_address,
             omni_master_port=omni_master_port,
@@ -1118,4 +1130,5 @@ def create_stage_runtime(
         diffusion_batch_size=diffusion_batch_size,
         async_chunk=async_chunk,
         tokenizer=tokenizer,
+        log_stats=log_stats,
     )

vllm_omni/entrypoints/cli/benchmark/cli_args.py

@@ -95,6 +95,23 @@ def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None:
         "'all': video + WAV together. Sets mm_processor_kwargs.use_audio_in_video=false "
         "and matches official separate video/audio streams.",
     )
+    group.add_argument(
+        "--daily-omni-pack-mode",
+        type=str,
+        choices=["qwen", "minicpm-interleave"],
+        default="qwen",
+        help="How to pack multimodal parts into OpenAI chat messages. "
+        "'qwen' (default): one video_url + one audio_url (Daily-Omni/Qwen protocol). "
+        "'minicpm-interleave': MiniCPM-o official recipe — 1fps frames interleaved with "
+        "matching 1s audio segments as image_url/audio_url pairs (needed to approach "
+        "OpenBMB ~80%% Daily-Omni; requires local Videos extract). By default this mode "
+        "writes JPEG/WAV segments under "
+        "<daily-omni-video-dir>/.minicpm_daily_omni_interleave and sends file:// URLs, "
+        "so start the server with --allowed-local-media-path <daily-omni-video-dir> "
+        "(or a parent). Alternatively pass --daily-omni-inline-local-video to embed "
+        "segments as data URLs. For MiniCPM-o 4.5 string chat templates also start the "
+        "server with --interleave-mm-strings.",
+    )
     group.add_argument(
         "--daily-omni-save-eval-items",
         action="store_true",

vllm_omni/entrypoints/omni_base.py

@@ -612,6 +612,7 @@ class OmniBase(PDDisaggregationMixin):
             request_id=req_id or "",
             finished=finished,
             stage_id=stage_id,
+            replica_id=result.replica_id,
             final_output_type=output_type,
             request_output=engine_outputs,
             images=images,

vllm_omni/entrypoints/openai/api_server.py

@@ -98,7 +98,6 @@ from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
     SUPPORTED_LAYERED_RESOLUTIONS,
-    encode_image_base64,
     encode_image_base64_with_compression,
     parse_size,
     validate_layered_layers,
@@ -1731,6 +1730,39 @@ async def show_available_models(raw_request: Request) -> JSONResponse:
 # Image generation API endpoints
 
 
+def _build_image_generation_response(
+    *,
+    images: list[Image.Image],
+    request: ImageGenerationRequest,
+    stage_durations: Any,
+    peak_memory_mb: Any,
+) -> ImageGenerationResponse | StreamingResponse:
+    """Encode generated images and apply the requested response format."""
+    output_format = _choose_output_format(request.output_format or "png", None)
+    image_data = [
+        ImageData(
+            b64_json=encode_image_base64_with_compression(image, format=output_format),
+            revised_prompt=None,
+        )
+        for image in images
+    ]
+    response_kwargs: dict[str, Any] = {
+        "created": int(time.time()),
+        "data": image_data,
+        "output_format": output_format,
+        "metrics": {
+            "stage_durations": stage_durations or None,
+            "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None,
+        },
+    }
+    if request.size is not None:
+        response_kwargs["size"] = request.size
+    response = ImageGenerationResponse(**response_kwargs)
+    if request.response_format == ResponseFormat.FILE:
+        return response.stream_response()
+    return response
+
+
 @router.post(
     "/v1/images/generations",
     dependencies=[Depends(validate_json_request)],
@@ -1832,15 +1864,11 @@ async def generate_images(
                     content=generation_result.model_dump(),
                 )
             flat_images, stage_durations, peak_memory_mb, _ = generation_result
-            image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images]
-
-            return ImageGenerationResponse(
-                created=int(time.time()),
-                data=image_data,
-                metrics={
-                    "stage_durations": stage_durations or None,
-                    "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None,
-                },
+            return _build_image_generation_response(
+                images=flat_images,
+                request=request,
+                stage_durations=stage_durations,
+                peak_memory_mb=peak_memory_mb,
             )
 
         # Build params - pass through user values directly
@@ -1928,32 +1956,14 @@ async def generate_images(
 
         logger.debug(f"Successfully generated {len(images)} image(s)")
 
-        # Determine output format (default to png)
-        output_format = _choose_output_format(request.output_format or "png", None)
-
-        # Encode images to base64 with the specified format
-        image_data = [
-            ImageData(b64_json=encode_image_base64_with_compression(img, format=output_format), revised_prompt=None)
-            for img in images
-        ]
-
         stage_durations = getattr(result, "stage_durations", None)
         peak_memory_mb = getattr(result, "peak_memory_mb", None)
-        response_kwargs = {
-            "created": int(time.time()),
-            "data": image_data,
-            "output_format": output_format,
-            "metrics": {
-                "stage_durations": stage_durations or None,
-                "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None,
-            },
-        }
-        if request.size:
-            response_kwargs["size"] = size_str
-        response = ImageGenerationResponse(**response_kwargs)
-        if request.response_format != ResponseFormat.FILE:
-            return response
-        return response.stream_response()
+        return _build_image_generation_response(
+            images=images,
+            request=request,
+            stage_durations=stage_durations,
+            peak_memory_mb=peak_memory_mb,
+        )
 
     except (EngineGenerateError, EngineDeadError) as exc:
         return _create_engine_error_json_response(raw_request, exc)

vllm_omni/entrypoints/openai/image_api_utils.py

@@ -101,13 +101,14 @@ def encode_image_base64_with_compression(
     """
     buffer = io.BytesIO()
     image = prepare_image_for_output_format(image, format)
+    pil_format = "jpeg" if format == "jpg" else format
     save_kwargs = {}
     if format in ("jpg", "jpeg", "webp"):
         save_kwargs["quality"] = output_compression
     elif format == "png":
         save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11))
 
-    image.save(buffer, format=format, **save_kwargs)
+    image.save(buffer, format=pil_format, **save_kwargs)
     buffer.seek(0)
     return base64.b64encode(buffer.read()).decode("utf-8")
 

vllm_omni/entrypoints/openai/protocol/images.py

@@ -21,6 +21,13 @@ from pydantic import BaseModel, Field, field_validator
 
 from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 
+_IMAGE_FILE_METADATA = {
+    "jpg": ("jpg", "image/jpeg"),
+    "jpeg": ("jpeg", "image/jpeg"),
+    "png": ("png", "image/png"),
+    "webp": ("webp", "image/webp"),
+}
+
 
 class ResponseFormat(str, Enum):
     """Image response format"""
@@ -183,7 +190,7 @@ class ImageGenerationResponse(BaseModel):
 
     created: int = Field(..., description="Unix timestamp of when the generation completed")
     data: list[ImageData] = Field(..., description="Array of generated images")
-    output_format: str = Field(None, description="The output format of the image generation")
+    output_format: str | None = Field(None, description="The output format of the image generation")
     size: str = Field(None, description="The size of the image generated")
     cot_output: str | None = Field(
         None,
@@ -201,12 +208,16 @@ class ImageGenerationResponse(BaseModel):
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
                 detail="No image data available for file response.",
             )
+        extension, media_type = _IMAGE_FILE_METADATA.get(
+            (self.output_format or "png").lower(),
+            _IMAGE_FILE_METADATA["png"],
+        )
         if len(self.data) == 1:
             image_bytes = base64.b64decode(self.data[0].b64_json)
-            filename = f"image_{uuid.uuid4().hex[:8]}.png"
+            filename = f"image_{uuid.uuid4().hex[:8]}.{extension}"
             return StreamingResponse(
                 io.BytesIO(image_bytes),
-                media_type="image/png",
+                media_type=media_type,
                 headers={
                     "Content-Disposition": f'attachment; filename="{filename}"',
                     "Content-Length": str(len(image_bytes)),
@@ -217,7 +228,7 @@ class ImageGenerationResponse(BaseModel):
             with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
                 for idx, item in enumerate(self.data):
                     if item.b64_json:
-                        zf.writestr(f"image_{idx}.png", base64.b64decode(item.b64_json))
+                        zf.writestr(f"image_{idx}.{extension}", base64.b64decode(item.b64_json))
             zip_bytes = zip_buffer.getvalue()
             filename = f"images_{uuid.uuid4().hex[:8]}.zip"
             return StreamingResponse(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2062,15 +2062,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     now_ts = time.time()
                     if req_state is not None and req_state.first_audio_ts is None:
                         req_state.first_audio_ts = now_ts
-                        stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
-                        # The orchestrator binds requests by their internal id,
-                        # not the user-visible external id, so look up the
-                        # replica with req_state.request_id (internal).
-                        replica_id = (
-                            stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id)
-                            if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools)
-                            else None
-                        )
+                        replica_id = getattr(omni_res, "replica_id", None)
+                        if replica_id is None:
+                            stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
+                            # Fallback for older outputs. The orchestrator binds
+                            # requests by internal id, but can release that
+                            # binding before the API sees the final audio chunk,
+                            # so prefer the OutputMessage replica copied onto
+                            # OmniRequestOutput above.
+                            replica_id = (
+                                stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id)
+                                if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools)
+                                else None
+                            )
                         req_state.audio_emit_stage_id = omni_res.stage_id
                         req_state.audio_emit_replica_id = replica_id
                         observe_audio_first_packet(

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -308,9 +308,6 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
     video_tensor = video_tensor.detach().cpu()
     if video_tensor.dim() == 5:
         raise ValueError("Batched video tensors are not supported for single-video encoding.")
-    elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
-        # [C, F, H, W] -> [F, H, W, C]
-        video_tensor = video_tensor.permute(1, 2, 3, 0)
 
     if video_tensor.is_floating_point():
         # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no
@@ -329,6 +326,8 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray:
 
     if video_array.ndim == 4:
         # Convert channel-first layouts to channel-last
+        # Prefer an explicit channel-last dimension for ambiguous 3/4-frame
+        # videos before interpreting a leading dimension as channels.
         if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4):
             video_array = np.transpose(video_array, (1, 2, 3, 0))
         elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4):

vllm_omni/inputs/data.py

@@ -263,7 +263,7 @@ class OmniDiffusionSamplingParams:
     # Scheduler parameters – ``None`` means "not explicitly set by the caller";
     # each pipeline's ``forward()`` decides its own model-specific default.
     num_inference_steps: int | None = None
-    guidance_scale: float = 0.0
+    guidance_scale: float | None = None
     guidance_scale_provided: bool = False
     guidance_scale_2: float | None = None
     guidance_scale_2_provided: bool = False

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -65,6 +65,27 @@ class BatchedToken2Wav(nn.Module):
         self._token2wav = token2wav
         self.flow = token2wav.flow
         self.hift = token2wav.hift
+        hift_parameter = next(self.hift.parameters(), None)
+        if hift_parameter is not None and hift_parameter.device.type == "cuda":
+            # Prime the CUDA state used by HiFT during backend construction.
+            # Otherwise, the first live audio chunk can fail when async stages
+            # share one GPU.
+            device = hift_parameter.device
+            dtype = hift_parameter.dtype
+            mel_channels = int(self.hift.conv_pre.in_channels)
+            with (
+                torch.inference_mode(),
+                torch.random.fork_rng(devices=[device]),
+                _autocast_disabled(device),
+            ):
+                # 50 mel frames match the default first streamed vocoder chunk.
+                speech, source = self.hift(
+                    torch.zeros((1, mel_channels, 50), device=device, dtype=dtype),
+                    torch.zeros((1, 1, 0), device=device, dtype=dtype),
+                )
+            torch.accelerator.synchronize(device)
+            del speech, source
+            torch.accelerator.empty_cache()
         self.float16 = bool(token2wav.float16)
         self.n_timesteps = int(token2wav.n_timesteps)
         self.mel_cache_len = int(token2wav.mel_cache_len)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -3412,10 +3412,23 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
     ) -> bool:
         return False
 
-    def get_image_prompt_texts(self, image_size: ImageSize, image_idx: int = 0) -> str:
+    def get_image_prompt_texts(
+        self,
+        image_size: ImageSize,
+        image_idx: int = 0,
+        max_slice_nums: int | None = None,
+        use_image_id: bool | None = None,
+    ) -> str:
+        # ``process_images`` forwards ``mm_processor_kwargs`` to the pixel path, so the
+        # placeholder grid must honor the same ``max_slice_nums``; otherwise a request that
+        # sets it (e.g. interleaved omni frames with max_slice_nums=1) produces fewer image
+        # features than placeholder ``<unk>`` slots. ``use_image_id`` is likewise per-request
+        # so callers can match MiniCPM-o's omni path, which emits no ``<image_id>`` prefix.
         return self.info.get_slice_image_placeholder(
             image_size,
             image_idx=image_idx,
+            max_slice_nums=max_slice_nums,
+            **({} if use_image_id is None else {"use_image_id": use_image_id}),
         )
 
     def get_video_prompt_texts(self, image_size: ImageSize, num_frames: int) -> str:
@@ -3605,13 +3618,21 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
                 additional_placeholders.append((modality, sub_pattern))
         placeholders += additional_placeholders
 
+        image_max_slice_nums = hf_processor_mm_kwargs.get("max_slice_nums")
+        image_use_image_id = hf_processor_mm_kwargs.get("use_image_id")
+
         def get_image_replacement(item_idx: int):
             images = mm_items.get_items("image", (MiniCPMVImageEmbeddingItems, ImageProcessorItems))
 
             image_size = images.get_image_size(item_idx)
 
             return PromptUpdateDetails.select_text(
-                self.get_image_prompt_texts(image_size, item_idx),
+                self.get_image_prompt_texts(
+                    image_size,
+                    item_idx,
+                    max_slice_nums=None if image_max_slice_nums is None else int(image_max_slice_nums),  # type: ignore[arg-type]
+                    use_image_id=None if image_use_image_id is None else bool(image_use_image_id),
+                ),
                 "<unk>",
             )
 

vllm_omni/outputs/__init__.py

@@ -83,6 +83,7 @@ class OmniRequestOutput:
         request_id: Unique identifier for this request
         finished: Whether generation is complete
         stage_id: Identifier of the stage that produced this output (pipeline mode)
+        replica_id: Identifier of the stage replica that produced this output
         final_output_type: Type of output ("text", "image", "audio", "latents")
         request_output: The underlying RequestOutput from the stage (pipeline mode)
         images: List of generated PIL images (diffusion mode)
@@ -96,6 +97,7 @@ class OmniRequestOutput:
 
     # Pipeline stage fields
     stage_id: int | None = None
+    replica_id: int | None = None
     final_output_type: str = "text"
     request_output: RequestOutput | None = None
 
@@ -154,6 +156,7 @@ class OmniRequestOutput:
         stage_id: int,
         final_output_type: str,
         request_output: RequestOutput,
+        replica_id: int | None = None,
     ) -> "OmniRequestOutput":
         """Create output from pipeline stage.
 
@@ -168,6 +171,7 @@ class OmniRequestOutput:
         return cls(
             request_id=getattr(request_output, "request_id", ""),
             stage_id=stage_id,
+            replica_id=replica_id,
             final_output_type=final_output_type,
             request_output=request_output,
             finished=True,

vllm_omni/outputs/output_processor.py

@@ -523,7 +523,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 upstream_outputs.append(eco)
 
         # Handle multimodal-only outputs (generation stages) locally.
-        mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs)
+        mm_request_outputs = self._process_mm_only_outputs(
+            mm_only_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
+        )
 
         # Delegate text/pooling outputs to upstream.
         processed = super().process_outputs(
@@ -537,6 +541,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_mm_only_outputs(
         self,
         engine_core_outputs: list[EngineCoreOutput],
+        *,
+        engine_core_timestamp: float | None = None,
+        iteration_stats: IterationStats | None = None,
     ) -> list[OmniRequestOutput | PoolingRequestOutput]:
         """Handle outputs from generation stages that have no detokenizer.
 
@@ -554,7 +561,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             stop_reason = eco.stop_reason
             kv_transfer_params = eco.kv_transfer_params
             routed_experts = eco.routed_experts
-            req_state.num_cached_tokens = eco.num_cached_tokens
+            self._update_stats_from_output(
+                req_state,
+                eco,
+                engine_core_timestamp,
+                iteration_stats,
+            )
+            prefill_stats = getattr(eco, "prefill_stats", None)
+            if prefill_stats is not None:
+                req_state.num_cached_tokens = prefill_stats.num_cached_tokens
             req_state.is_prefilling = False
 
             is_non_final_audio_chunk = (
@@ -578,6 +593,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             is_segment_finished = bool(getattr(eco, "is_segment_finished", False))
             if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk:
                 self._finish_request(req_state)
+                self._update_stats_from_finished(
+                    req_state,
+                    finish_reason,
+                    iteration_stats,
+                )
         return request_outputs
 
     def _update_stats_from_output(

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -122,6 +122,19 @@ class OmniConnectorModelRunnerMixin:
 
         # -- heterogeneous TP rank support --
         rank_cfg = self._parse_rank_mapping(model_config)
+        if self._kv_transfer_manager is not None:
+            topology = getattr(self._kv_transfer_manager, "tp_topology", None)
+            effective_mapping = (
+                getattr(topology, "source_tp_size", None),
+                getattr(topology, "target_tp_size", None),
+                getattr(topology, "local_rank", None),
+            )
+            if all(isinstance(value, int) for value in effective_mapping):
+                rank_cfg = {
+                    "from_tp": effective_mapping[0],
+                    "to_tp": effective_mapping[1],
+                    "local_rank": effective_mapping[2],
+                }
         self._from_tp: int = rank_cfg["from_tp"]
         self._to_tp: int = rank_cfg["to_tp"]
         self._local_rank: int = rank_cfg["local_rank"]
@@ -1275,6 +1288,10 @@ class OmniConnectorModelRunnerMixin:
         For heterogeneous TP receive, the local rank is the target rank and must
         fetch one or more source-rank shards keyed as ``from_rank -> to_rank``.
         """
+        if self._from_tp <= 1 and self._to_tp <= 1:
+            resolved_to_stage = self._next_stage_id if to_stage is None else to_stage
+            return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"]
+
         remote_ranks = self.get_kv_remote_ranks()
         return [
             self.get_kv_connector_key(
@@ -1307,6 +1324,10 @@ class OmniConnectorModelRunnerMixin:
         chunk_id: int = 0,
     ) -> list[str]:
         """Build send-side connector keys for this rank's KV shard(s)."""
+        if self._from_tp <= 1 and self._to_tp <= 1:
+            resolved_to_stage = self._next_stage_id if to_stage is None else to_stage
+            return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"]
+
         target_ranks = self.get_kv_target_ranks_for_send()
         return [
             self.get_kv_connector_key(