Update vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

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

.buildkite/cuda/test-merge.yml

@@ -169,7 +169,6 @@ steps:
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
           - vllm_omni/deploy/minicpmo_4_5.yaml
           - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
-          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
           - vllm_omni/experimental/fullduplex/
           - vllm_omni/model_executor/models/minicpmo_4_5/
           - vllm_omni/model_executor/models/step_audio2/

.buildkite/cuda/test-nightly.yml

@@ -83,7 +83,7 @@ steps:
           - export BENCHMARK_DIR=tests/dfx/perf/results
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and full_model"
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
             exit $$EXIT
@@ -375,7 +375,7 @@ steps:
       - label: ":full_moon: Diffusion · Quantization Test with L4"
         timeout_in_minutes: 60
         commands:
-          - pip install "vllm-gguf-plugin>=0.0.3"
+          - pip install "vllm-gguf-plugin==0.0.4"
           - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
         mirror_hardwares: l4_1
 

.buildkite/cuda/test-ready.yml

@@ -208,7 +208,6 @@ steps:
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
           - vllm_omni/deploy/minicpmo_4_5.yaml
           - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
-          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
           - vllm_omni/experimental/fullduplex/
           - vllm_omni/model_executor/models/minicpmo_4_5/
           - vllm_omni/model_executor/models/step_audio2/

.buildkite/cuda/test-weekly.yml

@@ -74,6 +74,19 @@ steps:
             exit $$EXIT
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Omni · Perf Test · Async Chunk · Random"
+        key: weekly-omni-performance-async-chunk-random
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and slow"
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Omni · Perf Test · Multi-Replica"
         key: weekly-omni-performance-multi-replicas
         timeout_in_minutes: 180

.claude/skills/add-diffusion-model/SKILL.md

@@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
 
 ## Path B: Custom/Private Repo Model
 
-For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
 
 ### B1. Understand the reference repo
 
@@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules:
 | **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
 | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
 | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
-| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
+| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone |
 | **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
-| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
+| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_<name>.py` |
 | **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
 
 ### B3. Handle external dependencies
@@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
 
 Custom models have two common patterns for weight loading:
 
-**Pattern 1: Bypass standard loader** (DreamID-Omni style)
+**Pattern 1: Bypass standard loader** (eager custom init)
 
 When the original model has complex custom init functions that load weights in `__init__`:
 

.claude/skills/add-diffusion-model/references/custom-model-patterns.md

@@ -16,20 +16,19 @@ vllm_omni/diffusion/models/wan2_2/
 
 The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
 
-### Custom model with external deps (e.g., DreamID-Omni)
+### Custom model with external deps
 
 ```
-vllm_omni/diffusion/models/dreamid_omni/
+vllm_omni/diffusion/models/<name>/
 ├── __init__.py                    # Exports pipeline only
-├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
-├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
-└── wan2_2.py                      # Re-implemented Wan backbone with split API
+├── pipeline_<name>.py             # Pipeline: loads ALL weights in __init__ via custom helpers
+└── <backbone>.py                  # Custom transformer / fusion architecture
 
-examples/offline_inference/x_to_video_audio/
-└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
+examples/offline_inference/<task>/
+└── download_<name>.py             # Downloads weights and/or clones code repo
 ```
 
-All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
+All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
 
 ### Custom model with ported code (e.g., BAGEL)
 
@@ -76,7 +75,7 @@ self.weights_sources = [
 ]
 ```
 
-### Pattern 3: Fully custom loading (DreamID-Omni)
+### Pattern 3: Fully custom loading
 
 ```
 init → load ALL weights eagerly via custom helpers → load_weights() = no-op
@@ -108,8 +107,8 @@ Standard diffusers `model_index.json`:
 Custom model `model_index.json` (minimal):
 ```json
 {
-    "_class_name": "DreamIDOmniPipeline",
-    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
+    "_class_name": "YourCustomPipeline",
+    "fusion": "weights/model.safetensors"
 }
 ```
 
@@ -117,7 +116,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS
 
 ## External Dependency Management
 
-### Git clone + .pth injection (DreamID-Omni pattern)
+### Git clone + .pth injection
 
 ```python
 def download_dependency():
@@ -163,7 +162,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config
 
 Diffusers models use `config.json` in each subfolder. Custom models often use:
 
-**Module-level config dicts** (DreamID-Omni):
+**Module-level config dicts**:
 ```python
 VIDEO_CONFIG = {
     "patch_size": [1, 2, 2], "model_type": "ti2v",
@@ -181,7 +180,7 @@ vae_cfg = bagel_cfg.get("vae_config", {})
 
 ## Custom Architecture Patterns
 
-### Split forward API (DreamID-Omni)
+### Split forward API
 
 When a fusion model needs to interleave blocks from two backbones:
 

benchmarks/tts/benchmark_cosyvoice3_trt_streams.py

@@ -0,0 +1,151 @@
+"""Benchmark CosyVoice3-style CUDA stream handoff strategies."""
+
+import argparse
+import json
+import statistics
+import time
+
+import torch
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def run_case(
+    *,
+    device: torch.device,
+    steps: int,
+    warmup_steps: int,
+    producer_cycles: int,
+    estimator_cycles: int,
+    use_stream_dependencies: bool,
+) -> dict[str, float | int | str]:
+    caller_stream = torch.cuda.Stream(device=device)
+    estimator_stream = torch.cuda.Stream(device=device)
+    input_tensor = torch.zeros(1024, device=device)
+    output_tensor = torch.empty_like(input_tensor)
+    submission_ms = []
+
+    def step(measure: bool) -> None:
+        with torch.cuda.stream(caller_stream):
+            torch.cuda._sleep(producer_cycles)
+            input_tensor.add_(1)
+
+        start = time.perf_counter_ns()
+        if use_stream_dependencies:
+            estimator_stream.wait_stream(caller_stream)
+        else:
+            caller_stream.synchronize()
+
+        with torch.cuda.stream(estimator_stream):
+            torch.cuda._sleep(estimator_cycles)
+            output_tensor.copy_(input_tensor)
+
+        if use_stream_dependencies:
+            caller_stream.wait_stream(estimator_stream)
+        else:
+            estimator_stream.synchronize()
+
+        if measure:
+            submission_ms.append((time.perf_counter_ns() - start) / 1e6)
+
+    for _ in range(warmup_steps):
+        step(measure=False)
+    caller_stream.synchronize()
+    estimator_stream.synchronize()
+    input_tensor.zero_()
+    torch.accelerator.reset_peak_memory_stats(device)
+    initial_allocated = torch.accelerator.memory_allocated(device)
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    wall_start = time.perf_counter()
+    with torch.cuda.stream(caller_stream):
+        start_event.record()
+    for _ in range(steps):
+        step(measure=True)
+    with torch.cuda.stream(caller_stream):
+        end_event.record()
+    end_event.synchronize()
+    wall_ms = (time.perf_counter() - wall_start) * 1e3
+    peak_delta_bytes = torch.accelerator.max_memory_allocated(device) - initial_allocated
+
+    expected = torch.full_like(output_tensor, steps)
+    torch.testing.assert_close(output_tensor, expected, rtol=0, atol=0)
+    return {
+        "mode": "wait_stream" if use_stream_dependencies else "synchronize",
+        "steps": steps,
+        "submission_median_ms": statistics.median(submission_ms),
+        "submission_p95_ms": statistics.quantiles(submission_ms, n=20)[18],
+        "submission_mean_ms": statistics.mean(submission_ms),
+        "gpu_elapsed_ms": start_event.elapsed_time(end_event),
+        "wall_elapsed_ms": wall_ms,
+        "peak_allocation_delta_bytes": peak_delta_bytes,
+        "final_value": int(output_tensor[0].item()),
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--steps", type=int, default=100)
+    parser.add_argument("--warmup-steps", type=int, default=10)
+    parser.add_argument("--repeats", type=int, default=7)
+    parser.add_argument("--producer-cycles", type=int, default=1_000_000)
+    parser.add_argument("--estimator-cycles", type=int, default=2_000_000)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = torch.device(args.device)
+    current_omni_platform.set_device(device)
+    torch.manual_seed(0)
+
+    old_results = []
+    new_results = []
+    for repeat in range(args.repeats):
+        common_args = {
+            "device": device,
+            "steps": args.steps,
+            "warmup_steps": args.warmup_steps,
+            "producer_cycles": args.producer_cycles,
+            "estimator_cycles": args.estimator_cycles,
+        }
+        if repeat % 2 == 0:
+            old_results.append(run_case(**common_args, use_stream_dependencies=False))
+            new_results.append(run_case(**common_args, use_stream_dependencies=True))
+        else:
+            new_results.append(run_case(**common_args, use_stream_dependencies=True))
+            old_results.append(run_case(**common_args, use_stream_dependencies=False))
+
+    old_submission_ms = statistics.median(result["submission_median_ms"] for result in old_results)
+    new_submission_ms = statistics.median(result["submission_median_ms"] for result in new_results)
+    old_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in old_results)
+    new_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in new_results)
+    print(
+        json.dumps(
+            {
+                "device": torch.cuda.get_device_name(device),
+                "torch": torch.__version__,
+                "cuda": torch.version.cuda,
+                "repeats": args.repeats,
+                "warmup_steps": args.warmup_steps,
+                "producer_cycles": args.producer_cycles,
+                "estimator_cycles": args.estimator_cycles,
+                "old_submission_median_ms": old_submission_ms,
+                "new_submission_median_ms": new_submission_ms,
+                "submission_median_speedup": old_submission_ms / new_submission_ms,
+                "old_gpu_elapsed_median_ms": old_gpu_ms,
+                "new_gpu_elapsed_median_ms": new_gpu_ms,
+                "gpu_elapsed_change_percent": (new_gpu_ms - old_gpu_ms) / old_gpu_ms * 100,
+                "parity": all(result["final_value"] == args.steps for result in old_results + new_results),
+                "old_runs": old_results,
+                "new_runs": new_results,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

collect_env.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 # ruff: noqa
 # code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
@@ -494,6 +494,16 @@ def is_xnnpack_available():
 
 
 def get_env_vars():
+    try:
+        from vllm_omni.config.environment_variable_inventory import ENVIRONMENT_VARIABLE_INVENTORY
+    except Exception:
+        # Keep environment collection usable when an incomplete/broken Omni
+        # installation is the reason the user is running this diagnostic.
+        # Importing an Omni submodule executes package initialization, which
+        # may fail with RuntimeError, AssertionError, or OSError as well as
+        # ImportError when vLLM or the accelerator stack is broken.
+        ENVIRONMENT_VARIABLE_INVENTORY = {}
+
     env_vars = ""
     secret_terms = ("secret", "token", "api", "access", "password")
     report_prefix = (
@@ -507,12 +517,17 @@ def get_env_vars():
         "MKL_",
         "NVIDIA",
     )
-    for k, v in os.environ.items():
-        if any(term in k.lower() for term in secret_terms):
+    for k, v in sorted(os.environ.items()):
+        omni_classification = ENVIRONMENT_VARIABLE_INVENTORY.get(k)
+        if (omni_classification is not None and omni_classification.redact_value) or any(
+            term in k.lower() for term in secret_terms
+        ):
             continue
-        if k in environment_variables:
-            env_vars = env_vars + "{}={}".format(k, v) + "\n"
-        if k.startswith(report_prefix):
+        if (
+            k in environment_variables
+            or (omni_classification is not None and omni_classification.is_public_omni)
+            or k.startswith(report_prefix)
+        ):
             env_vars = env_vars + "{}={}".format(k, v) + "\n"
 
     return env_vars

docs/.nav.yml

@@ -24,7 +24,6 @@ nav:
       - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
       - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
-      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md

docs/cli/README.md

@@ -2,6 +2,11 @@
 
 The CLI for vLLM-Omni inherits from vllm with some additional arguments.
 
+Environment variables are process configuration rather than CLI-only options.
+See the [Environment Variables](../configuration/environment_variables.md)
+reference for ownership, precedence, evaluation time, and per-stage `env`
+scoping.
+
 ## serve
 
 Starts the vLLM-Omni OpenAI Compatible API server.
@@ -39,4 +44,5 @@ vllm bench serve --omni \
     --num-prompts  5
 ```
 
-See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.
+See [vllm bench serve](./bench/serve.md) for serving benchmark arguments and
+dataset-specific examples, including OmniInteract native-duplex sessions.

docs/cli/bench/serve.md

@@ -2,7 +2,7 @@
 The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
 
 ## Notes
-vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
+vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` serving benchmark backends. It also adds the `omniinteract` dataset.
 
 ## Basic Parameter Description
 You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
   Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
 
 - `--backend`
-  Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices.
+  Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` to the upstream vLLM backend choices.
 
 - `--model`
   The model identifier to load, filled according to the models supported by vLLM-Omni.
@@ -19,7 +19,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
   The API endpoint exposed externally, to which clients send their requests.
 
 - `--dataset-name`
-  The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
+  The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio), while `omniinteract` replays official OmniInteract videos.
 
 - `--num-prompts`
   The total number of requests to send, an integer.
@@ -269,6 +269,44 @@ We use audio generation time / audio duration to calculate RTF.
 
 </details>
 
+### OmniInteract Realtime Benchmark
+
+OmniInteract runs each video as one native-duplex WebSocket sample in the standard serving-benchmark lifecycle:
+
+```bash
+vllm bench serve --omni \
+  --backend openai-realtime-duplex \
+  --dataset-name omniinteract \
+  --dataset-path /path/to/OmniInteract \
+  --model openbmb/MiniCPM-o-4_5 \
+  --base-url http://127.0.0.1:8000 \
+  --endpoint /v1/realtime \
+  --omniinteract-ref-audio /path/to/reference.wav \
+  --omniinteract-output-dir ./omniinteract-artifacts \
+  --num-prompts 3 \
+  --num-warmups 0
+```
+
+`--dataset-path` accepts an extracted directory, `data.tar[.gz]`, or a Hugging Face dataset ID; omitting it uses
+`lucky-lance/OmniInteract`. `--num-prompts` is the total across subsets and defaults to 3 for OmniInteract; explicit `0`
+selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the
+`/v1/realtime` endpoint.
+
+Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded
+before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use
+explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by
+`--omniinteract-media-timeout-s`, and concurrency defaults to 1. Standard request-rate, warmup, result-saving, and summary
+options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result.
+
+Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under
+`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write
+`.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts,
+not answer accuracy. Transcript timestamps are serialized playback-queue times. Clipped or cancelled outputs are ineligible and
+omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video horizon.
+
+TTFT, TTFP, and RTF start at client receipt of `response.created`. TPOT/ITL use engine stage-0 timing; ITL is emitted only when
+every token interval is present.
+
 ### Multi-Modal Benchmark
 
 <details class="admonition abstract" markdown="1">

docs/configuration/README.md

@@ -7,6 +7,9 @@ For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration g
 Each model defines fixed topology in a registered `PipelineConfig` and runtime
 overrides in a deploy YAML.
 
+For process-level settings shared by the CLI, deploy YAML, workers, and model
+integrations, see [Environment Variables](./environment_variables.md).
+
 For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py).
 
 For an introduction, see [Pipeline and deploy configurations](./stage_configs.md).

docs/configuration/environment_variables.md

@@ -0,0 +1,237 @@
+# Environment Variables
+
+Environment variables are process-level configuration shared by the CLI,
+deploy YAML, Python APIs, worker processes, and model integrations. They are
+therefore documented under **User Guide → Configuration**, rather than under
+CLI Reference.
+
+This page describes the public vLLM-Omni environment-variable contract on the
+`latest` documentation branch. Use the documentation for your installed
+vLLM-Omni release when behavior differs.
+
+!!! important "Public contract"
+    Only names in [Public vLLM-Omni variables](#public-vllm-omni-variables) are
+    public Omni configuration. A name found in source code is not automatically
+    supported: model-specific, benchmark, platform, and internal variables have
+    separate ownership and lifecycle rules.
+
+## Naming and implementation boundary
+
+New public variables owned by vLLM-Omni must use the `VLLM_OMNI_` prefix.
+Existing public names with older prefixes remain supported for compatibility,
+but they do not establish naming precedent for new settings. Inherited vLLM
+and third-party names retain the prefix chosen by their owning project.
+
+The source [environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py)
+records ownership, migration disposition, and redaction metadata. Unlike
+`vllm.envs`, it is not an executable value resolver: existing consumers still
+own parsing, defaults, precedence, and evaluation time. The public tables below
+describe those live consumer contracts. Stable model or stage behavior should
+move to typed configuration, and request-varying behavior should move to a
+request schema instead of adding another environment variable.
+
+## Precedence and evaluation time
+
+There is no repository-wide rule that makes environment variables override (or
+lose to) every CLI, YAML, or request option. The tables below state precedence
+for each public setting. In general:
+
+- request fields should control behavior that can vary per request;
+- typed deploy or stage configuration should control stable model and stage
+  behavior;
+- environment variables are appropriate for process launch, platform
+  integration, emergency diagnostics, and compatibility fallbacks.
+
+An environment variable can be read while a module is imported, when server or
+model configuration is constructed, when a worker is spawned, or whenever a
+property is accessed. Set import- and startup-time variables before invoking
+`vllm serve` or importing vLLM-Omni. Changing the parent shell after workers
+start does not update existing worker environments.
+
+The lifecycle labels used below are:
+
+- **Stable** — reviewed public configuration; changes require compatibility
+  consideration.
+- **Experimental** — public for evaluation, but its contract may still evolve.
+- **Diagnostic** — supported for investigation or temporary mitigation, not as
+  a normal tuning interface.
+- **Deprecated** — compatibility alias; migrate to the named replacement.
+
+## Public vLLM-Omni variables
+
+### Diffusion backend and cache selection
+
+| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
+| --- | --- | --- | --- | --- |
+| `DIFFUSION_ATTENTION_BACKEND` | Backend name or `auto`; default `auto` (platform selection) | Diffusion stages; read when `OmniDiffusionConfig` is constructed | `diffusion_attention_config.default` wins. An unknown backend fails during backend resolution. | Stable fallback |
+| `DIFFUSION_CACHE_BACKEND` | `none`, `cache_dit`, `tea_cache`, `mag_cache`, `step_cache`, `stepcache`, or `step_cache_dit`; default `none` | Diffusion runner startup | Explicit `cache_backend` in config wins. Otherwise this name wins over the deprecated alias. Unsupported values raise `ValueError` during runner setup. | Stable fallback |
+| `DIFFUSION_CACHE_ADAPTER` | Same values as `DIFFUSION_CACHE_BACKEND`; default `none` | Diffusion runner startup | Used only when neither explicit `cache_backend` nor `DIFFUSION_CACHE_BACKEND` is set. Unsupported values raise `ValueError`. | Deprecated; use `DIFFUSION_CACHE_BACKEND` |
+| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE` | Boolean: `1`, `true`, `yes`, `on`, `0`, `false`, `no`, or `off`; default disabled | Each diffusion runner; resolved during model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
+| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE` | Positive integer; default `32` entries | Each diffusion runner; resolved during model setup | A valid environment value overrides the explicit cache size. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
+| `OMNI_DIFFUSION_SESSION_STATE_MANAGER` | Boolean with the same accepted spellings as the prompt cache; default disabled | Experimental diffusion session manager; model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental |
+| `OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS` | Positive integer; default `64` | Experimental diffusion session manager; model setup | A valid environment value overrides the explicit maximum. A non-integer logs a warning; a non-positive value is ignored. | Experimental |
+
+Backend names for `DIFFUSION_ATTENTION_BACKEND` are the members of
+`DiffusionAttentionBackendEnum`, such as `FLASH_ATTN`, `TORCH_SDPA`,
+`SAGE_ATTN`, `FLASHINFER_ATTN`, and `TRTLLM_ATTN`. Platform support still
+depends on the installed kernels and model path.
+
+### Serving and runtime
+
+| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
+| --- | --- | --- | --- | --- |
+| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
+| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
+| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
+| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
+| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
+| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
+| `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
+| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |
+
+### Server storage
+
+Storage names use Pydantic's nested-settings delimiter (`__`). They are read
+when `ServerSettings` is initialized, normally during server import/startup.
+Invalid integer values raise a Pydantic validation error.
+
+| Name | Type and default | Meaning | Lifecycle |
+| --- | --- | --- | --- |
+| `VLLM_OMNI_SERVER_STORAGE__PATH` | Path; default `/tmp/storage` | Directory for completed server files. | Stable |
+| `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY` | Integer; default `4` | Maximum concurrent file operations. No additional range validation is currently applied. | Stable |
+| `VLLM_OMNI_SERVER_STORAGE__FILE_TTL` | Integer seconds or unset; default unset | Optional lifetime for locally stored files. | Stable |
+| `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL` | Integer seconds or unset | Sweep frequency. If file TTL is set and this is unset, the effective default is `300`. | Stable |
+| `VLLM_OMNI_STORAGE_PATH` | Path | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__PATH`. The new name wins when both are set. | Deprecated |
+| `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` | Integer | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`. The new name wins when both are set. | Deprecated |
+
+The deprecated aliases emit `DeprecationWarning`. They do not affect file TTL
+settings.
+
+### Quantization diagnostics and performance
+
+| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
+| --- | --- | --- | --- | --- |
+| `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary |
+| `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control |
+
+`QUACK_CACHE_DIR` is owned by the external quack library and is not an
+Omni-owned variable, even though vLLM-Omni supplies a persistent default when
+it is unset.
+
+## Per-stage environment
+
+Deploy configurations can set arbitrary environment keys for one stage:
+
+```yaml
+stages:
+  - stage_id: 0
+    env:
+      HCCL_IF_BASE_PORT: 23000
+      MY_MODEL_CACHE: /mnt/model-cache
+```
+
+The legacy nested form under `runtime.env` is also accepted. vLLM-Omni converts
+keys and values to strings, applies them while the stage engine or child process
+is launched, logs keys but not values, and restores the parent process's prior
+environment afterward. A child process inherits the scoped values. This keeps
+sibling stages from permanently changing one another's parent environment.
+
+This scoping has two consequences:
+
+1. A setting read inside stage/model initialization can see the stage value.
+2. A setting already cached by an import that occurred before stage launch
+   cannot be changed retroactively with stage `env`.
+
+Stage `env` is an escape hatch for process and third-party settings, not a
+replacement for typed stage configuration. Do not put request-varying behavior
+there. Secret values may reach child processes; logs and diagnostics must emit
+their keys only.
+
+## Inherited vLLM variables
+
+vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
+the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
+for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
+XPU, ModelScope, and FlashInfer workspace settings.
+
+An inherited name does not become Omni-owned because Omni reads or forwards it.
+Where behavior differs, the Omni-specific page should document only that
+difference. For example, device-visibility values can be further narrowed by a
+stage's `devices` setting.
+
+## Platform, external, and secret variables
+
+Hardware runtimes, distributed launchers, and third-party libraries retain
+ownership of their variables. Common groups include:
+
+- device visibility (`ASCEND_RT_VISIBLE_DEVICES`, `HIP_VISIBLE_DEVICES`, and
+  `MUSA_VISIBLE_DEVICES`);
+- distributed launch (`MASTER_ADDR`, `MASTER_PORT`, `RANK`, and `WORLD_SIZE`);
+- Hugging Face cache and authentication (`HF_HOME`, `HF_TOKEN`, and related
+  names);
+- accelerator libraries such as NCCL, FlashInfer, MindIE, MORI, and quack.
+
+Consult the owning project's documentation for accepted values and support.
+`HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, and `OPENAI_API_KEY` are explicitly marked
+for redaction in the reviewed inventory. Never print their values in environment
+collection, examples, bug reports, or logs.
+
+## Model-specific variables
+
+The audit found 56 variables read by a single model or pipeline family. They are
+not listed as public usage options here because doing so would turn implementation
+escape hatches into an accidental compatibility contract.
+
+Every audited model-specific name has a migration disposition in the
+[environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py):
+
+| Disposition | Count | Required outcome |
+| --- | ---: | --- |
+| Promote | 33 | Move a stable setting into typed stage or model configuration. |
+| Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
+| External | 0 | Retain only when a supported third-party library owns the contract. |
+| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
+
+The disposition is a migration target, not a statement that the existing
+environment switch is stable. Promote or request-scope work should land in
+model-owner-reviewed follow-up changes with validation and model-page updates.
+The family-sized migration packages and their ownership gate are tracked in
+[#6232](https://github.com/vllm-project/vllm-omni/issues/6232).
+
+## Environment collection and troubleshooting
+
+Run the repository collector when preparing a bug report:
+
+```bash
+python collect_env.py
+```
+
+The environment section includes the safe public Omni allowlist, registered
+vLLM variables, and selected platform prefixes. It omits names marked for
+redaction and names containing common secret terms. Review the output before
+sharing it. A key supplied through stage `env` is not collected merely because
+it appears in deploy configuration; the same allowlist and prefix rules apply.
+
+If a setting appears ineffective, first check its read time in the tables
+above. Restart processes after changing import- or startup-time values, verify
+that the variable reaches the intended worker, and check whether an explicit
+typed configuration field takes precedence.
+
+## Benchmark and internal variables
+
+The reviewed inventory marks 20 benchmark variables as transitional while
+benchmark configuration moves to CLI/request schemas in
+[issue #5376](https://github.com/vllm-project/vllm-omni/issues/5376). They are not
+part of the serving/runtime contract.
+
+Internal process-coordination and debug variables are deliberately excluded
+from this page. The CI scanner resolves string literals and module-level string
+constants used by direct `os` access, local wrapper helpers, aliased `os`
+imports, and environment membership checks. Server-storage names generated by
+Pydantic settings are checked separately against their model fields.
+
+Names assembled from cross-module metadata and arbitrary keys supplied through
+stage `env` cannot be inferred reliably. Contributors must classify those names
+explicitly. CI fails when a statically resolved name has no classification or
+when a reviewed public name is absent from this page.

docs/configuration/stage_configs.md

@@ -14,7 +14,7 @@ describe fixed model topology and are not accepted in deploy YAMLs.
 Common `PipelineConfig` fields include:
 
 | Field | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `model_type` | Pipeline identifier used during model and config resolution. |
 | `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
 | `model_arch` | Default Hugging Face architecture for the pipeline. |
@@ -26,7 +26,7 @@ Common `PipelineConfig` fields include:
 Common `StagePipelineConfig` fields include:
 
 | Field | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `stage_id` | Stable stage identifier. |
 | `model_stage` | Logical stage name used by runtime and strategy resolution. |
 | `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. |
@@ -49,10 +49,10 @@ deployment overrides.
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
 
 | Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
+| ------- | ------ | ---------- | --------- | ------------- |
 | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
 | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
-| `session_mode` | str | optional | `"turn"` | Session behavior. Use `"duplex"` only with a pipeline that enables duplex control. |
+| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
 | `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
 | `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
 | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
@@ -84,7 +84,7 @@ subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
 Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`.
 
 | Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
+| ------- | ------ | ---------- | --------- | ------------- |
 | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
 | `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. |
 | `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. |
@@ -113,9 +113,9 @@ connectors:
       # Additional connector-specific options
 ```
 
-| Connector class | Use case | `extra` keys |
-|-----------------|----------|--------------|
-| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
+| Connector class          | Use case                                                              | `extra` keys                                                                                                      |
+|--------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------|
+| `SharedMemoryConnector`  | Same-host KV transfer between stages (default for bundled YAMLs).     | None. All payloads use shared memory.                                                                             |
 | `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
 
 A stage references a connector by name in its `input_connectors` / `output_connectors`:
@@ -135,7 +135,7 @@ stages:
 ### CLI flags
 
 | Flag | Description |
-|------|-------------|
+| ------ | ------------- |
 | `--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. |
@@ -249,7 +249,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 Effective config per stage after the merge:
 
 | Stage | Field | Final value | Source |
-|-------|-------|-------------|--------|
+| ------- | ------- | ------------- | -------- |
 | 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
 | 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
 | 0 | `max_model_len` | `16384` | global CLI |
@@ -273,12 +273,14 @@ when loading a deploy configuration.
 Examples:
 
 For offline inference (assuming the necessary dependencies have been imported):
+
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
 omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
 ```
 
 For online serving:
+
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
 ```

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -4,8 +4,8 @@ When you want to add L4-level ***performance test*** cases, add entries to JSON
 
 | Model type | Runner | Example JSON files |
 | ---------- | ------ | ------------------ |
-| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` in `mark`) |
-| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
+| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` without `slow` in `mark`) |
+| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_async_chunk.json` (CUDA only), `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
 | TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
 | Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
 
@@ -20,7 +20,7 @@ Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-o
 
 #### Running perf cases
 
-Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"` or `-m "slow and omni and H100"` for weekly Omni configs).
+Pass **`--test-config-file`** to load one JSON file, or omit it for the bulk scan above. In either mode, pytest **`-m`** filters each case using its JSON `mark`. A single file can contain nightly, weekly, and platform-specific case objects. For example, the nightly object in `test_qwen3_omni_async_chunk.json` declares both H100 and A3, while its `slow` weekly object is CUDA-only; there is no NPU weekly counterpart. Use `-m "H100 and full_model and not slow"`, `-m "H100 and full_model and slow"`, or `-m "npu"` to select the intended schedule and platform.
 
 ```JSON
 {
@@ -75,7 +75,7 @@ Omit `mark` only for configs not meant to be filtered by `-m`. `server_type` is
 
 Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
 
-When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`.
+When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `slow`, `omni`, `tts`, `diffusion`, or `local_model`.
 
 Supported form:
 
@@ -142,11 +142,18 @@ Examples:
 pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion"
 pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
 
-# Single file (same as nightly CI Perf steps)
+# Single file (same selectors as the CI Perf steps)
 pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
   --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
 pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
-  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
+  -m "H100 and full_model and not slow"
+pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
+  -m "H100 and full_model and slow"
+pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \
+  -m "npu"
 ```
 
 <!-- markdownlint-disable-next-line MD051 -->

docs/design/feature/cfg_parallel.md

@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 ### N-Branch CFG (3+ branches)
 
-Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
+Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
 
 `predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
 

docs/design/feature/host_weight_runtime.md

@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
 
 ## Status
 
-The first implementation provides contracts and a CPU local-filesystem store.
-The initial diffusion consumer contract additionally defines typed,
+The implementation provides contracts and a CPU local-filesystem store. The
+diffusion consumer additionally defines typed,
 representation-independent final-layout identity/restoration mechanics plus a
-concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
-library-only: no loader or DLO path selects, publishes, restores, or transports
-that artifact yet.
+concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
+selects, publishes, restores, and transfers these artifacts.
 
 V1 includes:
 
@@ -33,13 +33,12 @@ V1 includes:
 
 V1 does not include:
 
-- a public CLI, loader activation, or default-on model integration;
-- lease handoff to DLO or another transport consumer;
-- FP8, quantized, merged-adaptation, or additional model producers;
-- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
+- default-on activation or consumers outside no-AllGather DLO;
+- online FP8, quantized, merged-adaptation, or additional model producers;
+- HWR interaction with DLO AllGather;
 - a remote artifact provider or cross-node coordination;
 - automatic eviction; or
-- a change to DLO AllGather or no-AllGather behavior.
+- a change to DLO collective or execution behavior.
 
 ## Motivation and use cases
 
@@ -255,7 +254,8 @@ The contract is intentionally separate from loader activation:
 - `FinalLayoutBF16Producer` accepts only the matching identity context and a
   finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
   coordinate. Its BF16 policy preserves model-declared FP32 parameters and
-  buffers and revalidates MiniMax H3 mixed-precision invariants.
+  buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
+  FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
 
 Other representations reuse source identity, typed parallel identity, tensor
 ownership, and exact restoration only when their policy proves those semantics.

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -196,16 +196,15 @@ adopt it.
 
 ### Final-layout Host Weight Runtime consumer
 
-> **Status:** PR2 implementation. The representation-independent identity,
-> producer, and restoration contracts landed in
-> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
-> wires the final-layout BF16 path into eligible no-AllGather DLO.
+> **Status:** PR2 landed in
+> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
+> registered direct H2D plus generic HWR publication/source-digest
+> optimizations without changing AllGather behavior.
 
 Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
 only to no-AllGather DLO. Enable it with
 `--host-weight-runtime-mode preferred` (or `required`) and
-`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
-and direct asynchronous H2D remain a separate PR3 transport change.
+`--host-weight-runtime-root <node-local-root>`.
 
 The modes express operator fallback policy, not different artifact formats:
 
@@ -221,9 +220,15 @@ model revision and parallel layout with `required`. TP coordinates have distinct
 identities, while equivalent DP replicas share them.
 
 The current producer/consumer boundary is the model-declared final-layout BF16
-contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
-and TP2 rank identities plus SP layout identities; online quantization, HSDP,
-LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
+contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
+contract covers both transformer block stacks, constructor-stable packed QKV
+mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
+buffers. The shared HWR machinery supports ordinary-loader final layouts for
+TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
+in this change is TP1-only; its TP2/SP layouts rely on the existing
+per-coordinate identity mechanics and remain unmeasured. Online quantization,
+HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
+HWR mode
 `disabled`, DLO-disabled, and DLO AllGather configurations stop before source
 identity or store construction and retain the existing checkpoint-mmap or
 ordinary-loader path.
@@ -254,6 +259,44 @@ outcomes are authoritative:
 Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
 not selected.
 
+Local source identity and publication avoid repeated full-payload passes while
+preserving the same semantic contract:
+
+- Canonical files without a trusted immutable Hub blob identity are hashed in
+  parallel. Their content IDs are cached per storage domain under the HWR root
+  and reused only when the path, inode, size, timestamps, symlink target, and
+  cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
+  cache-I/O failure falls back to hashing the canonical file directly.
+- Producers that emit payload bytes in canonical storage-key order may declare
+  `ordered=True`. The filesystem writer then computes file and tensor SHA256
+  values during the write and overlaps payload `fsync` with later producer
+  work. Producers without that guarantee retain unordered writes and use a
+  parallel readback checksum fallback.
+- Manifest and `READY` publication still wait for every payload `fsync` and
+  checksum to complete. These optimizations do not weaken artifact validation
+  or change lease ownership.
+
+#### Registered mmap transport
+
+After DLO takes an HWR lease, the no-AllGather backend may register the lease's
+complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
+Registration is transport state: it does not change artifact identity, the
+store, tensor ownership, or H2D payload. On success, each tensor view copies
+directly into the existing rotating HBM block buffers and the two private host
+staging slots are not allocated.
+
+`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
+over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
+memory policy, unsupported platform/capability, over-budget mapping, or fully
+rolled-back registration error selects the existing two-slot staging path. A
+partial registration that cannot be rolled back aborts startup because closing
+the lease would unmap memory still owned by the platform.
+
+Direct checkpoint mmap remains unchanged and continues to use staging. It may
+require loader-owned per-block transforms, while the HWR artifact already
+contains final runtime bytes. DLO AllGather never receives an HWR final-layout
+lease and therefore never enters this registration path.
+
 #### Pre-service transaction boundary
 
 The startup transaction does not end at restore commit. A lease-backed model is
@@ -317,6 +360,7 @@ loader owns
   -> backend takes once before asynchronous work
   -> backend drains pending H2D work
   -> backend releases hooks, staging, and model references
+  -> backend unregisters every mapped range
   -> backend closes lease
 ```
 
@@ -326,13 +370,15 @@ it. Once the backend takes ownership, only backend abort or teardown may drain
 work and close it.
 
 The implementation keeps the final-layout lease through backend setup and
-initial prefetch. The backend uses the existing two bounded rank-local staging
-slots for the immutable HWR tensors, and its transactional `enable()` cleanup
-removes partial hooks and closes the lease before reporting a setup failure.
-Preferred mode then uses the runner's fresh canonical retry; required mode
-propagates the failure.
+initial prefetch. Successful registration copies directly from immutable HWR
+views; otherwise the backend uses the existing two bounded rank-local staging
+slots. Transactional cleanup drains device work, removes source references,
+unregisters mappings, and only then closes the lease. An unregistration failure
+retains both registration and lease for retry/process teardown. Preferred mode
+then uses the runner's fresh canonical retry; required mode propagates the
+failure.
 
-#### PR2 promotion gates
+#### PR2 and PR3 promotion gates
 
 - Warm hit performs zero ordinary DiT materialization and zero producer calls.
 - Shared warm finalization changes neither restored bytes nor backing pointers.
@@ -344,6 +390,10 @@ propagates the failure.
 - Disabled mode and AllGather emit zero final-layout HWR interaction.
 - The lease carrier rejects duplicate take and serialization.
 - Backend setup or prefetch failure drains asynchronous work before lease close.
+- Registration is all-or-nothing; rollback/unregistration failure never closes
+  a mapping still owned by the platform.
+- Registered HWR transport bypasses host staging while preserving output and
+  H2D payload; unsupported registration retains the bounded staging path.
 - Compatible checkpoint mmap remains an unchanged benchmark and control path.
 - Prewarm uses one matching producer cohort per TP/SP topology and storage
   domain.
@@ -394,11 +444,12 @@ reconstructs those tensors with their recorded layouts. Other online methods
 must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
-The Host Weight Runtime representation and publication contracts are merged;
-see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
-[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
-no-AllGather consumer above is the PR2 implementation; registration and direct
-asynchronous H2D remain a separate transport follow-up.
+The Host Weight Runtime representation, publication, and no-AllGather consumer
+contracts are merged; see
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
+[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
+direct H2D is an optional transport layer over that merged lease contract.
 
 ## Validation coverage
 
@@ -411,8 +462,12 @@ Current source-level validation includes:
 - ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
   sharding of finalized weights and scales;
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
+- ordered publication hashing, overlapped durability, unordered parallel
+  checksum fallback, and node-local source-digest reuse/invalidation;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;
+- read-only registration preflight, direct source-to-device copies, safe
+  fallback, partial rollback, retryable unregistration, and lease ordering;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - bounded component allocator-cache retention, conservative/forced release,

docs/design/metrics.md

@@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage
 
 - Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
 - Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
-- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
+- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals.
 - Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
 - Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
 
@@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col
 
 ### The Problem
 
-vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
 
 Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
 
@@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
                                    |
         +--------+--------+--------+--------+--------+
         |                                            |
-   vllm:omni_*                                    vllm:*
+   vllm_omni:*                                    vllm:*
    collectors                                  collectors
         |                                            |
    +----+--------+   +-----------+   +----------+   +-----------+
@@ -58,17 +58,30 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
 
 ### Data Flow
 
-There are four independent paths for metric collection.
+There are five independent paths for metric collection.
 
-**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
+**Path 1: Pipeline-level metrics (`vllm_omni:*`)**
 
 `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
 
 - `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
-- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
-- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
+- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`.
+- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter.
+- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`.
 
-**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
+**Path 2: Image and diffusion metrics**
+
+Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice.
+
+- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages.
+- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations.
+- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample.
+- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output.
+- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits.
+
+Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active.
+
+**Path 3: Audio modality metrics (`vllm_omni:audio_*`)**
 
 `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
 
@@ -76,7 +89,7 @@ There are four independent paths for metric collection.
 - `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
 - `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
 
-**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
+**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)**
 
 `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
 
@@ -87,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf
 
 Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
 
-**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
+**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)**
 
 The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
 
@@ -97,9 +110,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in
 
 ### Metric Registration and Lifecycle
 
-All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
+All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
 
-The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
+The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
 
 ## OmniPrometheusStatLogger Wrap
 
@@ -130,11 +143,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr
 
 ## Gating: `--log-stats`
 
-All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
+All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
 
 Behavior with `--log-stats=off` (default):
 
-- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
+- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them.
 - `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
 - The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
 
@@ -162,10 +175,43 @@ second.
 
 | Metric | Type | Labels | Description |
 |--------|------|--------|-------------|
-| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
-| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
-| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
-| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
+| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
+| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+
+### Image and diffusion service-level metrics
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing |
+| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded |
+| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage |
+| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission |
+| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages |
+| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs |
+| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs |
+| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs |
+| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) |
+| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend |
+
+### Diffusion execution breakdown
+
+Labels: `{model_name, stage, replica}`.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request |
+| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count |
+| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time |
+| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time |
+| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it |
+| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled |
+| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used |
+| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output |
+| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages |
+
+The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained.
 
 ### Audio (7)
 
@@ -173,13 +219,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label.
 
 | Metric | Type | Extra label | Description |
 |--------|------|-------------|-------------|
-| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
-| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
-| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
+| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated |
+| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
 
 ### Cross-stage transfer (4)
 
@@ -187,10 +233,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
 
 | Metric | Type | Description |
 |--------|------|-------------|
-| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
-| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
 
 ### LLM stage-level (wrapped `vllm:*`)
 
@@ -203,7 +249,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
   - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
 - Counters use the `_total` suffix (auto-appended by `prometheus_client`).
 - Sizes use the `_bytes` suffix.
-- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
 
 ## Logging vs. Prometheus
 

docs/design/module/host_weight_runtime.md

@@ -16,7 +16,7 @@ validation_paths:
   - tests/host_weight_runtime/**
 upstream_refs:
   - https://github.com/vllm-project/vllm-omni/issues/6414
-last_reviewed: 2026-08-21
+last_reviewed: 2026-08-25
 ---
 
 # Host Weight Runtime
@@ -84,6 +84,13 @@ another. `WeightArtifactIdentity` includes:
 - producer version, implementation fingerprint, manifest schema, and restorer
   schema.
 
+Loader integrations may persist content digests for canonical source shards in
+the selected storage domain. Reuse requires the same path, inode, size,
+timestamps, symlink target, and cache-record checksum. Corrupt records are
+rebuilt, and coordination or cache-I/O failure falls back to hashing the
+canonical source directly. This is an identity-computation optimization, not a
+second source of truth or an artifact substitution rule.
+
 The identity excludes transfer policy such as registered mmap versus private
 pinned staging. Those paths move the same host representation and must not
 produce duplicate artifacts.
@@ -237,6 +244,12 @@ Every publication computes full file and tensor hashes. Local lookup supports:
   and
 - `full_checksum`, which additionally reads and hashes every payload file.
 
+An ordered producer may compute file and tensor digests while payload bytes are
+written and overlap each closed payload's `fsync` with later producer work. An
+unordered producer retains parallel readback hashing. In both cases manifest
+and `READY` publication wait for every payload digest and `fsync`; the atomic
+publication and validation contract is unchanged.
+
 Filesystem verity is modeled but unsupported. Invalid artifacts receive an
 external deny marker before their shared lookup lock is released; a build owner
 then quarantines the immutable directory under the exclusive lock.

docs/mkdocs/hooks/generate_examples.py

@@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
         "text_to_speech",
         "text_to_video",
         "x_to_text",
-        "x_to_video_audio",
     }
 )
 

docs/models/supported_models.md

@@ -58,7 +58,6 @@ th {
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -94,7 +93,6 @@ th {
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -65,6 +65,7 @@ omni = Omni(
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
 | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
+| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
 
 ## Host-weight loading
 
@@ -111,6 +112,12 @@ perform a collective and uses the same rank-local transfer behavior.
 HWR is an opt-in startup optimization for models that declare the final-layout
 BF16 restore contract. Use it only with no-AllGather DLO:
 
+The validated BF16 model contracts currently cover MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
+class but has not been validated against this contract. FLUX.2-dev, online
+FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
+the validated scope.
+
 ```bash
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
@@ -178,12 +185,48 @@ returning to `required`.
 
 On a cold start, the canonical loader remains authoritative and publishes a
 validated final-layout artifact for later workers. A warm start restores the
-DiT final tensors without ordinary DiT materialization, then DLO streams them
-through the same two bounded host staging slots. The artifact identity includes
-the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
-layouts do not alias one another. Publication failure remains separate from a
-`preferred` serving startup; the next `required` startup provides the explicit
-artifact-availability check.
+DiT final tensors without ordinary DiT materialization. DLO then attempts to
+register the immutable mapping for direct H2D; if registration is unavailable,
+it streams through the same two bounded host staging slots. The artifact
+identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
+and distinct SP layouts do not alias one another. Publication failure remains
+separate from a `preferred` serving startup; the next `required` startup
+provides the explicit artifact-availability check.
+
+For local canonical checkpoints, the first eligible worker may hash source
+shards to establish immutable identity. HWR caches those digests in the same
+node-local storage domain and validates file metadata before reuse, so later
+workers normally avoid repeating that read. Cold BF16 publication also hashes
+ordered payloads as they are written and overlaps payload durability work with
+later shards; this changes startup work only, not artifact contents or runtime
+H2D behavior.
+
+#### Registered direct H2D
+
+Registration is attempted automatically only for an eligible warm HWR hit when
+`pin_cpu_memory` is enabled. A successful path is:
+
+```text
+shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
+```
+
+It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
+the two host staging slots. It does not reduce the H2D payload or make GPU
+kernels access host memory directly.
+
+CUDA requires read-only host-registration capability for the immutable HWR
+mapping. Unsupported capability, a positive registration limit smaller than
+the complete page-aligned mapping, or a safely rolled-back registration error
+falls back to bounded staging. Programmatic configurations can set
+`pin_cpu_memory=False` to disable registration explicitly. A successful
+registration locks the complete mapped range in host memory for that worker's
+lifetime, so use
+`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
+required.
+
+Shutdown drains H2D work, releases hook/source references, unregisters every
+range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
+retain their existing transfer paths.
 
 When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
 not resolve HWR sources or construct its store.

docs/user_guide/diffusion_features.md

@@ -129,7 +129,6 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -147,7 +146,7 @@ The following tables show which models support each feature:
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
 > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
@@ -165,7 +164,6 @@ The following tables show which models support each feature:
 | **SANA-Video-2B T2V I2V**   |     ❓     |     ❓      |           ❓           |       ❓        |         ❓         |         ❓         |   ❓    |             ❓             |          ❓           |       ❓        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
-| **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -1,115 +0,0 @@
-# X-To-Video-Audio
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
-
-
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
-
-## Local CLI Usage
-### Download the Model locally
-Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
-
-```bash
-python download_dreamid_omni.py --output-dir ./dreamid_omni
-```
-After download, the model directory will look like this:
-
-```
-dreamid_omni/
-├── DreamID-Omni/
-│   ├── dreamid_omni.safetensors
-├── MMAudio/
-│   ├── ext_weights/
-│   │   ├── best_netG.pt
-│   │   ├── v1-16.pth
-├── Wan2.2-TI2V-5B/
-│   ├── google/*
-│   ├── models_t5_umt5-xxl-enc-bf16.pth
-│   ├── Wan2.2_VAE.pth
-│
-├── model_index.json # create by download_dreamid_omni.py
-```
-
-### Run the Inference
-```python
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "Two people walking together and singing happily" \
-  --image-path ./example0.png ./example1.png \
-  --audio-path ./example0.wav ./example1.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --output out_dreamid_omni_twoip.mp4
-```
-In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
-The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "Two people walking together and singing happily" \
-  --image-path ./example0.png ./example1.png \
-  --audio-path ./example0.wav ./example1.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --use-hsdp \
-  --hsdp-shard-size 2 \
-  --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
-  --image-path 9.png \
-  --audio-path 9.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --output out_dreamid_omni_oneip.mp4
-```
-
-
-Key arguments:
-- `--prompt`: text description (string).
-- `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
-- `--image-path`: path to the input image list.
-- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
-- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--num-inference-steps`: number of denoising steps (defaults 45).
-- `--video-negative-prompt`: negative prompt for video generation.
-- `--audio-negative-prompt`: negative prompt for audio generation.
-- `--enable-cpu-offload`: enable CPU offload (defaults False).
-
-## Example materials
-
-??? abstract "download_dreamid_omni.py"
-    ``````py
-    --8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
-    ``````
-??? abstract "x_to_video_audio.py"
-    ``````py
-    --8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
-    ``````

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -1,193 +0,0 @@
-# Text-To-Audio
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
-
-This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
-
-## Supported Models
-
-| Model | Description |
-|-------|-------------|
-| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
-
-## Start Server
-
-### Basic Start
-
-```bash
-vllm-omni serve stabilityai/stable-audio-open-1.0 \
-    --host 0.0.0.0 \
-    --port 8091 \
-    --gpu-memory-utilization 0.9 \
-    --trust-remote-code \
-    --enforce-eager \
-    --omni
-```
-
-## API Calls
-
-### Method 1: Using curl
-
-```bash
-# Run all curl examples
-bash curl_examples.sh
-
-# Or execute directly
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "The sound of a cat purring",
-        "audio_length": 10.0
-    }' --output cat.wav
-```
-
-### Method 2: Using Python Client
-
-```bash
-cd examples/online_serving/stable_audio
-
-# Simple generation
-python stable_audio_client.py \
-    --text "The sound of a cat purring"
-
-# With custom duration
-python stable_audio_client.py \
-    --text "A dog barking" \
-    --audio_length 5.0
-
-# With all parameters
-python stable_audio_client.py \
-    --text "Thunder and rain" \
-    --audio_length 15.0 \
-    --negative_prompt "Low quality" \
-    --guidance_scale 7.0 \
-    --num_inference_steps 100 \
-    --seed 42 \
-    --output thunder.wav
-```
-
-The Python client supports the following command-line arguments:
-
-- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
-- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
-- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
-- `--audio_start`: Audio start time in seconds (default: `0.0`)
-- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
-- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
-- `--num_inference_steps`: Number of inference steps (default: `100`)
-- `--seed`: Random seed for reproducibility (default: `None`)
-- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
-- `--output`: Output file path (default: `stable_audio_output.wav`)
-
-### Method 3: Using Python httpx
-
-```python
-import httpx
-
-response = httpx.post(
-    "http://localhost:8091/v1/audio/generate",
-    json={
-        "input": "The sound of ocean waves crashing on a beach",
-        "audio_length": 10.0,
-        "negative_prompt": "Low quality, distorted",
-        "guidance_scale": 7.0,
-        "num_inference_steps": 100,
-    },
-    timeout=300.0,
-)
-
-with open("ocean.wav", "wb") as f:
-    f.write(response.content)
-```
-
-## Request Format
-
-### Simple Generation
-
-```json
-{
-    "input": "The sound of ocean waves"
-}
-```
-
-### Generation with Parameters
-
-```json
-{
-    "input": "A piano playing a gentle melody",
-    "audio_length": 10.0,
-    "negative_prompt": "Low quality, distorted, noisy",
-    "guidance_scale": 8.0,
-    "num_inference_steps": 150,
-    "seed": 42,
-    "response_format": "wav"
-}
-```
-
-## API Reference
-
-### Endpoint
-
-```
-POST /v1/audio/generate
-Content-Type: application/json
-```
-
-### Generation Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text prompt describing the audio to generate |
-| `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
-| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
-| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
-| `audio_start` | float | 0.0 | Audio start time in seconds |
-| `negative_prompt` | string | null | Text describing what to avoid in generation |
-| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
-| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
-| `seed` | int | null | Random seed for reproducible generation |
-
-### Response Format
-
-Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
-
-## Tuning Tips
-
-1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
-2. **Quality vs Speed**:
-   - 50 steps: Fast, decent quality (quick previews)
-   - 100 steps: Good balance (general purpose)
-   - 150+ steps: High quality, slower (final / critical audio)
-3. **Guidance Scale**:
-   - Lower (3 - 5): More creative / varied output
-   - Default (7): Good balance
-   - Higher (10+): Strict adherence to the prompt
-4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
-5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
-
-## File Description
-
-| File | Description |
-|------|-------------|
-| `curl_examples.sh` | Curl examples covering common use cases |
-| `stable_audio_client.py` | Python client with full CLI argument support |
-
-## Troubleshooting
-
-1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
-2. **Connection refused**: Make sure the server is running on the correct port.
-3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
-4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
-5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
-
-## Example materials
-
-??? abstract "stable_audio_client.py"
-    ``````py
-    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
-    ``````
-??? abstract "curl_examples.sh"
-    ``````sh
-    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
-    ``````

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -1,108 +0,0 @@
-import argparse
-import fcntl
-import os
-import site
-import subprocess
-import tempfile
-import time
-from pathlib import Path
-
-from huggingface_hub import snapshot_download
-
-DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
-DEPENDENCY_BRANCH = "omni"
-CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
-LOCK_FILE = CACHE_DIR / ".install.lock"
-DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
-
-
-def download_dependency():
-    CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
-    with open(LOCK_FILE, "w") as f:
-        fcntl.flock(f, fcntl.LOCK_EX)
-        if not DEPENDENCY_DIR.exists():
-            print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
-            subprocess.run(
-                ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
-                check=True,
-            )
-            print("Download finished.")
-        fcntl.flock(f, fcntl.LOCK_UN)
-
-    # write .pth to site-packages
-    site_packages = Path(site.getsitepackages()[0])
-    pth_file = site_packages / "vllm_omni_dependency.pth"
-    pth_file.write_text(str(DEPENDENCY_DIR))
-    print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
-
-
-def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
-    """Download files from HF repo and log time + destination."""
-    if os.path.exists(local_dir):
-        print(f"Directory {local_dir} already exists. Skipping download.")
-        return
-    print(f"Starting download from {repo_id} into {local_dir}")
-    start_time = time.time()
-
-    snapshot_download(
-        repo_id=repo_id,
-        local_dir=local_dir,
-        local_dir_use_symlinks=False,
-        allow_patterns=allow_patterns,
-    )
-
-    elapsed = time.time() - start_time
-    print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
-
-
-def main(output_dir: str):
-    # Wan2.2
-    wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
-    timed_download(
-        repo_id="Wan-AI/Wan2.2-TI2V-5B",
-        local_dir=wan_dir,
-        allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
-    )
-
-    # MMAudio
-    mm_audio_dir = os.path.join(output_dir, "MMAudio")
-    timed_download(
-        repo_id="hkchengrex/MMAudio",
-        local_dir=mm_audio_dir,
-        allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
-    )
-
-    dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
-
-    timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
-
-    # Now we construct the config file
-    import json
-
-    data = {
-        "_class_name": "DreamIDOmniPipeline",
-    }
-
-    with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
-        json.dump(data, f, indent=2)
-
-    print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
-
-    transformer_dir = os.path.join(output_dir, "transformer")
-    os.makedirs(transformer_dir, exist_ok=True)
-    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
-        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
-    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
-
-    # now we download the dependency code
-    download_dependency()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Download models from Hugging Face")
-    parser.add_argument(
-        "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
-    )
-    args = parser.parse_args()
-    main(args.output_dir)

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -1,163 +0,0 @@
-# X-To-Video-Audio
-
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
-
-## Local CLI Usage
-### Download the Model locally
-Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
-
-```bash
-python download_dreamid_omni.py --output-dir ./dreamid_omni
-```
-After download, the model directory will look like this:
-
-```
-dreamid_omni/
-├── DreamID-Omni/
-│   ├── dreamid_omni.safetensors
-├── MMAudio/
-│   ├── ext_weights/
-│   │   ├── best_netG.pt
-│   │   ├── v1-16.pth
-├── Wan2.2-TI2V-5B/
-│   ├── google/*
-│   ├── models_t5_umt5-xxl-enc-bf16.pth
-│   ├── Wan2.2_VAE.pth
-│
-├── model_index.json
-└── transformer/
-    └── config.json   # create by download_dreamid_omni.py
-```
-
-### Run the Inference
-```python
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "Two people walking together and singing happily" \
-  --image-path ./example0.png ./example1.png \
-  --audio-path ./example0.wav ./example1.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --output out_dreamid_omni_twoip.mp4
-```
-In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
-The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "Two people walking together and singing happily" \
-  --image-path ./example0.png ./example1.png \
-  --audio-path ./example0.wav ./example1.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --use-hsdp \
-  --hsdp-shard-size 2 \
-  --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
-  --image-path 9.png \
-  --audio-path 9.wav \
-  --video-negative-prompt "jitter, bad hands, blur, distortion" \
-  --audio-negative-prompt "robotic, muffled, echo, distorted" \
-  --cfg-parallel-size 2 \
-  --num-inference-steps 45 \
-  --height 704 \
-  --width 1280 \
-  --output out_dreamid_omni_oneip.mp4
-```
-
-
-## MagiHuman (`--model-type magi-human`)
-
-MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
-T5Gemma text encoder. A detailed text prompt is the only required input; an optional
-image and/or audio file may be supplied for conditioning. Natively supports Tensor
-Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
-the MoE weights and the text encoder.
-
-> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
-> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
-
-Text-only generation:
-
-```bash
-python x_to_video_audio.py \
-  --model-type magi-human \
-  --model /path/to/daVinci-MagiHuman \
-  --prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
-  --tensor-parallel-size 4 \
-  --height 256 --width 448 \
-  --num-inference-steps 8 \
-  --seed 52 \
-  --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
-  --output output_magihuman.mp4
-```
-
-With optional image and audio conditioning:
-
-```bash
-python x_to_video_audio.py \
-  --model-type magi-human \
-  --model /path/to/daVinci-MagiHuman \
-  --prompt "A young woman..." \
-  --tensor-parallel-size 4 \
-  --height 256 --width 448 \
-  --num-inference-steps 8 \
-  --seed 52 \
-  --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
-  --output output_magihuman.mp4
-```
-
-MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
-in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
-
-- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
-  because the audio length then determines the number of frames).
-- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
-  the audio drives the frame count and is encoded as a latent condition into the DiT
-  (`is_a2v` mode). Omit for pure text-to-video+audio generation.
-- `image_path`: path to an image file for visual conditioning. Applied at both the
-  base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
-- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
-  when these are omitted.
-- `sr_num_inference_steps`: denoising steps for the SR stage.
-
-## DreamID-Omni arguments
-
-Key arguments:
-- `--prompt`: text description (string).
-- `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
-- `--image-path`: path to the input image list.
-- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
-- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--num-inference-steps`: number of denoising steps (defaults 45).
-- `--video-negative-prompt`: negative prompt for video generation.
-- `--audio-negative-prompt`: negative prompt for audio generation.
-- `--enable-cpu-offload`: enable CPU offload (defaults False).
-- `--cache-backend`: enable `cache_dit` for acceleration.
-- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
-  (VAEs, the T5 text encoder, norms and modulation stay bf16.)

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -1,279 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import json
-import re
-import time
-from typing import Any
-
-import numpy as np
-from PIL import Image
-from vllm.multimodal.media.audio import load_audio
-
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
-AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
-    parser.add_argument("--model", required=True, help="Model ckpt root directory.")
-    parser.add_argument(
-        "--model-type",
-        default="dreamid-omni",
-        choices=["dreamid-omni", "magi-human"],
-        help="Model type.",
-    )
-    parser.add_argument("--prompt", default=None, help="Text prompt.")
-
-    parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
-    parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
-    parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
-    parser.add_argument(
-        "--extra-body",
-        type=str,
-        default=None,
-        help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
-        'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
-    )
-
-    parser.add_argument("--height", type=int, default=704, help="Video height.")
-    parser.add_argument("--width", type=int, default=1280, help="Video width.")
-    parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
-    parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
-    parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
-    parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
-    parser.add_argument(
-        "--cfg-parallel-size",
-        type=int,
-        default=1,
-        choices=[1, 2, 3, 4],
-        help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
-    )
-    parser.add_argument(
-        "--tensor-parallel-size",
-        type=int,
-        default=1,
-        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
-    )
-    parser.add_argument(
-        "--video-negative-prompt",
-        default="jitter, bad hands, blur, distortion",
-        help="Negative prompt for video.",
-    )
-    parser.add_argument(
-        "--audio-negative-prompt",
-        default="robotic, muffled, echo, distorted",
-        help="Negative prompt for audio.",
-    )
-    parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "int8"],
-        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
-    )
-    parser.add_argument(
-        "--enable-cpu-offload",
-        action="store_true",
-        default=False,
-        help="Enable CPU offloading for diffusion models.",
-    )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        action="store_true",
-        help="Enable layerwise (blockwise) offloading on DiT modules.",
-    )
-    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
-    parser.add_argument(
-        "--use-hsdp",
-        action="store_true",
-        help="Enable HSDP for DreamID-Omni fused transformer blocks.",
-    )
-    parser.add_argument(
-        "--hsdp-shard-size",
-        type=int,
-        default=1,
-        help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
-    )
-    parser.add_argument(
-        "--hsdp-replicate-size",
-        type=int,
-        default=1,
-        help="Number of HSDP replica groups. Default 1 means pure sharding.",
-    )
-    return parser.parse_args()
-
-
-def load_image_and_audio(image_paths, audio_paths):
-    image = []
-    audio = []
-
-    for path in image_paths:
-        with Image.open(path) as img:
-            img = img.convert("RGB")
-            image.append(img)
-
-    for path in audio_paths:
-        audio_array, sr = load_audio(path, sr=16000)
-        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
-        audio.append(audio_array)
-    return image, audio
-
-
-def _extract_peak_memory_mb(result: Any) -> float:
-    """Pull worker-reported peak VRAM (MiB) generation result.
-
-    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
-    """
-    if isinstance(result, list):
-        result = result[0] if result else None
-    if result is None:
-        return 0.0
-    val = getattr(result, "peak_memory_mb", 0.0)
-    if not val:
-        inner = result
-        if isinstance(inner, list):
-            inner = inner[0] if inner else None
-        val = getattr(inner, "peak_memory_mb", 0.0)
-    try:
-        return float(val or 0.0)
-    except (TypeError, ValueError):
-        return 0.0
-
-
-def main() -> None:
-    args = parse_args()
-    if args.prompt is None and args.prompt_file is None:
-        raise ValueError("Either --prompt or --prompt-file must be provided.")
-
-    text_prompt = args.prompt
-    if args.prompt_file:
-        with open(args.prompt_file) as f:
-            text_prompt = json.load(f)
-            text_prompt = re.sub(
-                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
-            ).strip()
-            text_prompt = re.sub(
-                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
-            ).strip()
-            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
-            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
-
-    if args.model_type == "magi-human":
-        prompt = text_prompt
-        extra_args = json.loads(args.extra_body) if args.extra_body else {}
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args=extra_args,
-        )
-    else:
-        image, audio = load_image_and_audio(args.image_path, args.audio_path)
-
-        prompt = {
-            "prompt": text_prompt,
-            "video_negative_prompt": args.video_negative_prompt,
-            "audio_negative_prompt": args.audio_negative_prompt,
-            "multi_modal_data": {"image": image, "audio": audio},
-        }
-
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args={
-                "solver_name": args.solver_name,
-                "shift": args.shift,
-            },
-        )
-
-    cache_config = None
-    if args.cache_backend == "cache_dit":
-        cache_config = {
-            "Fn_compute_blocks": 1,
-            "Bn_compute_blocks": 0,
-            "max_warmup_steps": 4,
-            "max_cached_steps": 20,
-            "residual_diff_threshold": 0.24,
-            "max_continuous_cached_steps": 3,
-            "enable_taylorseer": False,
-            "taylorseer_order": 1,
-            "scm_steps_mask_policy": None,
-            "scm_steps_policy": "dynamic",
-        }
-
-    omni_kwargs = dict(
-        model=args.model,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-        model_type=args.model_type,
-        enable_cpu_offload=args.enable_cpu_offload,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        cache_backend=args.cache_backend,
-        cache_config=cache_config,
-    )
-    if args.quantization is not None:
-        omni_kwargs["quantization"] = args.quantization
-    omni = Omni(**omni_kwargs)
-    start = time.perf_counter()
-    outputs = omni.generate(prompt, sampling_params)
-    elapsed = time.perf_counter() - start
-
-    if not outputs:
-        raise RuntimeError("No output returned from the model.")
-    result = outputs[0]
-    if not result.images:
-        raise RuntimeError("No video frames found in model output.")
-    generated_video = result.images[0]
-    mm = result.multimodal_output or {}
-    generated_audio = mm.get("audio")
-
-    if args.model_type == "magi-human":
-        # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
-        fps = float(mm.get("fps", 25))
-        sample_rate = int(mm.get("audio_sample_rate", 24000))
-        frames = generated_video
-    else:
-        fps = float(mm.get("fps", 24))
-        sample_rate = int(mm.get("audio_sample_rate", 16000))
-        # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
-        # mux_video_audio_bytes expects (F, H, W, C) uint8.
-        if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
-            raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
-        frames = generated_video.transpose(1, 2, 3, 0)
-        frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
-
-    audio_np = None
-    if generated_audio is not None:
-        audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
-
-    output_path = args.output
-    video_bytes = mux_video_audio_bytes(
-        frames,
-        audio_np,
-        fps=float(fps),
-        audio_sample_rate=sample_rate,
-    )
-    with open(output_path, "wb") as f:
-        f.write(video_bytes)
-    print(f"Saved generated video to {output_path}")
-    print(f"Total time: {elapsed:.2f}s")
-    peak_mb = _extract_peak_memory_mb(outputs)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/minicpmo/README.md

@@ -32,12 +32,14 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
 profile admits at most four concurrent sequences per stage.
 
 | deploy config | GPUs | Notes |
-|---|---|---|
+| --- | --- | --- |
 | `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
 | `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
 | `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
 | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
-| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
+
+Every profile sets `session_mode: duplex`, so native full-duplex is served
+from the same process as `/v1/chat/completions`.
 
 The split pipeline preserves native-duplex epoch/turn identity, segment text,
 turn completion, reference voice, and terminal-audio metadata through
@@ -55,8 +57,8 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
 ```
 
 For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
-checkpoint path. To start the experimental native duplex backend, use
-`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`.
+checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on
+this same server.
 
 ### Per-stage overrides
 
@@ -145,7 +147,7 @@ assistant template, so it is not an apples-to-apples accuracy run.
 
 ## Run the Realtime duplex CLI demo
 
-After the duplex backend is running, stream one WAV through the Realtime
+After the server is running, stream one WAV through the Realtime
 WebSocket endpoint:
 
 ```bash

examples/online_serving/minicpmo/realtime_duplex_demo.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
 
 Run this after starting the duplex server. Strict lifecycle, overlap, and
@@ -19,8 +22,8 @@ if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
 from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
-    PCM16_BYTES_PER_SAMPLE,
-    PCM16_SAMPLE_RATE,
+    PCM16_BYTES_PER_SAMPLE,  # noqa: F401 - compatibility export
+    PCM16_SAMPLE_RATE,  # noqa: F401 - compatibility export
     RealtimeDuplexClient,
     RealtimeEventCollector,
     build_realtime_url,
@@ -28,6 +31,13 @@ from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
     wait_for,
     write_pcm16_wav,
 )
+from vllm_omni.experimental.fullduplex.client import chunk_period_ms as _chunk_period_ms  # noqa: E402
+from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
+    has_residual_model_unit as _has_residual_model_unit,
+)
+from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
+    reference_audio_data_url as _ref_audio_data_url,
+)
 
 
 class _StreamingOutputWriter:
@@ -133,25 +143,6 @@ def _latest_model_decision(
     return decision
 
 
-def _chunk_period_ms(events: list[dict[str, object]]) -> int:
-    for event in reversed(events):
-        session = event.get("session")
-        if not isinstance(session, dict):
-            continue
-        capabilities = session.get("capabilities")
-        if not isinstance(capabilities, dict):
-            continue
-        chunk_period_ms = capabilities.get("chunk_period_ms")
-        if isinstance(chunk_period_ms, int) and chunk_period_ms > 0:
-            return chunk_period_ms
-    return 1000
-
-
-def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
-    unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
-    return bool(unit_bytes > 0 and len(pcm16) % unit_bytes)
-
-
 def _response_in_progress(events: list[dict[str, object]]) -> bool:
     return sum(event.get("type") == "response.created" for event in events) > sum(
         event.get("type") == "response.done" for event in events
@@ -168,13 +159,6 @@ def _event_count_after(
     return sum(event.get("type") == event_type for event in events[index + 1 :])
 
 
-def _ref_audio_data_url(path: str | None) -> str | None:
-    if path is None:
-        return None
-    ref_path = Path(path).expanduser()
-    return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
-
-
 async def run_demo(args: argparse.Namespace) -> dict[str, object]:
     input_pcm16 = read_pcm16_wav(Path(args.input_wav))
     if not input_pcm16:

examples/online_serving/stable_audio/README.md

@@ -1,234 +0,0 @@
-# Stable Audio Online Serving
-
-Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
-
-## Features
-
-- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
-- **Flexible control**: Adjust audio length, guidance scale, inference steps
-- **Quality control**: Use negative prompts to avoid unwanted characteristics
-- **Reproducible**: Set random seed for deterministic generation
-
-## Quick Start
-
-### 1. Start the Server
-
-```bash
-vllm-omni serve stabilityai/stable-audio-open-1.0 \
-    --host 0.0.0.0 \
-    --port 8091 \
-    --gpu-memory-utilization 0.9 \
-    --trust-remote-code \
-    --enforce-eager \
-    --omni
-```
-
-### 2. Generate Audio
-
-#### Using curl
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "The sound of a cat purring",
-        "audio_length": 10.0
-    }' --output cat.wav
-```
-
-#### Using Python Client
-
-```bash
-python stable_audio_client.py \
-    --text "The sound of a cat purring" \
-    --audio_length 10.0 \
-    --output cat.wav
-```
-
-#### Using Bash Script
-
-```bash
-bash curl_examples.sh
-```
-
-## API Reference
-
-### Endpoint
-
-```
-POST /v1/audio/generate
-```
-
-### Request Body
-
-```json
-{
-    "input": "Text description of the audio",
-    "audio_length": 10.0,
-    "audio_start": 0.0,
-    "negative_prompt": "Low quality",
-    "guidance_scale": 7.0,
-    "num_inference_steps": 100,
-    "seed": 42,
-    "response_format": "wav"
-}
-```
-
-### Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text prompt describing the audio to generate |
-| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
-| `audio_start` | float | 0.0 | Audio start time in seconds |
-| `negative_prompt` | string | null | Text describing what to avoid in generation |
-| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
-| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
-| `seed` | int | null | Random seed for reproducibility |
-| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
-
-### Response
-
-Returns audio data in the requested format (default: WAV).
-
-## Usage Examples
-
-### Basic Generation
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "The sound of ocean waves"
-    }' --output ocean.wav
-```
-
-### Custom Duration
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "A dog barking",
-        "audio_length": 5.0
-    }' --output dog_5s.wav
-```
-
-### High Quality with Negative Prompt
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "A piano playing a gentle melody",
-        "audio_length": 10.0,
-        "negative_prompt": "Low quality, distorted, noisy",
-        "guidance_scale": 8.0,
-        "num_inference_steps": 150
-    }' --output piano_hq.wav
-```
-
-### Reproducible Generation
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Thunder and rain sounds",
-        "audio_length": 15.0,
-        "seed": 42
-    }' --output thunder.wav
-```
-
-### Quick Generation (Fewer Steps)
-
-For faster generation with slightly lower quality:
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Birds chirping in a forest",
-        "audio_length": 8.0,
-        "num_inference_steps": 50
-    }' --output birds_quick.wav
-```
-
-## Python Client Examples
-
-### Simple Generation
-
-```bash
-python stable_audio_client.py \
-    --text "The sound of a cat purring"
-```
-
-### Custom Parameters
-
-```bash
-python stable_audio_client.py \
-    --text "Thunder and rain" \
-    --audio_length 15.0 \
-    --negative_prompt "Low quality" \
-    --guidance_scale 7.0 \
-    --num_inference_steps 100 \
-    --seed 42 \
-    --output thunder.wav
-```
-
-### Different Output Format
-
-```bash
-python stable_audio_client.py \
-    --text "Guitar playing" \
-    --response_format mp3 \
-    --output guitar.mp3
-```
-
-## Tips
-
-1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
-2. **Quality vs Speed**:
-   - 50 steps: Fast, decent quality
-   - 100 steps: Good balance (default)
-   - 150+ steps: High quality, slower
-3. **Guidance Scale**:
-   - Lower (3-5): More creative/varied
-   - Default (7): Good balance
-   - Higher (10+): More literal to prompt
-4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
-5. **Seeds**: Use same seed for reproducible results
-
-## Performance
-
-| Inference Steps | Quality | Speed | Use Case |
-|----------------|---------|-------|----------|
-| 50 | Good | Fast | Quick previews |
-| 100 (default) | Very Good | Medium | Production |
-| 150+ | Excellent | Slow | Final/critical audio |
-
-## Troubleshooting
-
-### Server not responding
-- Check if server is running: `curl http://localhost:8091/health`
-- Check server logs for errors
-
-### Audio quality issues
-- Increase `num_inference_steps` (e.g., 150)
-- Add negative prompts: `"Low quality, distorted, noisy"`
-- Increase `guidance_scale` for more prompt adherence
-
-### Generation timeout
-- Reduce `num_inference_steps`
-- Reduce `audio_length`
-- Check GPU memory with `nvidia-smi`
-
-### Wrong audio length
-- Ensure `audio_length` is within model limits (~47s max)
-- Adjust `audio_start` if trimming is needed
-
-## See Also
-
-- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
-- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
-- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)

examples/online_serving/stable_audio/curl_examples.sh

@@ -1,54 +0,0 @@
-#!/bin/bash
-# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
-
-# Example 1: Simple request with default parameters
-echo "Example 1: Simple request with default parameters"
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "The sound audience clapping and cheering in a stadium"
-    }' --output stadium.wav
-
-# Example 2: Request with custom audio_length
-echo "Example 2: Custom audio length (5 seconds)"
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "The sound of a dog barking",
-        "audio_length": 5.0
-    }' --output dog_5s.wav
-
-# Example 3: Request with negative prompt for quality control
-echo "Example 3: With negative prompt"
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "A piano playing a gentle melody",
-        "audio_length": 10.0,
-        "negative_prompt": "Low quality, distorted, noisy"
-    }' --output piano.wav
-
-# Example 4: Full control with all parameters
-echo "Example 4: Full control (custom length, guidance, steps, seed)"
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Thunder and rain sounds",
-        "audio_length": 15.0,
-        "negative_prompt": "Low quality",
-        "guidance_scale": 7.0,
-        "num_inference_steps": 100,
-        "seed": 42
-    }' --output thunder_rain.wav
-
-# Example 5: Quick generation with fewer steps (faster but lower quality)
-echo "Example 5: Quick generation (fewer steps)"
-curl -X POST http://localhost:8091/v1/audio/generate \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Ocean waves crashing on a beach",
-        "audio_length": 8.0,
-        "num_inference_steps": 50
-    }' --output ocean.wav
-
-echo "All examples completed!"

examples/online_serving/stable_audio/stable_audio_client.py

@@ -1,170 +0,0 @@
-#!/usr/bin/env python3
-"""
-OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
-
-This script demonstrates how to use the OpenAI-compatible speech API
-to generate audio from text using Stable Audio models.
-
-Examples:
-    # Simple generation
-    python stable_audio_client.py --text "The sound of a cat purring"
-
-    # With custom duration
-    python stable_audio_client.py --text "A dog barking" --audio_length 5.0
-
-    # With all parameters
-    python stable_audio_client.py --text "Thunder and rain" \
-        --audio_length 15.0 \
-        --negative_prompt "Low quality" \
-        --guidance_scale 7.0 \
-        --num_inference_steps 100 \
-        --seed 42 \
-        --output thunder.wav
-"""
-
-import argparse
-import sys
-
-import requests
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
-    parser.add_argument(
-        "--api_url",
-        default="http://localhost:8091/v1/audio/generate",
-        help="API endpoint URL",
-    )
-    parser.add_argument(
-        "--text",
-        default="The sound of a cat purring",
-        help="Text prompt for audio generation",
-    )
-    parser.add_argument(
-        "--audio_length",
-        type=float,
-        default=10.0,
-        help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
-    )
-    parser.add_argument(
-        "--audio_start",
-        type=float,
-        default=0.0,
-        help="Audio start time in seconds",
-    )
-    parser.add_argument(
-        "--negative_prompt",
-        default="Low quality",
-        help="Negative prompt for classifier-free guidance",
-    )
-    parser.add_argument(
-        "--guidance_scale",
-        type=float,
-        default=7.0,
-        help="Guidance scale for diffusion (higher = more adherence to prompt)",
-    )
-    parser.add_argument(
-        "--num_inference_steps",
-        type=int,
-        default=100,
-        help="Number of inference steps (higher = better quality, slower)",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for reproducibility",
-    )
-    parser.add_argument(
-        "--output",
-        default="stable_audio_output.wav",
-        help="Output file path",
-    )
-    parser.add_argument(
-        "--response_format",
-        default="wav",
-        choices=["wav", "mp3", "flac", "pcm"],
-        help="Audio output format",
-    )
-    return parser.parse_args()
-
-
-def generate_audio(args):
-    """Generate audio using the API."""
-
-    # Build request payload
-    payload = {
-        "input": args.text,
-        "audio_length": args.audio_length,
-        "audio_start": args.audio_start,
-        "response_format": args.response_format,
-    }
-
-    # Add optional parameters
-    if args.negative_prompt:
-        payload["negative_prompt"] = args.negative_prompt
-    if args.guidance_scale:
-        payload["guidance_scale"] = args.guidance_scale
-    if args.num_inference_steps:
-        payload["num_inference_steps"] = args.num_inference_steps
-    if args.seed is not None:
-        payload["seed"] = args.seed
-
-    print(f"\n{'=' * 60}")
-    print("Stable Audio - Text-to-Audio Generation")
-    print(f"{'=' * 60}")
-    print(f"API URL: {args.api_url}")
-    print(f"Prompt: {args.text}")
-    print(f"Audio length: {args.audio_length}s")
-    print(f"Negative prompt: {args.negative_prompt}")
-    print(f"Guidance scale: {args.guidance_scale}")
-    print(f"Inference steps: {args.num_inference_steps}")
-    if args.seed is not None:
-        print(f"Seed: {args.seed}")
-    print(f"Output: {args.output}")
-    print(f"{'=' * 60}\n")
-
-    try:
-        # Make the API request
-        print("Generating audio...")
-        response = requests.post(
-            args.api_url,
-            json=payload,
-            headers={"Content-Type": "application/json"},
-            timeout=300,  # 5 minute timeout for long generations
-        )
-
-        # Check for errors
-        if response.status_code != 200:
-            print(f"Error: API returned status code {response.status_code}")
-            print(f"Response: {response.text}")
-            return False
-
-        # Save the audio
-        with open(args.output, "wb") as f:
-            f.write(response.content)
-
-        print(f"✓ Audio saved to {args.output}")
-        print(f"  File size: {len(response.content) / 1024:.1f} KB")
-        return True
-
-    except requests.exceptions.Timeout:
-        print("Error: Request timed out. Try reducing inference steps or audio length.")
-        return False
-    except requests.exceptions.ConnectionError:
-        print(f"Error: Could not connect to {args.api_url}")
-        print("Make sure the server is running.")
-        return False
-    except Exception as e:
-        print(f"Error: {e}")
-        return False
-
-
-def main():
-    args = parse_args()
-    success = generate_audio(args)
-    sys.exit(0 if success else 1)
-
-
-if __name__ == "__main__":
-    main()

recipes/MiniMaxAI/MiniMax-H3.md

@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
 export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
 ```
 
-Start from a non-offloaded FL2VA server command and add
+Start from a non-offloaded or DLO FL2VA server command and add
 `--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
 `--lora-path` preloads the adapter; each request still activates it and uses
 the published sampling settings:
@@ -853,9 +853,10 @@ the published sampling settings:
 
 For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
 ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
-dynamic-only and does not support prefusion, DLO, or LoRA composition.
-It also rejects model-level, layerwise, and distributed layerwise offload;
-the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
+dynamic-only and does not support prefusion or LoRA composition. DLO is
+supported by keeping the request-switchable LoRA A/B buffers resident on the
+accelerator while DLO streams only the base blocks; budget for this additional
+fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
 The five requested sigma points produce the four denoiser evaluations expected
 by the Turbo artifact.
 

recipes/README.md

@@ -85,6 +85,7 @@ recipes/
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
+| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |

recipes/StabilityAI/Stable-Audio-Open.md

@@ -23,8 +23,6 @@ offline validation sample with TeaCache and online serving through the
 - Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0>
 - Offline example:
   [`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
-- Online example:
-  [`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
 - Audio generation API:
   [`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
 - Related issue:
@@ -103,7 +101,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \
     "seed": 42,
     "response_format": "wav"
   }' \
-  --output examples/online_serving/stable_audio/piano_10s.wav
+  --output piano_10s.wav
 ```
 
 ##### Verification
@@ -137,12 +135,12 @@ PY
 Validate the online output:
 
 ```bash
-ls -lh examples/online_serving/stable_audio/piano_10s.wav
+ls -lh piano_10s.wav
 
 python - <<'PY'
 import soundfile as sf
 
-path = "examples/online_serving/stable_audio/piano_10s.wav"
+path = "piano_10s.wav"
 audio, sample_rate = sf.read(path)
 print("sample_rate:", sample_rate)
 print("shape:", audio.shape)

recipes/XiaomiMiMo/MiMo-Audio.md

@@ -0,0 +1,161 @@
+# MiMo-Audio for omni audio generation and understanding
+
+> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
+
+## Summary
+
+- Vendor: Xiaomi MiMo
+- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
+- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for running
+MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
+The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
+FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
+Stage-1 GPU pressure.
+
+## References
+
+- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
+- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
+- Offline example:
+  [`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
+- Online example:
+  [`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
+- Default deploy config (generic single GPU):
+  [`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
+- 5090D-tuned deploy config:
+  [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+- Pipeline / deploy docs:
+  [Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
+  [PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
+
+Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
+offline with `tts_sft`.
+
+#### Environment
+
+- OS: Ubuntu (Linux x86_64)
+- Python: 3.12
+- PyTorch: 2.11.0+cu129
+- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
+- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
+- vLLM: 0.26.0+cu129
+- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
+
+Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
+On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
+PyTorch 2.11+cu129.
+
+Download or cache:
+
+- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
+
+Export the runtime knobs below before offline or online runs:
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
+export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
+export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
+export VLLM_USE_FLASHINFER_SAMPLER=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+```
+
+Optional (mirror / download stability):
+
+```bash
+export HF_HUB_DISABLE_XET=1
+export HF_ENDPOINT=https://hf-mirror.com   # if you use a Hugging Face mirror
+```
+
+Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
+`attention_backend: TRITON_ATTN` in the deploy YAML instead.
+
+#### Command
+
+Offline TTS validation from the repository root:
+
+```bash
+cd examples/offline_inference/mimo_audio
+
+python3 -u end2end.py \
+  --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --text "The weather is so nice today."
+```
+
+Online serving (requires MiMo chat template):
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+
+vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
+  --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+  --served-model-name "MiMo-Audio-7B-Instruct" \
+  --port 18091 \
+  --chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
+```
+
+Pass the same environment exports to the server process when using a CPU
+tokenizer or disabling FlashInfer sampling.
+
+#### Verification
+
+Offline:
+
+```bash
+ls -lh output_audio/tts_sft/*.wav
+```
+
+Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
+with the decoded text. A successful run completes both stages without a
+Stage-1 `decode_vq` device mismatch.
+
+Online:
+
+```bash
+curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
+```
+
+Then use the client under
+[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
+
+#### Notes
+
+- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
+  with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
+- Key flags:
+  - `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
+    and `TRITON_ATTN`.
+  - `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
+  - `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
+    [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
+    fix) so code2wav decode uses the tokenizer device.
+  - `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
+    issues on consumer GPUs.
+- Known limitations:
+  - Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
+    tooling; the tested path uses TRITON attention and cu129 wheels.
+  - Without a working `flash-attn` build, audio quality may be slightly
+    metallic; install FlashAttention when a compatible wheel is available.
+  - First request triggers Triton JIT for rotary / attention kernels (latency
+    spike); subsequent requests are faster.
+  - For longer outputs or batching, you may need to raise `max_model_len` in
+    the deploy YAML and match `max_position_embeddings` in the model config
+    (see the offline example README).

tests/benchmarks/metrics/test_metrics.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for metrics.py
 """
 
+import math
+
 import pytest
 from vllm.benchmarks.serve import TaskType
 
@@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
     return output
 
 
+def _calculate_test_metrics(outputs, goodput=None):
+    return calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=1.0,
+        tokenizer=None,
+        selected_percentiles=[50.0],
+        goodput_config_dict=goodput or {},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=[],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=1.0,
+    )[0]
+
+
 # ============================================================================
 # total_input Tests
 # ============================================================================
@@ -126,6 +144,166 @@ def test_audio_continuity_aggregation():
     assert p99 is not None and p99 > 0.4
 
 
+def test_unmeasured_duplex_latency_does_not_add_zero_samples():
+    measured = _make_output(100, output_tokens=2)
+    measured.ttft = 0.1
+    measured.audio_ttfp = 0.2
+    measured.audio_rtf = 0.5
+    measured.duplex_session_metrics = {
+        "mean_ttft_ms": 100.0,
+        "mean_ttfp_ms": 200.0,
+        "mean_rtf": 0.5,
+    }
+    listen_only = _make_output(100, output_tokens=0)
+    listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0
+    listen_only.duplex_session_metrics = {
+        "mean_ttft_ms": None,
+        "mean_ttfp_ms": None,
+        "mean_rtf": None,
+    }
+
+    metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0})
+
+    assert metrics.mean_ttft_ms == 100.0
+    assert metrics.mean_audio_ttfp_ms == 200.0
+    assert metrics.mean_audio_rtf == 0.5
+    assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1)
+    assert metrics.request_goodput == 1.0
+
+
+def test_all_unmeasured_duplex_latency_is_not_reported_as_zero():
+    listen_only = _make_output(100, output_tokens=0)
+    listen_only.duplex_session_metrics = {
+        "mean_ttft_ms": None,
+        "mean_ttfp_ms": None,
+        "mean_rtf": None,
+    }
+
+    metrics = _calculate_test_metrics(
+        [listen_only],
+        {"ttft": float("inf"), "audio_ttft": float("inf")},
+    )
+
+    assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0)
+    assert math.isnan(metrics.mean_ttft_ms)
+    assert math.isnan(metrics.mean_audio_ttfp_ms)
+    assert math.isnan(metrics.mean_audio_rtf)
+    assert metrics.request_goodput == 0.0
+
+
+def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput():
+    missing_tpot = _make_output(100, output_tokens=5)
+    missing_tpot.itl = []
+    missing_tpot.text_latency = missing_tpot.ttft = 0.1
+    missing_tpot.tpot_measured = False
+    missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+    slow_ttft = _make_output(100, output_tokens=5)
+    slow_ttft.ttft = 1.0
+    slow_ttft.text_latency = 1.4
+    slow_ttft.itl = [0.1] * 4
+    slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0}
+
+    metrics = _calculate_test_metrics(
+        [missing_tpot, slow_ttft],
+        {"ttft": 500.0, "tpot": 200.0},
+    )
+
+    assert metrics.num_tpot_samples == 1
+    assert metrics.mean_tpot_ms == pytest.approx(100.0)
+    assert metrics.request_goodput == 0.0
+
+
+def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero():
+    output = _make_output(100, output_tokens=5)
+    output.itl = []
+    output.text_latency = output.ttft
+    output.tpot_measured = False
+    output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+    metrics = _calculate_test_metrics([output], {"tpot": float("inf")})
+
+    assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0)
+    assert math.isnan(metrics.mean_tpot_ms)
+    assert math.isnan(metrics.mean_itl_ms)
+    assert metrics.request_goodput == 0.0
+
+
+def test_duplex_response_timings_do_not_build_a_session_token_timeline():
+    output = _make_output(100, output_tokens=5)
+    output.latency = 101.0
+    output.duplex_request_metrics = [
+        {"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
+        {"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
+    ]
+    output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+    metrics = _calculate_test_metrics([output])
+
+    assert math.isnan(metrics.max_output_tokens_per_s)
+    assert metrics.max_concurrent_requests == 1
+    assert metrics.mean_tpot_ms == pytest.approx(100.0)
+    assert metrics.mean_itl_ms == pytest.approx(100.0)
+
+
+def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
+    output = _make_output(100, output_tokens=0)
+    output.generated_text = "timing metadata missing"
+    output.itl = []
+    output.text_latency = output.ttft
+    output.tpot_measured = False
+    output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+    def tokenizer(text, *, add_special_tokens):
+        assert text == output.generated_text
+        assert add_special_tokens is False
+        return type("Tokenized", (), {"input_ids": [1, 2, 3]})()
+
+    metrics, _ = calculate_metrics(
+        input_requests=[],
+        outputs=[output],
+        dur_s=1.0,
+        tokenizer=tokenizer,
+        selected_percentiles=[50.0],
+        goodput_config_dict={"tpot": float("inf")},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=[],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=1.0,
+    )
+
+    assert metrics.total_output == 3
+    assert metrics.num_tpot_samples == 0
+    assert math.isnan(metrics.mean_tpot_ms)
+    assert metrics.request_goodput == 0.0
+
+
+def test_measured_zero_itl_is_not_treated_as_missing():
+    output = _make_output(100, output_tokens=3)
+    output.itl = [0.0, 0.0]
+    output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+    metrics = _calculate_test_metrics([output], {"tpot": 1.0})
+
+    assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (1, 2)
+    assert metrics.mean_tpot_ms == 0.0
+    assert metrics.mean_itl_ms == 0.0
+    assert metrics.request_goodput == 1.0
+
+
+def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
+    text_only, audio_only = _make_output(100), _make_output(100)
+    text_only.ttft, text_only.audio_ttfp = 0.1, 0.0
+    text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None}
+    audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2
+    audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None}
+
+    metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0})
+
+    assert metrics.request_goodput == 0.0
+
+
 # ============================================================================
 # TTFT suppression for pure-audio (TTS) benchmarks
 # ============================================================================

tests/benchmarks/test_omniinteract.py

@@ -0,0 +1,778 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import io
+import json
+import subprocess
+import tarfile
+import threading
+import time
+import wave
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
+from vllm.benchmarks.serve import TaskType
+
+from vllm_omni.benchmarks import omniinteract as oi
+from vllm_omni.benchmarks import serve as benchmark_serve
+from vllm_omni.benchmarks.data_modules import omniinteract_dataset as data
+from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
+from vllm_omni.benchmarks.patch import patch as benchmark_patch
+from vllm_omni.entrypoints.cli.benchmark.cli_args import preprocess_serve_args
+from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
+from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+MODEL = "openbmb/MiniCPM-o-4_5"
+
+
+def _case(tmp_path: Path, *, subset: str = "1q1a", name: str = "video.mp4") -> data.OmniInteractCase:
+    video, annotation = tmp_path / name, tmp_path / f"{Path(name).stem}.json"
+    video.parent.mkdir(parents=True, exist_ok=True)
+    video.touch()
+    annotation.write_text("{}")
+    return data.OmniInteractCase(subset, name, video, annotation, "multi_turn")
+
+
+def _collector(*events: tuple[dict[str, object], float]) -> RealtimeEventCollector:
+    collector = RealtimeEventCollector()
+    for event, received_at in events:
+        collector.add(event, received_at_s=received_at)
+    return collector
+
+
+def _created(response_id: str = "r1") -> dict[str, object]:
+    return {"type": "response.created", "response": {"id": response_id}}
+
+
+def _done(response_id: str = "r1", status: str = "completed") -> dict[str, object]:
+    return {"type": "response.done", "response": {"id": response_id, "status": status}}
+
+
+def _audio(
+    response_id: str = "r1",
+    *,
+    samples: int = 2400,
+    rate: int | None = 24_000,
+    value: int = 1,
+) -> dict[str, object]:
+    event: dict[str, object] = {
+        "type": "response.audio.delta",
+        "response_id": response_id,
+        "format": "pcm16",
+        "delta": base64.b64encode(bytes((value, 0)) * samples).decode(),
+    }
+    if rate is not None:
+        event["sample_rate_hz"] = rate
+    return event
+
+
+def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]:
+    return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value}
+
+
+def _listen(*, buffering: bool = False) -> dict[str, object]:
+    return {
+        "type": "response.listen",
+        "response": {"metadata": {"model_listen": True, "buffering": buffering}},
+    }
+
+
+def _session_options(tmp_path: Path) -> data.OmniInteractSessionOptions:
+    return data.OmniInteractSessionOptions(tmp_path / "artifacts", 10.0, 5.0, "reference.wav", True)
+
+
+def _sample(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions):
+    return data.OmniInteractSampleRequest(
+        prompt="",
+        prompt_len=0,
+        expected_output_len=0,
+        multi_modal_data=None,
+        request_id="measured-0",
+        omniinteract_case=case,
+        omniinteract_options=options,
+    )
+
+
+def _write_dataset(root: Path) -> Path:
+    data_root = root / "data"
+    for subset in ("1q1a", "1q1a_math"):
+        subset_root = data_root / subset
+        (subset_root / "videos").mkdir(parents=True)
+        (subset_root / "annotations").mkdir()
+        (subset_root / "videos" / f"{subset}.mp4").touch()
+        (subset_root / "annotations" / f"{subset}.json").write_text("{}")
+        (subset_root / "video_json_map.json").write_text(
+            json.dumps(
+                {
+                    "entries": [
+                        {
+                            "video": f"videos/{subset}.mp4",
+                            "annotation": f"annotations/{subset}.json",
+                            "scene_type": "multi_turn",
+                        }
+                    ]
+                }
+            )
+        )
+    one_to_many = data_root / "1qna"
+    (one_to_many / "videos_bench" / "nested").mkdir(parents=True)
+    (one_to_many / "annotations" / "nested").mkdir(parents=True)
+    (one_to_many / "videos_bench" / "nested" / "guide.mp4").touch()
+    (one_to_many / "annotations" / "nested" / "guide.json").write_text("{}")
+    return data_root
+
+
+def _successful_output(
+    tmp_path: Path,
+    *,
+    audio_time: float = 10.0,
+    status: str = "completed",
+) -> tuple[data.OmniInteractCase, RealtimeEventCollector, oi.OmniInteractCaseResult]:
+    case = _case(tmp_path)
+    collector = _collector(
+        (_created(), 9.9),
+        (_audio(), audio_time),
+        (_text(), audio_time + 0.01),
+        (_done(status=status), audio_time + 0.1),
+    )
+    result = oi.OmniInteractCaseResult(case.subset, str(case.video_path), "")
+    return case, collector, result
+
+
+def _write_success(
+    root: Path,
+    case: data.OmniInteractCase,
+    collector: RealtimeEventCollector,
+    result: oi.OmniInteractCaseResult,
+    **kwargs,
+):
+    require_response = kwargs.pop("require_response", True)
+    return oi.prepare_success_artifacts(
+        root,
+        case,
+        collector,
+        stream_start=10.0,
+        video_duration_s=1.0,
+        require_response=require_response,
+        result=result,
+        **kwargs,
+    )
+
+
+def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
+    root = _write_dataset(tmp_path)
+    all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
+    selected = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=2, disable_shuffle=True)
+    assert [case.subset for case in all_cases] == ["1q1a", "1q1a_math", "1qna"]
+    assert all_cases[-1].video_rel == "videos_bench/nested/guide.mp4"
+    assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
+
+
+def _archive(path: Path, member: str) -> None:
+    with tarfile.open(path, "w") as archive:
+        content = b"{}"
+        item = tarfile.TarInfo(member)
+        item.size = len(content)
+        archive.addfile(item, io.BytesIO(content))
+
+
+def test_archive_is_safe_and_atomically_shared(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    bad = tmp_path / "bad.tar"
+    _archive(bad, "../escape")
+    with pytest.raises(ValueError, match="Unsafe path"):
+        data._extract_archive(bad, tmp_path / "bad-cache")
+
+    good = tmp_path / "good.tar"
+    _archive(good, "data/1q1a/video_json_map.json")
+    target = tmp_path / "cache"
+    barrier, safe_extract = threading.Barrier(2), data._safe_extract
+    monkeypatch.setattr(data, "_safe_extract", lambda *args: (barrier.wait(), safe_extract(*args)))
+    with ThreadPoolExecutor(2) as executor:
+        roots = list(executor.map(lambda _: data._extract_archive(good, target), range(2)))
+    assert roots[0] == roots[1]
+    assert (roots[0] / "1q1a" / "video_json_map.json").is_file()
+    assert not list(target.glob(".tmp-*"))
+
+
+def test_hub_archive_uses_vllm_filesystem(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    source = tmp_path / "source.tar"
+    _archive(source, "data/1q1a/video_json_map.json")
+
+    class FS:
+        def get_file(self, remote: str, local: str) -> None:
+            assert remote == "datasets/org/repo/data.tar.gz"
+            Path(local).write_bytes(source.read_bytes())
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
+    monkeypatch.setattr(data, "hf_fs", lambda: FS())
+    assert (data.resolve_omniinteract_root(None, "org/repo") / "1q1a").is_dir()
+
+
+def test_media_commands_are_bounded(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+    calls: list[float] = []
+    commands: list[list[str]] = []
+    bounded_run = oi._run_media_command
+
+    def run(command, *, timeout_s, text=False):
+        calls.append(timeout_s)
+        commands.append(command)
+        if command[0] == "ffprobe":
+            return subprocess.CompletedProcess(command, 0, "1.0", "")
+        return subprocess.CompletedProcess(command, 0, b"", b"")
+
+    monkeypatch.setattr(oi, "_run_media_command", run)
+    duration, pcm, frames = oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=1.0)
+    assert (duration, len(pcm), frames, calls) == (1.0, 32_000, [None], [3.0, 3.0, 3.0])
+    assert commands[1][commands[1].index("-t") + 1] == "1.0"
+    assert commands[1][commands[1].index("-fs") + 1] == "32000"
+
+    monkeypatch.setattr(subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(subprocess.TimeoutExpired("x", 3)))
+    with pytest.raises(TimeoutError, match="timed out"):
+        bounded_run(["ffmpeg"], timeout_s=3.0)
+
+
+def test_media_duration_is_bounded_before_decode(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+    calls = 0
+
+    def run(command, *, timeout_s, text=False):
+        nonlocal calls
+        calls += 1
+        return subprocess.CompletedProcess(command, 0, "601.0", "")
+
+    monkeypatch.setattr(oi, "_run_media_command", run)
+    with pytest.raises(ValueError, match="Invalid video duration"):
+        oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=600.0)
+    assert calls == 1
+
+
+def test_final_commit_keeps_a_partial_model_unit():
+    """Exact units retain a sample for the server's correlated final flush."""
+
+    events = [{"session": {"capabilities": {"chunk_period_ms": 1000}}}]
+    exact_unit = b"\0\0" * 16_000
+    assert oi._ensure_final_commit_tail(exact_unit, events) == exact_unit[:-2]
+    assert oi._ensure_final_commit_tail(exact_unit + b"\0\0", events) == exact_unit + b"\0\0"
+
+
+@pytest.mark.parametrize(
+    ("events", "match"),
+    [
+        (((_done(), 0.0),), "without response.created"),
+        (((_created(), 0.0), (_created(), 0.1)), "duplicate response.created"),
+        (((_created(), 0.0), (_done(), 0.1), (_done(), 0.2)), "duplicate response.done"),
+        (
+            (
+                (_created(), 0.0),
+                ({"type": "response.done", "response": {"id": "r1", "status_details": {"type": "failed"}}}, 0.1),
+            ),
+            "reports failure",
+        ),
+    ],
+)
+def test_response_ledger_rejects_identity_errors(events, match: str):
+    with pytest.raises(ValueError, match=match):
+        oi.response_ledger(_collector(*events))
+
+
+class _CompletionClient:
+    def __init__(self, collector: RealtimeEventCollector):
+        self.events, self.acks = collector, []
+
+    def raise_if_reader_stopped(self) -> None:
+        return None
+
+    async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
+        self.acks.append((response_id, played_ms))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("tail", "succeeds"),
+    [
+        (((_created(), 1.1), (_done(), 1.2)), True),
+        (((_listen(), 1.1),), True),
+        (((_listen(buffering=True), 1.1),), False),
+        ((({"type": "session.closed", "reason": "timeout"}, 1.1),), False),
+    ],
+)
+async def test_completion_requires_a_final_model_decision(tail, succeeds: bool):
+    events = _collector(({"type": "input_audio_buffer.committed"}, 1.0), *tail)
+    coroutine = oi.wait_for_session_completion(
+        _CompletionClient(events), oi._Playback(), commit_from=0, timeout_s=0.08, settle_s=0.0
+    )
+    if succeeds:
+        assert await coroutine == 0
+    else:
+        with pytest.raises((TimeoutError, RuntimeError)):
+            await coroutine
+
+
+@pytest.mark.asyncio
+async def test_deferred_commit_ignores_the_prior_response_terminal():
+    events = _collector(
+        (_created("old"), 0.0),
+        ({"type": "input_audio_buffer.committed", "event": {"overlap_deferred": True}}, 1.0),
+        (_done("old"), 1.1),
+        (_listen(), 1.2),
+    )
+    assert (
+        await oi.wait_for_session_completion(
+            _CompletionClient(events), oi._Playback(), commit_from=1, timeout_s=0.1, settle_s=0.0
+        )
+        == 1
+    )
+
+
+@pytest.mark.asyncio
+async def test_playback_acks_only_after_serial_audio_drain():
+    collector = _collector((_audio(samples=2400), 10.0), (_done(), 10.01))
+    client, playback = _CompletionClient(collector), oi._Playback()
+    await playback.acknowledge(client, now=10.05)
+    assert client.acks == []
+    await playback.acknowledge(client, now=10.1)
+    assert client.acks == [("r1", 100)]
+
+
+@pytest.mark.parametrize(
+    ("event", "match"),
+    [
+        (_audio(rate=16_000), "24000 Hz"),
+        ({**_audio(), "format": "opus"}, "pcm16"),
+        ({**_audio(), "delta": "!"}, "base64"),
+    ],
+)
+def test_playback_rejects_non_official_audio(event: dict[str, object], match: str):
+    playback = oi._Playback()
+    with pytest.raises(ValueError, match=match):
+        playback.ingest(_collector((event, 1.0)))
+
+
+def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: Path):
+    case, collector, result = _successful_output(tmp_path)
+    summary = _write_success(
+        tmp_path / "out",
+        case,
+        collector,
+        result,
+    )
+    context = result._artifact_context
+    assert result.success and result.eligible_for_official_eval
+    assert summary["scene_type"] == "multi_turn"
+    assert context and context.spans
+    assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta")
+    oi.publish_deferred_case_artifacts(tmp_path / "out", case, result)
+    directory = oi._output_dir(tmp_path / "out", case)
+    assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS)
+    with wave.open(str(directory / "output.wav")) as output:
+        assert (output.getframerate(), output.getnframes()) == (24_000, 24_000)
+    transcript = json.loads((directory / "wav_transcript.json").read_text())
+    assert transcript["text"] == "hello" and transcript["chunks"][0]["timestamp"] == [0.0, 0.1]
+    assert "serialized playback queue" in transcript["timestamp_semantics"]
+    assert result._artifact_context is None
+
+
+@pytest.mark.parametrize(
+    ("audio_time", "status", "reason"),
+    [(10.95, "completed", "audio_clipped"), (10.0, "cancelled", "cancelled_response")],
+)
+def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_time: float, status: str, reason: str):
+    case, collector, result = _successful_output(tmp_path, audio_time=audio_time, status=status)
+    _write_success(tmp_path / "out", case, collector, result, require_response=False)
+    assert reason in result.official_eval_ineligible_reasons
+    oi.write_batch_artifacts(tmp_path / "out", [case], [result])
+    assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
+
+
+def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):
+    case, _, result = _successful_output(tmp_path)
+    collector = _collector((_created(), 10.0), (_audio(), 10.0), (_done(), 10.1))
+    with pytest.raises(ValueError, match="audio and transcript"):
+        _write_success(tmp_path, case, collector, result)
+
+
+def test_publication_failure_removes_partial_success_bundle(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case, collector, result = _successful_output(tmp_path)
+    _write_success(tmp_path, case, collector, result)
+    original = oi._atomic_write_json
+
+    def fail(path: Path, value: object) -> None:
+        if path.name == "events.json":
+            raise OSError("disk full")
+        original(path, value)
+
+    monkeypatch.setattr(oi, "_atomic_write_json", fail)
+    with pytest.raises(OSError, match="disk full"):
+        oi.publish_deferred_case_artifacts(tmp_path, case, result)
+    directory = oi._output_dir(tmp_path, case)
+    assert not any((directory / name).exists() for name in oi.SUCCESS_ARTIFACTS)
+    assert not result.success and "artifact_write_failed" in result.official_eval_ineligible_reasons
+
+
+def test_atomic_write_preserves_destination_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    destination = tmp_path / "output.wav"
+    destination.write_bytes(b"old")
+    monkeypatch.setattr(oi, "write_pcm16_wav", lambda *a, **k: (_ for _ in ()).throw(OSError("write")))
+    with pytest.raises(OSError):
+        oi._atomic_write_wav(destination, b"\0\0", 24_000)
+    assert destination.read_bytes() == b"old"
+    assert not list(tmp_path.glob(".*.tmp"))
+
+
+class _RealtimeClient:
+    instances: list[_RealtimeClient] = []
+
+    def __init__(self, url: str, **kwargs):
+        self.url, self.events, self.acks, self.configure_kwargs = url, RealtimeEventCollector(), [], {}
+        self.instances.append(self)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        return None
+
+    async def configure(self, model: str, **kwargs) -> None:
+        self.configure_kwargs = kwargs
+        self.events.add({"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 1000}}})
+
+    async def send(self, event: dict[str, object]) -> None:
+        return None
+
+    async def commit(self) -> None:
+        now = time.monotonic()
+        for event, offset in (
+            ({"type": "input_audio_buffer.committed"}, 0.0),
+            (_created(), 0.001),
+            (_audio(samples=24), 0.002),
+            (_text(value="answer"), 0.003),
+            (_done(), 0.004),
+        ):
+            self.events.add(event, received_at_s=now + offset)
+
+    async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
+        self.acks.append((response_id, played_ms))
+
+    async def close_session(self, **kwargs) -> None:
+        self.events.add({"type": "session.closed"})
+
+    def raise_if_reader_stopped(self) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case = _case(tmp_path)
+    config = oi.OmniInteractBenchmarkConfig(
+        model=MODEL,
+        output_root=tmp_path / "out",
+        ref_audio="ref.wav",
+        require_response=True,
+        extra_body={"custom": "value"},
+    )
+    prepared = data.OmniInteractPreparedInput(1.0, b"\0\0\0\0", ("frame",), "data:audio/wav;base64,ref")
+    monkeypatch.setattr(oi, "RealtimeDuplexClient", _RealtimeClient)
+    monkeypatch.setattr(oi, "_COMPLETION_SETTLE_S", 0.0)
+    result = await oi.run_omniinteract_case(
+        case,
+        config,
+        request_index=0,
+        prepared_input=prepared,
+    )
+    assert result.success and result.transcript == "answer" and result._artifact_context is not None
+    assert not config.output_root.exists()
+    assert "autostart=0" in _RealtimeClient.instances[-1].url
+    assert _RealtimeClient.instances[-1].configure_kwargs["extra_body"] == {"custom": "value"}
+    assert _RealtimeClient.instances[-1].acks == [("r1", 1)]
+
+
+def test_standard_sample_loading_prepares_media_before_timing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    _write_dataset(tmp_path)
+    ref = tmp_path / "ref.wav"
+    ref.touch()
+    parser = TrackingArgumentParser()
+    OmniBenchmarkServingSubcommand.add_cli_args(parser)
+    args = parser.parse_args(
+        [
+            "--backend",
+            "openai-realtime-duplex",
+            "--dataset-name",
+            "omniinteract",
+            "--dataset-path",
+            str(tmp_path),
+            "--model",
+            MODEL,
+            "--endpoint",
+            "/v1/realtime",
+            "--num-prompts",
+            "0",
+            "--omniinteract-ref-audio",
+            str(ref),
+        ]
+    )
+    preprocess_serve_args(args)
+    assert args.endpoint == "/v1/realtime"
+    monkeypatch.setattr(benchmark_patch, "reference_audio_data_url", lambda _: "data:audio/wav;base64,ref")
+    monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
+    samples = benchmark_patch.get_samples(args, None)
+    assert len(samples) == args.num_prompts == 3
+    assert all(sample.omniinteract_prepared_input.video_frames == ("frame",) for sample in samples)
+
+
+def _request(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions) -> RequestFuncInput:
+    request = RequestFuncInput(
+        model=MODEL,
+        model_name=MODEL,
+        prompt="",
+        api_url="http://server:8000/v1/realtime",
+        prompt_len=0,
+        output_len=0,
+        logprobs=None,
+        multi_modal_content=None,
+        ignore_eos=False,
+        request_id="measured-0",
+    )
+    request.omniinteract_case, request.omniinteract_options = case, options
+    return request
+
+
+@pytest.mark.asyncio
+async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case, options = _case(tmp_path), _session_options(tmp_path)
+    request = _request(case, options)
+    output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+    assert not output.success and "prepared before benchmark request timing" in output.error
+
+    prepared = data.OmniInteractPreparedInput(1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref")
+    request.omniinteract_prepared_input = prepared
+    request.extra_headers = {"X-Proxy": "test"}
+    request.extra_body = {"custom": "value"}
+    monkeypatch.setenv("OPENAI_API_KEY", "secret")
+    seen = {}
+
+    async def run(_case_arg, config, **kwargs):
+        seen["headers"] = config.extra_headers
+        seen["extra_body"] = config.extra_body
+        seen.update(kwargs)
+        return oi.OmniInteractCaseResult(
+            case.subset,
+            str(case.video_path),
+            str(options.output_root),
+            success=True,
+            transcript="timing metadata missing",
+            output_tokens=0,
+            duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}],
+            duplex_session_metrics={"mean_ttft_ms": 10.0},
+        )
+
+    monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
+    output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+    assert output.success and seen["prepared_input"] is prepared
+    assert seen["headers"] == {"Authorization": "Bearer secret", "X-Proxy": "test", "x-request-id": "measured-0"}
+    assert seen["extra_body"] == {"custom": "value"}
+    assert output.itl == []
+    assert output.tpot_measured is False
+    assert seen["capture_artifacts"] is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("request_metrics", "expected_itl", "expected_text_latency"),
+    [
+        (
+            [
+                {"stage0_tokens": {"output_token_count": 3, "itls_ms": [10.0, 20.0]}},
+                {"stage0_tokens": {"output_token_count": 2, "itls_ms": [30.0]}},
+            ],
+            [0.01, 0.02, 0.03],
+            0.16,
+        ),
+        (
+            [
+                {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 10.0}},
+                {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 40.0}},
+            ],
+            [],
+            0.18,
+        ),
+        (
+            [
+                {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 0.0}},
+                {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 0.0}},
+            ],
+            [],
+            0.1,
+        ),
+    ],
+)
+async def test_adapter_reports_exact_or_weighted_token_timing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    request_metrics: list[dict[str, object]],
+    expected_itl: list[float],
+    expected_text_latency: float,
+):
+    case, options = _case(tmp_path), _session_options(tmp_path)
+    request = _request(case, options)
+    request.omniinteract_prepared_input = data.OmniInteractPreparedInput(
+        1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref"
+    )
+
+    async def run(*args, **kwargs):
+        return oi.OmniInteractCaseResult(
+            case.subset,
+            str(case.video_path),
+            str(options.output_root),
+            success=True,
+            output_tokens=5,
+            duplex_request_metrics=request_metrics,
+            duplex_session_metrics={"mean_ttft_ms": 100.0},
+        )
+
+    monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
+    output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+
+    assert output.success
+    assert output.itl == expected_itl
+    assert output.text_latency == pytest.approx(expected_text_latency)
+    assert output.tpot_measured is True
+
+
+def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case, options = _case(tmp_path), _session_options(tmp_path)
+    options.output_root.mkdir()
+    sample = _sample(case, options)
+    result = oi.OmniInteractCaseResult(
+        case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+    )
+    output = benchmark_patch.MixRequestFuncOutput(success=True)
+    output.omniinteract_case_result = result
+    published = []
+    monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: published.append(args))
+    summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+    assert summary and summary["total"] == summary["success"] == summary["eligible_for_official_eval"] == 1
+    assert summary["failed"] == summary["successful_but_ineligible"] == summary["audio_clipped_bytes"] == 0
+    assert published == [(options.output_root, case, result)]
+
+
+def test_case_artifact_failures_do_not_discard_serving_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case, options = _case(tmp_path), _session_options(tmp_path)
+    sample = _sample(case, options)
+    result = oi.OmniInteractCaseResult(
+        case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+    )
+    output = benchmark_patch.MixRequestFuncOutput(success=True, error="inference warning")
+    output.omniinteract_case_result = result
+    monkeypatch.setattr(
+        benchmark_patch,
+        "publish_deferred_case_artifacts",
+        lambda *args: (_ for _ in ()).throw(OSError("case disk full")),
+    )
+    monkeypatch.setattr(
+        benchmark_patch,
+        "write_failure_artifacts",
+        lambda *args: (_ for _ in ()).throw(OSError("failure disk full")),
+    )
+    monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", lambda *args: None)
+
+    summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+    metrics, _ = calculate_metrics(
+        [],
+        [output],
+        1.0,
+        None,
+        [50.0],
+        {},
+        TaskType.GENERATION,
+        [],
+        None,
+        float("inf"),
+        1.0,
+    )
+
+    assert summary and summary["failed"] == 1
+    assert summary["artifacts_complete"] is False
+    assert len(summary["artifact_errors"]) == 2
+    assert output.success and metrics.completed == 1
+    assert output.error.startswith("inference warning\n")
+    assert "Artifact publication failed: case disk full" in output.error
+
+
+def test_batch_artifact_failure_removes_partial_batch_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    case, options = _case(tmp_path), _session_options(tmp_path)
+    sample = _sample(case, options)
+    result = oi.OmniInteractCaseResult(
+        case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+    )
+    output = benchmark_patch.MixRequestFuncOutput(success=True)
+    output.omniinteract_case_result = result
+    monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: None)
+
+    def fail_batch(root, *args):
+        root.mkdir(parents=True, exist_ok=True)
+        (root / "batch_summary.json").write_text("partial")
+        raise OSError("batch disk full")
+
+    monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", fail_batch)
+
+    summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+
+    assert summary and summary["success"] == 1 and summary["artifacts_complete"] is False
+    assert summary["artifact_errors"] == ["Batch artifact publication failed: batch disk full"]
+    assert output.success
+    assert not (options.output_root / "batch_summary.json").exists()
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "match"),
+    [
+        ("omniinteract_ref_audio", None, "ref-audio"),
+        ("endpoint", "/v1/completions", "v1/realtime"),
+        ("max_concurrency", 0, "positive"),
+        ("skip_tokenizer_init", True, "skip-tokenizer-init"),
+    ],
+)
+def test_cli_rejects_incompatible_omniinteract_options(field: str, value, match: str):
+    args = argparse.Namespace(
+        dataset_name="omniinteract",
+        backend="openai-realtime-duplex",
+        endpoint="/v1/realtime",
+        omniinteract_ref_audio="ref.wav",
+        max_concurrency=1,
+        skip_tokenizer_init=False,
+    )
+    setattr(args, field, value)
+    with pytest.raises(ValueError, match=match):
+        preprocess_serve_args(args)
+
+
+def test_whole_benchmark_lock_serializes_a_shared_output_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    entered, release = threading.Event(), threading.Event()
+    calls: list[int] = []
+
+    async def fake_main_async(args):
+        calls.append(len(calls))
+        if len(calls) == 1:
+            entered.set()
+            await asyncio.to_thread(release.wait)
+        return {"total": 1}
+
+    monkeypatch.setattr(benchmark_serve, "main_async", fake_main_async)
+    args = argparse.Namespace(dataset_name="omniinteract", omniinteract_output_dir=tmp_path, extra_body=None)
+    with ThreadPoolExecutor(2) as executor:
+        first = executor.submit(benchmark_serve.main, args)
+        assert entered.wait(1)
+        second = executor.submit(benchmark_serve.main, args)
+        time.sleep(0.05)
+        assert calls == [0]
+        release.set()
+        assert first.result() == second.result() == {"total": 1}
+    assert calls == [0, 1]

tests/benchmarks/test_serve_cli.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import argparse
 import json
 import os
@@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
     "dataset_name",
     [
         "daily-omni",
+        "omniinteract",
         "seed-tts",
         "seed-tts-text",
         "seed-tts-design",
@@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
     [
         (["--print-stage"], "print_stage", True),
         (["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
+        (["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
         (["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
     ],
 )
@@ -80,12 +85,14 @@ def test_add_omni_args_preserves_implicit_defaults() -> None:
     args = parser.parse_args([])
     assert args.print_stage is False
     assert args.daily_omni_input_mode == "all"
+    assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
     assert args.seed_tts_locale == "en"
     assert args.explicit_keys == set()
 
 
 def test_update_omni_help_updates_upstream_actions() -> None:
     parser = TrackingArgumentParser()
+    parser.add_argument("--num-prompts", default=1000, help="Number of prompts.")
     parser.add_argument("--percentile-metrics", help="Upstream percentile help.")
     parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.")
     parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.")
@@ -93,6 +100,7 @@ def test_update_omni_help_updates_upstream_actions() -> None:
     update_omni_help(parser)
 
     actions = {action.dest: action for action in parser._actions}
+    assert "OmniInteract uses 3" in actions["num_prompts"].help
     assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf"))
     assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help
     assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help
@@ -120,6 +128,31 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
     assert args.extra_body == expected
 
 
+@pytest.mark.parametrize(
+    ("argv", "expected"),
+    [
+        ([], 3),
+        (["--num-prompts", "0"], 0),
+        (["--num-prompts", "12"], 12),
+    ],
+)
+def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
+    argv: list[str],
+    expected: int,
+) -> None:
+    parser = TrackingArgumentParser()
+    parser.add_argument("--dataset-name", default="omniinteract")
+    parser.add_argument("--backend", default="openai-realtime-duplex")
+    parser.add_argument("--endpoint", default="/v1/realtime")
+    parser.add_argument("--omniinteract-ref-audio", default="reference.wav")
+    parser.add_argument("--num-prompts", type=int, default=1000)
+    args = parser.parse_args(argv)
+
+    preprocess_serve_args(args)
+
+    assert args.num_prompts == expected
+
+
 @pytest.mark.parametrize(
     ("argv", "expected_extra_body", "expected_explicit"),
     [

tests/config/test_config_factory.py

@@ -1109,8 +1109,35 @@ class TestDeployConfigLoading:
         with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
             load_deploy_config(deploy_path)
 
-    def test_load_minicpmo_duplex_deploy_config(self):
-        deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
+    @pytest.mark.parametrize(
+        ("filename", "max_sessions"),
+        [
+            ("minicpmo_4_5.yaml", 4),
+            ("minicpmo_4_5_2gpu.yaml", 4),
+            ("minicpmo_4_5_3gpu.yaml", 4),
+            ("minicpmo_4_5_8x4090.yaml", 1),
+            ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
+            ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
+            ("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
+        ],
+    )
+    def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
+        """Every MiniCPM-o profile serves /v1/realtime, so the retired
+        minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
+        """
+        deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
+        pipeline = resolve_pipeline_config("minicpmo_4_5")
+        assert isinstance(pipeline, PipelineConfig)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert deploy.session_mode == "duplex"
+        assert deploy.active_stream_window == 1
+        assert deploy.duplex_session.max_sessions == max_sessions
+        assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
+        assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
+
+    def test_load_minicpmo_default_deploy_config(self):
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
 
         deploy = load_deploy_config(deploy_path)
         pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1118,17 +1145,15 @@ class TestDeployConfigLoading:
 
         stages = merge_pipeline_deploy(pipeline, deploy)
 
-        assert deploy.session_mode == "duplex"
         assert deploy.async_chunk is True
-        assert deploy.active_stream_window == 1
-        assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
-        assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
-        assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
-        assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
+        assert stages[0].yaml_engine_args["async_scheduling"] is True
+        assert stages[1].yaml_engine_args["async_scheduling"] is True
+        assert "Async" in (stages[0].scheduler_cls or "")
+        assert "Async" in (stages[1].scheduler_cls or "")
         assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
         assert deploy.stages[1].enforce_eager is False
         assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
-        assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
+        assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
         assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
         assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
         assert "codec_sampling_params" not in stages[1].yaml_engine_args
@@ -2116,7 +2141,8 @@ class TestBaseConfigInheritance:
     def test_minicpmo_overlays_inherit_talker_sampling_params(self):
         """Overlays must keep the base Talker codec Sampler knobs."""
         for filename in (
-            "minicpmo_4_5_duplex.yaml",
+            "minicpmo_4_5_2gpu.yaml",
+            "minicpmo_4_5_3gpu.yaml",
             "minicpmo_4_5_3gpu_stage1_replicas.yaml",
             "minicpmo_4_5_4gpu_stage1_replicas.yaml",
             "minicpmo_4_5_8x4090_stage1_replicas.yaml",
@@ -2128,9 +2154,6 @@ class TestBaseConfigInheritance:
             sampling = deploy.stages[1].default_sampling_params
             assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
             assert sampling["temperature"] == 0.8, filename
-            # Upstream utils.TTSSamplingParams. min_tokens is not checked here:
-            # the duplex overlay drops it to 0 because the model budgets each
-            # generate_chunk itself.
             assert sampling["top_k"] == 25, filename
             assert sampling["top_p"] == 0.85, filename
             assert sampling["repetition_penalty"] == 1.05, filename

tests/config/test_config_import_cycle.py

@@ -66,6 +66,8 @@ def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
     runtime_src = src[runtime_start:]
     assert "apply_minimax_h3_qwen3vl_patch" not in init_src
     assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
+    assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
+    assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
     assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
 
 

tests/config/test_environment_variables.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Drift checks for the reviewed environment-variable inventory."""
+
+import ast
+import builtins
+from collections import Counter
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.environment_variable_inventory import (
+    ENVIRONMENT_VARIABLE_INVENTORY,
+    EnvironmentVariableCategory,
+    ModelEnvironmentVariableDisposition,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).parents[2]
+_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni"
+_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py"
+_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md"
+
+
+def _module_string_constants(tree: ast.Module) -> dict[str, str]:
+    constants: dict[str, str] = {}
+    for statement in tree.body:
+        targets: list[ast.expr] = []
+        value: ast.expr | None = None
+        if isinstance(statement, ast.Assign):
+            targets = statement.targets
+            value = statement.value
+        elif isinstance(statement, ast.AnnAssign):
+            targets = [statement.target]
+            value = statement.value
+        if isinstance(value, ast.Constant) and isinstance(value.value, str):
+            for target in targets:
+                if isinstance(target, ast.Name):
+                    constants[target.id] = value.value
+    return constants
+
+
+def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]:
+    os_names: set[str] = set()
+    getenv_names: set[str] = set()
+    environ_names: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for imported in node.names:
+                if imported.name == "os":
+                    os_names.add(imported.asname or "os")
+        elif isinstance(node, ast.ImportFrom) and node.module == "os":
+            for imported in node.names:
+                if imported.name == "getenv":
+                    getenv_names.add(imported.asname or imported.name)
+                elif imported.name == "environ":
+                    environ_names.add(imported.asname or imported.name)
+    return os_names, getenv_names, environ_names
+
+
+def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool:
+    return (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id in os_names
+        and node.attr == "environ"
+    ) or (isinstance(node, ast.Name) and node.id in environ_names)
+
+
+def _environment_key_expressions(
+    node: ast.AST,
+    os_names: set[str],
+    getenv_names: set[str],
+    environ_names: set[str],
+) -> list[ast.expr]:
+    """Return key expressions used by one direct environment access."""
+    if isinstance(node, ast.Call) and node.args:
+        func = node.func
+        if (
+            isinstance(func, ast.Attribute)
+            and isinstance(func.value, ast.Name)
+            and func.value.id in os_names
+            and func.attr == "getenv"
+        ) or (isinstance(func, ast.Name) and func.id in getenv_names):
+            return [node.args[0]]
+        if (
+            isinstance(func, ast.Attribute)
+            and _is_environ(func.value, os_names, environ_names)
+            and func.attr in {"get", "pop", "setdefault"}
+        ):
+            return [node.args[0]]
+    elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names):
+        return [node.slice]
+    elif isinstance(node, ast.Compare):
+        expressions: list[ast.expr] = []
+        left_values = [node.left, *node.comparators[:-1]]
+        for left, operator, right in zip(left_values, node.ops, node.comparators):
+            if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names):
+                expressions.append(left)
+        return expressions
+    return []
+
+
+def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None:
+    value: str | None = None
+    if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
+        value = expression.value
+    elif isinstance(expression, ast.Name):
+        value = constants.get(expression.id)
+    return value or None
+
+
+def _environment_wrapper_parameters(
+    tree: ast.Module,
+    os_names: set[str],
+    getenv_names: set[str],
+    environ_names: set[str],
+) -> dict[str, set[tuple[str, int | None]]]:
+    """Find local helpers that forward one parameter as an environment key."""
+    wrappers: dict[str, set[tuple[str, int | None]]] = {}
+    for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
+        positional = [*function.args.posonlyargs, *function.args.args]
+        parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)}
+        parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs})
+        forwarded: set[tuple[str, int | None]] = set()
+        for node in ast.walk(function):
+            for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
+                if isinstance(expression, ast.Name) and expression.id in parameter_indexes:
+                    forwarded.add((expression.id, parameter_indexes[expression.id]))
+        if forwarded:
+            wrappers[function.name] = forwarded
+    return wrappers
+
+
+def _environment_accesses(path: Path) -> set[str]:
+    """Return statically resolvable names accessed through ``os``."""
+    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    constants = _module_string_constants(tree)
+    os_names, getenv_names, environ_names = _os_import_aliases(tree)
+    wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names)
+    names: set[str] = set()
+
+    for node in ast.walk(tree):
+        for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
+            if name := _resolve_environment_name(expression, constants):
+                names.add(name)
+
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            for parameter_name, position in wrappers.get(node.func.id, set()):
+                argument = node.args[position] if position is not None and position < len(node.args) else None
+                if argument is None:
+                    argument = next(
+                        (keyword.value for keyword in node.keywords if keyword.arg == parameter_name),
+                        None,
+                    )
+                if argument is not None and (name := _resolve_environment_name(argument, constants)):
+                    names.add(name)
+
+    return names
+
+
+def test_inventory_matches_reviewed_snapshot_counts():
+    """Make an inventory expansion an explicit review decision."""
+    category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
+    assert category_counts == {
+        EnvironmentVariableCategory.PUBLIC_OMNI: 23,
+        EnvironmentVariableCategory.INHERITED_VLLM: 20,
+        EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
+        EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
+        EnvironmentVariableCategory.INTERNAL: 2,
+    }
+
+    disposition_counts = Counter(
+        item.model_disposition
+        for item in ENVIRONMENT_VARIABLE_INVENTORY.values()
+        if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
+    )
+    assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
+        ModelEnvironmentVariableDisposition.PROMOTE: 33,
+        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
+        ModelEnvironmentVariableDisposition.EXTERNAL: 0,
+        ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
+        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
+    }
+
+
+def test_new_public_omni_names_use_project_prefix():
+    """Grandfather legacy names without allowing more prefix exceptions."""
+    legacy_public_names = {
+        "DIFFUSION_ATTENTION_BACKEND",
+        "DIFFUSION_CACHE_ADAPTER",
+        "DIFFUSION_CACHE_BACKEND",
+        "OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
+        "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
+        "OMNI_DIFFUSION_SESSION_STATE_MANAGER",
+        "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
+        "SPEAKER_MAX_UPLOADED",
+        "SPEAKER_SAMPLES_DIR",
+        "VLLM_VIDEO_ASYNC_CHUNK",
+        "VLLM_VIDEO_AUDIO_DELTA_MODE",
+    }
+    public_names = {
+        name
+        for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+        if item.category is EnvironmentVariableCategory.PUBLIC_OMNI
+    }
+
+    assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names
+
+
+def test_statically_resolvable_environment_accesses_are_classified():
+    discovered: set[str] = set()
+    for path in _PACKAGE_ROOT.rglob("*.py"):
+        discovered.update(_environment_accesses(path))
+
+    assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set()
+
+
+def test_model_and_benchmark_inventory_entries_are_still_referenced():
+    """Reject transitional rows after their implementation has disappeared."""
+    package_source = "\n".join(
+        path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE
+    )
+    transitional_categories = {
+        EnvironmentVariableCategory.MODEL_SPECIFIC,
+        EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL,
+    }
+
+    stale_names = {
+        name
+        for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+        if item.category in transitional_categories and name not in package_source
+    }
+
+    assert stale_names == set()
+
+
+def test_environment_scanner_covers_indirection_aliases_membership_and_casing():
+    expected_by_path = {
+        "metrics/definitions.py": {
+            "VLLM_OMNI_BENCH_AUDIO_CHANNELS",
+            "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
+        },
+        "entrypoints/openai/video_stream_envs.py": {
+            "VLLM_VIDEO_ASYNC_CHUNK",
+            "VLLM_VIDEO_AUDIO_DELTA_MODE",
+        },
+        "model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"},
+        "model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"},
+        "distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"},
+    }
+    for relative_path, expected in expected_by_path.items():
+        assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path)
+
+
+def test_generated_server_storage_environment_names_are_classified():
+    from vllm_omni.config.server_settings import FileBackend
+
+    generated_names = {
+        f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}"
+        for field_name in FileBackend.model_fields
+        if field_name != "type"
+    }
+    assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys()
+
+
+def test_public_omni_variables_are_in_the_reference_page():
+    reference = _REFERENCE_PAGE.read_text(encoding="utf-8")
+    missing = {
+        name
+        for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+        if item.is_public_omni and f"`{name}`" not in reference
+    }
+    assert missing == set()
+
+
+def test_secret_values_are_marked_for_redaction():
+    assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == {
+        "HF_TOKEN",
+        "HUGGINGFACE_HUB_TOKEN",
+        "OPENAI_API_KEY",
+    }
+
+
+def test_collect_env_reports_only_safe_public_omni_values(monkeypatch):
+    from collect_env import get_env_vars
+
+    monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
+    monkeypatch.setenv("HF_TOKEN", "hf-secret")
+    monkeypatch.setenv("OPENAI_API_KEY", "openai-secret")
+    monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica")
+
+    report = get_env_vars()
+
+    assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines()
+    assert "hf-secret" not in report
+    assert "openai-secret" not in report
+    assert "VLLM_OMNI_REPLICA_ID" not in report
+
+
+@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError])
+def test_collect_env_survives_inventory_import_error(monkeypatch, error_type):
+    from collect_env import get_env_vars
+
+    real_import = builtins.__import__
+
+    def fail_inventory_import(name, *args, **kwargs):
+        if name == "vllm_omni.config.environment_variable_inventory":
+            raise error_type("broken vllm-omni installation")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", fail_inventory_import)
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
+
+    assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines()

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -252,6 +252,37 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.chunk_transfer_adapter.save_async.assert_not_called()
 
 
+def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
+    """Native duplex + async: the next unit is usually already queued when the
+    current one stops, so _handle_stopped_request applies the update and the
+    stop site below it fences the same in-flight decode a second time. An
+    accumulated residue outlives the drain and swallows the next unit's
+    listen/speak, which the client sees as silence until its timeout.
+    """
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.resumable = True
+    session.append_output_token_ids([7])
+    session.num_computed_tokens = 4
+    session.num_output_placeholders = 1
+    # This step's frame (settled to 0 by update_from_output) plus one extra
+    # async decode still in flight.
+    session.num_in_flight_tokens = 2
+
+    sched = _make_scheduler(stage_id=0)
+    sched._enqueue_waiting_request = MagicMock()
+
+    def handle_stopped(request: Request) -> bool:
+        sched._update_request_as_session(request, _make_update([10, 20]))
+        return False
+
+    _run_resumable_segment_stop(session, handle_stopped=handle_stopped)
+
+    # Exactly the one unreported decode, so the drain reaches zero before the
+    # new segment's first frame arrives.
+    assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
+
+
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()

tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json

@@ -18,7 +18,7 @@
       "model": "openbmb/MiniCPM-o-4_5",
       "extra_cli_args": [
         "--deploy-config",
-        "vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
+        "vllm_omni/deploy/minicpmo_4_5.yaml",
         "--trust-remote-code"
       ]
     },

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -21,47 +21,6 @@
       ]
     },
     "benchmark_params": [
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "H100": {
-            "mean_e2el_ms": [
-              11701.3818,
-              27510.3956,
-              31907.9613,
-              42173.9832,
-              63970.0751
-            ],
-            "mean_audio_rtf": [
-              0.1175,
-              0.1647,
-              0.2003,
-              0.2726,
-              0.4195
-            ]
-          }
-        }
-      },
       {
         "dataset_name": "random-mm",
         "backend": "openai-chat-omni",
@@ -170,45 +129,6 @@
           }
         }
       },
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "H100": {
-            "mean_e2el_ms": [
-              4354.8945,
-              7672.4248,
-              9861.1687,
-              13086.8865,
-              18716.709
-            ]
-          }
-        }
-      },
       {
         "dataset_name": "random-mm",
         "backend": "openai-chat-omni",
@@ -324,5 +244,61 @@
         }
       }
     ]
+  },
+  {
+    "test_name": "test_qwen3_omni_chunk_cuda",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "omni",
+      "slow"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "extra_cli_args": [
+        "--async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          32,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          8,
+          32
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "H100": {
+            "mean_e2el_ms": [
+              4354.8945,
+              9861.1687,
+              18716.709
+            ]
+          }
+        }
+      }
+    ]
   }
 ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -29,16 +29,12 @@
         "endpoint": "/v1/completions",
         "num_prompts": [
           4,
-          16,
           32,
-          64,
           128
         ],
         "max_concurrency": [
           1,
-          4,
           8,
-          16,
           32
         ],
         "random_input_len": 2500,
@@ -54,9 +50,7 @@
           "H100": {
             "mean_e2el_ms": [
               4619.0068,
-              7868.3523,
               9760.6175,
-              13422.3491,
               18007.6918
             ]
           }
@@ -99,45 +93,6 @@
           }
         }
       },
-      {
-        "dataset_name": "random-mm",
-        "backend": "vllm",
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "endpoint": "/v1/completions",
-        "num_prompts": [
-          40
-        ],
-        "request_rate": [
-          0.5
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "H100": {
-            "mean_e2el_ms": [
-              524.5321
-            ]
-          }
-        }
-      },
       {
         "dataset_name": "random-mm",
         "backend": "vllm",

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -74,13 +74,14 @@
       "full_model",
       "diffusion"
     ],
-    "description": "Single-device baseline (no parallelism) with step execution",
+    "description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).",
     "server_type": "vllm-omni",
     "server_params": {
       "model": "Qwen/Qwen-Image",
       "serve_args": {
         "enable-diffusion-pipeline-profiler": true,
-        "step-execution": true
+        "step-execution": true,
+        "max-num-seqs": 8
       }
     },
     "benchmark_params": [
@@ -119,6 +120,27 @@
             "peak_memory_mb_mean": 64500.0
           }
         }
+      },
+      {
+        "name": "512x512_steps20_high_concurrency",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": [20, 40, 80, 160],
+        "max-concurrency": [1, 2, 4, 8],
+        "warmup-requests": 8,
+        "warmup-concurrency": 8,
+        "warmup-num-inference-steps": 20,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "H100": {
+            "throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
+            "latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
+            "peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
+          }
+        }
       }
     ]
   },
@@ -243,53 +265,5 @@
         }
       }
     ]
-  },
-  {
-    "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true,
-        "step-execution": true,
-        "max-num-seqs": 8
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_high_concurrency",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": [20, 40, 80, 160],
-        "max-concurrency": [1, 2, 4, 8],
-        "warmup-requests": 8,
-        "warmup-concurrency": 8,
-        "warmup-num-inference-steps": 20,
-        "enable-negative-prompt": true,
-        "baseline": {
-            "H100": {
-                "throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
-                "latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
-                "peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
-            }
-        }
-      }
-    ]
   }
 ]

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -157,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
                 "ref_audio": None,
                 "ref_text": None,
             },
-            ("CustomVoice", "does not support"),
-            id="customvoice_task_not_supported",
+            ("Invalid voice", "vivian"),
+            id="customvoice_invalid_voice",
         ),
         pytest.param(
             {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -181,13 +181,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
             ("mutually exclusive", "speaker_embedding", "ref_audio"),
             id="speaker_embedding_with_ref_audio_mutually_exclusive",
         ),
-        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
         pytest.param(
             {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
             ("max_new_tokens", "exceed 4096"),
             id="max_new_tokens_above_max",
         ),
-        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+        pytest.param(
+            {"seed": -(2**63) - 1},
+            ("seed", "greater_than_equal", "-9223372036854775808"),
+            id="seed_below_min",
+        ),
         pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
         pytest.param(
             {"initial_codec_chunk_frames": -1},
@@ -275,7 +279,9 @@ def test_speech_invalid_field_values(
         pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
         pytest.param({"input": "   ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
         pytest.param(
-            {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
+            {"input": "Hello world.", "max_new_tokens": 0},
+            ("max_new_tokens", "greater_than_equal"),
+            id="max_new_tokens_below_min",
         ),
         pytest.param(
             {
@@ -554,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, online_client: Onli
         pytest.param(
             "batch",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="batch_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),
         pytest.param(
             "item",
             {"max_new_tokens": 0},
-            ("max_new_tokens", "least 1"),
+            ("max_new_tokens", "greater_than_equal"),
             id="item_max_new_tokens_below_min",
             marks=_SKIP_ISSUE_3649,
         ),

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
         body["top_logprobs"] = 5
     elif case_id == "speaker_unknown":
         body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
+    elif case_id == "audio_format_unsupported":
+        body["modalities"] = ["text", "audio"]
+        body["audio"] = {"voice": "alloy", "format": "aac"}
     else:
         raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
     return {"json": body, "timeout": 120}
@@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
         pytest.param(
             "response_format_json_schema_incomplete",
             400,
-            ("response_format", "value_error", "json_schema"),
+            ("response_format", "BadRequestError", "json_schema"),
             id="invalid_response_format_json_schema",
         ),
         pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
@@ -135,6 +138,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
             ("Invalid speaker", "Supported"),
             id="speaker_unknown_preset",
         ),
+        pytest.param(
+            "audio_format_unsupported",
+            400,
+            ("Invalid audio format", "aac", "Supported formats"),
+            id="unsupported_audio_format",
+        ),
     ],
 )
 @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)

tests/diffusion/batching/test_diffusion_batching.py

@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
 ) -> None:
     """Run the same prompts sequentially then in parallel and print a comparison."""
 
-    omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
         single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
         await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
         return
 
-    omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
 
@@ -524,7 +524,12 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
     parser.add_argument("--model", type=str, required=True, help="Model name or path")
     parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
-    parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
+    )
     parser.add_argument(
         "--mode",
         choices=["batch", "single", "compare", "validate"],

tests/diffusion/cache/test_cache_backends.py

@@ -365,13 +365,10 @@ class TestCacheDiTBackend:
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
-    def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
-        """Test DreamID uses pipeline.transformer for cache enable/refresh.
-
-        NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
-        """
+    def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        """Generic cache path should pick up ``transformer.fused_blocks``."""
         mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
         mock_pipeline.transformer = Mock()
         mock_pipeline.transformer.fused_blocks = Mock()
         mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(

tests/diffusion/cache/test_cache_dit.py

@@ -6,7 +6,6 @@ Model specific tests for CacheDiT enablement.
 """
 
 import ast
-import sys
 from pathlib import Path
 from unittest.mock import Mock, patch
 
@@ -25,18 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.platforms import current_omni_platform
 
-# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
-# the class even though the dependency may not be set up correctly; this is ok for these
-# tests because we just inspect it and never initialize the model.
-for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
-    sys.modules.setdefault(mod, Mock())
-# isort: split
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel  # noqa: E402
-
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_TRANSFORMERS = [
-    DreamIdOmniModel,
     HeliosTransformer3DModel,
     LongCatImageTransformer2DModel,
     Cosmos3VFMTransformer,

tests/diffusion/layers/test_norm.py

@@ -340,6 +340,97 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo
     assert captured["epsilon"].item() == pytest.approx(1e-5)
 
 
+# ── RMSNorm residual tests ──
+
+
+def test_rmsnorm_matches_native_residual_contract():
+    """The non-NPU path preserves separate residual-add and RMSNorm math."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    eps = 1e-6
+    norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
+    norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
+    x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
+    residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)
+
+    expected_residual = residual + x
+    expected_fp32 = expected_residual.float()
+    expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
+    expected_output = (expected_output * norm.weight.float()).to(x.dtype)
+    output, updated_residual = norm.forward_native(x, residual)
+
+    torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
+    torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "forward_name",
+    ["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
+)
+def test_rmsnorm_non_npu_residual_delegates_to_native(
+    monkeypatch: pytest.MonkeyPatch,
+    forward_name: str,
+):
+    """Platform wrappers keep residual arithmetic in forward_native."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
+    x = torch.randn(2, 4, dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    expected_output = torch.randn_like(x)
+    expected_residual = torch.randn_like(x)
+    captured: dict[str, object] = {}
+
+    def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
+        captured.update(x=native_x, residual=native_residual)
+        return expected_output, expected_residual
+
+    monkeypatch.setattr(norm, "forward_native", forward_native)
+
+    output, updated_residual = getattr(norm, forward_name)(x, residual)
+
+    assert output is expected_output
+    assert updated_residual is expected_residual
+    assert captured["x"] is x
+    assert captured["residual"] is residual
+
+
+def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
+    """NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    captured: dict[str, object] = {}
+    fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
+    fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)
+
+    def npu_add_rms_norm(
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        gamma: torch.Tensor,
+        epsilon: float,
+    ):
+        captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
+        return fused_output, torch.empty(0), fused_residual
+
+    monkeypatch.setitem(
+        sys.modules,
+        "torch_npu",
+        types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
+    )
+    norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
+    x = torch.randn(2, 4, dtype=torch.bfloat16)
+    residual = torch.randn(2, 4, dtype=torch.bfloat16)
+
+    output, updated_residual = norm.forward_npu(x, residual)
+
+    assert output is fused_output
+    assert updated_residual is fused_residual
+    assert captured["x"] is x
+    assert captured["residual"] is residual
+    assert captured["gamma"] is norm.weight
+    assert captured["epsilon"] == 1e-5
+
+
 # ── RMSNorm compile-path regression tests ──
 
 

tests/diffusion/layers/test_norm_cuda.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
+]
+
+
+def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
+    eps = 1e-6
+    device = torch.device(current_omni_platform.device_type)
+    torch.manual_seed(0)
+    norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
+    x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    x_before = x.clone()
+    residual_before = residual.clone()
+
+    expected_output, expected_residual = norm.forward_native(x_before, residual_before)
+    output, updated_residual = norm.forward_cuda(x, residual)
+
+    torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
+    torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
+    torch.testing.assert_close(x, x_before, atol=0, rtol=0)
+    torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)

tests/diffusion/lora/test_base_linear.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
     assert torch.allclose(out, expected)
 
 
+def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
+    from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    torch.nn.Module.__init__(layer)
+    layer.base_layer = torch.nn.Module()
+    layer.n_slices = 1
+    layer.lora_a_stacked = (torch.zeros(2),)
+    layer.lora_b_stacked = (torch.zeros(3),)
+
+    layer._set_diffusion_lora_buffer_device(torch.device("meta"))
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+    def recreate_on_cpu(self, max_loras, lora_config, model_config):
+        self.lora_a_stacked = (torch.zeros(4),)
+        self.lora_b_stacked = (torch.zeros(5),)
+
+    monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
+    layer.create_lora_weights(max_loras=1, lora_config=object())
+
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+
 def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
     # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
     # LoRA tensors are still allocated and non-empty.

tests/diffusion/lora/test_lora_manager.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
@@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
     assert replace_calls == ["foo"]
 
 
+def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DLOLoRALayer(DummyBaseLayerWithLoRA):
+        lora_buffer_device = None
+
+        def _set_diffusion_lora_buffer_device(self, device):
+            self.lora_buffer_device = device
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
+    monkeypatch.setattr(
+        manager_mod,
+        "from_layer_diffusion",
+        lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
+    )
+
+    pipeline = torch.nn.Module()
+    pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _FakeLinearBase()
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
+
+    assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
+
+
 def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
     import vllm_omni.diffusion.lora.manager as manager_mod
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlanResult,
     TensorBinding,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.diffusion.registry import initialize_model
@@ -173,6 +174,15 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
         str(canonical_root / "model.safetensors"),
     )
     store_root = tmp_path / "hwr-store"
+    hash_calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal hash_calls
+        hash_calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
 
     def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
         loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
@@ -205,6 +215,8 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     assert warm_plan is not None
     assert warm_plan.lease_carrier is not None
     warm_plan.lease_carrier.close()
+    assert hash_calls == 1
+    assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
 
 
 def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
@@ -237,6 +249,30 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
     assert loader.take_host_weight_plan() is None
 
 
+def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
+    tmp_path: Path,
+    monkeypatch,
+):
+    canonical_root = tmp_path / "canonical"
+    canonical_root.mkdir()
+    save_file(
+        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+        str(canonical_root / "model.safetensors"),
+    )
+    loader = DiffusersPipelineLoader(
+        LoadConfig(),
+        _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
+    )
+    pipeline = _HWRPipeline(canonical_root)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+
+    with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
+        loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    assert pipeline.load_count == 0
+    assert loader.take_host_weight_plan() is None
+
+
 def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
@@ -249,10 +285,11 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
 
 
 @pytest.mark.parametrize(
-    ("dist_offload", "use_allgather"),
+    ("dist_offload", "use_allgather", "mode"),
     [
-        (False, False),
-        (True, True),
+        (False, False, "preferred"),
+        (True, True, "preferred"),
+        (True, False, "disabled"),
     ],
 )
 def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
@@ -260,12 +297,13 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
     tmp_path,
     dist_offload,
     use_allgather,
+    mode,
 ):
     """Disabled and AllGather paths must never construct or probe HWR."""
     from vllm_omni.host_weight_runtime import HostWeightRuntime
 
     root = tmp_path / "must-not-be-touched"
-    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
+    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
     model = _DummyPipelineModel(source_prefix="transformer.")
     modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
 

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -6,6 +6,8 @@ from __future__ import annotations
 
 import dataclasses
 import gc
+import json
+import threading
 from collections.abc import Sequence
 from pathlib import Path
 from types import SimpleNamespace
@@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import (
     FinalLayoutRequest,
     FinalLayoutTensorRestorer,
     ImplementationIdentity,
+    NodeSourceDigestCache,
     PreparedWeightSource,
     WeightSourceKind,
     build_final_layout_identity,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.model_loader.host_weights.contracts import (
     FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
     FinalLayoutArtifactSpec,
@@ -299,6 +303,83 @@ class _FakeH3Attention(nn.Module):
         super().__init__()
 
 
+class _FakeFluxLinear(_FakeH3Linear):
+    def __init__(
+        self,
+        *args: object,
+        bias: bool = False,
+        params_dtype: torch.dtype | None = None,
+        total_num_heads: int | None = None,
+        total_num_kv_heads: int | None = None,
+        **kwargs: object,
+    ) -> None:
+        super().__init__(
+            *args,
+            bias=bias,
+            params_dtype=params_dtype,
+            total_num_heads=total_num_heads,
+            total_num_kv_heads=total_num_kv_heads,
+            **kwargs,
+        )
+        self.weight.weight_loader = self._load_weight
+
+    @staticmethod
+    def _load_weight(
+        parameter: nn.Parameter,
+        loaded_weight: torch.Tensor,
+        shard_id: str | None = None,
+    ) -> None:
+        del shard_id
+        parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
+
+
+class _FakeFluxAttention(nn.Module):
+    def __init__(self, **kwargs: object) -> None:
+        del kwargs
+        super().__init__()
+        self.register_buffer("beta", torch.tensor([0.5]))
+        self.register_buffer("eps", torch.tensor([-1e-6]))
+
+
+class _Flux2KleinPipeline(nn.Module):
+    def __init__(self, transformer: nn.Module) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
+        self.vae = nn.Module()
+        self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
+
+
+def _reduced_flux2_klein(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    device: torch.device | str = "cpu",
+) -> nn.Module:
+    from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
+
+    monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
+
+    transformer = flux2.Flux2Transformer2DModel(
+        patch_size=1,
+        in_channels=8,
+        out_channels=8,
+        num_layers=1,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=1,
+        joint_attention_dim=8,
+        timestep_guidance_channels=8,
+        mlp_ratio=1.0,
+        axes_dims_rope=(2, 2, 2, 2),
+        guidance_embeds=False,
+    )
+    return transformer.to(device=device, dtype=torch.bfloat16)
+
+
 class _H3Pipeline(nn.Module):
     def __init__(self, transformer: nn.Module) -> None:
         super().__init__()
@@ -380,6 +461,7 @@ def _identity(
     *,
     request: FinalLayoutRequest | None = None,
     policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
+    source_digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutIdentityContext:
     transformer = model.get_submodule("transformer")
     return build_final_layout_identity(
@@ -388,6 +470,7 @@ def _identity(
         prepared_sources=(source,),
         request=request or _request(),
         policy=policy,
+        source_digest_cache=source_digest_cache,
     )
 
 
@@ -564,6 +647,131 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path
     assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
 
 
+def test_node_source_digest_cache_reuses_recovers_and_invalidates(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache_root = tmp_path / "store"
+    calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal calls
+        calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
+    first = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    second = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+
+    assert second.identity == first.identity
+    assert calls == 1
+
+    cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
+    cache_entry.write_text("not-json", encoding="utf-8")
+    recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert recovered.identity == first.identity
+    assert calls == 2
+
+    cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
+    cache_document["content_id"] = f"sha256:{'0' * 64}"
+    cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
+    checksum_recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert checksum_recovered.identity == first.identity
+    assert calls == 3
+
+    source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
+    changed = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert changed.identity != first.identity
+    assert calls == 4
+
+
+def test_source_identity_hashes_checkpoint_shards_in_parallel(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    second_file = source.resolved_root / "model-00002.safetensors"
+    second_file.write_bytes(b"second-checkpoint-shard")
+    source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
+    barrier = threading.Barrier(2)
+    thread_ids: set[int] = set()
+    thread_ids_lock = threading.Lock()
+    original_sha256 = source_identity_module._sha256_file
+
+    def synchronized_sha256(path: Path, state: object) -> str:
+        with thread_ids_lock:
+            thread_ids.add(threading.get_ident())
+        barrier.wait(timeout=5)
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
+    _identity(model, source)
+
+    assert len(thread_ids) == 2
+
+
+def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+
+    def unavailable_lock(*_args: object, **_kwargs: object) -> object:
+        raise source_identity_module.FileLockTimeoutError("busy")
+
+    monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+    assert not tuple(cache.entries.glob("*.json"))
+
+
+def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    original_mkdir = Path.mkdir
+
+    def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
+        if "source-digests-v1" in path.parts:
+            raise OSError("injected cache directory failure")
+        original_mkdir(path, *args, **kwargs)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+
+
 def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
     model = _TinyPipeline()
     source_root = tmp_path / "local-symlink-source"
@@ -1069,3 +1277,152 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
     assert contract.implementation_id == "minimax-h3-dit"
     assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
     assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
+
+
+def test_reduced_flux2_klein_declares_complete_final_layout_contract(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
+        Flux2Transformer2DModel,
+    )
+
+    transformer = _reduced_flux2_klein(monkeypatch)
+    pipeline = _Flux2KleinPipeline(transformer)
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    context = _identity(pipeline, source)
+    targets = collect_final_layout_targets(
+        pipeline,
+        (("transformer", transformer),),
+        policy=FINAL_LAYOUT_BF16_POLICY,
+        require_materialized=True,
+    )
+    target_names = {target.name for target in targets}
+
+    contract = Flux2Transformer2DModel.host_weight_restore_contract
+    assert isinstance(contract, FinalLayoutModelContract)
+    assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+    assert contract.implementation_id == "flux2-klein-dit"
+    assert contract.version == "1"
+    assert context.model_contracts == (contract,)
+    assert transformer.stacked_params_mapping == [
+        (".to_qkv.", ".to_q.", "q"),
+        (".to_qkv.", ".to_k.", "k"),
+        (".to_qkv.", ".to_v.", "v"),
+        (".add_kv_proj", ".add_q_proj", "q"),
+        (".add_kv_proj", ".add_k_proj", "k"),
+        (".add_kv_proj", ".add_v_proj", "v"),
+    ]
+    loaded = transformer.load_weights(
+        [
+            ("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
+            ("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
+            (
+                "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+                torch.tensor([3.0], dtype=torch.bfloat16),
+            ),
+        ]
+    )
+    assert loaded == {
+        "transformer_blocks.0.attn.to_qkv.weight",
+        "transformer_blocks.0.attn.add_kv_proj.weight",
+        "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+    }
+    assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
+    assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
+    assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
+    assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
+    assert all("text_encoder" not in name and "vae" not in name for name in target_names)
+    transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_restore_validator_rejects_non_ready_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    transformer = _reduced_flux2_klein(monkeypatch)
+
+    mapping = transformer.stacked_params_mapping
+    transformer.stacked_params_mapping = None
+    with pytest.raises(ValueError, match="packed QKV mapping"):
+        transformer.validate_restored_host_weights()
+    transformer.stacked_params_mapping = mapping
+
+    weight = transformer.x_embedder.weight
+    weight.data = weight.data.float()
+    with pytest.raises(ValueError, match="must stay bf16"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.to(torch.bfloat16)
+
+    weight.data = weight.data.T
+    with pytest.raises(ValueError, match="contiguous strided layout"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.contiguous()
+
+    attention = transformer.transformer_blocks[0].attn.attn
+    attention._non_persistent_buffers_set.add("beta")
+    with pytest.raises(ValueError, match="must be persistent"):
+        transformer.validate_restored_host_weights()
+    attention._non_persistent_buffers_set.remove("beta")
+
+    attention.beta.data = attention.beta.data.float()
+    with pytest.raises(ValueError, match="scalar bf16"):
+        transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_publication_restore_fallback_and_teardown(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    cold_transformer = _reduced_flux2_klein(monkeypatch)
+    cold_model = _Flux2KleinPipeline(cold_transformer)
+    with torch.no_grad():
+        for index, parameter in enumerate(cold_transformer.parameters(), start=1):
+            parameter.fill_(index)
+        for name, buffer in cold_transformer.named_buffers():
+            if name.endswith((".beta", ".eps")):
+                buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
+
+    context = _identity(cold_model, source)
+    runtime = _runtime(tmp_path / "flux2-klein-store")
+    miss = runtime.resolve(context.identity)
+    assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+    producer = FinalLayoutBF16Producer(
+        context,
+        cold_model,
+        (("transformer", cold_transformer),),
+        max_shard_bytes=64,
+    )
+    publication = runtime.publish_after_load(context.identity, producer=producer)
+    assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+
+    warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
+    warm_model = _Flux2KleinPipeline(warm_transformer)
+    warm_context = _identity(warm_model, source)
+    assert warm_context.identity == context.identity
+    hit = runtime.resolve(warm_context.identity)
+    assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert hit.lease is not None
+
+    plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+    assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
+    plan.commit()
+    warm_transformer.validate_restored_host_weights()
+
+    cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
+    warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
+    for name in warm_context.tensor_names:
+        assert torch.equal(warm_tensors[name], cold_tensors[name])
+        assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
+    assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
+
+    del warm_tensors, warm_model, warm_transformer, plan
+    gc.collect()
+    hit.lease.close()
+    assert hit.lease.closed
+    assert isinstance(runtime.store, FilesystemHostWeightStore)
+    assert runtime.store.cleanup(context.identity) is None

tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py

@@ -1,27 +0,0 @@
-import pytest
-import torch.nn as nn
-
-pytest.importorskip("dreamid_omni.modules.model")
-
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
-    model = object.__new__(FusionModel)
-    nn.Module.__init__(model)
-    model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
-    model.attn = nn.Identity()
-
-    conditions = getattr(model, "_hsdp_shard_conditions", None)
-
-    assert conditions is not None
-    assert len(conditions) == 1
-
-    matched = []
-    for name, module in model.named_modules():
-        if any(cond(name, module) for cond in conditions):
-            matched.append(name)
-
-    assert matched == ["fused_blocks.0", "fused_blocks.1"]

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -964,6 +964,80 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit():
     assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
 
 
+def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+    class FakeEncoderGroup:
+        rank_in_group = 0
+        world_size = 2
+
+        def all_reduce(self, tensor):
+            del tensor
+
+    config = SimpleNamespace(
+        hidden_size=8,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=2,
+        rms_norm_eps=1e-6,
+    )
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32)
+    # This CPU contract test only needs projection shapes and head reshaping.
+    # Avoid dispatching platform-specific norm and linear kernels (for example,
+    # rocm_unquantized_gemm when the test suite is collected on ROCm).
+    attention.q_norm = nn.Identity()
+    attention.k_norm = nn.Identity()
+    with torch.no_grad():
+        attention.qkv_proj.weight.zero_()
+    attention.o_proj = nn.Linear(
+        attention.qkv_proj.local_num_heads * attention.head_dim,
+        config.hidden_size,
+        bias=False,
+    )
+    attn_call = {}
+
+    def fake_attention(query, key, value):
+        attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
+        return query
+
+    monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention)
+    hidden_states = torch.randn(1, 3, config.hidden_size)
+    cos = torch.ones(1, 3, config.head_dim)
+    sin = torch.zeros_like(cos)
+
+    output = attention(hidden_states, (cos, sin))
+
+    assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim)
+    assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
+    assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
+    assert output.shape == hidden_states.shape
+
+
+def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+    sdpa_call = {}
+
+    def fake_sdpa(query, key, value, **kwargs):
+        sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
+        return query
+
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+    output = encoder_module._scaled_dot_product_attention(
+        torch.randn(1, 4, 3, 8),
+        torch.randn(1, 2, 3, 8),
+        torch.randn(1, 2, 3, 8),
+    )
+
+    assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8)
+    assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True}
+    assert output.shape == (1, 4, 3, 8)
+
+
 def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
     [
         "model-level CPU offload (--enable-cpu-offload)",
         "layerwise offload (--enable-layerwise-offload)",
-        "distributed layerwise offload (--enable-distributed-layerwise-offload)",
     ],
 )
 def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
@@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
         )
 
 
+def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+    loaded = pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+
+    assert loaded is not None
+    assert captured["unsupported_offload_mode"] is None
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+    from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+    def sdpa_sentinel(query, key, value):
+        del key, value
+        return query
+
+    def rope_sentinel(q, k, cos, sin):
+        del cos, sin
+        return q, k
+
+    monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
+    monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
+    monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+    monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+
+    npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
+
+    assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu
+    assert encoder_module._apply_rotary_pos_emb is rope_sentinel
+    assert not npu_minimax_h3._ROPE_PATCHED
+    assert npu_minimax_h3._SDPA_PATCHED
+
+    monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+    npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
+
+    assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize(
+    ("num_heads", "num_kv_heads", "expected_enable_gqa"),
+    [(2, 2, False), (4, 2, True)],
+)
+def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa(
+    monkeypatch,
+    num_heads: int,
+    num_kv_heads: int,
+    expected_enable_gqa: bool,
+) -> None:
+    from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+    captured: dict[str, object] = {}
+
+    def fake_sdpa(query, key, value, **kwargs):
+        captured.update(query=query, key=key, value=value, kwargs=kwargs)
+        return query
+
+    monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa)
+    query = torch.randn(1, num_heads, 3, 8)
+    key = torch.randn(1, num_kv_heads, 3, 8)
+    value = torch.randn_like(key)
+
+    output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+    assert captured["query"] is query
+    assert captured["key"] is key
+    assert captured["value"] is value
+    assert captured["kwargs"] == {
+        "dropout_p": 0.0,
+        "is_causal": True,
+        "enable_gqa": expected_enable_gqa,
+    }
+    assert output is query
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize(
+    ("num_heads", "num_key_heads", "num_value_heads", "message"),
+    [
+        (4, 2, 1, "key and value to have the same number of heads"),
+        (4, 0, 0, "at least one KV head"),
+        (3, 2, 2, "query heads to be a multiple of KV heads"),
+    ],
+)
+def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations(
+    num_heads: int,
+    num_key_heads: int,
+    num_value_heads: int,
+    message: str,
+) -> None:
+    from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+    query = torch.randn(1, num_heads, 3, 8)
+    key = torch.randn(1, num_key_heads, 3, 8)
+    value = torch.randn(1, num_value_heads, 3, 8)
+
+    with pytest.raises(ValueError, match=message):
+        npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None:
+    """Exercise the real torch-npu causal GQA operator."""
+    from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+    torch.manual_seed(0)
+    batch_size, seq_len = 1, 16
+    num_heads, num_kv_heads, head_size = 8, 2, 128
+    device = torch.device("npu")
+    dtype = torch.bfloat16
+
+    query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype)
+    key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype)
+    value = torch.randn_like(key)
+
+    actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+    repeat_num = num_heads // num_kv_heads
+    expected = F.scaled_dot_product_attention(
+        query,
+        key.repeat_interleave(repeat_num, dim=1),
+        value.repeat_interleave(repeat_num, dim=1),
+        dropout_p=0.0,
+        is_causal=True,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)

tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py

@@ -6,7 +6,38 @@ import torch
 
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+    from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+    def rope_sentinel(q, k, cos, sin):
+        del cos, sin
+        return q, k
+
+    def sdpa_sentinel(query, key, value):
+        del key, value
+        return query
+
+    monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
+    monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
+    monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+    monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+
+    npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
+
+    assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu
+    assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
+    assert npu_minimax_h3._ROPE_PATCHED
+    assert not npu_minimax_h3._SDPA_PATCHED
+
+    monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+    npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
+
+    assert encoder_module._apply_rotary_pos_emb is rope_sentinel
 
 
 @pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")

tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py

@@ -1,10 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
 import torch
+import torch.nn as nn
 
 from vllm_omni.diffusion.layers.norm import RMSNorm
-from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
+from vllm_omni.diffusion.models.minimax_h3.encoder import (
+    MiniMaxH3Qwen3VLRMSNorm,
+    MiniMaxH3Qwen3VLTextDecoderLayer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
@@ -24,3 +34,43 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
     assert norm.weight.dtype == torch.bfloat16
     assert set(norm.state_dict()) == {"weight"}
     torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
+
+
+def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None:
+    class ConstantAttention(nn.Module):
+        def forward(self, hidden_states, position_embeddings):
+            del position_embeddings
+            return torch.full_like(hidden_states, 2.0)
+
+    config = SimpleNamespace(
+        hidden_size=8,
+        intermediate_size=16,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=2,
+        rms_norm_eps=1e-6,
+    )
+    group = SimpleNamespace(rank_in_group=0, world_size=1)
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32)
+    layer.input_layernorm = nn.Identity()
+    layer.self_attn = ConstantAttention()
+    layer.mlp = nn.Identity()
+
+    hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8)
+    attention_output = torch.full_like(hidden_states, 2.0)
+    expected_residual = hidden_states + attention_output
+    expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual)
+
+    with patch.object(
+        layer.post_attention_layernorm,
+        "forward",
+        wraps=layer.post_attention_layernorm.forward_native,
+    ):
+        output = layer(hidden_states, (torch.empty(0), torch.empty(0)))
+
+    assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm)
+    torch.testing.assert_close(output, expected_residual + expected_normalized)

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import ctypes
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    _coalesce_ranges,
+    _mapped_regions,
+)
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationBudgetError,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+    register_host_mappings,
+)
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+        *,
+        read_only_supported: bool = True,
+        attribute_result: int = 0,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self._read_only_supported = read_only_supported
+        self._attribute_result = attribute_result
+        self.attribute_queries: list[tuple[int, int]] = []
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
+        self.attribute_queries.append((attribute, device))
+        ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported)
+        return self._attribute_result
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def _region(
+    file_name: str,
+    address: int,
+    length: int,
+    *,
+    read_only: bool = True,
+) -> MappedHostRegion:
+    return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only)
+
+
+@pytest.fixture(autouse=True)
+def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
+
+
+def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    regions = (_region("weights", 0x1000, 4096),)
+
+    def create(actual_regions, *, max_bytes):
+        assert actual_regions is regions
+        assert max_bytes == 4096
+        return sentinel
+
+    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+    assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel
+    with pytest.raises(HostRegistrationError, match="not supported on cpu"):
+        register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None)
+
+
+def test_ranges_align_per_file_without_merging_adjacent_files() -> None:
+    assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+    )
+    assert _mapped_regions(
+        (
+            _region("first.safetensors", 0x1000, 4096),
+            _region("second.safetensors", 0x2000, 4096),
+        )
+    ) == (
+        registration_module._AddressRange(0x1000, 0x2000),
+        registration_module._AddressRange(0x2000, 0x3000),
+    )
+
+
+def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(HostRegistrationError, match="is writable"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None)
+    with pytest.raises(HostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096)
+    assert runtime.registered == []
+
+
+def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], read_only_supported=False)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(HostRegistrationError, match="does not support read-only"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+    assert runtime.attribute_queries == [
+        (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
+    ]
+    assert runtime.registered == []
+
+
+def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], attribute_result=7)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    consumed: list[tuple[object, int]] = []
+    monkeypatch.setattr(
+        registration_module,
+        "_consume_last_cuda_error",
+        lambda actual_runtime, error: consumed.append((actual_runtime, error)),
+    )
+
+    with pytest.raises(HostRegistrationError, match="error-7"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+    assert consumed == [(runtime, 7)]
+
+
+def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+    with pytest.raises(HostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+    assert runtime.unregistered == [0x1000]
+
+
+def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0, 7], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+    with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error:
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+    registration = error.value.active_registration
+    assert registration is not None
+    assert registration.close() == ()
+    assert runtime.unregistered == [0x1000, 0x1000]
+
+
+def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+    registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None)
+
+    assert registration.total_bytes == 4096
+    assert registration.region_count == 1
+    assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",)
+    assert registration.close() == ()
+    assert runtime.unregistered == [0x1000, 0x1000]
+
+
+@pytest.mark.parametrize("pending_error", [0, 801])
+def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
+    monkeypatch: pytest.MonkeyPatch,
+    pending_error: int,
+) -> None:
+    runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error)
+    monkeypatch.setattr(
+        registration_module.ctypes,
+        "CDLL",
+        lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"),
+    )
+
+    registration_module._consume_last_cuda_error(runtime, 801)
+
+
+def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None:
+    global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801)
+    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime)
+
+    registration_module._consume_last_cuda_error(SimpleNamespace(), 801)
+
+
+def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    def missing_global_symbol(_name: object) -> None:
+        raise OSError("cudart symbols are local")
+
+    monkeypatch.setattr(
+        registration_module.ctypes,
+        "CDLL",
+        missing_global_symbol,
+    )
+
+    with pytest.raises(HostRegistrationError, match="cannot clear") as error:
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+
+    assert not isinstance(error.value, HostRegistrationCleanupError)
+    assert runtime.unregistered == [0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -5,6 +5,7 @@
 
 import gc
 import json
+import weakref
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -17,6 +18,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
 from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     TensorBinding,
@@ -33,12 +35,17 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadHook,
     PinnedResidentLayerGroup,
 )
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+)
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
+from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -703,11 +710,15 @@ class _HWRPipeline(nn.Module):
 
 
 class _FakeHostWeightLease:
-    def __init__(self, resolution_id: str):
+    def __init__(self, resolution_id: str, events: list[str] | None = None):
         self.closed = False
         self.provenance = SimpleNamespace(resolution_id=resolution_id)
+        self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
+        self._events = events
 
     def close(self):
+        if self._events is not None:
+            self._events.append("lease")
         self.closed = True
 
 
@@ -757,6 +768,206 @@ def _hwr_backend():
     return _HWRPipeline(), backend, carrier, lease
 
 
+def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
+    current_block = nn.Linear(2, 2, bias=False)
+    next_block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    next_block.weight.data.copy_(expected)
+    hook = DistributedLayerwiseOffloadHook(
+        next_block=next_block,
+        device=torch.device("cpu"),
+        dp_group=None,
+        dp_size=1,
+        rank=0,
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    hook.initialize_hook(current_block)
+    hook.registered_mmap = True
+    hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging")  # type: ignore[method-assign]
+
+    hook.prefetch_layer(slot=0, non_blocking=True)
+
+    assert torch.equal(next_block.weight, expected)
+
+
+def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
+    block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    block.weight.data.copy_(expected)
+    group = PinnedResidentLayerGroup(
+        [block],
+        device=torch.device("cpu"),
+        copy_stream=DummyStream(),
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    group.registered_mmap = True
+    group._cpu_staging_buffers.clear()
+
+    group.load()
+
+    assert torch.equal(block.weight, expected)
+
+
+def test_hwr_registration_uses_transport_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    plan, _, lease = _fake_hwr_plan()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+            dlo_host_registration_limit_gib=1.5,
+        ),
+        torch.device("cuda"),
+        host_weight_plan=plan,
+    )
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    calls: list[tuple[object, int | None]] = []
+
+    class Registration:
+        total_bytes = 4096
+        region_count = 1
+
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            return ()
+
+    def register(regions, *, device, max_bytes):
+        assert device == torch.device("cuda")
+        calls.append((regions, max_bytes))
+        return Registration()
+
+    monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+    pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
+
+    assert backend._try_register_hwr_mmap((pinned_source,))  # type: ignore[arg-type]
+    assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend._release_registered_mmap()
+
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+
+
+def test_hwr_registration_budget_validation_is_transport_scoped():
+    assert (
+        validate_dlo_host_registration_options(
+            limit_gib=1.5,
+            enable_dlo=True,
+            use_allgather=False,
+            hwr_mode="preferred",
+        )
+        == 1.5
+    )
+    invalid = (
+        dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
+    )
+    for options in invalid:
+        with pytest.raises(ValueError):
+            validate_dlo_host_registration_options(**options)
+
+
+def test_unregistration_precedes_lease_close_and_retries_failures(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    events: list[str] = []
+    responses = iter([("busy",), ()])
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return next(responses)
+
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retry", events)
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = Registration()
+    backend._using_rank_local_mmap = True
+    backend._using_registered_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    assert not lease.closed
+    assert backend._host_registration is not None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend.disable()
+
+    assert lease.closed
+    assert backend._host_registration is None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+    assert events == ["unregister", "unregister", "lease"]
+
+
+def test_failed_unregistration_retains_lease_after_backend_is_dropped(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    events: list[str] = []
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return ("busy",)
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retained", events)
+    registration = Registration()
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = registration
+    backend._using_rank_local_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    lease_ref = weakref.ref(lease)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    del backend
+    del lease
+    gc.collect()
+
+    retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
+    assert lease_ref() is retained_lease
+    assert not retained_lease.closed
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    retained_lease.close()
+
+
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -920,6 +1131,51 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+    def test_hwr_registration_failure_falls_back_to_bounded_staging(
+        self,
+        monkeypatch,
+        patched_offload_runtime,
+    ):
+        plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        pipeline = _HWRPipeline()
+        original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
+
+        def fail_registration(*args, **kwargs):
+            del args, kwargs
+            raise HostRegistrationError("registration unavailable in CPU test")
+
+        def allocate_unpinned_staging(hooks, resident_group=None):
+            # Registration selection still observes pin_cpu_memory=True. The
+            # CPU-only test avoids asking the host for CUDA-pinned allocations.
+            for hook in hooks:
+                hook.pin_memory = False
+            return original_staging_allocator(hooks, resident_group)
+
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
+        monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
+
+        backend.enable(pipeline)
+
+        assert carrier.taken
+        assert backend._using_rank_local_mmap
+        assert not backend._using_registered_mmap
+        assert backend._host_registration is None
+        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+        assert not lease.closed
+
+        backend.disable()
+        assert lease.closed
+
     def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
         self,
         monkeypatch,
@@ -1927,8 +2183,9 @@ class TestDynamicSlotTracking:
                 prev_idx = (i - 1) % num_blocks
 
                 # Dynamic slot tracking: read from prev's prefetched slot
-                if prefetched_slots[prev_idx] is not None:
-                    current_slots[i] = prefetched_slots[prev_idx]
+                prefetched_slot = prefetched_slots[prev_idx]
+                if prefetched_slot is not None:
+                    current_slots[i] = prefetched_slot
 
                 read_slot = current_slots[i]
 

tests/diffusion/quantization/test_wan_autoround_mxfp4.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests."""
+
+from unittest.mock import patch
+
+import pytest
+from torch import nn
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.model_loader.utils import configure_quant_config
+
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
+    WanTransformer3DModel,
+)
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.inc_config import OmniINCConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _autoround_mxfp4_config() -> dict:
+    return {
+        "bits": 4,
+        "act_bits": 4,
+        "data_type": "mx_fp",
+        "act_data_type": "mx_fp",
+        "group_size": 32,
+        "act_group_size": 32,
+        "sym": True,
+        "act_sym": True,
+        "act_dynamic": True,
+        "enable_quanted_input": False,
+        "block_name_to_quantize": "blocks",
+        "quant_method": "auto-round",
+        "packing_format": "auto_round:llm_compressor",
+        "autoround_version": "0.14.2",
+    }
+
+
+def test_build_wan_autoround_mxfp4_config():
+    config = build_quant_config(_autoround_mxfp4_config())
+
+    assert isinstance(config, OmniINCConfig)
+    assert config.get_name() == "inc"
+    assert config.weight_bits == 4
+    assert config.data_type == "mx_fp"
+    assert config.group_size == 32
+    assert config.packing_format == "auto_round:llm_compressor"
+    assert config.block_name_to_quantize == ["blocks"]
+
+
+def test_wan_autoround_config_uses_runtime_layer_names():
+    config = build_quant_config(_autoround_mxfp4_config())
+    configure_quant_config(config, WanTransformer3DModel)
+
+    assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping
+    assert config.block_name_to_quantize == ["blocks"]
+
+    layer = object.__new__(LinearBase)
+    nn.Module.__init__(layer)
+    for prefix in (
+        "blocks.0.ffn.net_0.proj",
+        "blocks.0.ffn.net_2",
+        "blocks.0.attn1.to_out",
+    ):
+        layer_config = config.config_parser.resolve(layer, prefix)
+        assert layer_config.quantized
+        assert layer_config.is_mxfp4
+
+
+def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4():
+    config = build_quant_config(_autoround_mxfp4_config())
+    configure_quant_config(config, WanTransformer3DModel)
+
+    layer = object.__new__(LinearBase)
+    nn.Module.__init__(layer)
+    layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj")
+    outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj")
+
+    assert layer_config.quantized
+    assert layer_config.is_mxfp4
+    assert not outside_config.quantized
+
+    fake_method = object()
+    with patch(
+        "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method",
+        return_value=fake_method,
+    ):
+        method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj")
+    assert method is fake_method

tests/diffusion/test_diffusion_engine_metrics.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Source-level regression tests for diffusion output/engine helpers.
 
 These tests verify naming conventions and patterns by inspecting source code
@@ -37,10 +37,50 @@ _FORMATTER_PATH = os.path.normpath(
         "output_formatter.py",
     )
 )
+_METRICS_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "metrics",
+        "utils.py",
+    )
+)
+_INLINE_CLIENT_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "diffusion",
+        "inline_stage_diffusion_client.py",
+    )
+)
+_STAGE_PROC_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "diffusion",
+        "stage_diffusion_proc.py",
+    )
+)
+_ORCHESTRATOR_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "engine",
+        "orchestrator.py",
+    )
+)
 
 
 def _read_source(path: str) -> str:
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         return f.read()
 
 
@@ -100,9 +140,99 @@ class TestMetricKeys:
                 raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
 
         assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
-        assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
+        # step_total_ms is no longer emitted as a metric (no family consumes
+        # it); it lives only in the debug log breakdown below.
+        assert '"diffusion_engine_total_time_ms"' not in step_streaming_source
         assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
 
+    def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None:
+        source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+        assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source
+        assert "_scheduler_num_waiting_reqs" in step_streaming_source
+
+    def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None:
+        source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+        assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source
+        assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source
+        assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index(
+            "self.postprocess_output("
+        )
+
+    def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None:
+        inline_source = _read_source(_INLINE_CLIENT_PATH)
+        abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async")
+        assert ".cancel(" not in abort_source
+
+        proc_source = _read_source(_STAGE_PROC_PATH)
+        run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop")
+        abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :]
+        abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')]
+        assert ".cancel(" not in abort_branch
+
+    def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
+        source = _read_source(_ORCHESTRATOR_PATH)
+        loop_source = _get_function_source(source, "Orchestrator", "_orchestration_loop")
+        snapshot_pos = loop_source.index("_update_stage_replica_waiting(")
+        sentinel_pos = loop_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
+        route_pos = loop_source.index("pool.record_output_timestamps([diffusion_output])")
+        assert snapshot_pos < sentinel_pos < route_pos
+
+
+class TestVaeDecodeEmit:
+    """VAE decode timing is sourced from ``DiffusionOutput.stage_durations``
+    (populated by the diffusion pipeline profiler) and emitted as
+    ``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up.
+    """
+
+    def test_extract_vae_decode_ms_helper_exists(self) -> None:
+        source = _read_source(_METRICS_PATH)
+        helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms")
+        assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix"
+
+    def test_step_streaming_emits_vae_decode_time_ms(self) -> None:
+        source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+        assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, (
+            "step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing"
+        )
+        assert '"vae_decode_time_ms"' in step_streaming_source, (
+            "step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator"
+        )
+
+
+class TestDiffuseForwardEmit:
+    """Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit."""
+
+    def test_extract_diffuse_ms_helper_exists(self) -> None:
+        source = _read_source(_METRICS_PATH)
+        helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms")
+        assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix"
+
+    def test_step_streaming_emits_forward_time_ms(self) -> None:
+        source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+        assert "extract_diffusion_denoise_ms(" in step_streaming_source, (
+            "step_streaming must call extract_diffusion_denoise_ms to source forward-only timing"
+        )
+        assert '"forward_time_ms"' in step_streaming_source, (
+            "step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator"
+        )
+
+
+class TestKvRecvEmit:
+    """KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms
+    (set by the runner's _prepare_request_for_forward, not the profiler)."""
+
+    def test_step_streaming_emits_kv_recv_time_ms(self) -> None:
+        source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+        assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms"
+        assert '"kv_recv_time_ms"' in step_streaming_source, (
+            "step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator"
+        )
+
 
 class TestDummyRunAllocation:
     """Verify _dummy_run generates exact-sized audio arrays."""

tests/diffusion/test_diffusion_ipc.py

@@ -181,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
 
 
 def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
-    # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
+    # LTX2 returns (video, audio) tuples as DiffusionOutput.output.
     video = torch.arange(300_000, dtype=torch.float32)
     audio = torch.arange(300_000, dtype=torch.float32) * 3
     output = DiffusionOutput(output=(video, audio))

tests/diffusion/test_diffusion_scheduler.py

@@ -773,6 +773,19 @@ class TestRequestScheduler:
         assert third.num_running_reqs == 1
         assert third.num_waiting_reqs == 0
 
+    def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None:
+        perf_counter = mocker.patch(
+            "vllm_omni.diffusion.sched.base_scheduler.time.perf_counter",
+            side_effect=[10.0, 10.125],
+        )
+        request = _make_step_request("queue-wait", num_inference_steps=2)
+
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+
+        assert request.scheduler_queue_wait_ms == pytest.approx(125.0)
+        assert perf_counter.call_count == 2
+
     def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
         scheduler = RequestScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=2))

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import multiprocessing as mp
@@ -1079,6 +1079,42 @@ class TestStageDiffusionClientErrorPropagation:
 
         assert client.get_diffusion_output_nowait() is None
 
+    def test_error_response_preserves_scheduler_metrics(self):
+        from vllm_omni.metrics import definitions as metric_defs
+
+        client = self._make_client()
+        client._response_socket.recv.side_effect = [b"message", zmq.Again()]
+        client._decoder.decode.return_value = {
+            "type": "error",
+            "request_id": "req-error",
+            "error": "gpu fault",
+            "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
+        }
+
+        output = client.get_diffusion_output_nowait()
+
+        assert output is not None
+        assert output.error == "gpu fault"
+        assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
+
+    def test_metrics_response_creates_metrics_only_output(self):
+        from vllm_omni.metrics import definitions as metric_defs
+        from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
+
+        client = self._make_client()
+        client._response_socket.recv.side_effect = [b"message", zmq.Again()]
+        client._decoder.decode.return_value = {
+            "type": "metrics",
+            "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
+        }
+
+        output = client.get_diffusion_output_nowait()
+
+        assert output is not None
+        assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID
+        assert output.error is None
+        assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
+
     def test_check_health_raises_when_proc_dead(self):
         """``check_health`` detects a dead subprocess via the manager's proc
         and raises ``EngineDeadError``, setting ``_engine_dead`` as a

tests/diffusion/test_worker_wrapper_base.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for WorkerWrapperBase class.
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+        events: list[str] = []
+        offload_backend = mocker.Mock()
+        offload_backend.disable.side_effect = lambda: events.append("offload")
+        kv_manager = mocker.Mock()
+        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+        destroy = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
+            side_effect=lambda: events.append("distributed"),
+        )
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_manager,
+        )
+
+        worker.shutdown()
+
+        assert events == ["offload", "kv", "distributed"]
+        destroy.assert_called_once_with()
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/e2e/features/fullduplex/test_client.py

@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
 import base64
 import wave
 from urllib.parse import parse_qs, urlsplit
@@ -20,7 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def test_realtime_client_builds_explicit_native_duplex_url():
     url = build_realtime_url(
-        "ws://localhost:8099/v1/realtime?custom=1",
+        "ws://localhost:8099/v1/realtime?custom=1&duplex=0&model=stale&minicpmo45_native_duplex=0&session_id=stale",
         "openbmb/MiniCPM-o-4_5",
         session_id="session-a",
     )
@@ -50,7 +54,7 @@ def test_seed_tts_initial_text_is_part_of_native_duplex_context():
 
 def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
     url = build_realtime_url(
-        "ws://localhost:8099/v1/realtime?duplex=1",
+        "ws://localhost:8099/v1/realtime?duplex=1&autostart=1",
         "openbmb/MiniCPM-o-4_5",
         autostart=False,
     )
@@ -73,10 +77,15 @@ async def test_realtime_client_configure_omits_ref_audio_by_default():
 
     client = Client()
 
-    await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
+    await client.configure(
+        "openbmb/MiniCPM-o-4_5",
+        idle_timeout_s=900,
+        timeout_s=1,
+    )
 
     session = client.sent[0]["session"]
     assert "ref_audio" not in session
+    assert session["idle_timeout_s"] == 900
 
 
 @pytest.mark.asyncio
@@ -100,6 +109,74 @@ async def test_realtime_client_configure_sends_explicit_ref_audio():
 
     session = client.sent[0]["session"]
     assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
+    assert "idle_timeout_s" not in session
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_reports_reader_exit_to_event_waiters():
+    client = RealtimeDuplexClient("ws://unused")
+
+    async def stopped():
+        return None
+
+    client._reader_task = asyncio.create_task(stopped())
+    await client._reader_task
+
+    with pytest.raises(ConnectionError, match="closed before"):
+        client.raise_if_reader_stopped()
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_configure_reports_reader_failure_immediately():
+    class Client(RealtimeDuplexClient):
+        async def send(self, event):
+            async def fail():
+                raise RuntimeError("reader failed")
+
+            self._reader_task = asyncio.create_task(fail())
+            await asyncio.sleep(0)
+
+    client = Client("ws://unused")
+    with pytest.raises(ConnectionError, match="reader failed"):
+        await client.configure("model", timeout_s=10.0)
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_close_reports_reader_cancellation_immediately():
+    class Client(RealtimeDuplexClient):
+        async def send(self, event):
+            async def wait_forever():
+                await asyncio.Future()
+
+            self._reader_task = asyncio.create_task(wait_forever())
+            self._reader_task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await self._reader_task
+
+    client = Client("ws://unused")
+    with pytest.raises(ConnectionError, match="reader was cancelled"):
+        await client.close_session(timeout_s=10.0)
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_close_waits_for_a_new_session_closed_event():
+    class Client(RealtimeDuplexClient):
+        async def send(self, event):
+            assert event["type"] == "session.close"
+
+            async def acknowledge_close():
+                await asyncio.sleep(0.02)
+                self.events.add({"type": "session.closed"})
+
+            asyncio.create_task(acknowledge_close())
+
+    client = Client("ws://unused")
+    client.events.add({"type": "session.closed", "reason": "old"})
+    started_at = asyncio.get_running_loop().time()
+
+    await client.close_session(timeout_s=0.2)
+
+    assert asyncio.get_running_loop().time() - started_at >= 0.01
 
 
 @pytest.mark.asyncio
@@ -235,6 +312,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
         "output_token_count": 4,
         "ttft_ms": 120.0,
         "tpot_ms": 15.0,
+        "itls_ms": [10.0, 14.0, 18.0],
         "inter_token_interval_ms": {
             "count": 3,
             "mean": 14.0,

tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py

@@ -1,290 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""E2E test for DreamID-Omni offline inference with Cache-Dit.
-
-Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
-synthetic inputs so the coverage matches the upstream demo structure.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-import os
-import re
-import time
-from functools import lru_cache
-from pathlib import Path
-from typing import NamedTuple
-
-import numpy as np
-import pytest
-import soundfile as sf
-from PIL import Image
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
-MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
-EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
-NUM_INFERENCE_STEPS = 45
-
-CACHE_CONFIG = {
-    "Fn_compute_blocks": 1,
-    "Bn_compute_blocks": 0,
-    "max_warmup_steps": 4,
-    "max_cached_steps": 20,
-    "residual_diff_threshold": 0.24,
-    "max_continuous_cached_steps": 3,
-    "enable_taylorseer": False,
-    "taylorseer_order": 1,
-    "scm_steps_mask_policy": None,
-    "scm_steps_policy": "dynamic",
-}
-
-
-class DreamIDGenerationResult(NamedTuple):
-    video: np.ndarray
-    audio: np.ndarray
-    fps: int
-    duration_s: float
-    elapsed_s: float
-
-
-class DreamIDExampleCase(NamedTuple):
-    prompt: str
-    image_paths: tuple[Path, ...]
-    audio_paths: tuple[Path, ...]
-
-
-def _require_model_path() -> str:
-    model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
-    if not Path(model_path).exists():
-        pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
-
-    if importlib.util.find_spec("dreamid_omni") is None:
-        pytest.skip(
-            "dreamid_omni dependency package is not importable. "
-            "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
-        )
-
-    return model_path
-
-
-def _dreamid_repo_root_candidates() -> list[Path]:
-    candidates: list[Path] = []
-    example_root = os.environ.get(EXAMPLE_ROOT_ENV)
-    if example_root:
-        candidates.append(Path(example_root).expanduser().resolve())
-
-    spec = importlib.util.find_spec("dreamid_omni")
-    if spec is not None and spec.origin:
-        candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
-
-    return candidates
-
-
-@lru_cache(maxsize=1)
-def _load_example_case() -> DreamIDExampleCase:
-    for candidate_root in _dreamid_repo_root_candidates():
-        caption_path = candidate_root / "captions" / "20.json"
-        image_paths = (
-            candidate_root / "imgs" / "20" / "0.png",
-            candidate_root / "imgs" / "20" / "1.png",
-        )
-        audio_paths = (
-            candidate_root / "audios" / "20" / "0.wav",
-            candidate_root / "audios" / "20" / "1.wav",
-        )
-        if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
-            prompt = _load_prompt_from_json(caption_path)
-            return DreamIDExampleCase(
-                prompt=prompt,
-                image_paths=image_paths,
-                audio_paths=audio_paths,
-            )
-
-    searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
-    pytest.skip(
-        "DreamID-Omni official example assets not found. "
-        f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
-        f"Searched: {searched}"
-    )
-
-
-def _load_prompt_from_json(path: Path) -> str:
-    raw_text = path.read_text(encoding="utf-8").strip()
-    parsed = json.loads(raw_text)
-    if isinstance(parsed, str):
-        prompt = parsed
-    elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
-        prompt = parsed["content"]
-    else:
-        raise AssertionError(f"Unexpected DreamID caption format in {path}")
-
-    prompt = re.sub(
-        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
-        "",
-        prompt,
-        flags=re.DOTALL,
-    ).strip()
-    prompt = re.sub(
-        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
-        "",
-        prompt,
-        flags=re.DOTALL,
-    ).strip()
-    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
-    prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
-    return prompt
-
-
-def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
-    images: list[Image.Image] = []
-    for path in example_case.image_paths:
-        with Image.open(path) as image:
-            images.append(image.convert("RGB"))
-    return images
-
-
-def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
-    audios: list[np.ndarray] = []
-    for path in example_case.audio_paths:
-        audio_array, sample_rate = sf.read(str(path), dtype="float32")
-        audio_np = np.asarray(audio_array, dtype=np.float32)
-        if audio_np.ndim == 2:
-            audio_np = np.mean(audio_np, axis=1)
-        audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
-        audios.append(audio_np.reshape(-1))
-    return audios
-
-
-def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
-    return OmniDiffusionSamplingParams(
-        height=704,
-        width=1280,
-        num_inference_steps=num_inference_steps,
-        seed=42,
-        extra_args={
-            "solver_name": "unipc",
-            "shift": 5.0,
-        },
-    )
-
-
-def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
-    return {
-        "prompt": example_case.prompt,
-        "video_negative_prompt": "jitter, bad hands, blur, distortion",
-        "audio_negative_prompt": "robotic, muffled, echo, distorted",
-        "multi_modal_data": {
-            "image": _load_example_images(example_case),
-            "audio": _load_example_audios(example_case),
-        },
-    }
-
-
-def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
-    """Extract video and audio from raw Omni outputs"""
-    assert outputs, "DreamID-Omni returned no outputs"
-
-    result = outputs[0]
-    assert hasattr(result, "images") and result.images, "not return video frames"
-
-    video = result.images[0]
-    assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
-    assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
-
-    multimodal_output = getattr(result, "multimodal_output", None) or {}
-    assert "audio" in multimodal_output, "not return audio"
-    audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
-    assert audio.size > 0, "returned empty audio"
-
-    audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
-    fps = multimodal_output.get("fps", 24)
-    duration_s = float(audio.size) / float(audio_sample_rate)
-
-    generation = DreamIDGenerationResult(
-        video=video.astype(np.float32, copy=False),
-        audio=audio,
-        fps=fps,
-        duration_s=duration_s,
-        elapsed_s=elapsed_s,
-    )
-    return generation
-
-
-def _run_generation(
-    *,
-    model_path: str,
-    example_case: DreamIDExampleCase,
-    use_cache_dit: bool,
-) -> DreamIDGenerationResult:
-    runner_kwargs = {
-        "model_type": "dreamid-omni",
-        "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
-    }
-    if use_cache_dit:
-        runner_kwargs["cache_backend"] = "cache_dit"
-        runner_kwargs["cache_config"] = CACHE_CONFIG
-
-    with OmniRunner(model_path, **runner_kwargs) as runner:
-        warmup_payload = _request_payload(example_case)
-        _ = runner.omni.generate(warmup_payload, _sampling_params(2))
-
-        elapsed_samples: list[float] = []
-        extracted_result: DreamIDGenerationResult | None = None
-
-        request_payload = _request_payload(example_case)
-        start = time.perf_counter()
-        outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
-        elapsed_s = time.perf_counter() - start
-        elapsed_samples.append(elapsed_s)
-        if extracted_result is None:
-            extracted_result = _extract_generation(outputs, elapsed_s)
-
-        assert extracted_result is not None
-        median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
-        return extracted_result._replace(elapsed_s=median_elapsed_s)
-
-
-def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
-    assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
-    return float(np.mean(np.abs(a - b)))
-
-
-def _audio_rms(audio: np.ndarray) -> float:
-    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
-    return float(np.sqrt(np.mean(audio**2)))
-
-
-def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
-    rms_a = _audio_rms(a.audio)
-    rms_b = _audio_rms(b.audio)
-    return rms_b / max(rms_a, 1e-6)
-
-
-@hardware_test(res={"cuda": "H100"})
-def test_dreamid_omni_cache_dit_e2e() -> None:
-    model_path = _require_model_path()
-    example_case = _load_example_case()
-
-    baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
-    cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
-
-    video_mae = _video_mae(baseline.video, cached.video)
-    audio_rms_ratio = _audio_rms_ratio(baseline, cached)
-    speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
-
-    assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
-    assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
-    assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -1,162 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
-
-import io
-
-import av
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
-
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
-_OMNI_RUNNER_PARAM = (
-    MODEL,
-    None,
-    {
-        "init_timeout": 1200,
-        "tensor_parallel_size": 2,
-        "vae_use_tiling": True,
-    },
-)
-
-_OMNI_RUNNER_PARAM_OFFLOAD = (
-    MODEL,
-    None,
-    {
-        "init_timeout": 1200,
-        "tensor_parallel_size": 2,
-        "enable_cpu_offload": True,
-        "vae_use_tiling": True,
-    },
-)
-
-pytestmark = [
-    pytest.mark.slow,
-    pytest.mark.diffusion,
-    pytest.mark.parametrize(
-        "omni_runner",
-        [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
-        indirect=True,
-    ),
-]
-
-
-def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
-    """Validate that the MP4 contains meaningful video and audio tracks."""
-    container = av.open(io.BytesIO(video_bytes))
-
-    v_streams = [s for s in container.streams if s.type == "video"]
-    assert len(v_streams) >= 1, "No video stream found in MP4"
-
-    a_streams = [s for s in container.streams if s.type == "audio"]
-    assert len(a_streams) >= 1, "No audio stream found in MP4"
-
-    v_stream = v_streams[0]
-    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
-    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
-
-    frame_count = 0
-    for frame in container.decode(video=0):
-        frame_count += 1
-        if frame_count >= min_frames:
-            break
-    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
-
-    container.close()
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_magi_human_e2e(omni_runner: OmniRunner):
-    """End-to-end test for MagiHuman generating video and audio."""
-    prompt = (
-        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
-        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
-        "stationary in front of a softly lit, blurred warm-toned interior. Her "
-        "overall disposition is warm, composed, and gently confident. The camera "
-        "holds a static medium close-up, framing her from the shoulders up, "
-        "with shallow depth of field keeping her face in sharp focus. Soft "
-        "directional key light falls from the upper left, casting a gentle "
-        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
-        "the levator labii superiors relaxing as her lips part. She speaks in "
-        "clear, warm, unhurried American English: "
-        "\"The most beautiful things in life aren't things at all — "
-        "they're moments, feelings, and the people who make you feel truly alive.\" "
-        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
-        "shapes each vowel with precision. A faint, genuine smile engages the "
-        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
-        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
-        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
-        "Dialogue:\n"
-        "<Young blonde woman, American English>: "
-        "\"The most beautiful things in life aren't things at all — "
-        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
-        "Background Sound:\n"
-        "<Soft, warm indoor ambience with a faint distant piano melody>"
-    )
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=256,
-        width=448,
-        num_inference_steps=8,
-        seed=52,
-        extra_args={
-            "seconds": 5,
-            "sr_height": 1080,
-            "sr_width": 1920,
-            "sr_num_inference_steps": 5,
-        },
-    )
-
-    omni = omni_runner.omni
-    outputs = list(
-        omni.generate(
-            prompts=[prompt],
-            sampling_params_list=[sampling_params],
-        )
-    )
-
-    assert len(outputs) > 0, "No outputs returned"
-    first = outputs[0]
-
-    assert hasattr(first, "images") and first.images, "No video frames in output"
-    video_frames = first.images[0]
-    assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
-    assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
-
-    mm = first.multimodal_output
-    assert mm, "multimodal_output is empty or missing"
-
-    audio_waveform = mm.get("audio")
-    assert audio_waveform is not None, "No audio waveform in multimodal_output"
-
-    audio_sample_rate = mm.get("audio_sample_rate")
-    assert audio_sample_rate is not None, (
-        "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
-    )
-    assert isinstance(audio_sample_rate, (int, float)), (
-        f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
-    )
-    assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
-
-    fps = mm.get("fps")
-    assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
-    assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
-    assert int(fps) > 0, f"fps must be positive, got {fps}"
-
-    video_bytes = mux_video_audio_bytes(
-        video_frames,
-        audio_waveform,
-        fps=float(fps),
-        audio_sample_rate=int(audio_sample_rate),
-    )
-    assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
-    assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
-
-    _validate_mp4(video_bytes)

tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E smoke test for Wan2.1 AutoRound MXFP4 inference."""
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+QUANTIZED_MODEL = os.environ.get(
+    "WAN21_AUTOROUND_MXFP4_MODEL",
+    "INCModel/Wan2.1-T2V-1.3B-Diffusers-MXFP4-AutoRound",
+)
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"xpu": "B60"})
+@pytest.mark.parametrize(
+    "omni_runner",
+    [(QUANTIZED_MODEL, None)],
+    indirect=True,
+)
+def test_wan21_autoround_mxfp4_generates_video(omni_runner, omni_runner_handler):
+    sampling_params = OmniDiffusionSamplingParams(
+        height=256,
+        width=256,
+        num_frames=5,
+        num_inference_steps=1,
+        guidance_scale=1.0,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+    response = omni_runner_handler.send_diffusion_request(
+        {
+            "prompt": "A red toy car driving on a table",
+            "sampling_params": sampling_params,
+        }
+    )
+
+    assert response.success, response.error_message
+    assert response.images is not None and len(response.images) == 1
+    frames = np.asarray(response.images[0])
+    assert frames.shape[1:4] == (5, 256, 256)
+    assert frames.std() > 0.01

tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py

@@ -15,20 +15,10 @@ import pytest
 from huggingface_hub import snapshot_download
 
 from tests.helpers.runtime import OmniServerParams, get_model_prefix
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "openbmb/MiniCPM-o-4_5"
-DEPLOY_CONFIG = modify_stage_config(
-    get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
-    updates={
-        "base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
-        # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
-        # is left automatic so duplex matches the simplex deploy profiles.
-        "stages": {
-            1: {"max_model_len": 4096},
-        },
-    },
-)
+DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
 ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
 RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
 RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"

tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py

@@ -58,18 +58,14 @@ from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
     read_pcm16_wav,
     summarize_session_request_metrics,
 )
+from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
+    reference_audio_data_url as _ref_audio_data_url,
+)
 
 _url_with_model = build_realtime_url
 _read_wav_pcm16 = read_pcm16_wav
 
 
-def _ref_audio_data_url(path: str | None) -> str | None:
-    if path is None:
-        return None
-    ref_path = Path(path).expanduser()
-    return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
-
-
 @dataclass
 class DemoState:
     events: list[dict[str, object]] = field(default_factory=list)

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -7,7 +7,6 @@ E2E Online tests for Qwen3-Omni model.
 import os
 
 import pytest
-from openai import BadRequestError
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -89,7 +88,7 @@ test_params = [
     ),
 ]
 
-test_token_params = [
+reduced_token_params = [
     pytest.param(
         OmniServerParams(
             model=model,
@@ -97,7 +96,7 @@ test_token_params = [
             use_stage_cli=True,
             server_args=["--async-chunk"],
         ),
-        id="batch_token_64",
+        id="batch_token_2048",
     )
 ]
 
@@ -138,29 +137,7 @@ def get_max_batch_size(size_type="few"):
 
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_audio_001(omni_server, online_client) -> None:
-    """
-    Input Modal: text
-    Output Modal: audio
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "modalities": ["audio"],
-        "stream": True,
-        "key_words": {"text": ["beijing"]},
-    }
-
-    online_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
 def test_text_to_text_audio_001(omni_server, online_client) -> None:
     """
     Input Modal: text
@@ -204,36 +181,13 @@ def test_text_video_to_text_001(omni_server, online_client) -> None:
 
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_audio_to_text_audio_001(omni_server, online_client) -> None:
-    """
-    Input Modal: text, audio
-    Output Modal: text, audio
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
-    messages = dummy_messages_from_mix_data(
-        audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "key_words": {"audio": AUDIO_KEY},
-    }
-
-    online_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
 def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
     """
     Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV)
     Output Modal: text, audio
     Input Setting: stream=False
-    Datasets: single request
+    Datasets: few requests
     """
     audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}"
     messages = dummy_messages_from_mix_data(
@@ -252,37 +206,13 @@ def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
 
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_image_to_text_audio_001(omni_server, online_client) -> None:
-    """
-    Input Modal: text, image
-    Output Modal: text, audio
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
-
-    messages = dummy_messages_from_mix_data(
-        image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image")
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "key_words": {"image": IMAGE_KEY},
-    }
-
-    online_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
     """
     Input Modal: text, high-resolution image (1080p-class JPEG)
     Output Modal: text, audio
     Input Setting: stream=False
-    Datasets: single request
+    Datasets: few requests
     """
     image_data_url = (
         f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}"
@@ -304,7 +234,7 @@ def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
 
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
     """
     Input Modal: text, video
@@ -328,36 +258,6 @@ def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config)
 
 
-@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_mix_to_text_audio_001(omni_server, online_client) -> None:
-    """
-    Input Modal: text, audio, image, video
-    Output Modal: text, audio
-    Input Setting: stream=True
-    Datasets: few requests
-    """
-    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
-    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
-    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        video_data_url=video_data_url,
-        image_data_url=image_data_url,
-        audio_data_url=audio_data_url,
-        content_text=get_prompt("mix"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
-    }
-    online_client.send_omni_request(request_config, request_num=get_max_batch_size())
-
-
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_in_video_001(omni_server, online_client) -> None:
@@ -507,22 +407,22 @@ def test_one_word_prompt_001(omni_server, online_client) -> None:
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_speaker_001(omni_server, online_client) -> None:
     """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
+    Input Modal: Chinese text.
+    Output Modal: Chinese text and audio using an uppercase preset speaker name.
     Input Setting: stream=True
     Datasets: single request
     """
     messages = dummy_messages_from_mix_data(
         system_prompt=get_system_prompt(),
-        content_text=get_prompt("text"),
+        content_text=get_prompt("text_chinese"),
     )
 
     request_config = {
         "model": omni_server.model,
         "messages": messages,
         "stream": True,
-        "speaker": "Chelsie",
-        "key_words": {"text": ["beijing"]},
+        "speaker": "CHELSIE",
+        "key_words": {"text": ["北京"]},
     }
 
     online_client.send_omni_request(request_config)
@@ -563,55 +463,6 @@ def test_speaker_002(omni_server, online_client) -> None:
             print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_003(omni_server, online_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "speaker": "CHELSIE",
-        "key_words": {"text": ["beijing"]},
-    }
-
-    online_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_language_001(omni_server, online_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text_chinese"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "key_words": {"text": ["北京"]},
-    }
-
-    online_client.send_omni_request(request_config)
-
-
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
@@ -619,7 +470,7 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
     Input Modal: text only (long-form generation prompt).
     Output Modal: text, audio (default ``modalities``);
     Input Setting: stream=True
-    Datasets: single request
+    Datasets: few requests
     """
     messages = dummy_messages_from_mix_data(
         system_prompt=get_system_prompt(),
@@ -631,21 +482,3 @@ def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
     text = responses[0].text_content if responses else ""
     word_count = len(text.split())
     assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
-def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
-    """Ensure invalid audio format is rejected with 400 before streaming starts."""
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt(),
-    )
-    with pytest.raises(BadRequestError, match="audio format"):
-        openai_client.client.chat.completions.create(
-            model=omni_server.model,
-            messages=messages,
-            modalities=["text", "audio"],
-            audio={"voice": "alloy", "format": "aac"},
-            stream=True,
-        )

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -8,9 +8,8 @@ and are supported by the following text-to-image models:
 - Qwen-Image-2512
 
 One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit).
-Supported features for Qwen-Image series: TeaCache, Cache-DiT, Ulysses-SP, Ring-Attention,
-CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise).
-See docs/user_guide/diffusion_acceleration.md.
+Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512).
+Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md.
 """
 
 import pytest
@@ -25,141 +24,116 @@ NEGATIVE_PROMPT = "blurry, low quality"
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
+MODEL_IMAGE = "Qwen/Qwen-Image"
+MODEL_2512 = "Qwen/Qwen-Image-2512"
 
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--enable-cpu-offload"],
-            ),
-            id="cpu_offload",
-            marks=SINGLE_CARD_FEATURE_MARKS,
+# One server per feature. Alternate models; keep feature pytest ids unchanged.
+FEATURE_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_IMAGE,
+            server_args=["--enable-cpu-offload"],
         ),
-        pytest.param(
-            OmniServerParams(model=model, server_args=["--step-execution"]),
-            id="step_execution",
-            marks=SINGLE_CARD_FEATURE_MARKS,
+        id="cpu_offload",
+        marks=SINGLE_CARD_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(model=MODEL_2512, server_args=["--step-execution"]),
+        id="step_execution",
+        marks=SINGLE_CARD_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(model=MODEL_IMAGE, server_args=["--cache-backend", "tea_cache"]),
+        id="cache_tea_cache",
+        marks=SINGLE_CARD_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_2512,
+            server_args=[
+                "--cache-backend",
+                "cache_dit",
+                "--enable-layerwise-offload",
+            ],
         ),
-        pytest.param(
-            OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
-            id="cache_tea_cache",
-            marks=SINGLE_CARD_FEATURE_MARKS,
+        id="layerwise_offload",
+        marks=SINGLE_CARD_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_IMAGE,
+            server_args=[
+                "--cache-backend",
+                "cache_dit",
+                "--ulysses-degree",
+                "2",
+            ],
         ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--enable-layerwise-offload",
-                ],
-            ),
-            id="layerwise_offload",
-            marks=SINGLE_CARD_FEATURE_MARKS,
+        id="ulysses_2",
+        marks=PARALLEL_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_2512,
+            server_args=[
+                "--cache-backend",
+                "cache_dit",
+                "--ring",
+                "2",
+            ],
         ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--ulysses-degree",
-                    "2",
-                ],
-            ),
-            id="ulysses_2",
-            marks=PARALLEL_FEATURE_MARKS,
+        id="ring_2",
+        marks=PARALLEL_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_IMAGE,
+            server_args=[
+                "--cache-backend",
+                "tea_cache",
+                "--cfg-parallel-size",
+                "2",
+            ],
         ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--ring",
-                    "2",
-                ],
-            ),
-            id="ring_2",
-            marks=PARALLEL_FEATURE_MARKS,
+        id="cfg_parallel_2",
+        marks=PARALLEL_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_2512,
+            server_args=[
+                "--cache-backend",
+                "cache_dit",
+                "--tensor-parallel-size",
+                "2",
+                "--vae-patch-parallel-size",
+                "2",
+                "--vae-use-tiling",
+                "--diffusion-quantization-config",
+                '{"method":"fp8"}',
+            ],
         ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "tea_cache",
-                    "--cfg-parallel-size",
-                    "2",
-                ],
-            ),
-            id="cfg_parallel_2",
-            marks=PARALLEL_FEATURE_MARKS,
+        id="vae_patch_parallel_2",
+        marks=PARALLEL_FEATURE_MARKS,
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_IMAGE,
+            server_args=[
+                "--use-hsdp",
+                "--hsdp-shard-size",
+                "2",
+            ],
         ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--tensor-parallel-size",
-                    "2",
-                    "--vae-patch-parallel-size",
-                    "2",
-                    "--vae-use-tiling",
-                    "--diffusion-quantization-config",
-                    '{"method":"fp8"}',
-                ],
-            ),
-            id="vae_patch_parallel_2",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--use-hsdp",
-                    "--hsdp-shard-size",
-                    "2",
-                ],
-            ),
-            id="parallel_hsdp",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-    ]
+        id="parallel_hsdp",
+        marks=PARALLEL_FEATURE_MARKS,
+    ),
+]
 
 
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image"),
-    indirect=True,
-)
+@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
 def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
-    """Test each diffusion feature with Qwen-Image (text-to-image), one feature per case."""
-    messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 2,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "true_cfg_scale": 4.0,
-            "seed": 42,
-        },
-    }
-    online_client.send_diffusion_request(request_config)
-
-
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-2512"),
-    indirect=True,
-)
-def test_qwen_image_2512(omni_server: OmniServer, online_client: OnlineOmniClient):
-    """Test each diffusion feature with Qwen-Image-2512 (text-to-image), one feature per case."""
+    """One diffusion feature per case; model is chosen in FEATURE_CASES."""
     messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
     request_config = {
         "model": omni_server.model,

tests/e2e/online_serving/test_stable_audio_online_expansion.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving test for Stable Audio Open text-to-audio diffusion.
+
+Stable Audio Open is served through the OpenAI-compatible
+`POST /v1/audio/generate` endpoint (JSON in, binary WAV out).
+"""
+
+import os
+
+import pytest
+
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+# Stable Audio Open is gated; override to a local bundle locally if needed.
+STABLE_AUDIO_TEST_MODEL = os.environ.get("STABLE_AUDIO_TEST_MODEL", "stabilityai/stable-audio-open-1.0")
+T2A_PROMPT = "A piano playing a gentle melody with soft room ambience."
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
+
+
+def _stable_audio_server_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(model=model),
+            id="t2a",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.fixture
+def _require_stable_audio_access() -> None:
+    """Skip cleanly (before the server boots) if the gated checkpoint is inaccessible."""
+    skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL)
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True)
+def test_stable_audio_t2a_online(
+    _require_stable_audio_access: None,
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    """Stable Audio Open text-to-audio: `/v1/audio/generate` returns a non-empty WAV.
+
+    Uses tiny steps / short duration to keep CI light, matching the offline smoke
+    test in `tests/e2e/offline_inference/test_stable_audio_expansion.py`.
+    """
+    request_config = {
+        "json": {
+            "model": omni_server.model,
+            "input": T2A_PROMPT,
+            "audio_length": 2.0,
+            "num_inference_steps": 4,
+            "guidance_scale": 7.0,
+            "negative_prompt": "Low quality.",
+            "seed": 42,
+            "response_format": "wav",
+        },
+        "timeout": 300,
+    }
+    responses = openai_client.send_audio_generate_http_request(request_config)
+    assert responses, "no response from /v1/audio/generate"
+    resp = responses[0]
+    assert resp.success, f"audio generate failed: {resp.status_code} {resp.error_message}"
+    assert resp.status_code == 200

tests/engine/test_async_omni_engine_stage_init.py

@@ -474,7 +474,7 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
     assert os.environ[runtime_env_var] == "parent-value"
 
 
-def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch):
+def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.engine.stage_init_utils as init_mod
 
@@ -506,7 +506,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
         use_inline=True,
     )
 
-    assert od_config.max_num_seqs == 4
+    assert od_config.max_num_seqs == 1
     assert captured == {
         "model": "dummy-model",
         "config": od_config,
@@ -517,7 +517,7 @@ def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypa
     }
 
 
-def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
+def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
     import vllm_omni.engine.stage_engine_startup as startup_mod
@@ -565,10 +565,84 @@ def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch
     )
 
     assert result is sentinel_client
-    assert od_config.max_num_seqs == 4
+    assert od_config.max_num_seqs == 1
     assert resources.manager is proc_manager
 
 
+def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.engine.stage_init_utils as init_mod
+
+    od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True)
+    monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+    monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object())
+    metadata = _make_diffusion_metadata(0)
+
+    init_mod.initialize_diffusion_stage(
+        0,
+        "dummy-model",
+        types.SimpleNamespace(),
+        metadata,
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=True,
+    )
+
+    assert od_config.max_num_seqs == 8
+
+
+def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch):
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
+    import vllm_omni.engine.stage_engine_startup as startup_mod
+
+    od_config = types.SimpleNamespace(
+        max_num_seqs=8,
+        step_execution=True,
+        parallel_config=types.SimpleNamespace(world_size=1),
+    )
+    monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
+    monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
+    monkeypatch.setattr(
+        startup_mod,
+        "register_stage_with_omni_master",
+        lambda **kwargs: types.SimpleNamespace(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+        ),
+    )
+    omni_master_server = types.SimpleNamespace(
+        address="127.0.0.1",
+        port=25000,
+        release_route_port_reservations=lambda *args, **kwargs: None,
+    )
+    proc_manager = types.SimpleNamespace(
+        addresses=types.SimpleNamespace(
+            inputs=["tcp://127.0.0.1:26002"],
+            outputs=["tcp://127.0.0.1:26003"],
+        )
+    )
+    monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
+    monkeypatch.setattr(
+        client_mod.StageDiffusionClient,
+        "from_addresses",
+        lambda metadata, **kwargs: object(),
+    )
+
+    startup_mod.launch_diffusion_stage_replica(
+        model="dummy-model",
+        stage_config=types.SimpleNamespace(),
+        metadata=types.SimpleNamespace(stage_id=0),
+        stage_init_timeout=12,
+        batch_size=4,
+        use_inline=False,
+        omni_master_server=omni_master_server,
+    )
+
+    assert od_config.max_num_seqs == 8
+
+
 def test_stage_runtime_initializes_stage_pools(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/engine/test_membership_controller.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -47,6 +47,7 @@ class FakePool:
         self.invalidated = []
         self.hub = None
         self.lb = None
+        self.replica_ids = {}
 
     def attach_hub(self, hub):
         self.hub = hub
@@ -57,7 +58,9 @@ class FakePool:
     def add_client(self, input_addr, client, *, replica_id=None):
         self.added.append((input_addr, client, replica_id))
         self.clients.append(client)
-        return len(self.clients) - 1 if replica_id is None else replica_id
+        resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id
+        self.replica_ids[input_addr] = resolved_replica_id
+        return resolved_replica_id
 
     def invalidate_addr(self, input_addr):
         self.invalidated.append(input_addr)
@@ -67,6 +70,12 @@ class FakePool:
         self.removed.append(input_addr)
         return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
 
+    def get_client_by_addr(self, input_addr):
+        return SimpleNamespace() if input_addr in self.replica_ids else None
+
+    def get_replica_id_by_addr(self, input_addr):
+        return self.replica_ids.get(input_addr)
+
 
 def _snapshot(*replicas):
     return SimpleNamespace(replicas=list(replicas))
@@ -209,6 +218,16 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
         return replacement
 
     controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory)
+    removed_replicas = []
+
+    async def cleanup(_request_ids):
+        return None
+
+    controller.install_unregister_handlers(
+        output_queue=asyncio.Queue(),
+        cleanup_callback=cleanup,
+        replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)),
+    )
     await controller.handle_unregister(0, "tcp://old-0")
     await controller.handle_unregister(0, "tcp://old-1")
 
@@ -219,6 +238,7 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
     assert pool.get_replica_id_by_addr("tcp://new-1") == 1
     assert pool.get_replica_id_by_addr("tcp://old-1") is None
     assert pool.available_replica_ids() == [1]
+    assert removed_replicas == [(0, 0), (0, 1)]
 
 
 @pytest.mark.asyncio

tests/engine/test_orchestrator.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import asyncio
@@ -2063,6 +2066,28 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
     assert metrics.output_unit_count == 3
 
 
+def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
+    stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    pool = StagePool(
+        1,
+        [stage],
+        output_processor=FakeOutputProcessor(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    output = SimpleNamespace(request_id="req-image", _custom_output={"image": "non-empty"})
+    pool.record_output_timestamps([output], output_ts=132.0)
+
+    metrics = pool.build_stage_metrics(
+        [output],
+        submit_ts=130.0,
+        request_timestamp=120.0,
+        replica_id=0,
+    )
+
+    assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0)
+    assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0)
+
+
 @pytest.mark.asyncio
 async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
     class FailingStageClient(FakeStageClient):

tests/engine/test_orchestrator_error_handling.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for error propagation paths within the Orchestrator.
 
 Covers:
@@ -16,7 +16,9 @@ from types import SimpleNamespace
 
 import janus
 import pytest
+from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.serial_utils import MsgpackEncoder
 
 from vllm_omni.engine.messages import (
     AddCompanionRequestMessage,
@@ -46,8 +48,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _sampling_params(max_tokens: int = 4):
-    from vllm.sampling_params import SamplingParams
-
     return SamplingParams(max_tokens=max_tokens)
 
 
@@ -188,6 +188,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
         raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
 
 
+class _OverflowOnEncodeStageClient(FakeStageClient):
+    """Msgpack-encodes the request's sampling params on dispatch, as the real
+    StageEngineCoreClient does. We use this to make sure overflow sampling params
+    do not break the orchestrator."""
+
+    async def add_request_async(self, request, *args, **kwargs) -> None:
+        MsgpackEncoder().encode(request.sampling_params)
+        await super().add_request_async(request, *args, **kwargs)  # unreachable
+
+
+@pytest.mark.asyncio
+async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
+    """Regression test for over-range seed handling."""
+    stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
+    orchestrator_fixture = orchestrator_factory([stage0])
+
+    # seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
+    overflow_params = SamplingParams(seed=2**70)
+
+    await _enqueue_add_request(
+        orchestrator_fixture,
+        request_id="req-overflow",
+        prompt=SimpleNamespace(
+            request_id="req-overflow",
+            prompt_token_ids=[1, 2],
+            sampling_params=overflow_params,
+        ),
+        original_prompt={"prompt": "hi"},
+        sampling_params_list=[overflow_params],
+        final_stage_id=0,
+    )
+
+    error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
+    assert error_msg.fatal is False
+    assert error_msg.status_code == 400
+    assert orchestrator_fixture.thread.is_alive()
+    assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
+
+
 @pytest.mark.asyncio
 async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
     """Multi-replica stage: one replica dying must fail only the requests bound

tests/engine/test_single_stage_mode.py

@@ -1175,7 +1175,7 @@ class TestSingleStageReplicaInitialization:
                 os.environ[device_env_var] = prev_device_env
 
         assert result is sentinel_client
-        assert od_config.max_num_seqs == 4
+        assert od_config.max_num_seqs is None
         mock_register.assert_called_once_with(
             omni_master_address="127.0.0.1",
             omni_master_port=25000,

tests/entrypoints/openai_api/test_duplex_handler.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import asyncio
@@ -105,16 +108,16 @@ class FakeEngineClient:
         self.close_result = close_result
         self.opened: list[str] = []
         self.opened_fences: list[DuplexFence | None] = []
-        self.appended: list[tuple[str, str, object, bool]] = []
+        self.appended: list[tuple[str, str, Any, bool]] = []
         self.append_operation_ids: list[str | None] = []
         self.appended_fences: list[DuplexFence | None] = []
-        self.opened_configs: list[dict[str, object]] = []
-        self.opened_runtime_configs: list[dict[str, object]] = []
+        self.opened_configs: list[dict[str, Any]] = []
+        self.opened_runtime_configs: list[dict[str, Any]] = []
         self.signals: list[tuple[str, str]] = []
         self.signal_fences: list[DuplexFence | None] = []
         self.signal_next_fences: list[DuplexFence | None] = []
-        self.signal_session_configs: list[dict[str, object] | None] = []
-        self.signal_runtime_configs: list[dict[str, object] | None] = []
+        self.signal_session_configs: list[dict[str, Any] | None] = []
+        self.signal_runtime_configs: list[dict[str, Any] | None] = []
         self.closed: list[tuple[str, str]] = []
         self.touched: list[tuple[str, DuplexLeaseActivity]] = []
         self.touch_fences: list[DuplexFence] = []
@@ -136,7 +139,7 @@ class FakeEngineClient:
         runtime_config: dict[str, object] | None = None,
         timeout: float | None = None,
         fence: DuplexFence | None = None,
-    ) -> None:
+    ) -> Any:
         del timeout
         if self.fail_open:
             raise RuntimeError("open failed")
@@ -157,7 +160,7 @@ class FakeEngineClient:
         timeout: float | None = None,
         collect_outputs: bool = True,
         fence: DuplexFence | None = None,
-    ) -> None:
+    ) -> Any:
         del timeout, collect_outputs
         self.appended.append((session_id, mode, payload, final))
         self.append_operation_ids.append(operation_id)
@@ -191,7 +194,7 @@ class FakeEngineClient:
         next_fence: DuplexFence | None = None,
         session_config: dict[str, object] | None = None,
         runtime_config: dict[str, object] | None = None,
-    ) -> None:
+    ) -> Any:
         del timeout
         if self.fail_signal or event in self.fail_signal_events:
             raise RuntimeError("signal failed")
@@ -209,7 +212,7 @@ class FakeEngineClient:
         reason: str = "client_close",
         timeout: float | None = None,
         fence: DuplexFence | None = None,
-    ) -> None:
+    ) -> Any:
         del timeout
         if self.fail_close:
             raise RuntimeError("close failed")
@@ -344,7 +347,7 @@ def _project_data_plane(
     result: object,
     *,
     session: DuplexSession | None = None,
-) -> list[dict[str, object]]:
+) -> list[dict[str, Any]]:
     return list(data_plane.project(result, context=_data_plane_context(session)))
 
 
@@ -4194,6 +4197,7 @@ async def test_realtime_resume_preserves_append_tail_order_across_connections():
         await asyncio.sleep(0.01)
     else:
         pytest.fail("resumable session did not open")
+    assert created is not None
     session = handler._registry.get("sid-resume-append-order")
     assert session is not None
     native_state = handler._minicpmo_session_state(session)
@@ -4323,6 +4327,7 @@ async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns():
         await asyncio.sleep(0.001)
     else:
         pytest.fail("first attachment did not create the session")
+    assert created is not None
 
     second = TimedWebSocket(receive_timeout_s=1.0)
     second.put(
@@ -4386,6 +4391,7 @@ async def test_realtime_takeover_drops_input_queued_by_replaced_attachment():
         await asyncio.sleep(0.001)
     else:
         pytest.fail("first attachment did not create the session")
+    assert created is not None
 
     second = TimedWebSocket(receive_timeout_s=1.0)
     second.put(
@@ -5002,11 +5008,12 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
             *,
             mode: str,
             payload: object,
+            operation_id: str | None = None,
             final: bool = False,
             timeout: float | None = None,
             collect_outputs: bool = True,
             fence: DuplexFence | None = None,
-        ):
+        ) -> Any:
             self.append_started.set()
             try:
                 await self.release_append.wait()
@@ -5017,6 +5024,7 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
                 session_id,
                 mode=mode,
                 payload=payload,
+                operation_id=operation_id,
                 final=final,
                 timeout=timeout,
                 collect_outputs=collect_outputs,
@@ -5128,7 +5136,9 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config()
     error = next(m for m in ws.sent if m.get("type") == "error")
     assert error["code"] == "runtime_signal_failed"
     assert error["runtime_control"]["error_count"] == 1
-    assert engine.signal_session_configs[-1]["instructions"] == "must roll back"
+    signaled_config = engine.signal_session_configs[-1]
+    assert signaled_config is not None
+    assert signaled_config["instructions"] == "must roll back"
 
 
 @pytest.mark.asyncio
@@ -5429,7 +5439,7 @@ async def test_minicpmo_native_duplex_open_unsupported_fails_session_create():
 @pytest.mark.asyncio
 async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
     class LegacyEngineClient(FakeEngineClient):
-        async def open_duplex_session_async(
+        async def open_duplex_session_async(  # type: ignore[override]
             self,
             session_id: str,
             *,
@@ -5437,7 +5447,7 @@ async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
             capabilities: dict[str, object] | None = None,
             session_config: dict[str, object] | None = None,
             timeout: float | None = None,
-        ) -> None:
+        ) -> Any:
             return await super().open_duplex_session_async(
                 session_id,
                 session_mode=session_mode,
@@ -5534,6 +5544,9 @@ async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkey
     assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192}
     assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20
     assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8
+    # The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex
+    # chunk is 26 codec samples, so check_stop would never release the request.
+    assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0}
     assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600))
 
 

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
 import base64
@@ -4147,16 +4150,15 @@ class TestCosyVoice3Serving:
         error = cosyvoice3_server._validate_tts_request(request)
         assert error is None
 
-    def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="hello",
-            max_new_tokens=0,
-        )
-        error = cosyvoice3_server._validate_tts_request(request)
-        assert error is not None
-        assert "max_new_tokens" in error
+    def test_validate_cosyvoice3_max_new_tokens_range(self):
+        """Ensure max_new_tokens below the minimum is rejected during request validation."""
+        with pytest.raises(ValidationError, match="max_new_tokens"):
+            OpenAICreateSpeechRequest(
+                input="Hello",
+                ref_audio="data:audio/wav;base64,abc",
+                ref_text="hello",
+                max_new_tokens=0,
+            )
 
     @pytest.mark.parametrize(
         ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -378,6 +378,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "preferred",
             "--host-weight-runtime-root",
             "/var/cache/vllm-omni/hwr",
+            "--dlo-host-registration-limit-gib",
+            "80",
         ]
     )
 
@@ -387,8 +389,10 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
 
     assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
     assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
     assert engine_args["host_weight_runtime_mode"] == "preferred"
     assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert engine_args["dlo_host_registration_limit_gib"] == 80
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

tests/entrypoints/test_omni_entrypoints.py

@@ -611,7 +611,13 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
     # external ID to avoid collisions, so this also tests mapping
     external_req_id = "req-1"
     req_id = "req-1-12345678"
+    recorded_failures = []
     try:
+        monkeypatch.setattr(
+            app,
+            "_record_request_failure_once",
+            lambda request_id, reason: recorded_failures.append((request_id, reason)),
+        )
         app.request_states[req_id] = ClientRequestState(
             request_id=req_id,
             external_request_id=external_req_id,
@@ -622,6 +628,7 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
 
     assert engine.aborted == [[req_id]]
     assert external_req_id not in app.request_states
+    assert recorded_failures == [(req_id, "client_abort")]
 
 
 @pytest.mark.asyncio
@@ -814,7 +821,13 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
     _patch_engine(monkeypatch, engine)
 
     app = Omni("dummy-model")
+    recorded_failures = []
     try:
+        monkeypatch.setattr(
+            app,
+            "_record_request_failure_once",
+            lambda request_id, reason: recorded_failures.append((request_id, reason)),
+        )
         app.request_states["req-1"] = object()
         app.abort("req-1")
     finally:
@@ -822,6 +835,7 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
 
     assert engine.aborted == [["req-1"]]
     assert "req-1" not in app.request_states
+    assert recorded_failures == [("req-1", "client_abort")]
 
 
 def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):

tests/host_weight_runtime/test_filesystem_store.py

@@ -57,7 +57,9 @@ from vllm_omni.host_weight_runtime.filesystem import (
     detect_storage_class,
 )
 from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
+from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module
 from vllm_omni.host_weight_runtime.filesystem.locks import FileLock
+from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -110,6 +112,150 @@ def _open_fd_targets() -> tuple[str, ...]:
     return tuple(targets)
 
 
+def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter:
+    staging = tmp_path / "staging"
+    staging.mkdir()
+    return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None)
+
+
+def test_ordered_writer_hashes_payload_during_writes(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    weight = _weight()
+    scale = torch.tensor([1.0, 2.0])
+
+    def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object:
+        raise AssertionError("ordered payload must not be rescanned")
+
+    monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan)
+    with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+        # The writer's ordered contract follows canonical storage-key order.
+        output.write_tensor("layer.scale", scale)
+        output.write_rows("layer.weight", 0, weight[:2])
+        output.write_rows("layer.weight", 2, weight[2:])
+
+    manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    payload = writer.staging_dir / "weights.safetensors"
+    assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest()
+    expected_tensor_digests = {
+        "layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(),
+        "layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(),
+    }
+    assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests
+
+
+def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None:
+    writer = _direct_writer(tmp_path)
+    output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+
+    with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"):
+        output.write_tensor("layer.weight", _weight())
+
+    output.abort()
+    assert not list(writer.staging_dir.iterdir())
+
+
+def test_payload_fsync_overlaps_producer_work(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    started = threading.Event()
+    release = threading.Event()
+    closed = threading.Event()
+    original_fsync = filesystem_writer_module._fsync_payload
+
+    def blocking_fsync(fd: int) -> None:
+        started.set()
+        if not release.wait(timeout=5):
+            raise TimeoutError("test did not release payload fsync")
+        original_fsync(fd)
+
+    monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync)
+    output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+    output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+    output.write_tensor("layer.weight", _weight())
+    close_errors: list[BaseException] = []
+
+    def close_output() -> None:
+        try:
+            output.close()
+        except BaseException as exc:
+            close_errors.append(exc)
+        finally:
+            closed.set()
+
+    close_thread = threading.Thread(target=close_output)
+    close_thread.start()
+    try:
+        assert started.wait(timeout=2)
+        assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it"
+    finally:
+        release.set()
+        close_thread.join(timeout=5)
+
+    assert not close_errors
+    writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+
+def test_payload_fsync_failure_prevents_publication(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+
+    def failing_fsync(fd: int) -> None:
+        try:
+            raise OSError(errno.EIO, "injected fsync failure")
+        finally:
+            os.close(fd)
+
+    monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync)
+    with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+        output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+        output.write_tensor("layer.weight", _weight())
+
+    with pytest.raises(OSError, match="injected fsync failure"):
+        writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    assert not (writer.staging_dir / "manifest.json").exists()
+    assert not (writer.staging_dir / "READY.json").exists()
+
+
+def test_unordered_payload_fallback_hashes_files_in_parallel(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    specs = (
+        TensorWriteSpec("a", (2,), torch.float32),
+        TensorWriteSpec("b", (2,), torch.float32),
+    )
+    for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True):
+        with writer.open_tensor_file(name, (spec,)) as output:
+            output.write_tensor(spec.name, torch.tensor([1.0, 2.0]))
+
+    original_scan = filesystem_writer_module._scan_payload
+    barrier = threading.Barrier(2)
+    thread_ids: set[int] = set()
+    thread_ids_lock = threading.Lock()
+
+    def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object:
+        with thread_ids_lock:
+            thread_ids.add(threading.get_ident())
+        barrier.wait(timeout=5)
+        return original_scan(path, file_specs)
+
+    monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan)
+    manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    assert len(manifest.files) == 2
+    assert len(thread_ids) == 2
+
+
 class FakeProducer:
     def __init__(
         self,
@@ -1026,7 +1172,9 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
 
     free_root = tmp_path / "free-limit"
     free_root.mkdir(mode=0o700)
-    unavailable_reserve = shutil.disk_usage(free_root).free + 1
+    # Leave enough margin for filesystem accounting to move between the
+    # preflight measurement and the capacity check.
+    unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30)
     free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
     free_result = free_limited.get_or_build(
         BuildRequest(identity),

tests/metrics/test_emit_calls.py

@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Verify emit-call wiring in production code paths.
+
+Uses behavior and Prometheus exposition to verify production call semantics.
+"""
+
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+import pytest
+
+from vllm_omni.metrics import OmniPrometheusMetrics
+from vllm_omni.metrics import definitions as defs
+
+if TYPE_CHECKING:
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODEL = "emit-calls-test"
+
+
+# ---------------------------------------------------------------------------
+# Behavioral pins — verify call semantics with mock OmniPrometheusMetrics.
+# ---------------------------------------------------------------------------
+
+
+def _make_omni_base_with_mock_prom(mocker):
+    """Build a minimal OmniBase shell wired with a mock prom_metrics.
+
+    Mirrors the pattern in ``test_prometheus.py::TestRequestLifecycleGauges``
+    — ``object.__new__`` skips ``__init__`` so we don't need a real engine.
+    """
+    from vllm_omni.entrypoints.omni_base import OmniBase
+
+    obj = object.__new__(OmniBase)
+    obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
+    obj.request_states = {}
+    obj._consumed_metric_messages = {}
+    obj.log_stats = True
+    return obj, obj.prom_metrics
+
+
+class TestFailureCounterWiring:
+    """All failure paths share one request-idempotent counter helper."""
+
+    def test_failure_reason_taxonomy_is_bounded(self) -> None:
+        from vllm_omni.metrics.utils import normalize_failure_reason
+
+        assert normalize_failure_reason("client_abort") == "client_abort"
+        assert normalize_failure_reason("client_disconnect") == "client_disconnect"
+        assert normalize_failure_reason("stage_error") == "stage_error"
+        assert normalize_failure_reason("unknown") == "unknown"
+        assert normalize_failure_reason("timeout for request req-1") == "unknown"
+        assert normalize_failure_reason("") == "unknown"
+        assert normalize_failure_reason(None) == "unknown"
+
+    def test_fire_failure_counter_passes_reason_through(self, mocker) -> None:
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+        obj.request_states["req-1"] = SimpleNamespace(
+            metrics=SimpleNamespace(e2e_done=set()),
+            failure_recorded=False,
+        )
+
+        obj._record_request_failure_once("req-1", reason="client_disconnect")
+
+        prom.request_failed.assert_called_once()
+        prom.inc_requests_failed.assert_called_once_with("client_disconnect")
+
+    def test_fire_failure_counter_normalizes_unrecognized_reason(self, mocker) -> None:
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+        obj.request_states["req-1"] = SimpleNamespace(
+            metrics=SimpleNamespace(e2e_done=set()),
+            failure_recorded=False,
+        )
+
+        obj._record_request_failure_once("req-1", reason="CUDA error for req-1")
+
+        prom.inc_requests_failed.assert_called_once_with("unknown")
+
+    def test_fire_failure_counter_skips_when_request_already_succeeded(self, mocker) -> None:
+        # When the request is in e2e_done (i.e. finalize already fired
+        # request_succeeded), the failure path must NOT double-count.
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+        obj.request_states["req-1"] = SimpleNamespace(
+            metrics=SimpleNamespace(e2e_done={"req-1"}),
+            failure_recorded=False,
+        )
+
+        obj._record_request_failure_once("req-1", reason="client_disconnect")
+
+        prom.request_failed.assert_not_called()
+        prom.inc_requests_failed.assert_not_called()
+
+    def test_fire_failure_counter_skips_when_request_state_missing(self, mocker) -> None:
+        # No request_states entry — already popped by abort path. Fail-safe
+        # must NOT raise.
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+
+        obj._record_request_failure_once("missing-req", reason="oom")
+
+        prom.request_failed.assert_not_called()
+        prom.inc_requests_failed.assert_not_called()
+
+    def test_failure_counter_is_idempotent_across_abort_and_cleanup(self, mocker) -> None:
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+        obj.request_states["req-1"] = SimpleNamespace(
+            metrics=SimpleNamespace(
+                e2e_done=set(),
+                build_and_log_summary=lambda: None,
+            ),
+            failure_recorded=False,
+        )
+
+        obj._record_request_failure_once("req-1", reason="client_disconnect")
+        obj._log_summary_and_cleanup("req-1", reason="stage_error")
+
+        prom.request_failed.assert_called_once()
+        prom.inc_requests_failed.assert_called_once_with("client_disconnect")
+
+    def test_log_summary_and_cleanup_default_reason(self, mocker) -> None:
+        obj, prom = _make_omni_base_with_mock_prom(mocker)
+        obj.request_states["req-1"] = SimpleNamespace(
+            metrics=SimpleNamespace(
+                e2e_done=set(),
+                build_and_log_summary=lambda: None,
+            ),
+            failure_recorded=False,
+        )
+
+        obj._log_summary_and_cleanup("req-1")
+
+        prom.inc_requests_failed.assert_called_once_with("stage_error")
+
+
+class TestEarlyReturnOnLogStatsOff:
+    """When --log-stats is off (default), observe methods must be silent no-ops.
+
+    This is the gating contract — emit sites can call unconditionally and the
+    helper short-circuits.
+    """
+
+    def test_observe_methods_silent_when_log_stats_false(self) -> None:
+        prom = OmniPrometheusMetrics(model_name=_MODEL, log_stats=False)
+        # None of these should raise or write to the registry.
+        prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=1.5)
+        prom.observe_stage_in_queue(stage=0, in_queue_s=0.2)
+        prom.observe_queue_wait(queue_wait_s=0.5)
+        prom.set_stage_waiting_requests(stage=0, n_waiting=3)
+        prom.observe_num_inference_steps(n_steps=20)
+        prom.inc_image_count(n_images=1)
+        prom.observe_image_pixels(n_pixels=512 * 512)
+        prom.set_peak_memory(stage=0, peak_memory_mb=1024.0)
+        prom.inc_requests_failed(reason="oom")
+        prom.observe_kv_wait(connector_type="shm", kv_wait_s=0.01)
+
+
+class TestStageGenerationLabels:
+    def test_stage_gen_time_exposes_bounded_stage_type(self) -> None:
+        from prometheus_client import REGISTRY, generate_latest
+
+        model = _MODEL + "-stage-type"
+        prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+        prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=0.25)
+        prom.observe_stage_gen_time(stage=1, stage_type="diffusion", gen_time_s=1.5)
+
+        out = generate_latest(REGISTRY).decode()
+        assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="0",stage_type="llm"}} 1.0' in out
+        assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="1",stage_type="diffusion"}} 1.0' in out
+
+
+class TestCounterExpositionNames:
+    def test_counter_families_have_exactly_one_total_suffix(self) -> None:
+        from prometheus_client import REGISTRY, generate_latest
+
+        model = _MODEL + "-counter-suffix"
+        prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+        prom.inc_image_count(2)
+        prom.inc_requests_failed("stage_error")
+
+        out = generate_latest(REGISTRY).decode()
+        assert f'{defs.IMAGE_COUNT_METRIC}_total{{model_name="{model}"}} 2.0' in out
+        assert f'{defs.REQUESTS_FAILED}_total{{model_name="{model}",reason="stage_error"}} 1.0' in out
+        assert f"{defs.IMAGE_COUNT_METRIC}_total_total" not in out
+        assert f"{defs.REQUESTS_FAILED}_total_total" not in out
+
+
+class TestQueueWaitExtraction:
+    def test_present_zero_is_a_valid_observation(self) -> None:
+        from vllm_omni.metrics.utils import extract_queue_wait_s
+
+        assert extract_queue_wait_s({"queue_wait_ms": 0.0}) == 0.0
+
+    def test_missing_queue_wait_is_not_synthetic_zero(self) -> None:
+        from vllm_omni.metrics.utils import extract_queue_wait_s
+
+        assert extract_queue_wait_s({"preprocess_ms": 2.0}) is None
+        assert extract_queue_wait_s(None) is None
+
+
+class TestStageInQueueObservation:
+    def test_prometheus_records_valid_zero_wait(self) -> None:
+        from prometheus_client import REGISTRY, generate_latest
+
+        model = _MODEL + "-zero-stage-wait"
+        prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+
+        prom.observe_stage_in_queue(stage=1, in_queue_s=0.0)
+
+        out = generate_latest(REGISTRY).decode()
+        sample = f'{defs.STAGE_IN_QUEUE_S}_count{{model_name="{model}",stage="1"}} 1.0'
+        assert sample in out
+
+
+class TestStageWorkloadMetricScope:
+    @staticmethod
+    def _observe(mocker, *, stage_type: str, output_unit_type: str):
+        from vllm_omni.metrics.utils import observe_stage_workload_metrics
+
+        prom = mocker.Mock(spec=OmniPrometheusMetrics)
+        stage_metrics = SimpleNamespace(
+            num_inference_steps=50,
+            output_unit_type=output_unit_type,
+            image_pixels=1024 * 1024,
+            output_unit_count=2,
+        )
+        observe_stage_workload_metrics(
+            prom,
+            stage_type=stage_type,
+            stage_metrics=stage_metrics,
+        )
+        return prom
+
+    def test_diffusion_image_observes_steps_and_image_workload(self, mocker) -> None:
+        prom = self._observe(mocker, stage_type="diffusion", output_unit_type="image")
+
+        prom.observe_num_inference_steps.assert_called_once_with(50)
+        prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
+        prom.inc_image_count.assert_called_once_with(2)
+
+    def test_diffusion_video_observes_steps_only(self, mocker) -> None:
+        prom = self._observe(mocker, stage_type="diffusion", output_unit_type="video")
+
+        prom.observe_num_inference_steps.assert_called_once_with(50)
+        prom.observe_image_pixels.assert_not_called()
+        prom.inc_image_count.assert_not_called()
+
+    def test_llm_text_observes_no_workload_metrics(self, mocker) -> None:
+        prom = self._observe(mocker, stage_type="llm", output_unit_type="text")
+
+        prom.observe_num_inference_steps.assert_not_called()
+        prom.observe_image_pixels.assert_not_called()
+        prom.inc_image_count.assert_not_called()
+
+    def test_non_diffusion_image_observes_image_workload_only(self, mocker) -> None:
+        prom = self._observe(mocker, stage_type="llm", output_unit_type="image")
+
+        prom.observe_num_inference_steps.assert_not_called()
+        prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
+        prom.inc_image_count.assert_called_once_with(2)
+
+    @pytest.mark.parametrize("peak_memory_mb", [2048.0, 0.0])
+    def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak_memory_mb: float) -> None:
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        obj = object.__new__(OmniBase)
+        obj._enable_ar_profiler = False
+        obj._consumed_metric_messages = {}
+        obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
+        obj.mod_metrics = mocker.Mock()
+        obj.engine = SimpleNamespace(
+            get_stage_metadata=lambda _stage_id: SimpleNamespace(
+                stage_type="diffusion",
+                final_output=False,
+                final_output_type="image",
+            )
+        )
+        stage_metrics = SimpleNamespace(
+            stage_gen_time_ms=1000.0,
+            diffusion_metrics={"diffusion_engine_exec_time_s": 0.8, "scheduler_queue_wait_s": 0.1},
+            num_inference_steps=20,
+            output_unit_type="image",
+            image_pixels=1024 * 1024,
+            output_unit_count=1,
+            serving_time_to_first_output_ms=250.0,
+            image_time_to_first_output_ms=250.0,
+            pipeline_timings={},
+        )
+        result = SimpleNamespace(
+            request_id="req-replay",
+            stage_id=1,
+            replica_id=0,
+            stage_submit_ts=10.0,
+            metrics=stage_metrics,
+            engine_outputs=SimpleNamespace(
+                stage_durations={},
+                peak_memory_mb=peak_memory_mb,
+                finished=True,
+                final_output_type="image",
+            ),
+        )
+        aggregator = mocker.Mock()
+        aggregator.stage_events = {}
+        aggregator.stage_first_ts = [None, None]
+        aggregator.stage_last_ts = [None, None]
+
+        for _ in range(2):
+            assert obj._process_single_result(result, 1, aggregator, {}, 0.0, 1) is None
+
+        aggregator.on_stage_metrics.assert_called_once_with(1, "req-replay", stage_metrics, "image")
+        obj.prom_metrics.inc_image_count.assert_called_once_with(1)
+        obj.prom_metrics.observe_image_pixels.assert_called_once_with(1024 * 1024)
+        obj.prom_metrics.observe_num_inference_steps.assert_called_once_with(20)
+        obj.prom_metrics.observe_stage_in_queue.assert_called_once_with(1, 0.1)
+        if peak_memory_mb > 0:
+            obj.prom_metrics.set_peak_memory.assert_called_once_with(1, peak_memory_mb)
+        else:
+            obj.prom_metrics.set_peak_memory.assert_not_called()
+        obj.mod_metrics.observe_image_ttfp.assert_called_once_with("1", "0", 0.25)
+
+
+class TestStageWaitingAggregation:
+    @staticmethod
+    def _make_orchestrator(mocker):
+        from vllm_omni.engine.orchestrator import Orchestrator
+
+        orchestrator = object.__new__(Orchestrator)
+        orchestrator._prom_metrics = mocker.Mock()
+        orchestrator._stage_replica_waiting = {}
+        return orchestrator
+
+    def test_multiple_replicas_are_summed(self, mocker) -> None:
+        orchestrator = self._make_orchestrator(mocker)
+
+        orchestrator._update_stage_replica_waiting(1, 0, 3)
+        orchestrator._update_stage_replica_waiting(1, 1, 5)
+
+        orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 8)
+
+    def test_zero_update_and_stage_isolation(self, mocker) -> None:
+        orchestrator = self._make_orchestrator(mocker)
+
+        orchestrator._update_stage_replica_waiting(1, 0, 3)
+        orchestrator._update_stage_replica_waiting(2, 0, 7)
+        orchestrator._update_stage_replica_waiting(1, 0, 0)
+
+        orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 0)
+        assert orchestrator._stage_replica_waiting[(2, 0)] == 7
+
+    def test_dead_replica_snapshot_is_removed(self, mocker) -> None:
+        orchestrator = self._make_orchestrator(mocker)
+        orchestrator._update_stage_replica_waiting(1, 0, 3)
+        orchestrator._update_stage_replica_waiting(1, 1, 5)
+
+        orchestrator._remove_stage_replica_waiting(1, 1)
+
+        orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 3)
+
+
+# ---------------------------------------------------------------------------
+# Behavioral pin: call sites now scope workload observations by stage/output
+# type. Keep the Prometheus wrapper's <= 0 guard as a second line of defense.
+# ---------------------------------------------------------------------------
+
+
+class TestZeroGuardContract:
+    """Zero-valued observations must not bump histogram ``_count``."""
+
+    def test_zero_image_pixels_not_observed(self) -> None:
+        from prometheus_client import REGISTRY, generate_latest
+
+        prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-guard")
+        prom.observe_image_pixels(n_pixels=0)
+        out = generate_latest(REGISTRY).decode()
+        # ``.labels()`` in OmniPrometheusMetrics.__init__ already creates the
+        # child sample, so the ``_count`` line exists with value 0.0 even
+        # without observations. The guard prevents ``_count`` from bumping
+        # to 1.0 — assert the value stays at 0.
+        needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-zero-guard"}}'
+        count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+        assert count_lines, "expected image_pixels_count line to exist after OmniPrometheusMetrics construction"
+        assert float(count_lines[0].split()[-1]) == 0.0, (
+            "zero-pixel observation leaked to registry — guard should early-return on <= 0"
+        )
+
+    def test_zero_num_inference_steps_not_observed(self) -> None:
+        from prometheus_client import REGISTRY, generate_latest
+
+        prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-steps")
+        prom.observe_num_inference_steps(n_steps=0)
+        out = generate_latest(REGISTRY).decode()
+        needle = f'vllm_omni:num_inference_steps_count{{model_name="{_MODEL}-zero-steps"}}'
+        count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+        assert count_lines, "expected num_inference_steps_count line to exist after construction"
+        assert float(count_lines[0].split()[-1]) == 0.0, (
+            "zero-step observation leaked to registry — guard should early-return on <= 0"
+        )
+
+    def test_positive_image_pixels_observed(self) -> None:
+        # Sanity check: positive values DO get observed.
+        from prometheus_client import REGISTRY, generate_latest
+
+        prom = OmniPrometheusMetrics(model_name=_MODEL + "-pos-guard")
+        prom.observe_image_pixels(n_pixels=512)
+        out = generate_latest(REGISTRY).decode()
+        needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-pos-guard"}}'
+        count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+        assert count_lines, "expected image_pixels_count line after positive observation"
+        assert float(count_lines[0].split()[-1]) == 1.0, "positive-pixel observation did not increment count to 1"
+
+
+# ---------------------------------------------------------------------------
+# KV-wait emit wiring — scheduler ENTER/EXIT lifecycle + orchestrator dispatch.
+# ---------------------------------------------------------------------------
+
+
+def _make_scheduler_shell() -> OmniARScheduler:
+    """Minimal OmniARScheduler shell — skips upstream __init__ via object.__new__."""
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+    obj = object.__new__(OmniARScheduler)
+    obj.requests = {}
+    obj.running = []
+    obj.waiting = []
+    obj._kv_wait_start_ts = {}
+    obj._omni_kv_config = None
+    return obj
+
+
+class TestKvWaitSchedulerEmit:
+    """Scheduler-side half of kv_wait_s: pops start ts, carries the wait across."""
+
+    def test_skips_when_no_start_ts_recorded(self) -> None:
+        sched = _make_scheduler_shell()
+        outputs: dict[int, list] = {}
+
+        sched._emit_kv_wait_output(outputs, "req-no-wait", req=SimpleNamespace(client_index=0))
+
+        assert outputs == {}, "emit must not append output when no start ts recorded"
+        assert "req-no-wait" not in sched._kv_wait_start_ts
+
+    def test_emits_wait_duration_and_connector_type(self) -> None:
+        sched = _make_scheduler_shell()
+
+        sched._kv_wait_start_ts["req-1"] = time.monotonic() - 0.25
+        outputs: dict[int, list] = {}
+        live_req = SimpleNamespace(client_index=3)
+
+        sched._emit_kv_wait_output(outputs, "req-1", req=live_req)
+
+        assert "req-1" not in sched._kv_wait_start_ts, "start ts must be popped after emit"
+        assert 3 in outputs and len(outputs[3]) == 1
+        params = outputs[3][0].kv_transfer_params
+        assert "kv_wait_s" in params and "connector_type" in params
+        assert 0.0 < params["kv_wait_s"] < 1.0
+        assert params["connector_type"] == "unknown"
+
+    def test_connector_type_resolved_from_omni_kv_config(self) -> None:
+        sched = _make_scheduler_shell()
+        sched._omni_kv_config = {"connector_config": {"type": "SharedMemoryConnector"}}
+
+        sched._kv_wait_start_ts["req-3"] = time.monotonic() - 0.01
+        outputs: dict[int, list] = {}
+
+        sched._emit_kv_wait_output(outputs, "req-3", req=SimpleNamespace(client_index=0))
+
+        eco = outputs[0][0]
+        assert eco.kv_transfer_params["connector_type"] == "SharedMemoryConnector"
+
+
+class TestKvWaitTerminalCleanup:
+    def test_finish_requests_clears_requested_wait_only(self) -> None:
+        from vllm.v1.request import RequestStatus
+
+        scheduler = _make_scheduler_shell()
+        scheduler._kv_wait_start_ts = {"req-1": 10.0, "req-2": 20.0}
+
+        finished = scheduler.finish_requests(
+            (request_id for request_id in ["req-1"]),
+            RequestStatus.FINISHED_ABORTED,
+        )
+
+        assert finished == []
+        assert "req-1" not in scheduler._kv_wait_start_ts
+        assert scheduler._kv_wait_start_ts["req-2"] == 20.0
+
+    def test_cleanup_is_idempotent_when_timestamp_is_already_gone(self) -> None:
+        from vllm.v1.request import RequestStatus
+
+        scheduler = _make_scheduler_shell()
+
+        scheduler.finish_requests("req-already-emitted", RequestStatus.FINISHED_ABORTED)
+
+        assert scheduler._kv_wait_start_ts == {}

tests/metrics/test_modality.py

@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 from prometheus_client import REGISTRY, generate_latest
 
@@ -10,6 +15,7 @@ from vllm_omni.metrics.modality import (
     observe_audio_streaming_finalize,
     observe_modality_at_finalize,
 )
+from vllm_omni.metrics.stats import OrchestratorAggregator
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -46,6 +52,11 @@ _EXPECTED_FAMILIES = [
     defs.DIFFUSION_EXEC_PER_STEP_S,
     defs.DIFFUSION_PREPROCESS_S,
     defs.DIFFUSION_POSTPROCESS_S,
+    defs.VAE_DECODE_S,
+    defs.DIFFUSION_FORWARD_S,
+    defs.DIFFUSION_KV_LOAD_S,
+    defs.IMAGE_TTFP_S,
+    defs.DENOISE_STEP_LATENCY_S,
 ]
 
 
@@ -63,6 +74,11 @@ class TestRegistration:
         mod.observe_diffusion_exec_per_step("s", "r", 0.01)
         mod.observe_diffusion_preprocess("s", "r", 0.01)
         mod.observe_diffusion_postprocess("s", "r", 0.2)
+        mod.observe_vae_decode("s", "r", 0.3)
+        mod.observe_diffusion_forward("s", "r", 0.8)
+        mod.observe_diffusion_kv_load("s", "r", 0.04)
+        mod.observe_image_ttfp("s", "r", 2.5)
+        mod.observe_denoise_step_latency("s", "r", 0.05)
 
         out = generate_latest(REGISTRY).decode()
         for name in _EXPECTED_FAMILIES:
@@ -247,6 +263,21 @@ class _StubModMetrics:
     def observe_diffusion_postprocess(self, s, r, v):
         self.calls.append(("observe_diffusion_postprocess", s, r, v))
 
+    def observe_vae_decode(self, s, r, seconds):
+        self.calls.append(("observe_vae_decode", s, r, seconds))
+
+    def observe_diffusion_forward(self, s, r, seconds):
+        self.calls.append(("observe_diffusion_forward", s, r, seconds))
+
+    def observe_diffusion_kv_load(self, s, r, seconds):
+        self.calls.append(("observe_diffusion_kv_load", s, r, seconds))
+
+    def observe_image_ttfp(self, s, r, seconds):
+        self.calls.append(("observe_image_ttfp", s, r, seconds))
+
+    def observe_denoise_step_latency(self, s, r, seconds):
+        self.calls.append(("observe_denoise_step_latency", s, r, seconds))
+
 
 class _Bag:
     def __init__(self, **kw):
@@ -306,11 +337,13 @@ class TestObserveModalityAtFinalize:
     def test_diffusion_path_full(self):
         stub = _StubModMetrics()
         stage_metrics = _Bag(
+            output_unit_type="image",
             stage_gen_time_ms=1000.0,
             diffusion_metrics={
                 "preprocess_time_s": 0.05,
                 "diffusion_engine_exec_time_s": 1.2,
                 "postprocess_time_s": 0.3,
+                "forward_time_s": 1.0,
                 "image_num": 1,
                 "resolution": 1024,
                 "num_inference_steps": 50,
@@ -328,6 +361,71 @@ class TestObserveModalityAtFinalize:
         assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
         assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
         assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
+        assert ("observe_denoise_step_latency", "2", "0", pytest.approx(0.02)) in stub.calls
+
+    def test_denoise_step_latency_requires_forward_time(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(
+            output_unit_type="image",
+            num_inference_steps=20,
+            diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
+        )
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=_Bag(),
+        )
+        assert not any(c[0] == "observe_denoise_step_latency" for c in stub.calls)
+
+    def test_stage_snapshot_keeps_forward_based_denoise_latency(self):
+        aggregator = OrchestratorAggregator(
+            num_stages=1,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=0,
+        )
+        aggregator.diffusion_metrics["req-image"]["forward_time_s"] = 0.8
+        stage_metrics = _Bag(
+            stage_id=None,
+            request_id=None,
+            final_output_type=None,
+            num_tokens_in=0,
+            num_tokens_out=0,
+            stage_gen_time_ms=1200.0,
+            rx_transfer_bytes=0,
+            rx_decode_time_ms=0.0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=_Bag(),
+            diffusion_metrics=None,
+            num_inference_steps=20,
+            output_unit_type="image",
+            denoise_step_latency_ms=0.0,
+        )
+
+        stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
+
+        assert stats.denoise_step_latency_ms == pytest.approx(40.0)
+
+    def test_stage_snapshot_falls_back_to_stage_time_without_profiler(self):
+        aggregator = object.__new__(OrchestratorAggregator)
+        stage_metrics = _Bag(
+            stage_id=None,
+            request_id=None,
+            final_output_type=None,
+            stage_gen_time_ms=1200.0,
+            diffusion_metrics=None,
+            num_inference_steps=20,
+            output_unit_type="image",
+            denoise_step_latency_ms=0.0,
+        )
+
+        stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
+
+        assert stats.diffusion_metrics is None
+        assert stats.denoise_step_latency_ms == pytest.approx(60.0)
 
     def test_diffusion_per_step_skipped_without_num_inference_steps(self):
         stub = _StubModMetrics()
@@ -422,6 +520,80 @@ class TestObserveModalityAtFinalize:
         assert stub.calls == []
 
 
+class TestDiffusionUnitConversion:
+    def test_accumulator_converts_all_engine_milliseconds_to_seconds(self):
+        aggregator = OrchestratorAggregator(
+            num_stages=1,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=0,
+        )
+        engine_output = SimpleNamespace(
+            metrics={
+                "preprocess_time_ms": 50.0,
+                "diffusion_engine_exec_time_ms": 1500.0,
+                "postprocess_time_ms": 20.0,
+                "vae_decode_time_ms": 300.0,
+                "forward_time_ms": 800.0,
+                "scheduler_queue_wait_ms": 125.0,
+                "kv_recv_time_ms": 40.0,
+            }
+        )
+
+        aggregator.accumulate_diffusion_metrics("diffusion", "req-units", engine_output)
+
+        assert dict(aggregator.diffusion_metrics["req-units"]) == {
+            "preprocess_time_s": pytest.approx(0.05),
+            "diffusion_engine_exec_time_s": pytest.approx(1.5),
+            "postprocess_time_s": pytest.approx(0.02),
+            "vae_decode_time_s": pytest.approx(0.3),
+            "forward_time_s": pytest.approx(0.8),
+            "scheduler_queue_wait_s": pytest.approx(0.125),
+            "kv_recv_time_s": pytest.approx(0.04),
+        }
+
+
+class TestObserveDiffusionFinalize:
+    @pytest.mark.parametrize(
+        ("metric_key", "call_name", "value"),
+        [
+            ("vae_decode_time_s", "observe_vae_decode", 0.3),
+            ("forward_time_s", "observe_diffusion_forward", 0.8),
+            ("kv_recv_time_s", "observe_diffusion_kv_load", 0.04),
+        ],
+    )
+    def test_emits_optional_breakdown_metric(self, metric_key, call_name, value):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="image",
+            stage_id=2,
+            replica_id=0,
+            stage_metrics=_Bag(diffusion_metrics={metric_key: value}),
+            engine_outputs=_Bag(),
+        )
+
+        assert (call_name, "2", "0", value) in stub.calls
+
+    def test_audio_diffusion_stage_fires_both_paths(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=_Bag(
+                stage_gen_time_ms=1500.0,
+                audio_generated_frames=24000,
+                diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
+            ),
+            engine_outputs=_Bag(multimodal_output={"audio_sample_rate": 24000}),
+        )
+
+        assert ("observe_diffusion_exec", "1", "0", 1.2) in stub.calls
+        assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+
+
 class TestObserveAudioFirstPacket:
     def test_observes_with_valid_inputs(self):
         stub = _StubModMetrics()
@@ -541,3 +713,39 @@ class TestBucketSelection:
         # SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS.
         fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"'
         assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001"
+
+    def test_vae_decode_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+        # VAE decode can stretch into seconds for high-resolution latents;
+        # stays on the wide bucket so the 300 s ceiling covers video VAE.
+        stage, replica = "vae_decode_sec", "0"
+        mod.observe_vae_decode(stage, replica, 1.5)
+        out = generate_latest(REGISTRY).decode()
+        sec_marker = f'{defs.VAE_DECODE_S}_bucket{{le="300.0"'
+        assert sec_marker in out, "vae_decode_s should use SECONDS_BUCKETS containing le=300.0"
+
+    def test_diffusion_forward_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+        # Denoise loop can exceed the 60s fast-bucket ceiling for video gen.
+        stage, replica = "diff_forward_sec", "0"
+        mod.observe_diffusion_forward(stage, replica, 2.0)
+        out = generate_latest(REGISTRY).decode()
+        sec_marker = f'{defs.DIFFUSION_FORWARD_S}_bucket{{le="300.0"'
+        assert sec_marker in out, "diffusion_forward_s should use SECONDS_BUCKETS containing le=300.0"
+
+    def test_denoise_step_latency_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None:
+        # Per-step denoise is sub-second; stays on the fast bucket for ms-level
+        # resolution.
+        stage, replica = "denoise_fast", "0"
+        mod.observe_denoise_step_latency(stage, replica, 0.075)
+        out = generate_latest(REGISTRY).decode()
+        fast_marker = f'{defs.DENOISE_STEP_LATENCY_S}_bucket{{le="0.001"'
+        assert fast_marker in out, "denoise_step_latency_s should use SECONDS_FAST_BUCKETS containing le=0.001"
+
+
+class TestDenoiseStepLatencyGuard:
+    @pytest.mark.parametrize(("value", "expected_count"), [(0.0, None), (-0.1, None), (0.05, 1.0)])
+    def test_nonpositive_values_are_skipped(self, mod: OmniModalityMetrics, value, expected_count) -> None:
+        stage, replica = f"denoise_{value}", "0"
+        mod.observe_denoise_step_latency(stage, replica, value)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.DENOISE_STEP_LATENCY_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == expected_count

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -281,6 +281,153 @@ class TestCFM:
 
         assert out.shape == mu.shape
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
+    def test_trt_estimator_uses_stream_dependencies(self, monkeypatch):
+        from omegaconf import DictConfig
+
+        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ConditionalCFM
+
+        class FakeStream:
+            def __init__(self, name, state):
+                self.name = name
+                self.state = state
+                self.cuda_stream = hash(name)
+                self.synchronize_calls = 0
+                self.waited_on = []
+
+            def synchronize(self):
+                self.synchronize_calls += 1
+
+            def wait_stream(self, stream):
+                self.waited_on.append(stream)
+
+        class FakeStreamContext:
+            def __init__(self, stream):
+                self.stream = stream
+
+            def __enter__(self):
+                self.previous = self.stream.state.current
+                self.stream.state.current = self.stream
+                return self.stream
+
+            def __exit__(self, exc_type, exc, traceback):
+                self.stream.state.current = self.previous
+
+        state = SimpleNamespace(current=None)
+        caller_stream = FakeStream("caller", state)
+        estimator_stream = FakeStream("estimator", state)
+        state.current = caller_stream
+        monkeypatch.setattr(torch.cuda, "current_stream", lambda *args, **kwargs: state.current)
+        stream_contexts = []
+
+        def stream_context(stream):
+            stream_contexts.append(stream)
+            return FakeStreamContext(stream)
+
+        monkeypatch.setattr(torch.cuda, "stream", stream_context)
+
+        class FakeContext:
+            def __init__(self):
+                self.execute_stream = None
+
+            def set_input_shape(self, name, shape):
+                pass
+
+            def set_tensor_address(self, name, address):
+                pass
+
+            def execute_async_v3(self, stream):
+                self.execute_stream = stream
+                return True
+
+        class FakeEngine:
+            @staticmethod
+            def get_tensor_name(index):
+                return f"tensor_{index}"
+
+        context = FakeContext()
+
+        class FakeEstimatorPool:
+            io_dtype = torch.float32
+
+            def __init__(self):
+                self.released = []
+
+            def acquire_estimator(self):
+                return [context, estimator_stream], FakeEngine()
+
+            def release_estimator(self, released_context, released_stream):
+                self.released.append((released_context, released_stream))
+
+        estimator_pool = FakeEstimatorPool()
+        cfm = ConditionalCFM(
+            in_channels=80,
+            cfm_params=DictConfig(
+                {
+                    "sigma_min": 1e-6,
+                    "solver": "euler",
+                    "t_scheduler": "cosine",
+                    "training_cfg_rate": 0.2,
+                    "inference_cfg_rate": 0.7,
+                }
+            ),
+            n_spks=1,
+            spk_emb_dim=80,
+            estimator=estimator_pool,
+        )
+
+        x = torch.randn(2, 80, 4)
+        mask = torch.ones(2, 1, 4)
+        mu = torch.randn(2, 80, 4)
+        timestep = torch.randn(2)
+        speakers = torch.randn(2, 80)
+        condition = torch.randn(2, 80, 4)
+
+        output = cfm.forward_estimator(x, mask, mu, timestep, speakers, condition)
+
+        assert output.shape == x.shape
+        assert caller_stream.synchronize_calls == 0
+        assert estimator_stream.synchronize_calls == 0
+        assert estimator_stream.waited_on == [caller_stream]
+        assert caller_stream.waited_on == [estimator_stream]
+        assert stream_contexts == [estimator_stream]
+        assert context.execute_stream == estimator_stream.cuda_stream
+        assert estimator_pool.released == [(context, estimator_stream)]
+
+    @pytest.mark.core_model
+    @pytest.mark.cpu
+    def test_trt_context_pool_stores_cuda_stream(self, monkeypatch):
+        from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import TrtContextWrapper
+
+        execution_context = object()
+
+        class FakeEngine:
+            @staticmethod
+            def create_execution_context():
+                return execution_context
+
+        cuda_stream = object()
+        monkeypatch.setattr(torch.cuda, "Stream", lambda device: cuda_stream)
+
+        def reject_stream_context(stream):
+            pytest.fail("TrtContextWrapper must store the CUDA stream, not a StreamContext")
+
+        monkeypatch.setattr(torch.cuda, "stream", reject_stream_context)
+
+        engine = FakeEngine()
+        wrapper = TrtContextWrapper(engine, device="cuda:0")
+        [context, stream], acquired_engine = wrapper.acquire_estimator()
+
+        assert context is execution_context
+        assert stream is cuda_stream
+        assert acquired_engine is engine
+
+        wrapper.release_estimator(context, stream)
+        [reused_context, reused_stream], _ = wrapper.acquire_estimator()
+        assert reused_context is context
+        assert reused_stream is stream
+
 
 class TestSDPAFallback:
     """Test SDPA fallback for float32 inputs."""

tests/model_extras/test_model_extras.py

@@ -177,23 +177,6 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
     assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("MagiHumanPipeline") == frozenset(
-        {
-            "seconds",
-            "audio_path",
-            "image_path",
-            "sr_height",
-            "sr_width",
-            "sr_num_inference_steps",
-        }
-    )
-    assert get_extra_output_params("MagiHumanPipeline") == frozenset()
-    assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_ltx_extra_registry_declares_official_guidance_params() -> None:

tests/model_tests/diffusion/test_alignment.py

@@ -57,7 +57,6 @@ EXCLUDED_MODELS = [
     "OmniGen2Pipeline",
     "HeliosPipeline",
     "HeliosPyramidPipeline",
-    "DreamIDOmniPipeline",
     "SenseNovaU1Pipeline",
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
@@ -71,7 +70,6 @@ EXCLUDED_MODELS = [
     # Requires camera/action assets and persistent AR state; covered by its
     # dedicated contract tests and real-checkpoint E2E smoke instead.
     "LingBotWorldCausalDMDPipeline",
-    "MagiHumanPipeline",
     "OmniVoicePipeline",
     "OmniVoice",
     "Cosmos3OmniDiffusersPipeline",

vllm_omni/benchmarks/data_modules/omniinteract_dataset.py

@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniInteract dataset discovery for the standard serving benchmark."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import random
+import shutil
+import tarfile
+import tempfile
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest
+from vllm.tokenizers import TokenizerLike
+from vllm.transformers_utils.repo_utils import hf_fs
+
+OMNIINTERACT_SUBSETS = ("1q1a", "1q1a_math", "1qna")
+DEFAULT_OMNIINTERACT_REPO = "lucky-lance/OmniInteract"
+DEFAULT_MAX_VIDEO_DURATION_S = 3600.0
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniInteractCase:
+    subset: str
+    video_rel: str
+    video_path: Path
+    annotation_path: Path
+    scene_type: str
+
+
+@dataclass(frozen=True)
+class OmniInteractSessionOptions:
+    output_root: Path
+    timeout_s: float
+    media_timeout_s: float
+    ref_audio: str
+    require_response: bool = False
+    max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
+
+
+@dataclass(frozen=True)
+class OmniInteractPreparedInput:
+    duration_s: float
+    pcm16: bytes = field(repr=False)
+    video_frames: tuple[str | None, ...] = field(repr=False)
+    ref_audio_data_url: str = field(repr=False)
+
+
+@dataclass
+class OmniInteractSampleRequest(SampleRequest):
+    omniinteract_case: OmniInteractCase | None = None
+    omniinteract_options: OmniInteractSessionOptions | None = None
+    omniinteract_prepared_input: OmniInteractPreparedInput | None = field(default=None, repr=False)
+
+
+class OmniInteractDataset(BenchmarkDataset):
+    def __init__(
+        self,
+        *,
+        data_root: str | None,
+        dataset_repo: str = DEFAULT_OMNIINTERACT_REPO,
+        subsets: Sequence[str] = OMNIINTERACT_SUBSETS,
+        random_seed: int = 0,
+        disable_shuffle: bool = False,
+    ) -> None:
+        super().__init__(
+            dataset_path=data_root or dataset_repo,
+            random_seed=random_seed,
+            disable_shuffle=disable_shuffle,
+        )
+        self.root = resolve_omniinteract_root(data_root, dataset_repo)
+        self.subsets = tuple(subsets)
+
+    def sample(
+        self,
+        tokenizer: TokenizerLike | None,
+        num_requests: int,
+        *,
+        request_id_prefix: str = "",
+        options: OmniInteractSessionOptions,
+        **_: Any,
+    ) -> list[SampleRequest]:
+        del tokenizer
+        cases = discover_omniinteract_cases(
+            self.root,
+            self.subsets,
+            num_prompts=num_requests,
+            seed=self.random_seed,
+            disable_shuffle=self.disable_shuffle,
+        )
+        return [
+            OmniInteractSampleRequest(
+                prompt="",
+                prompt_len=0,
+                expected_output_len=0,
+                multi_modal_data=None,
+                request_id=f"{request_id_prefix}{index}",
+                omniinteract_case=case,
+                omniinteract_options=options,
+            )
+            for index, case in enumerate(cases)
+        ]
+
+
+def _data_dir(root: Path) -> Path:
+    for candidate in (root, root / "data"):
+        if any((candidate / subset).is_dir() for subset in OMNIINTERACT_SUBSETS):
+            return candidate
+    raise FileNotFoundError(f"OmniInteract data not found under {root}")
+
+
+def _confined_path(root: Path, relative: object, *, field: str) -> Path:
+    text, resolved_root = str(relative or ""), root.resolve()
+    path = Path(text)
+    destination = (root / path).resolve()
+    if not text or path.is_absolute() or ".." in path.parts or not destination.is_relative_to(resolved_root):
+        raise ValueError(f"Unsafe OmniInteract {field} path: {text!r}")
+    return destination
+
+
+def _safe_extract(archive: tarfile.TarFile, target: Path) -> None:
+    root = target.resolve()
+    members = archive.getmembers()
+    for member in members:
+        path = Path(member.name)
+        if (
+            path.is_absolute()
+            or ".." in path.parts
+            or not (root / path).resolve().is_relative_to(root)
+            or not (member.isdir() or member.isfile())
+        ):
+            raise ValueError(f"Unsafe path in OmniInteract archive: {member.name!r}")
+    archive.extractall(target, members=members, filter="data")
+
+
+def _extract_archive(archive: Path, target: Path) -> Path:
+    if target.is_symlink():
+        raise ValueError(f"Refusing to extract through symlink: {target}")
+    stat = archive.stat()
+    published = target / f"{stat.st_size}-{stat.st_mtime_ns}"
+    if published.is_symlink():
+        raise ValueError(f"Refusing to use symlinked extraction: {published}")
+    if published.is_dir():
+        try:
+            return _data_dir(published)
+        except FileNotFoundError:
+            pass
+    target.mkdir(parents=True, exist_ok=True)
+    staging = Path(tempfile.mkdtemp(prefix=".tmp-", dir=target))
+    try:
+        with tarfile.open(archive, "r:*") as handle:
+            _safe_extract(handle, staging)
+        root = _data_dir(staging)
+        try:
+            staging.replace(published)
+            return published / root.relative_to(staging)
+        except OSError:
+            if not published.is_dir():
+                raise
+            return _data_dir(published)
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+
+def resolve_omniinteract_root(data_root: str | None, dataset_repo: str = DEFAULT_OMNIINTERACT_REPO) -> Path:
+    """Resolve a local tree/archive or download the official archive."""
+
+    if data_root:
+        local = Path(data_root).expanduser().resolve()
+        if local.is_file():
+            return _extract_archive(local, local.parent / f".{local.stem}.vllm_omni_extracted")
+        if not local.is_dir():
+            raise FileNotFoundError(f"--dataset-path does not exist: {local}")
+        try:
+            return _data_dir(local)
+        except FileNotFoundError:
+            for name in ("data.tar.gz", "data.tar"):
+                if (archive := local / name).is_file():
+                    return _extract_archive(archive, local / ".vllm_omni_extracted")
+            raise
+
+    cache = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
+    target = cache / "vllm_omni" / "omniinteract" / dataset_repo.replace("/", "__")
+    downloads = target / ".downloads"
+    downloads.mkdir(parents=True, exist_ok=True)
+    filesystem, errors = hf_fs(), []
+    for name in ("data.tar.gz", "data.tar"):
+        archive = downloads / name
+        try:
+            if not archive.is_file():
+                descriptor, temporary_name = tempfile.mkstemp(prefix=f".{name}.", dir=downloads)
+                os.close(descriptor)
+                temporary = Path(temporary_name)
+                try:
+                    filesystem.get_file(f"datasets/{dataset_repo}/{name}", str(temporary))
+                    if not temporary.stat().st_size:
+                        raise ValueError("downloaded archive is empty")
+                    temporary.replace(archive)
+                finally:
+                    temporary.unlink(missing_ok=True)
+            return _extract_archive(archive, target)
+        except Exception as exc:  # noqa: BLE001, PERF203 - Hub errors vary by version
+            errors.append(f"{name}: {exc}")
+    raise FileNotFoundError(f"Could not download OmniInteract from {dataset_repo!r}: {'; '.join(errors)}")
+
+
+def _mapping_cases(root: Path, subset: str) -> list[OmniInteractCase]:
+    mapping_path = root / "video_json_map.json"
+    if not mapping_path.is_file():
+        raise FileNotFoundError(f"Missing OmniInteract mapping: {mapping_path}")
+    mapping = json.loads(mapping_path.read_text())
+    entries = mapping.get("entries") if isinstance(mapping, dict) else None
+    if not isinstance(entries, list):
+        raise ValueError(f"Invalid OmniInteract mapping: {mapping_path}")
+    cases = []
+    for row in entries:
+        if not isinstance(row, dict):
+            raise ValueError(f"Invalid OmniInteract mapping row in {mapping_path}")
+        video_rel, annotation_rel = str(row.get("video") or ""), str(row.get("annotation") or "")
+        video = _confined_path(root, video_rel, field="video")
+        annotation = _confined_path(root, annotation_rel, field="annotation")
+        if not video.is_file() or not annotation.is_file():
+            raise FileNotFoundError(f"OmniInteract mapping references missing files: {video_rel!r}, {annotation_rel!r}")
+        cases.append(
+            OmniInteractCase(
+                subset,
+                video_rel,
+                video,
+                annotation,
+                str(row.get("scene_type") or "multi_turn").lower(),
+            )
+        )
+    return cases
+
+
+def _one_to_many_cases(root: Path) -> list[OmniInteractCase]:
+    videos, annotations = root / "videos_bench", root / "annotations"
+    if not videos.is_dir() or not annotations.is_dir():
+        raise FileNotFoundError(f"Invalid OmniInteract 1qna layout under {root}")
+    cases = []
+    for video in sorted(videos.rglob("*.mp4")):
+        if not video.resolve().is_relative_to(videos.resolve()):
+            raise ValueError(f"Unsafe OmniInteract video path: {video}")
+        relative = video.relative_to(videos)
+        annotation = (annotations / relative).with_suffix(".json").resolve()
+        if not annotation.is_relative_to(annotations.resolve()):
+            raise ValueError(f"Unsafe OmniInteract annotation path: {relative}")
+        if not annotation.is_file():
+            raise FileNotFoundError(f"Missing OmniInteract annotation: {annotation}")
+        cases.append(OmniInteractCase("1qna", f"videos_bench/{relative.as_posix()}", video, annotation, "1qna"))
+    return cases
+
+
+def discover_omniinteract_cases(
+    root: Path,
+    subsets: Sequence[str],
+    *,
+    num_prompts: int,
+    seed: int = 0,
+    disable_shuffle: bool = False,
+) -> list[OmniInteractCase]:
+    invalid = set(subsets) - set(OMNIINTERACT_SUBSETS)
+    if invalid:
+        raise ValueError(f"Unsupported OmniInteract subsets: {sorted(invalid)}")
+    if not subsets:
+        raise ValueError("At least one OmniInteract subset is required")
+    if len(set(subsets)) != len(subsets):
+        raise ValueError("OmniInteract subsets must not contain duplicates")
+    if num_prompts < 0:
+        raise ValueError("num_prompts must be non-negative")
+    data_root, cases = _data_dir(root.resolve()), []
+    for subset in subsets:
+        subset_root = data_root / subset
+        selected = _one_to_many_cases(subset_root) if subset == "1qna" else _mapping_cases(subset_root, subset)
+        if not selected:
+            raise ValueError(f"No OmniInteract sessions found for requested subset {subset!r}")
+        cases.extend(selected)
+    if len({case.video_path for case in cases}) != len(cases):
+        raise ValueError("OmniInteract dataset contains duplicate video paths")
+    if not disable_shuffle:
+        random.Random(seed).shuffle(cases)
+    if num_prompts:
+        if num_prompts > len(cases):
+            logger.warning(
+                "Requested %d OmniInteract prompts but only %d are available; using all cases", num_prompts, len(cases)
+            )
+        cases = cases[:num_prompts]
+    return cases
+
+
+def case_manifest(case: OmniInteractCase, output_dir: Path) -> dict[str, Any]:
+    return {
+        "sample_id": f"{case.subset}__{output_dir.name}",
+        "video": str(case.video_path),
+        "gt_json": str(case.annotation_path.resolve()),
+        "model_json": str((output_dir / "wav_transcript.json").resolve()),
+        "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
+    }

vllm_omni/benchmarks/metrics/metrics.py

@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import warnings
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import field, make_dataclass
 
 import numpy as np
@@ -15,6 +19,11 @@ _FLOAT_LIST_TYPE = list[float]
 _INT_LIST_TYPE = list[int]
 
 _MULTIMODAL_BENCHMARK_FIELDS = [
+    ("num_ttft_samples", int, field(default=0)),
+    ("num_tpot_samples", int, field(default=0)),
+    ("num_itl_samples", int, field(default=0)),
+    ("num_audio_ttfp_samples", int, field(default=0)),
+    ("num_audio_rtf_samples", int, field(default=0)),
     (defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
     (defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
     (defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
@@ -176,6 +185,17 @@ def _stage_modality_flags(
     )
 
 
+def has_metric_samples(metrics: object, metric_name: str) -> bool:
+    sample_count_attr = {
+        "ttft": "num_ttft_samples",
+        "tpot": "num_tpot_samples",
+        "itl": "num_itl_samples",
+        defs.AUDIO_TTFP: "num_audio_ttfp_samples",
+        defs.AUDIO_RTF: "num_audio_rtf_samples",
+    }.get(metric_name)
+    return sample_count_attr is None or getattr(metrics, sample_count_attr, 0) > 0
+
+
 def print_metrics(
     task_type,
     selected_percentile_metrics,
@@ -308,6 +328,9 @@ def process_one_metric(
     metric_attribute_name: str,
     metrics: MultiModalsBenchmarkMetrics,
 ):
+    if not has_metric_samples(metrics, metric_attribute_name):
+        return
+
     metric_header_map = {
         "ttft": "Time to First Token",
         "tpot": "Time per Output Token (excl. 1st token)",
@@ -718,10 +741,12 @@ def calculate_metrics(
     good_completed = 0
     itls: list[float] = []
     tpots: list[float] = []
-    all_tpots: list[float] = []
+    all_tpots: list[float | None] = []
     ttfts: list[float] = []
+    goodput_ttfts: list[float | None] = []
     e2els: list[float] = []
     audio_ttfps: list[float] = []
+    goodput_audio_ttfps: list[float | None] = []
     audio_rtfs: list[float] = []
     audio_duration: list[float] = []
     audio_frames: list[int] = []
@@ -748,7 +773,7 @@ def calculate_metrics(
                     output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
             actual_output_lens.append(output_len)
             total_input += outputs[i].prompt_len
-            tpot = 0
+            tpot: float | None = 0.0
             if output_len > 1:
                 if outputs[i].itl:
                     # Use mean(ITL) directly so per-request TPOT == mean(ITL).
@@ -756,19 +781,33 @@ def calculate_metrics(
                     # bundle multiple tokens per chunk, so len(itl)+1 != output_len.
                     # Using mean(itl) keeps TPOT and ITL on the same footing.
                     tpot = sum(outputs[i].itl) / len(outputs[i].itl)
-                else:
+                elif getattr(outputs[i], "tpot_measured", True):
                     try:
                         latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
                     except Exception:
                         latency_minus_ttft = outputs[i].latency - outputs[i].ttft
                     tpot = latency_minus_ttft / (output_len - 1)
-                tpots.append(tpot)
+                else:
+                    tpot = None
+                if tpot is not None:
+                    tpots.append(tpot)
             # Note: if output_len <= 1, we regard tpot as 0 for goodput
             all_tpots.append(tpot)
             itls += outputs[i].itl
-            ttfts.append(outputs[i].ttft)
-            audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
-            audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
+            session_metrics = getattr(outputs[i], "duplex_session_metrics", None)
+            ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None
+            audio_ttfp_measured = (
+                not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None
+            )
+            audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None
+            if ttft_measured:
+                ttfts.append(outputs[i].ttft)
+            if audio_ttfp_measured:
+                audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
+            if audio_rtf_measured:
+                audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
+            goodput_ttfts.append(outputs[i].ttft if ttft_measured else None)
+            goodput_audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0) if audio_ttfp_measured else None)
             audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
             audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
             image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
@@ -789,14 +828,14 @@ def calculate_metrics(
             actual_output_lens.append(0)
 
     if goodput_config_dict:
-        valid_metrics = []
+        valid_metrics: list[Sequence[float | None]] = []
         slo_values = []
 
         if "ttft" in goodput_config_dict:
-            valid_metrics.append(ttfts)
+            valid_metrics.append(goodput_ttfts)
             slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
         if "audio_ttft" in goodput_config_dict:
-            valid_metrics.append(audio_ttfps)
+            valid_metrics.append(goodput_audio_ttfps)
             slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
         if "tpot" in goodput_config_dict:
             valid_metrics.append(all_tpots)
@@ -806,7 +845,7 @@ def calculate_metrics(
             slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
 
         for req_metric in zip(*valid_metrics):
-            is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
+            is_good_req = all(r is not None and s >= r for s, r in zip(slo_values, req_metric))
             if is_good_req:
                 good_completed += 1
 
@@ -822,6 +861,7 @@ def calculate_metrics(
     # Calculate max output tokens per second metric
     max_output_tokens_per_s = 0.0
     max_concurrent_requests = 0
+    token_timeline_available = True
 
     # Find the time range across all successful requests
     successful_outputs = [output for output in outputs if output.success]
@@ -836,19 +876,26 @@ def calculate_metrics(
         concurrent_requests_per_second = np.zeros(duration_seconds)
 
         for i, output in enumerate(successful_outputs):
-            # Calculate token generation timestamp using
-            # start_time, ttft, and itl
-            token_times = [output.start_time + output.ttft]
-            current_time = token_times[0]
-            for itl_value in output.itl:
-                current_time += itl_value
-                token_times.append(current_time)
+            session_metrics = getattr(output, "duplex_session_metrics", None)
+            if isinstance(session_metrics, dict):
+                # Duplex TTFT/ITL samples are response-local. Without each
+                # response.created offset they cannot form a session-global
+                # token timeline, so do not publish a misleading peak rate.
+                token_timeline_available = False
+            else:
+                # Calculate token generation timestamp using
+                # start_time, ttft, and itl
+                token_times = [output.start_time + output.ttft]
+                current_time = token_times[0]
+                for itl_value in output.itl:
+                    current_time += itl_value
+                    token_times.append(current_time)
 
-            # Add tokens to second buckets
-            for token_time in token_times:
-                second_bucket = int(token_time - min_start_time)
-                if 0 <= second_bucket < duration_seconds:
-                    tokens_per_second[second_bucket] += 1
+                # Add tokens to second buckets
+                for token_time in token_times:
+                    second_bucket = int(token_time - min_start_time)
+                    if 0 <= second_bucket < duration_seconds:
+                        tokens_per_second[second_bucket] += 1
 
             # Track concurrent requests for each second this request was active
             request_start_second = int(output.start_time - min_start_time)
@@ -860,18 +907,22 @@ def calculate_metrics(
         # Find the maximum tokens per second and corresponding
         # concurrent requests
         if len(tokens_per_second) > 0:
-            max_output_tokens_per_s = float(np.max(tokens_per_second))
+            if token_timeline_available:
+                max_output_tokens_per_s = float(np.max(tokens_per_second))
+            else:
+                max_output_tokens_per_s = float("nan")
             max_concurrent_requests = int(np.max(concurrent_requests_per_second))
 
         if TERM_PLOTLIB_AVAILABLE:
             import termplotlib as tpl
 
             fig = tpl.figure()
-            fig.plot(
-                np.arange(len(tokens_per_second)),
-                tokens_per_second,
-                title="Output tokens per second",
-            )
+            if token_timeline_available:
+                fig.plot(
+                    np.arange(len(tokens_per_second)),
+                    tokens_per_second,
+                    title="Output tokens per second",
+                )
             fig.plot(
                 np.arange(len(concurrent_requests_per_second)),
                 concurrent_requests_per_second,
@@ -881,6 +932,10 @@ def calculate_metrics(
         else:
             print("tip: install termplotlib and gnuplot to plot the metrics")
 
+    duplex_metrics_present = any(
+        isinstance(getattr(output, "duplex_session_metrics", None), dict) for output in outputs
+    )
+    missing_duplex_value = float("nan") if duplex_metrics_present else 0
     metrics = MultiModalsBenchmarkMetrics(
         completed=completed,
         failed=len(failed_outputs),
@@ -890,18 +945,18 @@ def calculate_metrics(
         request_goodput=good_completed / dur_s,
         output_throughput=sum(actual_output_lens) / dur_s,
         total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
-        mean_ttft_ms=np.mean(ttfts or 0) * 1000,  # ttfts is empty if streaming is not supported by the endpoint
-        std_ttft_ms=np.std(ttfts or 0) * 1000,
-        median_ttft_ms=np.median(ttfts or 0) * 1000,
-        percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
-        mean_tpot_ms=np.mean(tpots or 0) * 1000,
-        std_tpot_ms=np.std(tpots or 0) * 1000,
-        median_tpot_ms=np.median(tpots or 0) * 1000,
-        percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
-        mean_itl_ms=np.mean(itls or 0) * 1000,
-        std_itl_ms=np.std(itls or 0) * 1000,
-        median_itl_ms=np.median(itls or 0) * 1000,
-        percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
+        mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000,
+        std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000,
+        median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000,
+        percentiles_ttft_ms=[(p, np.percentile(ttfts or missing_duplex_value, p) * 1000) for p in selected_percentiles],
+        mean_tpot_ms=np.mean(tpots or missing_duplex_value) * 1000,
+        std_tpot_ms=np.std(tpots or missing_duplex_value) * 1000,
+        median_tpot_ms=np.median(tpots or missing_duplex_value) * 1000,
+        percentiles_tpot_ms=[(p, np.percentile(tpots or missing_duplex_value, p) * 1000) for p in selected_percentiles],
+        mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000,
+        std_itl_ms=np.std(itls or missing_duplex_value) * 1000,
+        median_itl_ms=np.median(itls or missing_duplex_value) * 1000,
+        percentiles_itl_ms=[(p, np.percentile(itls or missing_duplex_value, p) * 1000) for p in selected_percentiles],
         mean_e2el_ms=np.mean(e2els or 0) * 1000,
         std_e2el_ms=np.std(e2els or 0) * 1000,
         median_e2el_ms=np.median(e2els or 0) * 1000,
@@ -910,11 +965,16 @@ def calculate_metrics(
         max_concurrent_requests=max_concurrent_requests,
         rtfx=input_audio_duration / dur_s,
         **{
-            defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
-            defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
-            defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
+            "num_ttft_samples": len(ttfts),
+            "num_tpot_samples": len(tpots),
+            "num_itl_samples": len(itls),
+            "num_audio_ttfp_samples": len(audio_ttfps),
+            "num_audio_rtf_samples": len(audio_rtfs),
+            defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or missing_duplex_value) * 1000,
+            defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or missing_duplex_value) * 1000,
+            defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or missing_duplex_value) * 1000,
             defs.PERCENTILES_AUDIO_TTFP_MS: [
-                (p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles
+                (p, np.percentile(audio_ttfps or missing_duplex_value, p) * 1000) for p in selected_percentiles
             ],
             defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
             defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0),
@@ -925,10 +985,12 @@ def calculate_metrics(
             defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration),
             defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
             defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
-            defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
-            defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
-            defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
-            defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
+            defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value),
+            defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value),
+            defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value),
+            defs.PERCENTILES_AUDIO_RTF: [
+                (p, np.percentile(audio_rtfs or missing_duplex_value, p)) for p in selected_percentiles
+            ],
             defs.TOTAL_IMAGES: total_images,
             defs.IMAGE_THROUGHPUT: total_images / dur_s,
             defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,