[Doc] Address README review feedback

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/cuda/test-nightly.yml

@@ -457,14 +457,20 @@ steps:
   - group: ":card_index_dividers: Diffusion Test"
     key: nightly-diffusion-test-group
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     steps:
-      - label: ":full_moon: Diffusion · Distributed Test"
+      - label: ":full_moon: Diffusion · Distributed Test with L4"
         timeout_in_minutes: 60
         commands:
-          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
         mirror_hardwares: l4_4
 
+      - label: ":full_moon: Diffusion · Distributed Test with H100"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
+        mirror_hardwares: h100_4
+
       - label: ":full_moon: Diffusion · Offloader Test"
         timeout_in_minutes: 60
         commands:

.buildkite/npu/test-npu-nightly.yml

@@ -95,6 +95,80 @@ steps:
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
             exit $$EXIT
+  - group: ":card_index_dividers: Diffusion I2I Model Test"
+    key: nightly-HunyuanImage3-x2i-test-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
+    steps:
+      - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
+          HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model
+  - group: ":card_index_dividers: Diffusion Benchmark Perf Test"
+    key: nightly-diffusion-benchmark-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
+    steps:
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4"
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
     depends_on: upload-nightly-pipeline
@@ -103,6 +177,14 @@ steps:
       build.pull_request.labels includes "nightly-test" ||
       build.pull_request.labels includes "diffusion-x2v-test"
     steps:
+      - label: ":full_moon: Diffusion X2V · Function Test"
+        mirror_hardwares: a2b3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
       - label: ":full_moon: Diffusion X2V · Perf Test"
         key: nightly-diffusion-x2v-performance
         timeout_in_minutes: 180

README.md

@@ -16,14 +16,10 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
 - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
-- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
-- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
-- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
+- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
-- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
-- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
-- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---
@@ -55,12 +51,13 @@ vLLM-Omni is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
+- Full-duplex realtime serving with streaming audio input and output (experimental)
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
 - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
-- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
 - **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 ## Getting Started

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" \
@@ -44,6 +45,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "hyperpyyaml>=1.2.0" \
       "onnx>=1.16.0" \
       "librosa>=0.10.0" \
+      "torchmetrics>=1.9.0" \
       --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60 && \

docs/README.md

@@ -52,12 +52,13 @@ vLLM-Omni is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
+- Full-duplex realtime serving with streaming audio input and output (experimental)
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
 - TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
-- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
 - Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 For more information, checkout the following:

docs/cli/README.md

@@ -18,18 +18,12 @@ Specify the port:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-For a migrated model, load a custom deploy configuration with `--deploy-config`:
+Load a custom deploy configuration with `--deploy-config`:
 
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
 ```
 
-The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
-```
-
 ## bench
 
 Run benchmark tests for online serving throughput.

docs/cli/serve.md

@@ -8,7 +8,6 @@ The stage-based CLI is designed for deployments that require launching each pipe
 - For **migrated models** that utilize the bundled deployment YAML configurations located in
   `vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
   automatically loads the bundled deployment configuration.
-- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`.
 
 Example: Initializing Stage 0 (Orchestrator and API Server):
 The commands below show a common device mapping where Stage 0 uses GPU 0 and
@@ -32,7 +31,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --omni-master-port 26000
 ```
 
-When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
+When utilizing a custom deployment YAML, append `--deploy-config /path/to/override.yaml` to each command execution.
 
 In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
 However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
+For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/configuration/composable_parallel.md

@@ -34,12 +34,8 @@ Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs`
 | Flag | Description |
 |------|-------------|
 | `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. |
-| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). |
 | `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. |
 
-!!! warning
-    `--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy.
-
 ## Sub-schemas
 
 ### `axis`
@@ -216,17 +212,6 @@ Strategy keys MUST be `model_stage` *names* (strings), the same labels stage con
 !!! warning "Mixing CLI overrides with a strategy on the same axis"
     A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
 
-!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
-    Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
-
-    ```text
-    --strategy-config (<path>) was provided but model '<model>' resolves via the
-    legacy stage_configs YAML path, which does not support composable-parallel
-    strategies; the strategy is ignored. Use a registry-based model to apply it.
-    ```
-
-    Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy.
-
 !!! warning "Conflicting `omni_lb_policy` across stages"
     `omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises:
 

docs/configuration/pd_disaggregation.md

@@ -158,7 +158,7 @@ runtime:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+    --deploy-config /path/to/qwen3_omni_pd.yaml
 ```
 
 ## Operational Notes

docs/configuration/stage_configs.md

@@ -1,9 +1,9 @@
-# Deploy and legacy stage configurations
+# Pipeline and deploy configurations
 
-In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
+In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run.
 
 !!! note
-    Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`.
+    Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when `--deploy-config` is omitted. The model registry resolves the right pipeline and deploy YAML by `model_type`.
 
 ## Deploy configuration schema
 
@@ -86,10 +86,9 @@ stages:
 
 | Flag | Description |
 |------|-------------|
-| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
+| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
 | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
 | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
-| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
 
 ### Stage-Based CLI Paradigm
 
@@ -119,8 +118,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --omni-master-port 26000
 ```
 
-When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
-to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
+When instantiating a custom deployment YAML, append the `--deploy-config /path/to/override.yaml` directive to all node invocations.
 
 In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
 for delineating stage-specific tuning from the CLI interface:
@@ -220,8 +218,7 @@ Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment config
 
 To override specific parameters, explicitly inject the customized configuration schema
 in both online and offline instantiation flows. Use the `--deploy-config` flag
-when loading a deploy configuration, reserving the `--stage-configs-path` parameter
-exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
+when loading a deploy configuration.
 
 Examples:
 
@@ -236,11 +233,6 @@ For online serving:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```
 
-Legacy online serving:
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
 !!! important
     We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
 

docs/contributing/model/adding_diffusion_model.md

@@ -653,7 +653,7 @@ python end2end.py --model your-org/your-model-name --modality text2img --prompts
 Mirror BAGEL’s online serving setup:
 
 - Server launcher: `examples/online_serving/your_model_name/run_server.sh`
-  - Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed)
+  - Wrap `vllm serve ... --omni --port ...` (and `--deploy-config ...` if needed)
 - Client: `examples/online_serving/your_model_name/openai_chat_client.py`
   - Send requests to `POST /v1/chat/completions`
   - Support multimodal inputs (e.g., base64 image) if your model needs it

docs/contributing/model/adding_omni_model.md

@@ -48,8 +48,8 @@ vllm_omni/deploy/
 ```
 
 New in-tree models should register their `PipelineConfig` in
-`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for
-custom legacy `stage_args` YAMLs.
+`vllm_omni/config/pipeline_registry.py`; use `--deploy-config` for custom
+deployment YAMLs.
 
 ## Step-by-Step Implementation
 

docs/contributing/model/adding_tts_model.md

@@ -449,9 +449,8 @@ Key configuration fields:
 !!! note
     New in-tree models should define frozen topology in `models/<model>/pipeline.py`,
     register it in `vllm_omni/config/pipeline_registry.py`, and put deployment
-    knobs in `vllm_omni/deploy/<model>.yaml`. Legacy `stage_args` YAMLs are still
-    accepted for custom configs via `--stage-configs-path`, but should not be used
-    for new bundled defaults.
+    knobs in `vllm_omni/deploy/<model>.yaml`. Use `--deploy-config` to load a
+    custom deployment YAML.
 
 ### Batch mode
 

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
 

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -419,8 +419,7 @@ Notes:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
   --omni \
-  --port 8091 \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+  --port 8091
 ```
 
 #### 3) Key config knobs
@@ -494,7 +493,7 @@ The default config (`qwen3_tts.yaml`) enables the full optimization stack:
 vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
   --omni \
   --port 8000 \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
+  --no-async-chunk
 ```
 
 #### 3) Key config knobs

docs/getting_started/installation/README.md

@@ -1,5 +1,8 @@
 # Installation
 
+!!! important
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
+
 vLLM-Omni supports the following hardware platforms:
 
 - [GPU](gpu.md)

docs/getting_started/installation/gpu/cuda.inc.md

@@ -5,7 +5,7 @@
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends vLLM. So please follow instructions below mainly for vLLM.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.1 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
   -t vllm-omni-cuda .
 ```
 

docs/getting_started/installation/npu/npu.inc.md

@@ -53,16 +53,20 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.18.0
-git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
+# Pin vLLM and vLLM-Ascend to the v0.26 release line
+git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
+cd vllm
 VLLM_TARGET_DEVICE=empty pip install -v -e .
+cd ..
 
-git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
+git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
+cd vllm-ascend
 pip install -v -e .
+cd ..
 
 # Install vLLM-Omni from the latest main branch
 git clone https://github.com/vllm-project/vllm-omni.git
-cd /vllm-workspace/vllm-omni
+cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn

docs/models/supported_models.md

@@ -44,6 +44,7 @@ th {
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
+| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion/attention_backends.md

@@ -15,7 +15,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 | Value | Notes |
 |---|---|
 | `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
-| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
+| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
@@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T
 Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
 `TRTLLM_ATTN` raises rather than silently degrading.
 
+## TRTLLM_ATTN SAGE Quantization
+
+On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
+FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
+in low precision. Enable it through the typed `quant` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
+| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
+| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
+
+V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
+neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
+
+The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
+fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
+`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
+`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
+    ),
+)
+```
+
+
+**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
+availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
+(B300); `int8` QK kernels are compiled for **SM100 only**.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
@@ -193,9 +235,17 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/te
     --model Lightricks/LTX-2 ...
 ```
 
-### FA4 not yet integrated
+### FlashAttention-4 on Blackwell
 
-FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
+Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
+
+```bash
+pip install 'vllm-omni[fa4]'
+```
+
+`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
+when it is unavailable. Version `4.0.0b18` is required; earlier beta
+wheels had known JIT failures on Blackwell.
 
 ## Choosing a Backend Manually
 
@@ -353,64 +403,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t
     --tensor-parallel-size 2 \
     --output outputs/hv15_sage3.mp4
 ```
-
-### Mixed backends across roles
-
-Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --diffusion-attention-config.per_role.self.backend FLASH_ATTN \
-    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_mixed.mp4
-```
-
-### Compare against FlashAttention
-
-Unset the backend override, or explicitly use `FLASH_ATTN`:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 704 --width 1280 --num-frames 49 \
-    --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_fa3.mp4
-```
-
-## Validation Guidance
-
-Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
-
-Always compare:
-
-- End-to-end runtime
-- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
-- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
-
-At minimum, keep the same:
-
-- model
-- prompt
-- seed
-- resolution
-- frame count / step count
-- parallel config (TP / CFG-parallel / Ulysses degrees)
-
-## Reproducing the Benchmark Table
-
-The end-to-end numbers above were collected by running `text_to_video.py` /
-`text_to_image.py` with the same prompt and seed while varying
-`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
-backends without loading a model:
-
-```bash
-python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-```
-
-It runs all three BF16 backends on representative DiT attention shapes and
-prints a ranking table at the end.

docs/user_guide/examples/online_serving/bagel.md

@@ -40,7 +40,7 @@ bash run_server_stage_cli.sh --stage 0
 bash run_server_stage_cli.sh --stage 1
 ```
 
-To use a custom deploy YAML (note: `--stage-configs-path` is deprecated in favor of `--deploy-config`):
+To use a custom deploy YAML:
 
 ```bash
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \

docs/user_guide/examples/online_serving/glm_image.md

@@ -11,8 +11,8 @@ memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
 vllm serve zai-org/GLM-Image --omni --port 8091
 ```
 
-The config system auto-detects the pipeline from the model's `model_index.json` — no
-manual `--stage-configs-path` or `--deploy-config` needed.
+The config system auto-detects the pipeline from the model's `model_index.json`, so
+no manual `--deploy-config` is needed.
 
 By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
 both stages on a single GPU, override per stage:

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -15,9 +15,9 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file, launch the server with command below
+To use a custom deploy config, launch the server with the command below:
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```
 
 ### Send Multi-modal Request
@@ -182,7 +182,7 @@ The script supports the following arguments:
 - `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
 - `--server-port`: Port for vLLM server (default: 8091)
 - `--gradio-port`: Port for Gradio demo (default: 7861)
-- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--deploy-config`: Path to a custom deploy config YAML file (optional)
 - `--server-host`: Host for vLLM server (default: 0.0.0.0)
 - `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
 - `--share`: Share Gradio demo publicly (creates a public link)
@@ -195,9 +195,9 @@ The script supports the following arguments:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file:
+To use a custom deploy config:
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```
 
 **Step 2: Run the Gradio demo**

docs/user_guide/fault_injection_reliability_matrix.md

@@ -69,19 +69,19 @@ The behaviors and conclusions above are summarized from current fault injection
 
 ### Example: `SIGTERM` fault injection log (Qwen3-Omni)
 
-The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
+The following log is based on a real run where a `SIGTERM` signal was sent to the serve root process. Its command line has been normalized to the current `--deploy-config` option.
 
 ```text
-Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
 Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
 OmniServer started successfully
-[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
 [reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
 [reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
 [reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
 [reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
 
-[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
 [0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
 [0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
 [0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal

examples/online_serving/audex/run_server.sh

@@ -53,5 +53,5 @@ vllm-omni serve "$MODEL" \
     --host 0.0.0.0 \
     --port "$PORT" \
     --trust-remote-code \
-    --stage-configs-path "$DEPLOY_YAML" \
+    --deploy-config "$DEPLOY_YAML" \
     --omni

examples/online_serving/chart-helm/templates/_helpers.tpl

@@ -28,8 +28,11 @@ If image.command is set, uses that as a full override.
 - {{ .Values.omniArgs.numGpus | quote }}
 {{- end }}
 {{- if .Values.omniArgs.stageConfigsPath }}
-- "--stage-configs-path"
-- {{ .Values.omniArgs.stageConfigsPath | quote }}
+{{- fail "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath" }}
+{{- end }}
+{{- if .Values.omniArgs.deployConfigPath }}
+- "--deploy-config"
+- {{ .Values.omniArgs.deployConfigPath | quote }}
 {{- end }}
 {{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
 - "--cache-backend"

examples/online_serving/chart-helm/tests/deployment_test.yaml

@@ -78,6 +78,26 @@ tests:
           path: spec.template.spec.containers[0].command
           content: "--enable-layerwise-offload"
 
+  - it: should include deploy config in command
+    set:
+      omniArgs:
+        deployConfigPath: "/configs/deploy.yaml"
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command[8]
+          value: "--deploy-config"
+      - equal:
+          path: spec.template.spec.containers[0].command[9]
+          value: "/configs/deploy.yaml"
+
+  - it: should reject removed stage config value
+    set:
+      omniArgs:
+        stageConfigsPath: "/configs/legacy.yaml"
+    asserts:
+      - failedTemplate:
+          errorMessage: "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath"
+
   - it: should use full command override when image.command is set
     set:
       image:

examples/online_serving/chart-helm/values.yaml

@@ -67,8 +67,8 @@ omniArgs:
   enableCpuOffload: false
   # -- Number of GPUs for diffusion inference (empty = auto)
   numGpus:
-  # -- Path to stage configs file (empty = auto-detected from model)
-  stageConfigsPath:
+  # -- Path to deploy config file (empty = auto-detected from model)
+  deployConfigPath:
   # -- Cache backend: none, tea_cache, or cache_dit
   cacheBackend: "none"
   # -- Default sampling params as a JSON string

examples/online_serving/minicpmo/README.md

@@ -23,9 +23,11 @@ including `librosa`.
 
 The deploy config auto-loads via `--omni`.
 The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on
-logical device 0 with memory budgets of 55%, 22%, and 22%. The profile admits
-at most four sequences per stage and bounds startup video profiling to 32
-frames per video. For throughput,`minicpmo_4_5_2gpu.yaml` gives the Thinker
+logical device 0 with memory budgets of 55%, 15%, and 15%, leaving headroom
+for runtime kernels such as the HiFi-GAN vocoder's cuDNN workspace. The
+profile admits at most four sequences per stage and bounds startup video
+profiling to 32 frames per video. For throughput,
+`minicpmo_4_5_2gpu.yaml` gives the Thinker
 GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
 profile admits at most four concurrent sequences per stage.
 

examples/online_serving/minicpmo/run_gradio_demo.sh

@@ -3,7 +3,7 @@
 #
 # Prereq:
 #   Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
-#   8x4090 stage config under vllm_omni/model_executor/stage_configs).
+#   8x4090 deploy config at vllm_omni/deploy/minicpmo_4_5_8x4090.yaml).
 set -e
 
 HERE="$(cd "$(dirname "$0")" && pwd)"

examples/online_serving/qwen2_5_omni/README.md

@@ -12,9 +12,9 @@ Please refer to [README.md](../../../README.md)
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file, launch the server with command below
+To use a custom deploy config, launch the server with the command below:
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```
 
 ### Send Multi-modal Request
@@ -173,7 +173,7 @@ The script supports the following arguments:
 - `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
 - `--server-port`: Port for vLLM server (default: 8091)
 - `--gradio-port`: Port for Gradio demo (default: 7861)
-- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--deploy-config`: Path to a custom deploy config YAML file (optional)
 - `--server-host`: Host for vLLM server (default: 0.0.0.0)
 - `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
 - `--share`: Share Gradio demo publicly (creates a public link)
@@ -186,9 +186,9 @@ The script supports the following arguments:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file:
+To use a custom deploy config:
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```
 
 **Step 2: Run the Gradio demo**

examples/online_serving/qwen2_5_omni/run_gradio_demo.sh

@@ -13,7 +13,7 @@ set -e
 MODEL="Qwen/Qwen2.5-Omni-7B"
 SERVER_PORT=8091
 GRADIO_PORT=7861
-STAGE_CONFIGS_PATH=""
+DEPLOY_CONFIG_PATH=""
 SERVER_HOST="0.0.0.0"
 GRADIO_IP="127.0.0.1"
 GRADIO_SHARE=false
@@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do
             GRADIO_PORT="$2"
             shift 2
             ;;
-        --stage-configs-path)
-            STAGE_CONFIGS_PATH="$2"
+        --deploy-config)
+            DEPLOY_CONFIG_PATH="$2"
             shift 2
             ;;
         --server-host)
@@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do
             echo "  --model MODEL                 Model name/path (default: Qwen/Qwen2.5-Omni-7B)"
             echo "  --server-port PORT            Port for vLLM server (default: 8091)"
             echo "  --gradio-port PORT            Port for Gradio demo (default: 7861)"
-            echo "  --stage-configs-path PATH     Path to custom stage configs YAML file (optional)"
+            echo "  --deploy-config PATH          Path to custom deploy config YAML file (optional)"
             echo "  --server-host HOST            Host for vLLM server (default: 0.0.0.0)"
             echo "  --gradio-ip IP                IP for Gradio demo (default: 127.0.0.1)"
             echo "  --share                       Share Gradio demo publicly"
@@ -87,8 +87,8 @@ echo "=========================================="
 
 # Build vLLM server command
 SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST")
-if [ -n "$STAGE_CONFIGS_PATH" ]; then
-    SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
+if [ -n "$DEPLOY_CONFIG_PATH" ]; then
+    SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG_PATH")
 fi
 
 # Function to cleanup on exit

examples/online_serving/text_to_speech/README.md

@@ -336,7 +336,7 @@ vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
 # or:
 ./moss_tts_nano/run_server.sh
 ```
-The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
+The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
 
 ### Sending requests
 ```bash

pyproject.toml

@@ -75,6 +75,7 @@ dev = [
 ]
 
 demo = [
+    "gradio>=6.7.0",
     "opencv-python>=4.12.0.88",  # cli/stream_client frame reading
     "requests>=2.28.0",  # cli/stream_client HTTP
     "cn2an>=0.5.22",  # Chinese number normalization for TTS (e.g. "1次" -> "一次")
@@ -95,6 +96,12 @@ quack = [
     "quack-kernels>=0.3.11",
 ]
 
+# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because
+# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific.
+fa4 = [
+    "flash-attn-4[cu13]==4.0.0b18",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -0,0 +1,162 @@
+# MiniMax H3 on Ascend NPU
+
+> Joint video and audio generation with text, first-frame, image/audio, or
+> multi-video conditions — Ascend NPU deployment guide
+
+## Summary
+
+- Vendor: MiniMaxAI
+- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
+- Tasks: T2VA, FL2VA, and Ref2VA
+- Mode: OpenAI-compatible `/v1/videos` HTTP serving
+- Hardware: Atlas 800I A3
+- Maintainer: Community
+
+This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU
+environments. Differences from the GPU path:
+
+- Runs on a CPU-only (aarch64) PyTorch build with `torch_npu`; no CUDA
+  runtime is present.
+- Audio loading does **not** require TorchCodec (whose aarch64 wheels are
+  built against CUDA torch and fail to load on CPU-only builds). vLLM-Omni
+  automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio
+  inputs.
+
+## Prerequisites
+
+### Checkpoint
+
+Same as the GPU recipe — Hugging Face access approval is required:
+
+```bash
+hf auth login
+export MODEL_ROOT=/path/to/MiniMax-H3
+hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+```
+
+### Environment
+
+- Ascend driver & firmware: 25.5.1
+- CANN toolkit: 9.0.1
+- Python: 3.12
+- PyTorch: 2.10.0+cpu
+- torch_npu: 2.10.0.post2
+- Install vLLM-Omni from a checkout with MiniMax H3 support:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e .
+```
+
+Install the **mindie-sd** operator library to enable Ascend-optimized fused
+operators (`adalayernorm`, etc.):
+
+```bash
+git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
+
+# Comment out the tik_ops build step (not needed for this use case)
+sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
+
+python setup.py bdist_wheel
+cd dist
+pip install mindiesd-*.whl
+```
+
+- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
+  reference-video preparation and MP4 output.
+- Reference-video decoding uses `decord` when available and falls back to
+  PyAV otherwise.
+- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded
+  at native sample rate through the soundfile / ffmpeg fallback.
+
+## Start a server
+
+One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
+and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
+
+### Multi-NPU: 768P validated configuration
+
+Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
+parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
+VAE patch parallelism degree 8, and layerwise offload:
+
+```bash
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export PORT=9098
+export MODEL="${MODEL_ROOT}/FL2VA"
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
+export PYTHONDONTWRITEBYTECODE=1
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 8 \
+  --usp 8 \
+  --ring 1 \
+  --text-encoder-tp-size 8 \
+  --enable-layerwise-offload \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --vae-patch-parallel-size 8 \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+Do not add `--enforce-eager`. The first request includes regional
+compilation; warm the server once before measuring steady-state latency.
+H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
+
+To serve Ref2VA, stop the FL2VA server and restart with:
+
+```bash
+export MODEL="${MODEL_ROOT}/Ref2VA"
+```
+
+## HTTP API examples
+
+The request format is identical to the GPU recipe; see
+[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples).
+Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU.
+
+## Key parameters
+
+Same as the GPU recipe; see
+[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters).
+The validated resolution on NPU is 768P (e.g. 1344x768).
+
+## Validated NPU evidence
+
+Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1,
+PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU
+configuration above:
+
+| Workload | Configuration |
+|----------|---------------|
+| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
+
+These measurements describe the validated shapes rather than a general
+throughput guarantee.
+
+## Known limitations
+
+- Each server process loads only one checkpoint partition.
+- H3 currently executes one generation request per diffusion batch.
+- The first regional-compile request is a warmup and should not be included
+  in steady-state performance measurements.
+- Image+audio Ref2VA accepts exactly one image and one audio reference.
+- Video Ref2VA accepts one or more video files, but not an additional
+  standalone audio reference.
+- VAE patch parallelism requires size 1 or the full DiT group size and
+  supports the H3 native `tile` mode only.
+
+## Additional resources
+
+- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)

recipes/MiniMaxAI/MiniMax-H3.md

@@ -0,0 +1,329 @@
+# MiniMax H3
+
+> Joint video and audio generation with text, first-frame, image/audio, or
+> multi-video conditions
+
+## Summary
+
+- Vendor: MiniMaxAI
+- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
+- Tasks: T2VA, FL2VA, and Ref2VA
+- Mode: OpenAI-compatible `/v1/videos` HTTP serving
+- Maintainer: Community
+
+MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
+checkpoint has two independently served partitions:
+
+- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
+  (`fl2va`)
+- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
+  (`ref2va`)
+
+The generated MP4 contains H.264 video and synchronized stereo audio.
+
+## Prerequisites
+
+The checkpoint requires Hugging Face access approval. Download it and point
+`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
+subdirectories:
+
+```bash
+hf auth login
+export MODEL_ROOT=/path/to/MiniMax-H3
+hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
+```
+
+Install vLLM-Omni from the checkout containing MiniMax H3 support. On
+Blackwell, install the optional FlashAttention-4 dependency:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e '.[fa4]'
+```
+
+`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
+reference-video preparation and MP4 output.
+
+## Start a server
+
+One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
+and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
+
+### Single GPU: accuracy and memory first
+
+The single-GPU configuration uses model-level CPU offload.
+This matches the accuracy-qualified reference path and prevents the Qwen3-VL
+encoder and DiT from being resident on the GPU at the same time.
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+export PORT=8091
+
+CUDA_VISIBLE_DEVICES=0 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 1 \
+  --enable-cpu-offload \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+Use a GPU with enough memory for the active H3 component and enough system RAM
+for the offloaded components. CPU offload reduces GPU memory pressure but adds
+PCIe/NVLink transfer latency.
+
+### Four GPUs: measured best-practice throughput
+
+The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
+
+- no CPU or layerwise offload;
+- Ulysses sequence parallelism degree 4;
+- native tiled VAE patch parallelism degree 4;
+- regional `torch.compile` for the repeated DiT blocks;
+- FlashAttention, with Ring and TP left at 1.
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+export PORT=8091
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 4 \
+  --usp 4 \
+  --ring 1 \
+  --vae-patch-parallel-size 4 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+Do not add `--enforce-eager` to this performance configuration. The first
+request includes regional compilation; warm the server once before measuring
+steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
+1. The H3 VAE supports its native `tile` mode, not
+`spatial_shard_height` or `spatial_shard_width`.
+
+### Text encoder tensor parallelism
+
+The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
+default fully resident on the DiT main rank.  On multi-GPU no-offload runs that
+rank becomes the peak-memory hotspot.  Add `--text-encoder-tp-size N` to shard
+the encoder across the first `N` DiT ranks (the encoder is implemented with
+vLLM-style tensor-parallel layers and runs with distributed collectives over
+its own encoder process group):
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+export PORT=8091
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 4 \
+  --usp 4 \
+  --ring 1 \
+  --text-encoder-tp-size 4 \
+  --vae-patch-parallel-size 4 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --diffusion-attention-backend FLASH_ATTN
+```
+
+`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
+valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs).  The
+encoder TP rank set is the first `N` DiT ranks; on 4 GPUs
+`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main
+rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the
+other ranks each gain `~51.5/N` GB.  The encoder output remains identical to
+the reference path within bf16 rounding: every encoder rank all-reduces the
+row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
+replicated on every rank.
+
+To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
+command with:
+
+```bash
+export MODEL="${MODEL_ROOT}/Ref2VA"
+```
+
+## HTTP API examples
+
+The following requests use the synchronous endpoint so the returned body can
+be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can
+also be used when job polling is preferred.
+
+All four tasks use 24 FPS, 50 sigma points, seed values from the validated
+workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3.
+Decimal durations are passed through `extra_params`.
+
+Set the endpoint once:
+
+```bash
+export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
+```
+
+### 1. T2VA: text to video and audio
+
+Run this request against the `FL2VA` partition:
+
+```bash
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=1101' \
+  -F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \
+  -o t2va.mp4
+```
+
+### 2. FL2VA: first frame to video and audio
+
+Run this request against the `FL2VA` partition. When width and height are
+omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
+edge.
+
+```bash
+export FIRST_FRAME=/path/to/fl2va_first_frame.png
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=2101' \
+  -F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \
+  -F "input_reference=@${FIRST_FRAME};type=image/png" \
+  -o fl2va.mp4
+```
+
+### 3. Ref2VA: image and audio reference
+
+Run this request against the `Ref2VA` partition. `audio_reference` accepts an
+HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference
+assets to the serving host:
+
+```bash
+python -m http.server 8092 \
+  --bind 127.0.0.1 \
+  --directory /path/to/reference_assets
+```
+
+Then submit the image and audio request from another terminal:
+
+```bash
+export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
+export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=3101' \
+  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
+  -F "input_reference=@${REF_IMAGE};type=image/png" \
+  -F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \
+  -o ref2va_image_audio.mp4
+```
+
+The requested duration should cover the complete audio. If `duration` is
+shorter, the reference soundtrack is truncated to the generated clip.
+
+### 4. Ref2VA: two video references
+
+Run this request against the `Ref2VA` partition. Repeat the
+`input_references` multipart field once per source video. H3 consumes the
+videos in form order and preserves their original soundtracks during
+conditioning.
+
+```bash
+export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4
+export BACKGROUND_VIDEO=/path/to/fairytale_background.mov
+
+curl -sS -X POST "${API_URL}" \
+  -F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \
+  -F 'width=1344' \
+  -F 'height=768' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=50' \
+  -F 'flow_shift=12' \
+  -F 'seed=3101' \
+  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
+  -F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \
+  -F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \
+  -o ref2va_video_video.mp4
+```
+
+The server stores uploaded multi-video references only for the lifetime of the
+request and deletes the temporary files after generation. Video Ref2VA uses
+the source-video soundtracks and does not accept a separate
+`audio_reference`.
+
+## Key parameters
+
+| Parameter | Recommended value | Notes |
+|-----------|-------------------|-------|
+| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
+| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
+| `fps` | `24` | H3 output FPS is fixed |
+| `num_inference_steps` | `50` | Matches the reference accuracy workloads |
+| `flow_shift` | `12` | Video sigma shift |
+| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
+| `seed` | Task-specific | Use a fixed value for reproducibility |
+| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
+
+## Validated four-GPU evidence
+
+The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
+excluded warmup followed by three requests.
+
+| Workload | Configuration | Observed result |
+|----------|---------------|-----------------|
+| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
+| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
+
+These measurements describe the validated shapes rather than a general
+throughput guarantee. Multi-video Ref2VA is much slower because the two
+reference videos expand both the Qwen3-VL vision sequence and the packed DiT
+attention sequence.
+
+## Known limitations
+
+- Each server process loads only one checkpoint partition.
+- H3 currently executes one generation request per diffusion batch.
+- The first regional-compile request is a warmup and should not be included in
+  steady-state performance measurements.
+- Image+audio Ref2VA accepts exactly one image and one audio reference.
+- Video Ref2VA accepts one or more video files, but not an additional standalone
+  audio reference.
+- VAE patch parallelism requires size 1 or the full DiT group size and supports
+  the H3 native `tile` mode only.
+
+## Additional resources
+
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)
+- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -77,7 +77,9 @@ a fallback because it would duplicate state machines and correctness paths.
 The default
 [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
 co-locates Thinker, codec-only Talker, and Code2Wav on GPU 0. Their
-`gpu_memory_utilization` budgets are 0.55, 0.22, and 0.22. All stages admit
+`gpu_memory_utilization` budgets are 0.55, 0.15, and 0.15. The remaining
+device memory is left available for runtime workspaces such as the HiFi-GAN
+vocoder's cuDNN kernels. All stages admit
 up to four sequences. Startup video profiling is bounded to 32 frames per
 video. Use the two-GPU profile for production throughput.
 
@@ -202,9 +204,10 @@ speech output (TTS)"** checkbox on / off.
 
 #### Notes
 
-- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.22, and
-  0.22 of GPU 0. The larger Thinker share protects its multimodal KV cache;
-  all three model processes still share one CUDA device.
+- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and
+  0.15 of GPU 0. The larger Thinker share protects its multimodal KV cache,
+  while the unreserved memory remains available to runtime kernels; all three
+  model processes still share one CUDA device.
 - `--trust-remote-code` is required — the HF repo ships a custom
   `MiniCPMO` config / model class.
 - Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains

recipes/README.md

@@ -45,6 +45,7 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
+| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |

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/config/test_omni_config.py

@@ -440,6 +440,7 @@ def test_sub_config_fields_match_rfc_scopes():
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
+        "text_encoder_tp_size",
         "vae_parallel_mode",
         "use_hsdp",
         "mask_sp_padding",

tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py

@@ -0,0 +1,74 @@
+"""Regression tests for vllm-project/vllm-omni#5349 (P1): normal completion
+goes through _free_request(), not finish_requests() (the external
+abort/cancel entry point). Without a cleanup_receiver() call there,
+_active_streams never releases on ordinary completion, so the bounded-K
+window fills with stale entries after K completions.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / RequestStatus are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_scheduler(*, chunk_transfer_adapter=None) -> OmniARScheduler:
+    """Minimal OmniARScheduler exercising _free_request()'s no-KV-transfer
+    happy path."""
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._omits_kv_transfer_cache = {}
+    sched._connector_finished = lambda request: (False, None)
+    sched.encoder_cache_manager = MagicMock()
+    sched.finished_req_ids = set()
+    sched._new_prompt_len_snapshot = {}
+    sched.finished_req_ids_dict = None
+    sched._should_transfer_kv_for_request = lambda req_id: False
+    sched._free_blocks = MagicMock()
+    sched._free_input_coordinator_request = MagicMock()
+    sched.chunk_transfer_adapter = chunk_transfer_adapter
+    return sched
+
+
+class _FakeFinishedRequest:
+    """Not a SimpleNamespace: SimpleNamespace defines __eq__, which makes it
+    unhashable, and _free_request() puts the request in a set."""
+
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+
+    def is_finished(self) -> bool:
+        return True
+
+
+def _make_finished_request(request_id: str = "req-1"):
+    return _FakeFinishedRequest(request_id)
+
+
+def test_free_request_releases_chunk_transfer_adapter_receiver_state():
+    adapter = MagicMock()
+    sched = _make_scheduler(chunk_transfer_adapter=adapter)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)
+
+    adapter.cleanup_receiver.assert_called_once_with("req-1")
+
+
+def test_free_request_is_safe_without_a_chunk_transfer_adapter():
+    """Most stages have no chunk transfer adapter configured."""
+    sched = _make_scheduler(chunk_transfer_adapter=None)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)  # must not raise
+
+    sched._free_input_coordinator_request.assert_called_once_with("req-1")

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/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -75,6 +75,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     coordinator = FakeInputCoordinator()
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
+    scheduler.chunk_transfer_adapter = None
     scheduler._omits_kv_transfer_cache = {"req-free": True}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()

tests/dfx/perf/scripts/run_benchmark.py

@@ -1,6 +1,8 @@
 import json
 import os
 import threading
+from collections.abc import Callable
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from pathlib import Path
 from typing import Any
 
@@ -63,35 +65,78 @@ paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_C
 _omni_server_lock = threading.Lock()
 
 
+class _SingleActiveContext:
+    """Reuse one active context while its configuration key is unchanged."""
+
+    def __init__(self) -> None:
+        self._key: Any = None
+        self._stack: ExitStack | None = None
+        self._value: Any = None
+
+    def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any:
+        if self._stack is not None and key == self._key:
+            return self._value
+
+        self.close()
+        stack = ExitStack()
+        value = stack.enter_context(factory())
+        self._key = key
+        self._stack = stack
+        self._value = value
+        return value
+
+    def close(self) -> None:
+        stack = self._stack
+        self._key = None
+        self._stack = None
+        self._value = None
+        if stack is not None:
+            stack.close()
+
+
+@contextmanager
+def _start_omni_server(server_param):
+    test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param
+
+    print(f"Starting OmniServer with test: {test_name}, model: {model}")
+
+    server_args: list[str] = []
+    if use_omni:
+        server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
+    # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
+    # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
+    if stage_config_path:
+        server_args = ["--deploy-config", stage_config_path] + server_args
+    if stage_overrides:
+        server_args = ["--stage-overrides", stage_overrides] + server_args
+    if extra_cli_args:
+        server_args = list(extra_cli_args) + server_args
+    with OmniServer(model, server_args, use_omni=use_omni) as server:
+        server.test_name = test_name
+        print("OmniServer started successfully")
+        yield server
+        print("OmniServer stopping...")
+
+    print("OmniServer stopped")
+
+
 @pytest.fixture(scope="module")
-def omni_server(request):
+def omni_server_context():
     """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
+    Reuse it for adjacent benchmark cases with the same server configuration.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param
+        active_context = _SingleActiveContext()
+        try:
+            yield active_context
+        finally:
+            active_context.close()
 
-        print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args: list[str] = []
-        if use_omni:
-            server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
-        # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
-        # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
-        if stage_config_path:
-            server_args = ["--deploy-config", stage_config_path] + server_args
-        if stage_overrides:
-            server_args = ["--stage-overrides", stage_overrides] + server_args
-        if extra_cli_args:
-            server_args = list(extra_cli_args) + server_args
-        with OmniServer(model, server_args, use_omni=use_omni) as server:
-            server.test_name = test_name
-            print("OmniServer started successfully")
-            yield server
-            print("OmniServer stopping...")
-
-        print("OmniServer stopped")
+@pytest.fixture
+def omni_server(request, omni_server_context):
+    return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))
 
 
 @pytest.fixture

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }
@@ -27,16 +28,9 @@
             "async_chunk": false,
             "trust_remote_code": true,
             "connectors": {
-              "rdma_connector": {
-                "name": "MooncakeTransferEngineConnector",
-                "extra": {
-                  "host": "auto",
-                  "zmq_port": 50051,
-                  "protocol": "rdma",
-                  "device_name": "",
-                  "memory_pool_size": 17179869184,
-                  "memory_pool_device": "cpu"
-                }
+              "shared_mem_connector": {
+                "name": "SharedMemoryConnector",
+                "extra": {}
               }
             },
             "stages": [
@@ -66,7 +60,7 @@
                   "need_send_cache": true
                 },
                 "output_connectors": {
-                  "to_stage_1": "rdma_connector"
+                  "to_stage_1": "shared_mem_connector"
                 },
                 "default_sampling_params": {
                   "temperature": 0.0,
@@ -96,7 +90,7 @@
                 },
                 "vae_use_tiling": true,
                 "input_connectors": {
-                  "from_stage_0": "rdma_connector"
+                  "from_stage_0": "shared_mem_connector"
                 },
                 "default_sampling_params": {
                   "num_inference_steps": 50,

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_runner_metadata.py

@@ -1,6 +1,8 @@
 """Tests for DFX runner metadata field exclusion."""
 
 import json
+from contextlib import contextmanager
+from types import SimpleNamespace
 
 import pytest
 
@@ -254,8 +256,9 @@ def test_benchmark_param_id_suffix_from_task_eval_phase():
     ]
 
 
-def test_create_paired_omni_benchmark_pytest_params(tmp_path):
+def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_path, monkeypatch):
     from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
+    from tests.dfx.perf.scripts import run_benchmark
 
     configs = [
         {
@@ -285,3 +288,51 @@ def test_create_paired_omni_benchmark_pytest_params(tmp_path):
     assert bench_row == ("test_tts", 0)
     assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
     assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)
+
+    omni_p0_server = by_id["test_omni-p0"].values[0]
+    omni_p1_server = by_id["test_omni-p1"].values[0]
+    tts_server = by_id["test_tts-p0"].values[0]
+    assert omni_p0_server == omni_p1_server
+    assert omni_p0_server != tts_server
+
+    events = []
+
+    @contextmanager
+    def fake_start(server_param):
+        events.append(("start", server_param))
+        yield object()
+        events.append(("stop", server_param))
+
+    monkeypatch.setattr(run_benchmark, "_start_omni_server", fake_start)
+    active_context = run_benchmark._SingleActiveContext()
+    try:
+        first = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p0_server),
+            active_context,
+        )
+        second = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p1_server),
+            active_context,
+        )
+        assert first is second
+        assert events == [("start", omni_p0_server)]
+
+        third = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=tts_server),
+            active_context,
+        )
+        assert third is not first
+        assert events == [
+            ("start", omni_p0_server),
+            ("stop", omni_p0_server),
+            ("start", tts_server),
+        ]
+    finally:
+        active_context.close()
+
+    assert events == [
+        ("start", omni_p0_server),
+        ("stop", omni_p0_server),
+        ("start", tts_server),
+        ("stop", tts_server),
+    ]

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -241,12 +241,12 @@ def _fault_keywords_match_response(*, status: int, body: bytes) -> bool:
     return any(key in text for key in body_fault_hints)
 
 
-def _stage_config_path_from_omni_server(omni_server: _HasServeArgs) -> str | None:
+def _deploy_config_path_from_omni_server(omni_server: _HasServeArgs) -> str | None:
     args: list[str] = omni_server.serve_args
     for i, arg in enumerate(args):
-        if arg == "--stage-configs-path" and i + 1 < len(args):
+        if arg == "--deploy-config" and i + 1 < len(args):
             return args[i + 1]
-        if arg.startswith("--stage-configs-path="):
+        if arg.startswith("--deploy-config="):
             return arg.split("=", 1)[1]
     return None
 
@@ -331,7 +331,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
     """
     device_spec = resolve_oom_device_spec(
         OOM_INJECTION_CONFIG,
-        _stage_config_path_from_omni_server(omni_server_function),
+        _deploy_config_path_from_omni_server(omni_server_function),
     )
     handle = inject_gpu_oom(
         device=device_spec,
@@ -394,7 +394,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
 def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
     device_spec = resolve_oom_device_spec(
         OOM_INJECTION_CONFIG,
-        _stage_config_path_from_omni_server(omni_server_function),
+        _deploy_config_path_from_omni_server(omni_server_function),
     )
     handle = inject_gpu_oom(
         device=device_spec,
@@ -437,7 +437,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
 def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
     device_spec = resolve_oom_device_spec(
         OOM_INJECTION_CONFIG,
-        _stage_config_path_from_omni_server(omni_server_function),
+        _deploy_config_path_from_omni_server(omni_server_function),
     )
     handle = inject_gpu_oom(
         device=device_spec,
@@ -791,7 +791,7 @@ def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
     """
     device_spec = resolve_oom_device_spec(
         OOM_RECOVER_INJECTION_CONFIG,
-        _stage_config_path_from_omni_server(omni_server_function),
+        _deploy_config_path_from_omni_server(omni_server_function),
     )
     handle = inject_gpu_oom(
         device=device_spec,

tests/diffusion/attention/test_attention_config.py

@@ -52,6 +52,44 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"):
             AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5})
 
+    def test_quant_serialized_with_defaults_and_overrides(self):
+        assert AttentionSpec(backend="TRTLLM_ATTN", quant={"dtype_qk": "fp8_e4m3"}).backend_kwargs()["quant"] == {
+            "dtype_qk": "fp8_e4m3",
+            "q_block_size": 1,
+            "k_block_size": 16,
+        }
+        assert AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+        ).backend_kwargs()["quant"] == {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+
+    def test_quant_superset_fields_passed_through(self):
+        # FLASHINFER_ATTN-style config (dtype_vo / flashinfer_backend) round-trips on the shared spec.
+        bk = AttentionSpec(
+            backend="FLASHINFER_ATTN",
+            quant={"dtype_qk": "bfloat16", "dtype_vo": "fp8_e4m3", "flashinfer_backend": "trtllm-gen"},
+        ).backend_kwargs()["quant"]
+        assert bk["dtype_qk"] == "bfloat16"
+        assert bk["dtype_vo"] == "fp8_e4m3"
+        assert bk["flashinfer_backend"] == "trtllm-gen"
+
+    def test_quant_and_skip_softmax_coexist(self):
+        bk = AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
+        ).backend_kwargs()
+        assert bk["target_sparsity"] == 0.5 and bk["quant"]["dtype_qk"] == "int8"
+
+    @pytest.mark.parametrize(
+        "spec, match",
+        [
+            ({"backend": "TORCH_SDPA", "quant": {"dtype_qk": "int8"}}, "only supported by the TRTLLM_ATTN"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int4"}}, "quant.dtype_qk"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int8", "k_block_size": 8}}, "quant.k_block_size"),
+        ],
+    )
+    def test_quant_validation_rejects(self, spec, match):
+        with pytest.raises(ValueError, match=match):
+            AttentionSpec(**spec)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):

tests/diffusion/attention/test_flash_attn.py

@@ -20,7 +20,6 @@ from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
-
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813
@@ -330,6 +329,48 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     assert calls[0]["softmax_scale"] == 0.5
 
 
+def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
+    calls = []
+
+    def fake_varlen_func(q, k, v, **kwargs):
+        calls.append((q.shape, k.shape, v.shape, kwargs))
+        return q
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True] * 6 + [False] * 2]),
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert len(calls) == 1
+    assert calls[0][0] == torch.Size([8, 2, 4])
+    assert calls[0][3]["cu_seqlens_q"] is cu_seqlens
+    assert calls[0][3]["max_seqlen_q"] == 6
+
+
+def test_packed_varlen_metadata_must_be_complete(monkeypatch):
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 8, 2, 4)
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_ring_fa4.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.ring import ring_kernels
+from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
+from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch):
+    expected_out = object()
+    expected_lse = object()
+    captured = {}
+
+    def fake_fa4(q, k, v, **kwargs):
+        captured.update(q=q, k=k, v=v, **kwargs)
+        return expected_out, expected_lse
+
+    monkeypatch.setattr(ring_kernels, "HAS_FA4", True)
+    monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4)
+
+    out, lse = ring_kernels.flash_attn4_func_forward(
+        "q",
+        "k",
+        "v",
+        softmax_scale=0.125,
+        causal=True,
+        window_size=(-1, -1),
+    )
+
+    assert (out, lse) == (expected_out, expected_lse)
+    assert captured == {
+        "q": "q",
+        "k": "k",
+        "v": "v",
+        "softmax_scale": 0.125,
+        "causal": True,
+        "window_size": (None, None),
+        "softcap": 0.0,
+        "return_lse": True,
+    }
+
+
+def test_ring_prefers_fa4_on_blackwell(monkeypatch):
+    captured = {}
+
+    def fake_ring_flash_attn_func(*args, **kwargs):
+        captured.update(kwargs)
+        return "output"
+
+    monkeypatch.setattr(ring_parallel, "HAS_FA4", True)
+    monkeypatch.setattr(ring_parallel, "HAS_FA3", True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+
+    from vllm_omni.diffusion.attention.backends import ring_flash_attn
+
+    monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func)
+
+    query = SimpleNamespace(
+        dtype=torch.bfloat16,
+        shape=(1, 16, 4, 128),
+        is_cuda=True,
+        device=torch.device("cuda:0"),
+    )
+    strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object()))
+
+    assert strategy.run_attention(query, "key", "value", None) == "output"
+    assert captured["attn_type"] is AttnType.FA4

tests/diffusion/attention/test_trtllm_attn.py

@@ -47,6 +47,44 @@ def test_skip_config_pure_resolution():
     assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096)
 
 
+def test_quant_config_from_backend_kwargs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    assert not QuantConfig.from_backend_kwargs(None).enabled
+    assert not QuantConfig.from_backend_kwargs({}).enabled
+    assert QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "fp8_e4m3"}}) == QuantConfig("fp8_e4m3", 1, 16)
+    qc = QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}})
+    assert qc.enabled and (qc.dtype_qk, qc.q_block_size, qc.k_block_size) == ("int8", 4, 16)
+
+
+def test_quant_rejects_non_sage_dtype():
+    with pytest.raises(RuntimeError, match="supports dtype_qk"):
+        _impl(quant={"dtype_qk": "bfloat16"})
+
+
+def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
+    import vllm_omni.diffusion.attention.backends.trtllm_attn as mod
+
+    monkeypatch.setattr(mod, "_sage_kernel_available", lambda: True)
+    monkeypatch.setattr(mod, "_sage_quantize_fn", lambda: None)
+    with pytest.raises(RuntimeError, match="trtllm_sage_attention_quantize"):
+        _impl(quant={"dtype_qk": "int8"})
+
+
+def test_quant_quantize_calls_routine_and_shapes_sfs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    captured = {}
+
+    def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
+        captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
+        return "qq", "kq", "vq", "qsfs", "ksfs", "vsfs"
+
+    q_q, k_q, v_q, sfs, blk = QuantConfig(dtype_qk="int8").quantize(object(), object(), object(), fake_quantize)
+    assert (q_q, k_q, v_q, sfs, blk) == ("qq", "kq", "vq", ("qsfs", "ksfs", None, "vsfs"), (1, 16, 0, 1))
+    assert captured == {"q_block_size": 1, "k_block_size": 16, "qk_quant_dtype": torch.int8}
+
+
 def test_skip_factor_none_without_curve():
     assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
 
@@ -257,6 +295,32 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+requires_sage = pytest.mark.skipif(
+    not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
+    reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
+)
+
+
+@requires_sage
+@pytest.mark.parametrize("dtype_qk", ["fp8_e4m3", "int8"])
+def test_sage_quant_matches_sdpa(dtype_qk):
+    torch.manual_seed(0)
+    B, S, H, D = 2, 256, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    impl = _impl(quant={"dtype_qk": dtype_qk, "q_block_size": 1, "k_block_size": 16})
+    try:
+        out = impl.forward_cuda(q, k, v, None)
+    except RuntimeError as e:  # pragma: no cover - arch-dependent (int8 is SM100-only)
+        if "Missing TRTLLM-GEN kernel" in str(e):
+            pytest.skip(f"no trtllm-gen SAGE kernel for dtype_qk={dtype_qk} on this GPU arch")
+        raise
+    assert out.shape == (B, S, H, D) and torch.isfinite(out).all()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)

tests/diffusion/cache/test_magcache_hook.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for the MagCache hook forward numerics.
+
+The cached MagCache residual is the residual of the whole block stack
+(tail output - head input). On a skipped step the head hook applies it
+once; every other hooked block must pass its inputs through unchanged,
+matching the reference implementation in diffusers.hooks.mag_cache.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    _MAG_CACHE_LEADER_BLOCK_HOOK,
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.hooks.base import HookRegistry
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _CountingBlock(nn.Module):
+    """Toy dual-stream block: adds a constant delta to hidden_states.
+
+    Returns (encoder_hidden_states, hidden_states) to match the default
+    TransformerBlockMetadata(return_hidden_states_index=1,
+    return_encoder_hidden_states_index=0) registered by the hook.
+    """
+
+    def __init__(self, delta: float):
+        super().__init__()
+        self.delta = delta
+        self.forward_calls = 0
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        self.forward_calls += 1
+        return encoder_hidden_states, hidden_states + self.delta
+
+
+class _ToyTransformer(nn.Module):
+    def __init__(self, n_blocks: int, delta: float = 1.0):
+        super().__init__()
+        self.blocks = nn.ModuleList(_CountingBlock(delta) for _ in range(n_blocks))
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        for block in self.blocks:
+            encoder_hidden_states, hidden_states = block(hidden_states, encoder_hidden_states)
+        return hidden_states
+
+    @property
+    def total_forward_calls(self) -> int:
+        return sum(block.forward_calls for block in self.blocks)
+
+
+def _make_config(num_inference_steps: int = 2) -> MagCacheConfig:
+    # retention_ratio=0.5 with 2 steps -> retention_step = 1: step 0 always
+    # computes, step 1 with mag_ratio 1.0 accumulates zero error, which is a
+    # guaranteed cache hit.
+    return MagCacheConfig(
+        transformer_type="Toy",
+        threshold=0.24,
+        max_skip_steps=5,
+        retention_ratio=0.5,
+        num_inference_steps=num_inference_steps,
+        mag_ratios=torch.ones(num_inference_steps),
+    )
+
+
+def test_skipped_step_applies_full_stack_residual_once():
+    n_blocks, delta = 3, 1.0
+    model = _ToyTransformer(n_blocks=n_blocks, delta=delta)
+    apply_mag_cache_hook(model, _make_config())
+
+    x = torch.zeros(1, 4, 8)
+    e = torch.zeros(1, 2, 8)
+
+    # Step 0 computes and caches residual = tail_output - head_input.
+    out_computed = model(x, e)
+    assert torch.allclose(out_computed, x + n_blocks * delta)
+    assert model.total_forward_calls == n_blocks
+
+    # Step 1 is a cache hit: no block computes, and the cached residual is
+    # applied exactly once, so the output matches the computed step.
+    out_skipped = model(x, e)
+    assert model.total_forward_calls == n_blocks, "step 1 was not skipped"
+    assert torch.allclose(out_skipped, out_computed), (
+        f"cached residual over-applied on skip: expected "
+        f"{out_computed.flatten()[0].item()}, got {out_skipped.flatten()[0].item()}"
+    )
+
+
+def test_single_block_advances_schedule_on_skip():
+    model = _ToyTransformer(n_blocks=1)
+    apply_mag_cache_hook(model, _make_config(num_inference_steps=2))
+
+    x = torch.zeros(1, 4, 8)
+    e = torch.zeros(1, 2, 8)
+
+    model(x, e)  # step 0: computed
+    assert model.total_forward_calls == 1
+    model(x, e)  # step 1: cache hit
+    assert model.total_forward_calls == 1, "step 1 was not skipped"
+
+    # The head hook is also the tail here, so a skipped step must still
+    # advance the schedule; step 1 was the last step, which resets the
+    # state for the next inference run.
+    registry = HookRegistry.check_if_exists_or_initialize(model.blocks[0])
+    hook = registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
+    state = hook.state_manager.get_state()
+    assert state.step_index == 0
+    assert state._is_first_step
+    assert state.previous_residual is None
+
+    # A new run therefore starts by computing, not by replaying the stale
+    # residual from the previous run.
+    model(x, e)
+    assert model.total_forward_calls == 2, "new run reused stale skip state"

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -39,7 +39,7 @@ MODEL_CASES = [
             "needs_image": False,
             "is_moe": False,
             "mean_threshold": 3e-2,
-            "p99_threshold": 3e-2,
+            "p99_threshold": 4e-2,
         },
         id="wan22_t2v_a14b",
     ),
@@ -146,7 +146,7 @@ def _run_generate(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4})
 @pytest.mark.parametrize("model_case", MODEL_CASES)
 def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
     if current_omni_platform.is_npu():

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -16,7 +16,6 @@ from PIL import Image
 from torch import nn
 
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -858,10 +857,6 @@ def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos
     assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls)
 
 
-@pytest.mark.skipif(
-    current_omni_platform.is_rocm(),
-    reason="Pytorch shipped with vLLM does not built with LAPACK support",
-)
 def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None:
     from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
         FlowUniPCMultistepScheduler,

tests/diffusion/models/minimax_h3/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_pipeline_import_registry_and_component_discovery():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+    )
+
+    assert _DIFFUSION_MODELS["MiniMaxH3Pipeline"] == (
+        "minimax_h3",
+        "pipeline_minimax_h3",
+        "MiniMaxH3Pipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
+    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
+    assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
+    assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
+
+
+def test_joint_postprocess_is_multiprocessing_picklable():
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        get_minimax_h3_post_process_func,
+    )
+
+    postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
+    postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
+    video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+    audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
+
+    result = postprocess((video, audio), output_type="np")
+
+    assert isinstance(result["video"], list)
+    assert result["video"][0].shape == (2, 4, 5, 3)
+    np.testing.assert_array_equal(result["audio"], audio.numpy())
+    assert result["audio_sample_rate"] == 32000
+    assert result["fps"] == 24
+
+
+def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    od_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(cfg_parallel_size=2),
+    )
+    with pytest.raises(ValueError, match="CFG-distilled"):
+        MiniMaxH3Pipeline(od_config=od_config)
+
+
+def test_shape_contract_matches_three_reference_tasks():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        MINIMAX_H3_SHAPE_PLANNER,
+        minimax_h3_align_frame_count,
+    )
+
+    assert minimax_h3_align_frame_count(round(8.7 * 24)) == 209
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(209) == 62
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(8.7) == 348
+    assert minimax_h3_align_frame_count(round(5.0 * 24)) == 124
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(124) == 37
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(5.0) == 200
+    assert MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(62) == 209
+
+
+def test_shifted_sigma_schedule_matches_reference_values():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        minimax_h3_time_shift_sigmas,
+    )
+
+    sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0)
+
+    assert sigmas == pytest.approx(
+        [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0],
+        abs=1e-7,
+    )
+
+
+def test_reference_image_resize_contract():
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _reference_image_shape,
+    )
+
+    assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == (
+        2048,
+        2720,
+    )
+    with pytest.raises(ValueError, match="aspect ratio"):
+        _reference_image_shape(Image.new("RGB", (100, 501)))
+
+
+def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values = torch.ones(1, 4)
+    image_grid_thw = torch.tensor([[1, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+
+def test_encoder_forward_forwards_video_inputs():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values_videos = torch.ones(1, 4)
+    video_grid_thw = torch.tensor([[2, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=None,
+        image_grid_thw=None,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+
+def test_reference_video_shape_uses_h3_adapt_shape_policy():
+    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+        _reference_video_shape,
+    )
+
+    assert _reference_video_shape(1280, 720) == (1344, 768)
+    assert _reference_video_shape(3844, 2160) == (1344, 768)
+
+
+def test_text_encoder_stub_constructs_without_group_or_weights():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = MiniMaxH3Qwen3VLEncoder(
+        "/nonexistent/text_encoder",
+        device=torch.device("cpu"),
+        load_model=False,
+        encoder_group=None,
+    )
+    assert not encoder.is_loaded
+    assert encoder.tp_size == 1
+    # The stub has no parameters, so it never contributes to the runner's
+    # strict missing-parameter check on non-encoder ranks.
+    assert list(encoder.named_parameters()) == []
+
+
+def test_no_offload_keeps_text_encoder_resident():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    input_ids = torch.tensor([1, 2])
+
+    actual = pipeline._encode_text_hidden(input_ids, {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
+
+
+def test_model_offload_uses_hooked_text_encoder_call():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+    )
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    vision_kwargs = {"pixel_values": torch.ones(1, 4)}
+
+    actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
+
+    assert actual is expected
+    pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
+    pipeline.text_encoder.load_to_device.assert_not_called()
+
+
+def test_layerwise_offload_releases_text_encoder():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=True,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+
+def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = torch.nn.Linear(1, 1).half()
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("cpu"),
+    )
+
+    assert next(video_vae.parameters()).dtype == torch.float32
+
+
+def test_video_vae_encode_uses_configured_parallel_tiling():
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        MiniMaxH3VideoVAE,
+    )
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.parallel_tiling = True
+            self.encode_calls = []
+
+        def encode_videos(self, frames, *, use_fp16_latent):
+            assert self.parallel_tiling
+            self.encode_calls.append((frames, use_fp16_latent))
+            return [torch.ones(1, 1, 2, 2, 2)]
+
+    video_vae = object.__new__(MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(video_vae)
+    video_vae.model = FakeModel()
+    video_vae.config_dict = {
+        "latent_channels": 1,
+        "latents_mean": [0.0],
+        "latents_std": [1.0],
+    }
+
+    rows, shape = video_vae.encode_video("frames")
+
+    assert video_vae.model.parallel_tiling
+    assert video_vae.model.encode_calls == [("frames", True)]
+    assert rows.shape == (2, 4)
+    assert shape == (2, 2, 2)
+
+
+def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    prepared = [
+        {"prepared_path": "video-1.mp4"},
+        {"prepared_path": "video-2.mp4"},
+    ]
+
+    class FakeVideoVAE:
+        def __init__(self):
+            self.calls = []
+
+        def is_distributed_enabled(self):
+            return True
+
+        def encode_video(self, frames):
+            self.calls.append(frames)
+            index = len(self.calls)
+            return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3)
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = FakeVideoVAE()
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "_dit_rank_world",
+        lambda: ("dit-group", 1, 4),
+    )
+
+    def fake_broadcast_object_list(values, *, src, group, device):
+        assert values == [None]
+        assert (src, group, device) == (0, "dit-group", torch.device("cpu"))
+        values[0] = prepared
+
+    monkeypatch.setattr(
+        pipeline_module.dist,
+        "broadcast_object_list",
+        fake_broadcast_object_list,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "load_video_frames",
+        lambda path: f"frames:{path}",
+    )
+
+    rows, shapes = pipeline._encode_video_conditions(None, count=2)
+
+    assert pipeline.video_vae.calls == [
+        "frames:video-1.mp4",
+        "frames:video-2.mp4",
+    ]
+    torch.testing.assert_close(
+        rows,
+        torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
+    )
+    assert shapes == [(1, 2, 3), (2, 2, 3)]

tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import numpy as np
+import pytest
+
+pytestmark = [
+    pytest.mark.local_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
+
+
+@pytest.mark.skipif(
+    not os.environ.get(MODEL_ENV),
+    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
+)
+def test_minimax_h3_t2va_ulysses8_smoke():
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    if torch.accelerator.device_count() < 8:
+        pytest.skip("MiniMax H3 Ulysses smoke requires eight GPUs")
+
+    engine = Omni(
+        model=os.environ[MODEL_ENV],
+        parallel_config=DiffusionParallelConfig(ulysses_degree=8),
+        trust_remote_code=True,
+        enable_cpu_offload=True,
+        enforce_eager=True,
+    )
+    try:
+        outputs = engine.generate(
+            "A quiet cinematic night scene with matching ambient sound.",
+            OmniDiffusionSamplingParams(
+                height=256,
+                width=448,
+                num_frames=29,
+                fps=24,
+                num_inference_steps=2,
+                seed=42,
+                output_type="np",
+                extra_args={
+                    "task": "t2va",
+                    "flow_shift": 12.0,
+                    "audio_flow_shift": 3.0,
+                },
+            ),
+            use_tqdm=False,
+        )
+    finally:
+        engine.close()
+
+    assert len(outputs) == 1
+    frames = np.asarray(outputs[0].images[0])
+    assert frames.shape == (39, 256, 448, 3)
+    multimodal = outputs[0].multimodal_output
+    assert multimodal is not None
+    assert np.asarray(multimodal["audio"]).shape[1] == 2
+    assert multimodal["audio_sample_rate"] == 32000
+    assert multimodal["fps"] == 24

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_video_patchify_round_trip_preserves_values():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_patchify_video_latent,
+        minimax_h3_unpatchify_video_tokens,
+    )
+
+    latent = torch.arange(2 * 3 * 2 * 4 * 6).reshape(2, 3, 2, 4, 6)
+    rows = minimax_h3_patchify_video_latent(
+        latent,
+        patch_size=(1, 2, 2),
+    )
+    restored = minimax_h3_unpatchify_video_tokens(
+        rows,
+        latent_shape=(2, 2, 3, 3),
+        patch_size=(1, 2, 2),
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_audio_pack_round_trip_preserves_channel_major_order():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_pack_audio_latent,
+        minimax_h3_unpack_audio_tokens,
+    )
+
+    latent = torch.arange(2 * 4 * 5).reshape(2, 4, 5)
+    rows = minimax_h3_pack_audio_latent(latent)
+    restored = minimax_h3_unpack_audio_tokens(
+        rows,
+        audio_t=10,
+        audio_channel=2,
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_t2va_and_fl2va_packing_keep_update_rows_separate():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence,
+    )
+
+    common = dict(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+    )
+    t2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=False,
+    )
+    fl2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=True,
+        keyframe_frame_indices=[0],
+        frame_count=5,
+    )
+
+    assert int(t2va["seq_len"]) == 64
+    assert t2va["img_pos"].numel() == 12
+    assert t2va["update_mask"].all()
+    assert fl2va["img_pos"].numel() == 18
+    assert fl2va["update_mask"].sum().item() == 12
+    assert (~fl2va["update_mask"]).sum().item() == 6
+    assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
+
+
+def test_ref2va_packing_tracks_video_and_audio_update_masks():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+
+    packed = minimax_h3_packed_sequence_ref2va_blocks(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 2},
+        ],
+    )
+
+    assert int(packed["seq_len"]) == 64
+    assert packed["img_pos"].numel() == 16
+    assert packed["update_mask"].sum().item() == 12
+    assert packed["audio_pos"].numel() == 10
+    assert packed["audio_update_mask"].sum().item() == 6
+    assert packed["cu_seqlens"].tolist() == [0, 30, 64]
+
+
+def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
+    from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
+        minimax_h3_audio_cond_noise_aug_rows,
+        minimax_h3_imgvid_cond_noise_aug_rows,
+    )
+
+    image_rows = torch.randn(4, 96)
+    audio_rows = torch.randn(6, 32)
+
+    torch.testing.assert_close(
+        minimax_h3_imgvid_cond_noise_aug_rows(
+            image_rows,
+            condition_shapes=[(1, 4, 4)],
+            target_latent_t=2,
+            imgvid_cond_num_frames=1,
+            seed=42,
+            noise_aug=1.0,
+        ),
+        image_rows,
+    )
+    torch.testing.assert_close(
+        minimax_h3_audio_cond_noise_aug_rows(
+            audio_rows,
+            condition_audio_t=[3],
+            seed=42,
+            noise_aug=1.0,
+        ),
+        audio_rows,
+    )
+
+    first = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    second = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    torch.testing.assert_close(first, second)

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_grouped_qkv_checkpoint_reorder():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        _reorder_grouped_qkv_to_qkv,
+    )
+
+    # Two groups with rows [q, k, v] become [q0, q1, k0, k1, v0, v1].
+    grouped = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+    reordered = _reorder_grouped_qkv_to_qkv(
+        grouped,
+        num_query_groups=2,
+        heads_per_group=1,
+        head_dim=1,
+    )
+
+    assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+
+
+def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
+    from cache_dit import ForwardPattern
+
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    assert MiniMaxH3DiTModel._repeated_blocks == ["MiniMaxH3DiTBlock"]
+    assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
+    assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
+    assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
+    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.token_refiner = nn.Module()
+    model.token_refiner.blocks = nn.ModuleList([nn.Linear(4, 4)])
+    model.final_layer = nn.Linear(4, 4)
+
+    matched = [
+        name
+        for name, module in model.named_modules()
+        if any(condition(name, module) for condition in MiniMaxH3DiTModel._hsdp_shard_conditions)
+    ]
+    assert matched == ["blocks.0", "blocks.1"]
+
+
+def test_packed_attention_is_a_regional_compile_boundary():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
+
+
+@pytest.mark.parametrize(
+    ("tp_size", "message"),
+    [
+        (3, "num_attention_heads"),
+        (5, "num_attention_heads"),
+    ],
+)
+def test_tp_rejects_non_divisible_head_counts(tp_size, message):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    with pytest.raises(ValueError, match=message):
+        model._validate_tp_config(
+            arch=MiniMaxH3DiTArchConfig(),
+            tp_size=tp_size,
+        )
+
+
+def test_tp_accepts_checkpoint_supported_sizes():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    arch = MiniMaxH3DiTArchConfig()
+    for tp_size in (1, 2, 4, 7):
+        model._validate_tp_config(arch=arch, tp_size=tp_size)

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/diffusion/test_omni_diffusion_config_offline_model_path.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OmniDiffusionConfig HF-offline model-path resolution."""
+
+import os
+
+import huggingface_hub
+import pytest
+
+import vllm_omni.diffusion.data as diffusion_data
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_LOCAL_SNAPSHOT = "/root/.cache/huggingface/hub/models--foo--bar/snapshots/deadbeef"
+
+
+@pytest.fixture
+def record_get_model_path(monkeypatch: pytest.MonkeyPatch):
+    calls: list[tuple[str, str | None]] = []
+
+    def fake_get_model_path(model, revision=None):
+        calls.append((model, revision))
+        if os.path.exists(model):
+            return model
+        return _LOCAL_SNAPSHOT
+
+    monkeypatch.setattr(diffusion_data, "get_model_path", fake_get_model_path)
+    return calls
+
+
+class TestOmniDiffusionConfigOfflineModelPath:
+    def test_offline_repo_id_replaced_with_local_path(
+        self, monkeypatch: pytest.MonkeyPatch, record_get_model_path
+    ) -> None:
+        monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+        config = OmniDiffusionConfig(model="foo/bar")
+        assert config.model == _LOCAL_SNAPSHOT
+        assert record_get_model_path == [("foo/bar", None)]
+
+    def test_online_leaves_model_untouched(self, monkeypatch: pytest.MonkeyPatch, record_get_model_path) -> None:
+        monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", False)
+        config = OmniDiffusionConfig(model="foo/bar")
+        assert config.model == "foo/bar"
+        assert record_get_model_path == []
+
+    def test_offline_local_path_unchanged(self, monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
+        monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+        local = str(tmp_path)
+        config = OmniDiffusionConfig(model=local)
+        assert config.model == local
+
+    @pytest.mark.parametrize("model", [None, ""])
+    def test_offline_empty_model_unchanged(
+        self, monkeypatch: pytest.MonkeyPatch, record_get_model_path, model: str | None
+    ) -> None:
+        monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+        config = OmniDiffusionConfig(model=model)
+        assert config.model == model
+        assert record_get_model_path == []

tests/diffusion/utils/test_flow_matching.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.utils import flow_matching
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_safe_linalg_solve_uses_torch_native_path(monkeypatch: pytest.MonkeyPatch):
+    matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32)
+    rhs = torch.tensor([9.0, 8.0], dtype=torch.float32)
+    expected = torch.tensor([2.0, 3.0], dtype=torch.float32)
+
+    def fake_torch_solve(solve_matrix: torch.Tensor, solve_rhs: torch.Tensor) -> torch.Tensor:
+        assert torch.equal(solve_matrix, matrix)
+        assert torch.equal(solve_rhs, rhs)
+        return expected
+
+    def fail_numpy(*_args, **_kwargs):
+        raise AssertionError("NumPy fallback should not run")
+
+    monkeypatch.setattr(flow_matching.torch.linalg, "solve", fake_torch_solve)
+    monkeypatch.setattr(flow_matching.np.linalg, "solve", fail_numpy)
+
+    result = flow_matching.safe_linalg_solve(matrix, rhs)
+
+    assert result is expected
+
+
+def test_safe_linalg_solve_falls_back_on_missing_cpu_lapack(monkeypatch: pytest.MonkeyPatch):
+    matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32)
+    rhs = torch.tensor([9.0, 8.0], dtype=torch.float32)
+    expected = np.array([2.0, 3.0], dtype=np.float32)
+
+    def fail_with_missing_lapack(*_args, **_kwargs):
+        raise RuntimeError(
+            "Calling torch.linalg.lu_factor on a CPU tensor requires compiling PyTorch with LAPACK. "
+            "Please use PyTorch built with LAPACK support."
+        )
+
+    monkeypatch.setattr(flow_matching.torch.linalg, "solve", fail_with_missing_lapack)
+    monkeypatch.setattr(flow_matching.np.linalg, "solve", lambda *_args, **_kwargs: expected)
+
+    result = flow_matching.safe_linalg_solve(matrix, rhs)
+
+    assert torch.equal(result, torch.from_numpy(expected))
+    assert result.dtype == matrix.dtype
+    assert result.device == matrix.device
+
+
+def test_safe_linalg_solve_reraises_unrelated_runtime_error(monkeypatch: pytest.MonkeyPatch):
+    matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32)
+    rhs = torch.tensor([9.0, 8.0], dtype=torch.float32)
+
+    def fail_with_unrelated_error(*_args, **_kwargs):
+        raise RuntimeError("solver exploded for another reason")
+
+    def fail_numpy(*_args, **_kwargs):
+        raise AssertionError("NumPy fallback should not run")
+
+    monkeypatch.setattr(flow_matching.torch.linalg, "solve", fail_with_unrelated_error)
+    monkeypatch.setattr(flow_matching.np.linalg, "solve", fail_numpy)
+
+    with pytest.raises(RuntimeError, match="solver exploded for another reason"):
+        flow_matching.safe_linalg_solve(matrix, rhs)

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,
     )
@@ -192,7 +196,7 @@ def test_load_poll(build_adapter):
     assert request.additional_information == payload
     assert adapter.get_req_chunk["req-1"] == 1
     assert "req-1" in adapter._finished_load_reqs
-    assert "req-1" in adapter.finished_requests
+    assert "req-1" in adapter.upstream_exhausted_requests
     assert "req-1" not in adapter._pending_load_reqs
 
 
@@ -459,7 +463,7 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap
     assert request.additional_information["meta"]["phase"] == "decode"
     assert request.additional_information["kv_metadata"] == {"foo": "bar"}
     assert "req-non-ar" in adapter._finished_load_reqs
-    assert "req-non-ar" in adapter.finished_requests
+    assert "req-non-ar" in adapter.upstream_exhausted_requests
 
 
 def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
@@ -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)
@@ -527,7 +573,7 @@ def test_process_and_restore_queues(build_adapter):
 
 
 def test_fifo_promotion(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, max_num_seqs=2, active_stream_window=2)
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
     reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
     waiting_queue = DummyWaitingQueue(reqs)
     running_queue = []
@@ -540,11 +586,12 @@ def test_fifo_promotion(build_adapter):
     assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
     assert reqs[2].status == RequestStatus.WAITING
 
-    adapter.finished_requests.add("req-1")
+    # Ordinary completion calls cleanup_receiver(), not finish_requests()
+    # (the abort path) -- see test_omni_ar_scheduler_free_request_cleanup.py.
     # Eviction is deferred to postprocess_scheduler_output in the runtime path
-    # (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2).
-    # Simulate it here so promotion can pick up the freed slot.
-    adapter._evict_finished_active_streams({"req-1"})
+    # (commit c4d95fd9 -- otherwise the terminal chunk deadlocks at c=8 K=2).
+    # Simulate the restore step here so promotion can pick up the freed slot.
+    adapter.cleanup_receiver("req-1")
     adapter.restore_queues(waiting_queue, running_queue)
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -590,7 +637,7 @@ def test_legacy_k0(build_adapter):
 
 
 def test_finished_releases_slot(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1)
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=1, active_stream_window=1)
     req_1 = _req("req-1", RequestStatus.WAITING)
     req_2 = _req("req-2", RequestStatus.WAITING)
     waiting_queue = DummyWaitingQueue([req_1, req_2])
@@ -600,10 +647,11 @@ def test_finished_releases_slot(build_adapter):
     assert list(adapter._active_streams) == ["req-1"]
     assert waiting_queue == []
 
-    adapter.finished_requests.add("req-1")
-    # Eviction is deferred to postprocess_scheduler_output in the runtime path
-    # (commit c4d95fd9). Simulate it so promotion can pick up the freed slot.
-    adapter._evict_finished_active_streams({"req-1"})
+    # req-1 finishes ordinarily (see test_fifo_promotion for why this is
+    # cleanup_receiver(), not finish_requests()). Eviction is deferred to
+    # postprocess_scheduler_output in the runtime path (commit c4d95fd9);
+    # simulate the restore step so promotion can pick up the freed slot.
+    adapter.cleanup_receiver("req-1")
     adapter.restore_queues(waiting_queue, running_queue)
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -631,6 +679,79 @@ def test_postprocess_scheduler_output(build_adapter):
     assert adapter.requests_with_ready_chunks == {"leftover"}
 
 
+@pytest.mark.parametrize("model_mode", ["ar", "generation"])
+def test_active_stream_window_stalls_lone_running_request_after_upstream_finishes(build_adapter, model_mode):
+    """Regression test for vllm-project/vllm-omni#5349.
+
+    upstream_exhausted_requests means the upstream sent its terminal chunk,
+    not that this stage's own generation is done. A downstream stage can
+    still have a long decode ahead after its upstream finishes. Evicting on
+    that signal and then permanently denying re-admission wedges the
+    request out of running_queue forever, with no error.
+
+    requests_with_ready_chunks is set up front so the request stays in
+    running_queue this tick instead of legitimately parking in
+    waiting_for_chunk_running_requests (a separate mechanism this test must
+    not conflate with the bug).
+
+    Parametrized over both worker types: the real issue hit an "ar" stage
+    (Qwen3-Omni's talker); must not regress "generation" (Qwen3-TTS's
+    Code2Wav), which this feature was designed for.
+    """
+    adapter, _ = build_adapter(stage_id=1, model_mode=model_mode, max_num_seqs=2, active_stream_window=2)
+    running_req = _req("req-1", RequestStatus.RUNNING)
+    running_queue = [running_req]
+    waiting_queue = DummyWaitingQueue([])
+
+    adapter._active_streams["req-1"] = running_req
+    adapter.requests_with_ready_chunks.add("req-1")
+
+    # Upstream terminal chunk arrives; this stage hasn't finished.
+    adapter.upstream_exhausted_requests.add("req-1")
+    scheduler_output = SimpleNamespace(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=SimpleNamespace(req_ids=["req-1"]),
+    )
+    adapter.postprocess_scheduler_output(scheduler_output)
+    assert not running_req.is_finished()  # sanity: not done
+
+    # No competition for the slot; req-1 must stay schedulable.
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert running_queue == [running_req], "req-1 must still be schedulable -- losing it here is #5349"
+    assert not adapter.waiting_for_chunk_running_requests
+    assert running_req not in adapter._held_non_active
+    assert list(adapter._active_streams) == ["req-1"]
+
+
+def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter):
+    """Companion to test_fifo_promotion: after EACH of the K active streams
+    finishes (not just one), promotion of a new K-sized batch must still
+    work -- the window must not stay exhausted by stale entries. See
+    test_omni_ar_scheduler_free_request_cleanup.py for the scheduler-level
+    proof that ordinary completion calls cleanup_receiver().
+    """
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
+    reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
+    waiting_queue = DummyWaitingQueue(reqs)
+    running_queue = []
+
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+    assert list(adapter._active_streams) == ["req-1", "req-2"]
+
+    adapter.cleanup_receiver("req-1")
+    adapter.cleanup_receiver("req-2")
+    # Eviction is deferred to postprocess_scheduler_output in the runtime
+    # path; simulate the restore step so promotion can pick up the freed
+    # slots (see test_fifo_promotion).
+    adapter.restore_queues(waiting_queue, running_queue)
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert list(adapter._active_streams) == ["req-3", "req-4"], (
+        "both freed slots must be reusable after ordinary completion"
+    )
+
+
 # ---------------------------------------------------------------
 # Cleanup tests
 # ---------------------------------------------------------------
@@ -638,7 +759,7 @@ def test_postprocess_scheduler_output(build_adapter):
 
 def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
     """Fill every per-request structure so cleanup can be verified."""
-    adapter.finished_requests.add(req_id)
+    adapter.upstream_exhausted_requests.add(req_id)
     adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id)
     adapter.get_req_chunk[req_id] = 3
     adapter.requests_with_ready_chunks.add(req_id)
@@ -659,7 +780,7 @@ def test_cleanup_clears_all_state(build_adapter):
 
     adapter.cleanup(req_id, ext_id)
 
-    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.upstream_exhausted_requests
     assert req_id not in adapter._active_streams
     assert req_id not in adapter.get_req_chunk
     assert req_id not in adapter.requests_with_ready_chunks
@@ -712,7 +833,7 @@ def test_cleanup_request_id_reuse_not_polluted(build_adapter):
 
     adapter.cleanup(req_id, ext_id)
 
-    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.upstream_exhausted_requests
     assert req_id not in adapter.get_req_chunk
 
 
@@ -738,7 +859,7 @@ def test_cleanup_only_affects_target_request(build_adapter):
 
     adapter.cleanup("req-a", "ext-a")
 
-    assert "req-b" in adapter.finished_requests
+    assert "req-b" in adapter.upstream_exhausted_requests
     assert "req-b" in adapter.get_req_chunk
     assert "ext-b" in adapter.put_req_chunk
     assert "ext-b" in adapter.request_payload
@@ -761,13 +882,13 @@ def test_cleanup_after_poll_flow(build_adapter):
     connector.get.return_value = (payload, 8)
     adapter._poll_single_request(request)
 
-    assert "req-flow" in adapter.finished_requests
+    assert "req-flow" in adapter.upstream_exhausted_requests
     assert adapter.get_req_chunk["req-flow"] == 1
     assert "req-flow" in adapter.request_ids_mapping
 
     adapter.cleanup("req-flow", "ext-flow")
 
-    assert "req-flow" not in adapter.finished_requests
+    assert "req-flow" not in adapter.upstream_exhausted_requests
     assert "req-flow" not in adapter.get_req_chunk
     assert "req-flow" not in adapter.request_ids_mapping
     assert "ext-flow" not in adapter.request_payload
@@ -801,7 +922,7 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter
     adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
     adapter.waiting_for_chunk_running_requests = deque([other, zombie])
     adapter.requests_with_ready_chunks.add("req-zombie")
-    adapter.finished_requests.add("req-zombie")
+    adapter.upstream_exhausted_requests.add("req-zombie")
     requests_map = {
         "req-zombie": zombie,
         "req-live": other,
@@ -816,7 +937,7 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter
     assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
     assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
     assert "req-zombie" not in adapter.requests_with_ready_chunks
-    assert "req-zombie" not in adapter.finished_requests
+    assert "req-zombie" not in adapter.upstream_exhausted_requests
 
 
 def test_finish_requests_releases_active_stream_slot(build_adapter):
@@ -886,7 +1007,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     cleanup_calls = []
 
     adapter_mock = mocker.MagicMock()
-    adapter_mock.finished_requests = {"req-s1"}
+    adapter_mock.upstream_exhausted_requests = {"req-s1"}
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
 
     from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -1122,7 +1243,7 @@ def test_wire_round_trip_struct_to_dict_contract():
 def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     """Build a mock scheduler with requests queued for deferred finish."""
     adapter_mock = mocker.MagicMock()
-    adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs}
+    adapter_mock.upstream_exhausted_requests = {r.request_id for r in pending_finish_reqs}
     cleanup_calls = []
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
 

tests/distributed/omni_connectors/test_kv_recv_tp_consensus.py

@@ -0,0 +1,257 @@
+"""Pure-TP KV receive consensus (#5627) — CPU-testable mechanism.
+
+In pure TP (LOCAL role) each rank fetches its own KV shard independently and
+``distribute_kv_cache`` is a no-op, so one rank's silent receive timeout used
+to send only that rank into the pipeline's no-KV prefill — desyncing TP
+collectives until the NCCL watchdog killed the worker. These tests cover the
+identical-state verdict exchange: skip guards, both rank roles, KV dropping
+on any cross-rank divergence (primary or CFG companion roles), that the
+exchange runs on the TP group and never on ``world`` (which also holds DP/PP
+ranks working on other requests), and the wiring from both distributed-receive
+entry points.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+    ReceiveRole,
+    _TransferTopoConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeGroup:
+    """Scripted stand-in for one rank's view of a GroupCoordinator."""
+
+    def __init__(self, rank_in_group, world_size, peer_signatures=None, verdict=None):
+        self.rank_in_group = rank_in_group
+        self.world_size = world_size
+        self._peer_signatures = list(peer_signatures or [])
+        self._verdict = verdict
+        self.sent = []
+        self.broadcasts = []
+
+    def recv_object(self, src):
+        return self._peer_signatures.pop(0)
+
+    def send_object(self, obj, dst):
+        self.sent.append((obj, dst))
+
+    def broadcast_object(self, obj=None, src=0):
+        self.broadcasts.append(obj)
+        return obj if obj is not None else self._verdict
+
+    @property
+    def touched(self):
+        return bool(self.sent or self.broadcasts)
+
+
+def _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None, world=None):
+    tp_size = tp_group.world_size if tp_group is not None else 1
+    return _TransferTopoConfig(
+        role=role,
+        tp_active=tp_active,
+        cfg_size=1,
+        cfg_rank=0,
+        cfg_group=None,
+        sp_size=1,
+        sp_rank=0,
+        sp_group=None,
+        world=world,
+        tp_group=tp_group,
+        tp_size=tp_size,
+    )
+
+
+def _manager(topo, need_recv_cache=True):
+    mgr = OmniKVTransferManager(OmniKVCacheConfig(need_recv_cache=need_recv_cache))
+    mgr._topo_config = topo
+    return mgr
+
+
+def _req_with_kv(rid="r1", cfg_text=False):
+    kv = SimpleNamespace(key_cache=[], value_cache=[])
+    sp = SimpleNamespace(past_key_values=kv, kv_metadata={"ropes": [3]})
+    if cfg_text:
+        sp.cfg_text_past_key_values = SimpleNamespace(key_cache=[], value_cache=[])
+    return SimpleNamespace(
+        request_id=rid,
+        past_key_values=kv,
+        kv_metadata={"ropes": [3]},
+        sampling_params=sp,
+    )
+
+
+def _signature_of(req, received=True):
+    return (received, tuple(sorted(OmniKVTransferManager._collect_request_kv_payload(req))))
+
+
+# --------------------------------------------------------------------------- #
+#  Skip guards — no exchange may happen (a one-sided exchange would deadlock)
+# --------------------------------------------------------------------------- #
+
+
+def test_pass_through_outside_pure_tp():
+    group = _FakeGroup(0, 2, peer_signatures=[(False, ())])
+    for topo in (
+        _topo(role=ReceiveRole.LEADER, tp_active=True, tp_group=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=False, tp_group=group),
+        _topo(role=ReceiveRole.LOCAL, tp_active=True, tp_group=None),
+    ):
+        mgr = _manager(topo)
+        assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+        assert not group.touched
+
+
+def test_pass_through_single_rank_group_no_recv_config_and_dummy():
+    group = _FakeGroup(0, 1)
+    assert _manager(_topo(tp_group=group))._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    group2 = _FakeGroup(0, 2, peer_signatures=[(False, ())])
+    mgr = _manager(_topo(tp_group=group2), need_recv_cache=False)
+    assert mgr._tp_local_receive_consensus(_req_with_kv(), True) is True
+
+    mgr3 = _manager(_topo(tp_group=group2))
+    dummy = _req_with_kv(rid=DUMMY_DIFFUSION_REQUEST_ID)
+    assert mgr3._tp_local_receive_consensus(dummy, False) is False
+    assert not group2.touched
+
+
+def test_world_group_is_never_used_for_the_exchange():
+    """DP replicas / PP stages live in ``world`` and work on other requests —
+    the exchange must be scoped to the TP group (review #5636, P1)."""
+    tp_group = _FakeGroup(0, 2, peer_signatures=[(True, ())])
+    world = _FakeGroup(0, 4, peer_signatures=[(False, ()), (False, ()), (False, ())])
+    mgr = _manager(_topo(tp_group=tp_group, world=world))
+    req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace())
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert tp_group.touched
+    assert not world.touched
+
+
+# --------------------------------------------------------------------------- #
+#  Rank 0: collects signatures, decides, broadcasts
+# --------------------------------------------------------------------------- #
+
+
+def test_rank0_drops_kv_when_peer_missed_primary():
+    req = _req_with_kv()
+    peer_sig = (False, ())
+    group = _FakeGroup(0, 2, peer_signatures=[peer_sig])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.broadcasts == [False]  # peers were told to fall back too
+    assert req.past_key_values is None
+    assert req.kv_metadata is None
+    assert req.sampling_params.past_key_values is None
+    assert req.sampling_params.kv_metadata is None
+
+
+def test_rank0_keeps_kv_when_all_ranks_match():
+    req = _req_with_kv()
+    group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.broadcasts == [True]
+    assert req.past_key_values is not None
+    assert req.sampling_params.past_key_values is not None
+
+
+def test_rank0_drops_kv_on_companion_role_divergence():
+    """Primary received everywhere but cfg_text only landed on this rank —
+    divergent branch caches would desync TP collectives (review #5636, P1)."""
+    req = _req_with_kv(cfg_text=True)
+    peer_sig = _signature_of(_req_with_kv(cfg_text=False))  # peer lacks cfg_text
+    group = _FakeGroup(0, 2, peer_signatures=[peer_sig])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.broadcasts == [False]
+    assert req.past_key_values is None
+    assert req.sampling_params.cfg_text_past_key_values is None
+
+
+def test_rank0_keeps_symmetric_partial_state():
+    """cfg_text absent on every rank is consistent — Bagel's empty
+    unconditional branch — and must NOT be cleared or failed."""
+    req = _req_with_kv(cfg_text=False)
+    group = _FakeGroup(0, 2, peer_signatures=[_signature_of(req)])
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert req.past_key_values is not None
+
+
+def test_rank0_own_miss_fails_verdict_even_if_peers_received():
+    req = SimpleNamespace(request_id="r1", sampling_params=SimpleNamespace())
+    group = _FakeGroup(0, 3, peer_signatures=[(True, ("past_key_values",)), (True, ("past_key_values",))])
+    mgr = _manager(_topo(tp_group=group))
+    assert mgr._tp_local_receive_consensus(req, False) is False
+    assert group.broadcasts == [False]
+
+
+# --------------------------------------------------------------------------- #
+#  Rank > 0: reports signature, honors verdict
+# --------------------------------------------------------------------------- #
+
+
+def test_follower_rank_drops_kv_on_negative_verdict():
+    req = _req_with_kv()
+    sig_before_clear = _signature_of(req)
+    group = _FakeGroup(1, 2, verdict=False)
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is False
+    assert group.sent == [(sig_before_clear, 0)]
+    assert req.past_key_values is None
+    assert req.sampling_params.past_key_values is None
+
+
+def test_follower_rank_keeps_kv_on_positive_verdict():
+    req = _req_with_kv()
+    group = _FakeGroup(1, 2, verdict=True)
+    mgr = _manager(_topo(tp_group=group))
+
+    assert mgr._tp_local_receive_consensus(req, True) is True
+    assert group.sent == [(_signature_of(req), 0)]
+    assert req.past_key_values is not None
+
+
+# --------------------------------------------------------------------------- #
+#  Wiring — both distributed-receive entry points run the consensus
+# --------------------------------------------------------------------------- #
+
+
+def _wired_manager(monkeypatch, consensus_calls):
+    mgr = _manager(_topo(tp_group=_FakeGroup(0, 2)))
+
+    def fake_consensus(req, received):
+        consensus_calls.append(received)
+        return False
+
+    monkeypatch.setattr(mgr, "_tp_local_receive_consensus", fake_consensus)
+    monkeypatch.setattr(mgr, "receive_multi_kv_cache", lambda *a, **k: True)
+    return mgr
+
+
+def test_receive_multi_kv_cache_distributed_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.receive_multi_kv_cache_distributed(_req_with_kv()) is False
+    assert calls == [True]
+
+
+def test_consume_and_distribute_kv_cache_runs_consensus(monkeypatch):
+    calls = []
+    mgr = _wired_manager(monkeypatch, calls)
+    assert mgr.consume_and_distribute_kv_cache(_req_with_kv()) is False
+    assert calls == [True]

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -682,6 +682,10 @@ class TestDistributedReceive:
         mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
         req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
         world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
+        # broadcast_value=True: rank 0's verdict says every TP rank has the
+        # same KV state, so the per-rank receive result stands (see #5627
+        # consensus). The exchange runs on the TP group, not world.
+        tp_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, broadcast_value=True)
         mgr.receive_multi_kv_cache = MagicMock(return_value=True)
 
         with (
@@ -704,10 +708,19 @@ class TestDistributedReceive:
                 return_value=0,
             ),
             patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None),
+            patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group),
         ):
             assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
 
         mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
+        # Pure TP still fetches per-rank (no leader→follower distribution),
+        # but each rank now reports its KV-state signature on the TP group
+        # for the identical-state verdict instead of silently diverging on a
+        # miss (#5627). World must stay untouched — it holds DP/PP ranks.
+        assert len(tp_group.send_calls) == 1
+        dst, sig = tp_group.send_calls[0]
+        assert dst == 0 and sig[0] is True
+        assert world_group.send_calls == []
 
     # ── SP-only scenarios ────────────────────────────────────────────
 

tests/e2e/accuracy/helpers.py

@@ -696,7 +696,7 @@ def _dreamzero_omni_server_from_params(params):
 
     server_args = list(params.server_args or [])
     if params.stage_config_path is not None:
-        server_args += ["--stage-configs-path", params.stage_config_path]
+        server_args += ["--deploy-config", params.stage_config_path]
     return OmniServer(
         params.model,
         server_args,

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -359,12 +359,12 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
     return image, cot_text, elapsed
 
 
-def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+def _run_online(deploy_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
 
     server_args = [
-        "--stage-configs-path",
-        stage_configs_path,
+        "--deploy-config",
+        deploy_config_path,
         "--stage-init-timeout",
         "300",
         "--init-timeout",

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -32,10 +32,28 @@ PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
 SSIM_THRESHOLD = 0.97
-PSNR_THRESHOLD = 30.0
+PSNR_THRESHOLD_CUDA = 30.0
+PSNR_THRESHOLD_NPU = 26.0
+
+
+def _psnr_threshold() -> float:
+    from vllm_omni.platforms import current_omni_platform
+
+    return PSNR_THRESHOLD_NPU if current_omni_platform.is_npu() else PSNR_THRESHOLD_CUDA
+
 
-BASELINE_PATH = get_asset_path("hunyuan/hunyuan_baseline.png")
 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
+
+
+def _baseline_path() -> Path:
+    from vllm_omni.platforms import current_omni_platform
+
+    suffix = "_npu" if current_omni_platform.is_npu() else ""
+    path = get_asset_path(f"hunyuan/hunyuan_baseline{suffix}.png")
+    assert path.exists(), f"Baseline image not found at {path}"
+    return path
+
+
 _OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
 
 # DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
@@ -74,6 +92,23 @@ _DEPLOY_CONFIG = {
             },
         },
     ],
+    "platforms": {
+        "npu": {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "gpu_memory_utilization": 0.65,
+                    "devices": "0,1,2,3",
+                    "moe_backend": "auto",
+                    "max_num_batched_tokens": 32768,
+                    "parallel_config": {
+                        "tensor_parallel_size": 4,
+                        "enable_expert_parallel": True,
+                    },
+                },
+            ],
+        },
+    },
 }
 
 
@@ -216,8 +251,8 @@ def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, out
 
 
 def _assert_against_baseline(image: Image.Image, label: str) -> None:
-    assert BASELINE_PATH.exists(), f"Baseline image not found at {BASELINE_PATH}"
-    baseline_image = Image.open(BASELINE_PATH).convert("RGB")
+    baseline_path = _baseline_path()
+    baseline_image = Image.open(baseline_path).convert("RGB")
 
     assert_images_pixel_close(
         model_name=f"{MODEL_NAME} ({label} vs baseline)",
@@ -231,11 +266,11 @@ def _assert_against_baseline(image: Image.Image, label: str) -> None:
         vllm_image=image,
         diffusers_image=baseline_image,
         ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
+        psnr_threshold=_psnr_threshold(),
     )
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
     model = _model_name()
     output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
@@ -249,7 +284,7 @@ def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> N
     _assert_against_baseline(image, "online")
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
     model = _model_name()
     output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)

tests/e2e/features/fullduplex/openai/test_runtime_control_redaction.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Regression tests for duplex ``runtime_control`` payload redaction.
+
+The redacted result is emitted on the WebSocket (for example
+``session.created.runtime_control``), so it must always survive
+``json.dumps``. Control results carry typed values straight from the engine --
+``fence`` and ``accepted_fence`` are ``DuplexFence`` instances added by
+``DuplexControlClient.execute`` -- which previously reached
+``WebSocket.send_json`` unconverted and killed the session at handshake.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from enum import Enum
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.openai.runtime_bridge import NativeRuntimeBridgeMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_redact = NativeRuntimeBridgeMixin._redact_runtime_control_result
+
+
+def _control_result() -> dict[str, object]:
+    """The shape ``DuplexControlClient.execute`` returns for an open."""
+    return {
+        "ok": True,
+        "operation": "open",
+        "session_id": "sess-1",
+        "fence": DuplexFence("sess-1", epoch=0, turn_id=0, response_seq=0, incarnation=0),
+        "accepted_fence": DuplexFence("sess-1", epoch=1, turn_id=2, response_seq=3, incarnation=4),
+        "stage_results": [{"stage_id": -1, "replica_id": -1, "result": {"supported": True}}],
+    }
+
+
+def test_redacted_control_result_is_json_serializable() -> None:
+    """The whole payload must serialize the way send_json does."""
+    redacted = _redact(_control_result())
+    # Must not raise.
+    json.dumps({"type": "session.created", "runtime_control": redacted})
+
+
+def test_fences_are_converted_to_plain_dicts() -> None:
+    redacted = _redact(_control_result())
+    assert redacted["fence"] == {
+        "session_id": "sess-1",
+        "epoch": 0,
+        "turn_id": 0,
+        "response_seq": 0,
+        "incarnation": 0,
+    }
+    # Field values are preserved, not just the type.
+    assert redacted["accepted_fence"]["epoch"] == 1
+    assert redacted["accepted_fence"]["turn_id"] == 2
+    assert redacted["accepted_fence"]["response_seq"] == 3
+    assert redacted["accepted_fence"]["incarnation"] == 4
+
+
+def test_nested_dataclasses_are_converted() -> None:
+    """Fences can appear inside stage results, not only at the top level."""
+    redacted = _redact(
+        {
+            "stage_results": [
+                {"stage_id": 0, "result": {"fence": DuplexFence("sess-2", epoch=7)}},
+            ]
+        }
+    )
+    json.dumps(redacted)
+    assert redacted["stage_results"][0]["result"]["fence"]["epoch"] == 7
+
+
+def test_enums_are_reduced_to_values() -> None:
+    class _Mode(str, Enum):
+        APPEND = "append_audio_chunk"
+
+    redacted = _redact({"mode": _Mode.APPEND})
+    json.dumps(redacted)
+    assert redacted["mode"] == "append_audio_chunk"
+
+
+def test_redaction_still_drops_heavy_payload_keys() -> None:
+    """The pre-existing contract must be preserved."""
+    redacted = _redact(
+        {
+            "ok": True,
+            "stage_handoff": object(),
+            "tts_handoff": object(),
+            "omni_payload": object(),
+            "tts_hidden_states": object(),
+            "tts_token_ids": object(),
+            "traceback": "long trace",
+        }
+    )
+    assert redacted == {"ok": True}
+
+
+def test_plain_values_pass_through_unchanged() -> None:
+    payload = {"ok": True, "count": 3, "name": "x", "missing": None, "items": [1, "two"]}
+    assert _redact(payload) == payload
+
+
+def test_dataclass_types_are_not_treated_as_instances() -> None:
+    """A class object is not a dataclass instance and must not be asdict'd."""
+
+    @dataclass
+    class _Shape:
+        x: int = 1
+
+    assert _redact(_Shape) is _Shape

tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py

@@ -265,14 +265,26 @@ class DemoState:
                 indices[event_type] = index
         return indices
 
-    def event_order_ok(self, *, require_input_commit: bool = True) -> bool:
+    def event_order_ok(
+        self,
+        *,
+        require_input_commit: bool = True,
+        require_commit_before_response: bool = True,
+    ) -> bool:
         if not self.events or self.events[0].get("type") != "session.created":
             return False
         first_commit_index = self.first_index("input_audio_buffer.committed")
         first_response_index = self.first_index("response.created")
         if first_response_index is None:
             return False
-        if require_input_commit and (first_commit_index is None or first_commit_index > first_response_index):
+        if require_input_commit and first_commit_index is None:
+            return False
+        if (
+            require_input_commit
+            and require_commit_before_response
+            and first_commit_index is not None
+            and first_commit_index > first_response_index
+        ):
             return False
         indices_by_type = self.first_response_lifecycle_indices()
         if not indices_by_type:
@@ -1409,7 +1421,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
             require_input_commit=not continuous_input,
         )
         if validation_mode == "model-policy"
-        else state.event_order_ok(require_input_commit=not continuous_input)
+        else state.event_order_ok(
+            require_input_commit=not continuous_input,
+            require_commit_before_response=False,
+        )
     )
     input_transcription_ok = _input_transcription_ok(
         state.input_transcription_count,

tests/e2e/online_serving/test_minicpmo_4_5_duplex.py

@@ -114,7 +114,7 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
             )
         )
     )
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
     assert result["audio_delta_count"] > 0
     assert result["done_count"] == 1
     assert result["error_count"] == 0
@@ -137,7 +137,7 @@ def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str,
             )
         )
     )
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
     assert result["session_count"] == 2
     assert result["resume"]["ok"] is True
     assert result["takeover"]["ok"] is True

tests/e2e/online_serving/test_minicpmo_4_5_expansion.py

@@ -150,7 +150,16 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     request_config = {
         "model": omni_server.model,
         "messages": messages,
+        "modalities": ["text", "audio"],
         "stream": True,
+        # Delimit the assistant answer with the TTS template so the Talker
+        # does not spend its codec-token budget speaking hidden reasoning.
+        "extra_body": {
+            "chat_template_kwargs": {
+                "use_tts_template": True,
+                "enable_thinking": False,
+            }
+        },
         "key_words": {"audio": ["Beijing"]},
     }
 

tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py

@@ -874,6 +874,39 @@ def test_realtime_duplex_demo_model_policy_accepts_continuous_input_without_comm
     assert state.model_policy_event_order_ok(expected_turns=3, require_input_commit=False)
 
 
+def test_realtime_duplex_demo_distinguishes_late_and_missing_commit():
+    demo = _load_demo_module()
+    state = demo.DemoState()
+    response_id = "resp-1"
+    item_id = f"item_{response_id}"
+    for event in (
+        {"type": "session.created"},
+        {"type": "response.created", "response_id": response_id},
+        {"type": "conversation.item.added", "item": {"id": item_id}},
+        {"type": "response.output_item.added", "response_id": response_id},
+        {"type": "response.content_part.added", "response_id": response_id},
+        {"type": "response.speak", "response_id": response_id},
+        {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA"},
+        {"type": "response.audio.done", "response_id": response_id},
+        {"type": "response.content_part.done", "response_id": response_id},
+        {"type": "response.output_item.done", "response_id": response_id},
+        {"type": "response.done", "response_id": response_id},
+        {"type": "input_audio_buffer.committed"},
+    ):
+        state.add(event)
+
+    assert not state.event_order_ok(require_input_commit=True)
+    assert state.event_order_ok(
+        require_input_commit=True,
+        require_commit_before_response=False,
+    )
+    state.events.pop()
+    assert not state.event_order_ok(
+        require_input_commit=True,
+        require_commit_before_response=False,
+    )
+
+
 def test_realtime_duplex_demo_accepts_overlap_unit_terminating_continuous_response():
     demo = _load_demo_module()
     state = demo.DemoState()

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/e2e/online_serving/test_qwen_image.py

@@ -58,7 +58,6 @@ def _get_default_case(model: str):
     ]
 
 
-@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)

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

@@ -1,3 +1,4 @@
+import concurrent.futures
 import contextlib
 import importlib
 import os
@@ -7,6 +8,7 @@ import types
 import pytest
 
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+from vllm_omni.engine import async_omni_engine as async_omni_engine_module
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_init_utils import (
     LogicalStageInitPlan,
@@ -19,6 +21,30 @@ from vllm_omni.engine.stage_runtime import StageRuntime
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+    engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
+    monkeypatch.setattr(engine, "_try_shutdown", lambda _message: None)
+
+    ticks = iter((0.0, 1.0))
+    monkeypatch.setattr(async_omni_engine_module.time, "monotonic", lambda: next(ticks))
+
+    warnings: list[tuple[object, ...]] = []
+    monkeypatch.setattr(
+        async_omni_engine_module.logger,
+        "warning",
+        lambda *args: warnings.append(args),
+    )
+
+    with pytest.raises(TimeoutError, match="did not become ready within 1s"):
+        engine._wait_for_orchestrator_init(concurrent.futures.Future(), startup_timeout=1)
+
+    assert len(warnings) == 1
+    message = str(warnings[0][0])
+    assert "--init-timeout" in message
+    assert "--stage-init-timeout" in message
+
+
 def _make_llm_metadata(
     stage_id: int,
     *,
@@ -146,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,
@@ -316,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
 
@@ -511,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)
@@ -561,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_cfg_companion_lifecycle.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Deterministic tests for the CFG-companion output lifecycle.
+
+Regression suite for the companion-output race behind the nightly Bagel
+failures: ``kv_ready`` used to mark a companion done before its processed
+output was stashed, letting the parent dispatch to the diffusion stage with
+0/N companion outputs (degraded CFG conditioning, and after an abort a TP
+rank desync). The invariant under test: **dispatch to the diffusion stage
+never happens with an incomplete companion bundle** — the parent is
+re-deferred while outputs are pending and failed when the bundle can no
+longer complete.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.messages import AbortRequestMessage, ErrorMessage
+from vllm_omni.engine.orchestrator import Orchestrator
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ───────── tracker-level ordering semantics ─────────
+
+
+def test_companion_first_then_parent_is_not_deferred():
+    tracker = CfgCompanionTracker()
+    tracker.register_companion("p", "neg", "p__neg")
+    tracker.set_companion_output("p__neg", {"o": 1})
+    assert tracker.on_companion_completed("p__neg") is None  # parent not deferred yet
+    # Parent arrives afterwards: bundle already complete, no defer needed.
+    assert tracker.all_companions_done("p")
+    assert tracker.get_companion_outputs("p") == [{"o": 1}]
+
+
+def test_parent_first_deferred_then_companion_flushes():
+    tracker = CfgCompanionTracker()
+    tracker.register_companion("p", "neg", "p__neg")
+    tracker.defer_parent("p", {"parent": True}, stage_id=0)
+    tracker.set_companion_output("p__neg", {"o": 1})
+    assert tracker.on_companion_completed("p__neg") == "p"
+    deferred = tracker.pop_pending_parent("p")
+    assert deferred is not None and deferred["engine_outputs"] == {"parent": True}
+
+
+def test_duplicate_companion_completion_is_idempotent():
+    tracker = CfgCompanionTracker()
+    tracker.register_companion("p", "neg", "p__neg")
+    tracker.defer_parent("p", {}, stage_id=0)
+    tracker.set_companion_output("p__neg", {"o": 1})
+    assert tracker.on_companion_completed("p__neg") == "p"
+    assert tracker.on_companion_completed("p__neg") is None  # duplicate: no re-flush
+    assert tracker.get_companion_outputs("p") == [{"o": 1}]
+
+
+def test_bundle_survives_resubmission_until_cleanup():
+    """get_companion_outputs is non-destructive: a streaming re-submission of
+    the parent bundles the complete set again (never a partial/empty one).
+    Only cleanup_parent releases the outputs."""
+    tracker = CfgCompanionTracker()
+    tracker.register_companion("p", "neg", "p__neg")
+    tracker.set_companion_output("p__neg", {"o": 1})
+    assert tracker.get_companion_outputs("p") == [{"o": 1}]
+    assert tracker.get_companion_outputs("p") == [{"o": 1}]  # re-submit: full bundle again
+    tracker.cleanup_parent("p")
+    assert tracker.get_companion_outputs("p") == []
+
+
+# ───────── orchestrator-method-level lifecycle ─────────
+
+
+class _FakePool:
+    """Minimal stage pool satisfying _forward_to_next_stage's preamble and cleanup."""
+
+    def __init__(self, stage_type: str = "diffusion") -> None:
+        self.stage_type = stage_type
+        self.stage_client = SimpleNamespace(requires_multimodal_data=False, custom_process_input_func=None)
+        self.aborted: list[list[str]] = []
+
+    async def abort_requests(self, request_ids):
+        self.aborted.append(list(request_ids))
+
+    def release_bindings(self, request_ids):
+        pass
+
+
+def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
+    orch = Orchestrator.__new__(Orchestrator)
+    orch._cfg_tracker = CfgCompanionTracker()
+    orch.request_states = {}
+    orch.stage_pools = [_FakePool("llm"), _FakePool("diffusion")][:num_stages]
+    orch.output_async_queue = asyncio.Queue()
+    orch.async_chunk = False
+    orch.duplex_control_plane = None
+    orch._pd_kv_params = {}
+    orch._pd_pair = None
+    orch._running_counter = None
+    orch._transfer_emitter = None
+    return orch
+
+
+def _req_state(final_stage_id: int = 1) -> SimpleNamespace:
+    return SimpleNamespace(
+        final_stage_id=final_stage_id,
+        stage_submit_ts={},
+        sampling_params_list=[SimpleNamespace(output_kind=None)] * (final_stage_id + 1),
+        streaming=SimpleNamespace(enabled=False, segment_finished=False),
+        running_counter_registered=False,
+        prompt={"prompt": "x"},
+        pipeline_timings={},
+    )
+
+
+@pytest.mark.asyncio
+async def test_kv_ready_does_not_mark_companion_done():
+    """kv_ready alone must not count a companion as done — only the processed
+    output (set_companion_output) may. Guards the original race."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch.request_states["p__neg"] = _req_state()
+
+    raw = SimpleNamespace(request_id="p__neg", kv_transfer_params={"kv_ready": True})
+    await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw]))
+
+    assert not orch._cfg_tracker.all_companions_done("p")
+    assert orch._cfg_tracker.pop_pending_parent("p") is None
+
+
+@pytest.mark.asyncio
+async def test_parent_kv_ready_defers_until_companion_outputs_stashed():
+    """Parent kv_ready with an incomplete bundle defers; the companion's
+    processed output later flushes it (companion-first vs parent-first order
+    both funnel through the same defer/flush path)."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch.request_states["p"] = _req_state()
+
+    raw = SimpleNamespace(request_id="p", kv_transfer_params={"kv_ready": True})
+    await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw]))
+
+    deferred = orch._cfg_tracker.pop_pending_parent("p")
+    assert deferred is not None and deferred["stage_id"] == 0
+
+
+@pytest.mark.asyncio
+async def test_forward_re_defers_when_bundle_incomplete():
+    """The forward gate re-defers (non-destructively) instead of dispatching
+    with 0/N companion outputs."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})  # stashed but NOT done-marked
+    state = _req_state()
+    orch.request_states["p"] = state
+
+    await orch._forward_to_next_stage("p", 0, {"parent": True}, state)
+
+    # Not dispatched: re-deferred with the stashed output left intact.
+    deferred = orch._cfg_tracker.pop_pending_parent("p")
+    assert deferred is not None
+    assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}]
+    assert orch.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_forward_fails_request_when_bundle_lost():
+    """Companions done but outputs gone without a prior bundle: the request is
+    failed (terminal ErrorMessage + abort) instead of dispatching degraded CFG."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})
+    orch._cfg_tracker.on_companion_completed("p__neg")
+    orch._cfg_tracker._companion_outputs.pop("p__neg")  # simulate lost output (unreachable normally)
+    state = _req_state()
+    orch.request_states["p"] = state
+
+    await orch._forward_to_next_stage("p", 0, {"parent": True}, state)
+
+    msg = orch.output_async_queue.get_nowait()
+    assert isinstance(msg, ErrorMessage)
+    assert msg.request_id == "p"
+    assert "incomplete" in msg.error
+    assert "p" not in orch.request_states
+    assert any("p" in batch for pool in orch.stage_pools for batch in pool.aborted)
+
+
+@pytest.mark.asyncio
+async def test_companion_abort_fails_deferred_parent():
+    """A companion aborted directly (parent not in the batch) can never
+    complete the bundle — the deferred parent is failed immediately."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0)
+    orch.request_states["p"] = _req_state()
+
+    await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"]))
+
+    msg = orch.output_async_queue.get_nowait()
+    assert isinstance(msg, ErrorMessage)
+    assert msg.request_id == "p"
+    assert "p__neg" in msg.error
+    assert "p" not in orch.request_states
+    assert not orch._cfg_tracker.is_companion("p__neg")
+    # Idempotent: a repeated abort of the same companion is a no-op.
+    await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"]))
+    assert orch.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_companion_abort_after_completion_does_not_fail_parent():
+    """A late/duplicate abort of a companion whose output is already stashed
+    must not cancel the (possibly already dispatched) parent."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})
+    orch._cfg_tracker.on_companion_completed("p__neg")
+    orch.request_states["p"] = _req_state()
+
+    await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"]))
+
+    assert orch.output_async_queue.empty()  # no parent failure emitted
+    assert "p" in orch.request_states  # parent untouched
+    assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}]  # bundle intact
+
+
+class _CapturedDispatchError(Exception):
+    """Sentinel raised by the capture hook to stop the forward after bundling."""
+
+
+@pytest.mark.asyncio
+async def test_repeated_forward_bundles_full_set_again():
+    """A second forward of the same parent (streaming terminal update) must
+    dispatch the complete [parent, *companions] bundle again — never [parent]
+    alone. Captured at the diffusion input hook, i.e. actual dispatch content."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})
+    orch._cfg_tracker.on_companion_completed("p__neg")
+    state = _req_state()
+    orch.request_states["p"] = state
+
+    captured: list[list] = []
+
+    def _capture(source_outputs, prompt, requires_multimodal_data):
+        captured.append(list(source_outputs))
+        raise _CapturedDispatchError
+
+    orch.stage_pools[1].stage_client.custom_process_input_func = _capture
+
+    for _ in range(2):
+        with pytest.raises(_CapturedDispatchError):
+            await orch._forward_to_next_stage("p", 0, {"parent": True}, state)
+
+    assert captured == [[{"parent": True}, {"o": 1}], [{"parent": True}, {"o": 1}]]
+
+
+@pytest.mark.asyncio
+async def test_replica_loss_cleanup_releases_parent_tracker_state():
+    """Replica-loss/membership teardown calls _cleanup_request_ids directly
+    (no cleanup_parent at the call site): the parent's tracker state must be
+    released and its companions folded into the same cleanup batch."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})
+    orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0)
+    orch.request_states["p"] = _req_state()
+
+    await orch._cleanup_request_ids(["p"], abort=True)
+
+    assert not orch._cfg_tracker.has_companions("p")
+    assert not orch._cfg_tracker.is_companion("p__neg")
+    assert orch._cfg_tracker.get_companion_outputs("p") == []
+    assert orch._cfg_tracker.pop_pending_parent("p") is None
+    aborted = {rid for pool in orch.stage_pools for b in pool.aborted for rid in b}
+    assert {"p", "p__neg"} <= aborted
+
+
+@pytest.mark.asyncio
+async def test_replica_loss_of_companion_fails_deferred_parent():
+    """A companion torn down by replica loss before its output arrived means
+    the parent can never complete its bundle — the cleanup fails the parent."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0)
+    orch.request_states["p"] = _req_state()
+
+    await orch._cleanup_request_ids(["p__neg"])
+
+    msg = orch.output_async_queue.get_nowait()
+    assert isinstance(msg, ErrorMessage)
+    assert msg.request_id == "p"
+    assert "p__neg" in msg.error
+    assert "p" not in orch.request_states
+    assert not orch._cfg_tracker.has_companions("p")
+
+
+@pytest.mark.asyncio
+async def test_finished_companion_cleanup_preserves_parent_bundle():
+    """The normal companion-finished cleanup (_route_output cleans just the
+    companion id) must not disturb the parent's stashed bundle."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.set_companion_output("p__neg", {"o": 1})
+    orch._cfg_tracker.on_companion_completed("p__neg")
+    orch.request_states["p"] = _req_state()
+
+    await orch._cleanup_request_ids(["p__neg"])
+
+    assert orch.output_async_queue.empty()
+    assert "p" in orch.request_states
+    assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}]
+
+
+@pytest.mark.asyncio
+async def test_parent_abort_still_expands_to_companions():
+    """Regression-guard the existing behavior: aborting the parent aborts its
+    companions and clears tracker state, without emitting an error."""
+    orch = _make_orchestrator()
+    orch._cfg_tracker.register_companion("p", "neg", "p__neg")
+    orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0)
+    orch.request_states["p"] = _req_state()
+
+    await orch._handle_abort(AbortRequestMessage(request_ids=["p"]))
+
+    assert orch.output_async_queue.empty()
+    assert "p" not in orch.request_states
+    assert not orch._cfg_tracker.is_companion("p__neg")
+    aborted = {rid for pool in orch.stage_pools for batch in pool.aborted for rid in batch}
+    assert {"p", "p__neg"} <= aborted

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
@@ -2282,12 +2418,15 @@ async def test_request_cleanup_failure_is_deferred_to_control_plane():
         def finalize_closed_sessions(self, session_ids: list[str]) -> None:
             self.finalized.extend(session_ids)
 
+    from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+
     orchestrator = object.__new__(Orchestrator)
     orchestrator.stage_pools = [_FailingPool()]
     orchestrator.duplex_control_plane = _Plane()
     orchestrator._pd_kv_params = {}
     orchestrator.request_states = {}
     orchestrator._running_counter = None
+    orchestrator._cfg_tracker = CfgCompanionTracker()
 
     with pytest.raises(RuntimeError, match="stage abort failed"):
         await orchestrator._cleanup_request_ids(

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

@@ -12,6 +12,7 @@ import pytest
 from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
@@ -476,6 +477,89 @@ class TestSingleStageModeDetection:
         assert not hasattr(engine, "_omni_master_server")
 
 
+# ---------------------------------------------------------------------------
+# Endpoint-restriction resolution vs. tri-state trust_remote_code (#5495)
+# ---------------------------------------------------------------------------
+
+
+class TestEndpointRestrictionsTrustRemoteCode:
+    """Regression tests for issue #5495.
+
+    ``AsyncOmniEngine.__init__`` takes ``trust_remote_code: bool | None`` where
+    ``None`` means "not specified" (the CLI default, since ``--trust-remote-code``
+    is ``store_true`` and its absence is mapped to ``None``). Endpoint-restriction
+    resolution eagerly loads the HF config via vLLM's ``get_config``, which does
+    ``trust_remote_code |= ...`` internally and raises ``TypeError`` on ``None``.
+    That used to be swallowed, leaving ``endpoint_restrictions`` empty so
+    ``/v1/completions`` was never shut down and the request crashed stage-1.
+
+    ``__init__`` must collapse the ``None`` "not specified" case to ``False`` at
+    the ``get_pipeline_config`` call site so the restriction is still computed.
+    """
+
+    @pytest.fixture(autouse=True)
+    def clear_config_factory_caches(self):
+        """Isolate the process-wide ``functools.cache`` on StageConfigFactory.
+
+        ``get_hf_config`` / ``try_infer_model_type`` cache per-model results for
+        the whole process. These tests patch ``get_config`` with a synthetic
+        Qwen3-Omni config for ``"fake-model"``; without clearing on teardown the
+        synthetic entry would leak to later tests in the same worker (after the
+        mock is restored), making the suite order-dependent. Clear both before
+        and after.
+        """
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+        yield
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+
+    def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine:
+        mocker.patch.object(
+            AsyncOmniEngine,
+            "_resolve_stage_configs",
+            return_value=("/fake/path", [_make_stage_cfg(0)]),
+        )
+        mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
+        mocker.patch("threading.Thread")
+        mocker.patch("concurrent.futures.Future")
+        return AsyncOmniEngine(model="fake-model", **kwargs)
+
+    @staticmethod
+    def _install_qwen3_omni_config(mocker: MockerFixture) -> None:
+        """Patch get_config to a thinker+talker config, mirroring vLLM's
+        ``trust_remote_code |= ...`` (raises TypeError on None)."""
+        from transformers import Qwen3OmniMoeConfig
+
+        def fake_get_config(model, trust_remote_code, **_):
+            trust_remote_code |= False  # TypeError if None reaches here
+            return Qwen3OmniMoeConfig(enable_audio_output=True)
+
+        mocker.patch("vllm_omni.config.config_factory.get_config", side_effect=fake_get_config)
+
+    @pytest.mark.parametrize("trc_kwargs", [{}, {"trust_remote_code": False}, {"trust_remote_code": True}])
+    def test_completions_restriction_survives(self, mocker: MockerFixture, trc_kwargs: dict):
+        """COMPLETIONS restriction is computed regardless of how (or whether)
+        trust_remote_code is passed — including the unset case that resolves to
+        the engine's ``None`` default."""
+        from vllm_omni.config.endpoint_policy import OmniServingCapability
+
+        self._install_qwen3_omni_config(mocker)
+        engine = self._make_engine_no_thread(mocker, **trc_kwargs)
+        assert any(r.capability is OmniServingCapability.COMPLETIONS for r in engine.endpoint_restrictions), (
+            f"COMPLETIONS restriction dropped for {trc_kwargs or 'unset (None default)'}"
+        )
+
+    def test_unset_trust_remote_code_defaults_to_none(self, mocker: MockerFixture):
+        """Guard the tri-state premise: the engine keeps ``None`` (not ``False``)
+        when trust_remote_code is not passed, so the call-site coercion is what
+        prevents the crash."""
+        import inspect
+
+        default = inspect.signature(AsyncOmniEngine.__init__).parameters["trust_remote_code"].default
+        assert default is None
+
+
 # ---------------------------------------------------------------------------
 # AsyncOmniEngine single-stage initialization paths
 # ---------------------------------------------------------------------------
@@ -778,6 +862,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 +896,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 +916,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/engine/test_stage_metadata.py

@@ -158,3 +158,16 @@ def test_extract_stage_metadata_preserves_legacy_one_argument_api():
     assert metadata.stage_id == 0
     assert metadata.model_stage == "thinker"
     assert metadata.custom_process_input_func is operator.add
+
+
+def test_extract_stage_metadata_defaults_missing_engine_input_source():
+    pipeline, deploy = _metadata_inputs()
+    legacy_config = merge_pipeline_deploy(
+        pipeline,
+        deploy,
+    )[0].to_omegaconf()
+    del legacy_config["engine_input_source"]
+
+    metadata = extract_stage_metadata(legacy_config)
+
+    assert metadata.engine_input_source == []

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_serving_speech.py

@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
 
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
 from vllm_omni.entrypoints.openai import api_server as api_server_module
+from vllm_omni.entrypoints.openai import serving_speech as serving_speech_module
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.protocol.audio import (
     BatchSpeechRequest,
@@ -1390,6 +1391,20 @@ class TestTTSMethods:
         # Verify speakers are normalized to lowercase
         assert server.supported_speakers == {"ryan", "vivian", "aiden"}
 
+    def test_load_supported_speakers_skips_non_tts_omni_model(
+        self,
+        speech_server,
+        mocker: MockerFixture,
+    ):
+        warning = mocker.patch.object(serving_speech_module.logger, "warning")
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(tts_config=SimpleNamespace())
+        )
+
+        assert speech_server._is_tts is False
+        assert speech_server._load_supported_speakers() == set()
+        warning.assert_not_called()
+
     def test_load_supported_languages_from_config(self, speech_server):
         """Languages/dialects from codec_language_id are loaded title-cased; 'Auto' is added."""
         speech_server._tts_model_type = "qwen3_tts"

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/openai_api/test_video_server.py

@@ -12,6 +12,7 @@ import os
 import threading
 import time
 from contextlib import asynccontextmanager
+from pathlib import Path
 from types import SimpleNamespace
 
 import numpy as np
@@ -62,11 +63,19 @@ class FakeAsyncOmni:
         self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
         self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
         self.captured_prompt = None
+        self.captured_reference_video_bytes = None
         self.captured_sampling_params_list = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_prompt = prompt
         self.captured_sampling_params_list = sampling_params_list
+        reference_videos = prompt.get("multi_modal_data", {}).get("video")
+        if (
+            isinstance(reference_videos, list)
+            and reference_videos
+            and all(isinstance(item, str) for item in reference_videos)
+        ):
+            self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos]
         num_outputs = sampling_params_list[0].num_outputs_per_prompt
         videos = [object() for _ in range(num_outputs)]
         yield MockVideoResult(videos)
@@ -501,6 +510,44 @@ def test_v2v_video_generation_with_video_reference_form(test_client, mocker: Moc
     assert input_video[0].size == (32, 24)
 
 
+@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"])
+def test_multi_video_generation_preserves_uploaded_files_until_generation(
+    endpoint,
+    test_client,
+    mocker: MockerFixture,
+):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        endpoint,
+        data={
+            "prompt": "Composite the subject from the first video into the second.",
+            "extra_params": json.dumps({"task": "ref2va", "duration": 15.0}),
+        },
+        files=[
+            ("input_references", ("subject.mp4", b"subject-video", "video/mp4")),
+            ("input_references", ("background.mov", b"background-video", "video/quicktime")),
+        ],
+    )
+
+    assert response.status_code == 200
+    if endpoint.endswith("/sync"):
+        assert response.content == b"fake-video"
+    else:
+        video_id = response.json()["id"]
+        _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_videos = engine.captured_prompt["multi_modal_data"]["video"]
+    assert engine.captured_reference_video_bytes == [b"subject-video", b"background-video"]
+    assert len(input_videos) == 2
+    assert all(not Path(path).exists() for path in input_videos)
+    assert engine.captured_sampling_params_list[0].extra_args["task"] == "ref2va"
+    assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
+
+
 def test_decode_video_bytes_can_keep_last_frames():
     from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
 

tests/entrypoints/test_serve.py

@@ -46,6 +46,26 @@ def test_serve_parser_accepts_strategy_config() -> None:
     assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml"
 
 
+def test_serve_parser_accepts_deploy_config() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--deploy-config", "/tmp/deploy.yaml"])
+
+    assert args.deploy_config == "/tmp/deploy.yaml"
+    assert args.get_explicit_kwargs_dict()["deploy_config"] == "/tmp/deploy.yaml"
+
+
+def test_serve_parser_rejects_stage_configs_path() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    with pytest.raises(SystemExit):
+        parser.parse_args(["serve", "fake-model", "--omni", "--stage-configs-path", "/tmp/stages.yaml"])
+
+
 def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="subcommand")
@@ -66,7 +86,6 @@ def _make_headless_args(**kwargs) -> TrackingNamespace:
         "omni_replica_address": None,
         "omni_dp_size_local": 1,
         "worker_backend": "multi_process",
-        "stage_configs_path": None,
         "deploy_config": None,
         "log_stats": False,
         "disable_log_stats": False,
@@ -154,6 +173,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
     captured: dict = {}
 
     def _fake_resolve(*args, **kwargs):
+        captured["args"] = args
         captured.update(kwargs)
         # Return a stage that does NOT match stage_id=0 so run_headless stops
         # right after the resolver call (we only care about how it was called).
@@ -166,12 +186,15 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
 
     args = _make_headless_args(
         stage_id=0,
+        deploy_config="/tmp/deploy.yaml",
         strategy_config="/tmp/strategy.yaml",
         stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
     )
     with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
         run_headless(args)
 
+    assert captured["args"][1] is None
+    assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
     assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
     assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
 

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
 # ---------------------------------------------------------------------------