Revert "[Bugfix] Restore name-based model detection for HF cache snapshot pat…"

This reverts commit 7ba5e8a146.

.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-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/

.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()

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/configuration/environment_variables.md

@@ -82,6 +82,7 @@ depends on the installed kernels and model path.
 | --- | --- | --- | --- | --- |
 | `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 |
@@ -177,7 +178,7 @@ collection, examples, bug reports, or logs.
 
 ## Model-specific variables
 
-The audit found 57 variables read by a single model or pipeline family. They are
+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.
 
@@ -190,7 +191,7 @@ Every audited model-specific name has a migration disposition in the
 | 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 | 7 | Remove a compatibility escape hatch that has no continuing contract. |
+| 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

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/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/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/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_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"
-    ``````

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

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

@@ -164,10 +164,10 @@ 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: 22,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 23,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
-        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
+        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
     }
 
 

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

@@ -63,6 +63,7 @@ _PUBLIC_OMNI = (
     "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
     "SPEAKER_MAX_UPLOADED",
     "SPEAKER_SAMPLES_DIR",
+    "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT",
     "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
     "VLLM_OMNI_ORCH_MONITOR_PATH",
     "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
@@ -197,7 +198,6 @@ _MODEL_INTERNALIZE = (
 
 _MODEL_DEPRECATE_REMOVE = (
     "HIGGS_AUDIO_V2_TOKENIZER_PATH",
-    "ROSVOT_SOURCE_DIR",
     "VLLM_OMNI_LINGBOT_ACTION_ROOT",
     "VLLM_OMNI_VOXCPM_CODE_PATH",
     "VLLM_USE_FLASHINFER_MOE_FP16",

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -3,8 +3,9 @@
 
 from __future__ import annotations
 
+import time
 from collections import defaultdict
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 from typing import Any
 
 import numpy as np
@@ -107,6 +108,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # Cache per-request flag to avoid repeated deserialization of additional_information
         self._omits_kv_transfer_cache: dict[str, bool] = {}
+
+        # KV-wait start ts for the vllm_omni:kv_wait_s metric; see
+        # _emit_kv_wait_output for the engine-core → orchestrator carry.
+        self._kv_wait_start_ts: dict[str, float] = {}
         self._init_omni_io_scheduling_state()
         # Snapshot prompt length for each streaming input update
         self._new_prompt_len_snapshot: dict[str, int] = {}
@@ -116,6 +121,70 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Output placeholders are zero when async scheduling isn't used
         return request.num_computed_tokens - request.num_output_placeholders
 
+    def _resolve_kv_connector_type(self) -> str:
+        """Connector backend name for the ``kv_wait_s`` label, or ``unknown``."""
+        connector_cfg = self._get_omni_kv_config_value("connector_config")
+        if isinstance(connector_cfg, dict):
+            c_type = connector_cfg.get("type")
+            if isinstance(c_type, str) and c_type:
+                return c_type
+        return "unknown"
+
+    def _emit_kv_wait_output(
+        self,
+        outputs: dict[int, list[OmniEngineCoreOutput]],
+        req_id: str,
+        req: Request,
+    ) -> None:
+        """Carry the KV-wait duration for ``req_id`` across the process boundary.
+
+        Pops the start ts recorded at ENTER (``_free_request``); skips silently
+        when none was recorded. The wait rides ``kv_transfer_params`` to the
+        orchestrator, which calls ``observe_kv_wait``.
+        """
+        start_ts = self._kv_wait_start_ts.pop(req_id, None)
+        if start_ts is None:
+            return
+        kv_wait_s = max(time.monotonic() - start_ts, 0.0)
+        outputs.setdefault(req.client_index, []).append(
+            OmniEngineCoreOutput(
+                request_id=req_id,
+                new_token_ids=[],
+                kv_transfer_params={
+                    "kv_wait_s": kv_wait_s,
+                    "connector_type": self._resolve_kv_connector_type(),
+                },
+            )
+        )
+
+    def _clear_kv_wait_starts(self, request_ids: Iterable[str]) -> None:
+        """Drop incomplete KV-wait measurements for terminal requests."""
+        wait_starts = getattr(self, "_kv_wait_start_ts", None)
+        if wait_starts is None:
+            return
+        for request_id in request_ids:
+            wait_starts.pop(request_id, None)
+
+    def finish_requests(
+        self,
+        request_ids: str | Iterable[str] | None,
+        finished_status: RequestStatus,
+    ) -> list[Request]:
+        """Finish requests and discard any incomplete KV-wait timing."""
+        if isinstance(request_ids, str):
+            cleanup_ids = (request_ids,)
+            finish_request_ids: str | tuple[str, ...] | None = request_ids
+        elif request_ids is None:
+            cleanup_ids = ()
+            finish_request_ids = None
+        else:
+            finish_request_ids = list(request_ids) if isinstance(request_ids, Iterator) else request_ids
+            cleanup_ids = tuple(finish_request_ids)
+
+        finished = super().finish_requests(finish_request_ids, finished_status)
+        self._clear_kv_wait_starts(cleanup_ids)
+        return finished
+
     def _get_kv_transfer_criteria(self) -> dict | None:
         return self._get_omni_kv_config_value("kv_transfer_criteria")
 
@@ -226,7 +295,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
-        for queue in (self.waiting, self.running):
+        waiting = getattr(self, "waiting")
+        running = getattr(self, "running")
+        for queue in (waiting, running):
             for req in list(queue):
                 if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
                     queue.remove(req)
@@ -234,7 +305,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         original_waiting = None
         if self._should_defer_waiting_admission():
-            original_waiting = self.waiting
+            original_waiting = waiting
             self.waiting = create_request_queue(self.policy)
 
         try:
@@ -309,6 +380,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                                 kv_transfer_params={"kv_ready": True},
                             )
                         )
+                    # EXIT: extraction ack ends the wait — carry the duration.
+                    # req is still alive here (the del happens in the cleanup
+                    # loop below); skip if it was already freed upstream.
+                    if req is not None:
+                        self._emit_kv_wait_output(outputs, req_id, req)
                 except Exception:
                     init_logger(__name__).exception("Failed to pre-process KV extraction for %s", req_id)
 
@@ -505,8 +581,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     # valid new-segment frames or underflowed the drain
                     # assert whenever placeholder counts diverged from
                     # scheduled counts (spec drafts, in-flight prefill).
+                    # Assign, never accumulate: _handle_stopped_request above
+                    # may already have fenced the same in-flight tokens through
+                    # a queued streaming update, and adding twice swallows the
+                    # next duplex unit's listen/speak under async scheduling.
                     if request.num_in_flight_tokens > 0:
-                        request.num_stale_output_tokens += request.num_in_flight_tokens
+                        request.num_stale_output_tokens = request.num_in_flight_tokens
                     if outstanding_async_tokens > 0:
                         # Discard only outputs that are already in flight and
                         # roll back their optimistic computed-token accounting.
@@ -642,10 +722,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # pre-replacement frame will drain, so the counter reaches exactly
         # zero and the new segment's frames are never swallowed. This also
         # covers an in-flight prefill chunk, which carries no placeholders
-        # but must still have its late output dropped.
+        # but must still have its late output dropped. Assign, never
+        # accumulate: update_from_output fences the same in-flight tokens when
+        # a resumable stop applies a queued update through this helper.
         in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
         if in_flight_tokens > 0:
-            session.num_stale_output_tokens += in_flight_tokens
+            session.num_stale_output_tokens = in_flight_tokens
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous
             # segment's next token. Drop that late token instead of
@@ -733,6 +815,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         # It triggered but hasn't finished yet. We MUST wait.
                         logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
                         self.waiting_for_transfer_free.add(request_id)
+                        self._kv_wait_start_ts[request_id] = time.monotonic()
                         kv_xfer_params = None
                         return kv_xfer_params, None
                     elif request_id in self.waiting_for_transfer_free:
@@ -745,6 +828,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         )
                 else:
                     self.waiting_for_transfer_free.add(request_id)
+                    self._kv_wait_start_ts[request_id] = time.monotonic()
                     confirmed_computed = self._get_confirmed_num_computed_tokens(request)
                     self._mark_request_for_kv_transfer(request_id, confirmed_computed)
                     # Return KV transfer metadata so it propagates to RequestOutput

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -636,8 +636,8 @@ class OmniSchedulerMixin:
             )
         }
 
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
+        if chunk_transfer_adapter := getattr(self, "chunk_transfer_adapter", None):
+            chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
         self._realign_request_status_to_queues(
             request_ids,

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -1,6 +1,13 @@
 # MiniCPM-o 4.5 three-stage deploy on one large-memory GPU.
 pipeline: minicpmo_4_5
 async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+active_stream_window: 1
+duplex_session:
+  # Matches Stage 0/1 max_num_seqs below.
+  max_sessions: 4
 
 trust_remote_code: true
 enable_prefix_caching: false

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -4,6 +4,13 @@
 # isolated from codec/vocoder contention while preserving stage-2 batching.
 pipeline: minicpmo_4_5
 async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+active_stream_window: 1
+duplex_session:
+  # Matches Stage 0/1 max_num_seqs below.
+  max_sessions: 4
 
 trust_remote_code: true
 enable_prefix_caching: false

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -4,6 +4,13 @@
 # Code2Wav run independently on GPUs 0, 1, and 2.
 pipeline: minicpmo_4_5
 async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+active_stream_window: 1
+duplex_session:
+  # Matches Stage 0/1 max_num_seqs below.
+  max_sessions: 4
 
 trust_remote_code: true
 enable_prefix_caching: false

vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml

@@ -5,7 +5,6 @@
 # Stage 1: replicas on GPU 1 and 2
 base_config: minicpmo_4_5.yaml
 pipeline: minicpmo_4_5
-session_mode: duplex
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml

@@ -5,7 +5,6 @@
 # Stage 1: replicas on GPU 1, 2, and 3
 base_config: minicpmo_4_5.yaml
 pipeline: minicpmo_4_5
-session_mode: duplex
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_8x4090.yaml

@@ -15,6 +15,10 @@
 # so 4096 is the empirically-verified safe ceiling.
 pipeline: minicpmo_4_5
 async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. The default
+# duplex_session caps sessions at 1, matching Stage 0/1 max_num_seqs below.
+session_mode: duplex
+active_stream_window: 1
 
 trust_remote_code: true
 enable_prefix_caching: false

vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml

@@ -5,7 +5,6 @@
 # Stage 1: replicas on GPU 4, 5, 6, and 7
 base_config: minicpmo_4_5_8x4090.yaml
 pipeline: minicpmo_4_5
-session_mode: duplex
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_duplex.yaml

@@ -1,42 +0,0 @@
-# MiniCPM-o 4.5 native duplex deployment.
-#
-# The registered minicpmo_4_5 pipeline owns the fixed
-# Thinker -> Talker -> Code2Wav topology. This overlay changes session
-# lifetime and bounds all three stages to two live duplex sessions.
-base_config: minicpmo_4_5.yaml
-pipeline: minicpmo_4_5
-session_mode: duplex
-active_stream_window: 1
-duplex_session:
-  idle_ttl_s: 300
-  disconnect_grace_s: 30
-  reaper_interval_s: 5
-  resume_replay_ttl_s: 60
-  resume_replay_max_bytes_per_session: 8388608
-  max_pending_input_bytes_per_session: 16777216
-  max_pending_turns_per_session: 4
-  # Matches the configured Stage0/Stage1 sequence capacity below.
-  max_sessions: 2
-  completed_append_cache_size: 256
-
-stages:
-  - stage_id: 0
-    max_num_seqs: 2
-    async_scheduling: false
-
-  - stage_id: 1
-    max_num_seqs: 2
-    async_scheduling: false
-    # Talker context is 4096 (tts_config.max_position_embeddings).
-    max_model_len: 4096
-    # Overlay only the knobs duplex changes. Codec temperature / top_k stay
-    # on the base YAML. Native duplex still ends each Talker request after
-    # one generate_chunk (26 samples, last is codec EOS) inside the model;
-    # YAML max_tokens is only the hard ceiling.
-    default_sampling_params:
-      max_tokens: 4096
-      min_tokens: 0
-
-  - stage_id: 2
-    max_num_seqs: 2
-    async_scheduling: false

vllm_omni/diffusion/data.py

@@ -1528,6 +1528,9 @@ class DiffusionOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
+    # KV-recv wall-clock (ms) from the runner; feeds vllm_omni:diffusion_kv_load_s.
+    kv_recv_ms: float = 0.0
+
     # When True, move tensor fields to CPU at construction time. Useful when
     # the output is shipped across process boundaries (e.g. step-execution
     # mode) and the receiving side must not initialise a stray CUDA context.

vllm_omni/diffusion/diffusion_engine.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
 
@@ -56,6 +56,11 @@ from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
 from vllm_omni.errors import client_error_from_metadata, is_client_error_status
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.metrics.utils import (
+    diffusion_scheduler_waiting_metrics,
+    extract_diffusion_denoise_ms,
+    extract_diffusion_vae_decode_ms,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.outputs import OmniRequestOutput
@@ -375,6 +380,9 @@ class DiffusionEngine:
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
+        # Copied onto the existing output metrics payload so queue monitoring
+        # reuses the normal diffusion result path without additional IPC.
+        self._scheduler_num_waiting_reqs = 0
 
     def _init_execute_fn(self) -> None:
         if self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
@@ -444,7 +452,14 @@ class DiffusionEngine:
                     )
                     raise
             postprocess_start_time = time.perf_counter()
-            formatted_outputs = self.postprocess_output(request, output)
+            scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0))
+            try:
+                formatted_outputs = self.postprocess_output(request, output)
+            except Exception as exc:
+                # Preserve the latest scheduler snapshot across terminal
+                # abort/error paths, which do not produce formatted outputs.
+                setattr(exc, "diffusion_metrics", scheduler_metrics)
+                raise
             postprocess_time = time.perf_counter() - postprocess_start_time
             step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
             logger.debug(
@@ -456,14 +471,24 @@ class DiffusionEngine:
                 step_total_ms,
             )
             for request_output in formatted_outputs:
-                request_output.metrics.update(
-                    {
-                        "preprocess_time_ms": preprocess_time * 1000,
-                        "diffusion_engine_exec_time_ms": exec_total_time * 1000,
-                        "diffusion_engine_total_time_ms": step_total_ms,
-                        "postprocess_time_ms": postprocess_time * 1000,
-                    }
-                )
+                metrics_update = {
+                    "preprocess_time_ms": preprocess_time * 1000,
+                    "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+                    "postprocess_time_ms": postprocess_time * 1000,
+                    **scheduler_metrics,
+                }
+                if request.scheduler_queue_wait_ms is not None:
+                    metrics_update["scheduler_queue_wait_ms"] = request.scheduler_queue_wait_ms
+                vae_decode_ms = extract_diffusion_vae_decode_ms(output)
+                if vae_decode_ms is not None:
+                    metrics_update["vae_decode_time_ms"] = vae_decode_ms
+                forward_ms = extract_diffusion_denoise_ms(output)
+                if forward_ms is not None:
+                    metrics_update["forward_time_ms"] = forward_ms
+                kv_recv_ms = getattr(output, "kv_recv_ms", 0.0)
+                if kv_recv_ms > 0:
+                    metrics_update["kv_recv_time_ms"] = kv_recv_ms
+                request_output.metrics.update(metrics_update)
             yield formatted_outputs
 
     async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
@@ -554,6 +579,7 @@ class DiffusionEngine:
                 self._wait_for_admission_if_needed_locked()
 
                 sched_output = self.scheduler.schedule()
+                self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0)
 
             if sched_output.is_empty:
                 self._emit_finished_outputs(sched_output.finished_req_ids, None)

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -235,8 +235,8 @@ class CFGParallelMixin(metaclass=ABCMeta):
         Predict noise with N-branch CFG dispatch across M GPUs.
 
         This is the multi-branch counterpart of predict_noise_maybe_with_cfg().
-        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel,
-        DreamID). Existing 2-branch models should continue using
+        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel).
+        Existing 2-branch models should continue using
         predict_noise_maybe_with_cfg().
 
         Args:

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime.
 
 Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process
@@ -23,6 +26,7 @@ from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.errors import client_error_metadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
+from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID, diffusion_exception_metrics
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -158,6 +162,15 @@ class InlineStageDiffusionClient(StageClientBase):
                 self._output_queue.put_nowait(result)
         except DiffusionRequestAbortedError as e:
             logger.info("request_id: %s aborted: %s", request_id, str(e))
+            metrics = diffusion_exception_metrics(e)
+            if metrics:
+                self._output_queue.put_nowait(
+                    OmniRequestOutput(
+                        request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID,
+                        finished=True,
+                        metrics=metrics,
+                    )
+                )
         except Exception as e:
             logger.exception("Diffusion request %s failed: %s", request_id, e)
             status_code, error_type = client_error_metadata(e)
@@ -167,6 +180,7 @@ class InlineStageDiffusionClient(StageClientBase):
                 status_code=status_code,
                 error_type=error_type,
             )
+            error_output.metrics.update(diffusion_exception_metrics(e))
             self._output_queue.put_nowait(error_output)
         finally:
             self._tasks.pop(request_id, None)
@@ -181,9 +195,8 @@ class InlineStageDiffusionClient(StageClientBase):
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
         for rid in request_ids:
-            task = self._tasks.pop(rid, None)
-            if task:
-                task.cancel()
+            # Keep the consumer alive until the engine emits its terminal
+            # abort output. That path carries the refreshed scheduler gauge.
             self._engine.abort(rid)
 
     async def submit_interaction_async(

vllm_omni/diffusion/ipc.py

@@ -158,7 +158,7 @@ def _pack_value_if_large(
 
     Walks the container shapes pipelines return as ``DiffusionOutput.output``:
     bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and
-    tuples/lists (e.g. LTX2 and DreamID ``(video, audio)``). Other values pass
+    tuples/lists (e.g. LTX2 ``(video, audio)``). Other values pass
     through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep
     the two in sync.
     """

vllm_omni/diffusion/layers/norm.py

@@ -85,7 +85,10 @@ class RMSNorm(CustomOp):
     def forward_cuda(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            return self.forward_native(x, residual)
         # During torch.compile tracing, fused_rms_norm writes to `out` in-place
         # (returns None) and accesses self.weight.data, which is a DTensor under
         # HSDP. Both patterns confuse inductor's compute_ancestors scheduler.
@@ -101,7 +104,10 @@ class RMSNorm(CustomOp):
     def forward_hip(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            return self.forward_native(x, residual)
         if torch.compiler.is_compiling():
             return self.forward_native(x)
         try:
@@ -112,7 +118,10 @@ class RMSNorm(CustomOp):
     def forward_musa(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            return self.forward_native(x, residual)
         # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
         # H3 Q/K norm with its inline RoPE path on MUSA.
         return F.rms_norm(
@@ -125,9 +134,19 @@ class RMSNorm(CustomOp):
     def forward_npu(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         import torch_npu
 
+        if residual is not None:
+            output, _, updated_residual = torch_npu.npu_add_rms_norm(
+                x,
+                residual,
+                self.weight,
+                self.variance_epsilon,
+            )
+            return output, updated_residual
+
         output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0]
 
         return output
@@ -135,19 +154,27 @@ class RMSNorm(CustomOp):
     def forward_xpu(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.forward_native(x)
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            return self.forward_native(x)
+        return self.forward_native(x, residual)
 
     def forward_native(
         self,
         x: torch.Tensor,
-    ) -> torch.Tensor:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            residual = residual + x
+            x = residual
         input_dtype = x.dtype
         x = x.to(torch.float32)
         variance = x.pow(2).mean(-1, keepdim=True)
         out = x * torch.rsqrt(variance + self.variance_epsilon)
         out = self.weight.to(torch.float32) * out
-        return out.to(input_dtype)
+        out = out.to(input_dtype)
+        return (out, residual) if residual is not None else out
 
 
 class RMSNormVAE(CustomOp):

vllm_omni/diffusion/models/dreamid_omni/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-__all__ = [
-    "DreamIDOmniPipeline",
-]
-
-
-def __getattr__(name: str):
-    if name == "DreamIDOmniPipeline":
-        from .pipeline_dreamid_omni import DreamIDOmniPipeline
-
-        return DreamIDOmniPipeline
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -1,460 +0,0 @@
-import re
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-from cache_dit import ForwardPattern
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import ColumnParallelLinear
-
-from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-try:
-    from dreamid_omni.modules.model import WanLayerNorm
-except ImportError:
-    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
-
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import (
-    DistributedRMSNorm,
-    WanModel,
-    rope_apply,
-)
-
-logger = init_logger(__name__)
-
-
-class FusedBlock(nn.Module):
-    """Wrapper pairing a video block and audio block for layerwise offloading.
-
-    Registers both blocks as submodules so their parameters are visible to the offload hooks.
-    """
-
-    def __init__(
-        self,
-        vid_block: nn.Module,
-        audio_block: nn.Module,
-        device: torch.device,
-    ):
-        super().__init__()
-        self.vid_block = vid_block
-        self.audio_block = audio_block
-        self.device = device
-
-    def _cross_attention_forward(
-        self,
-        attn: Attention,
-        cross_attn_block,
-        src_seq,
-        src_grid_sizes,
-        src_freqs,
-        target_seq,
-        target_seq_lens,
-        target_grid_sizes,
-        target_freqs,
-        context,
-        context_lens,
-        src_ref_lengths=None,
-        src_freqs_scaling=None,
-        target_ref_lengths=None,
-        target_freqs_scaling=None,
-    ):
-        b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim
-        if hasattr(cross_attn_block, "k_img"):
-            q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context)
-        else:
-            q, k, v = cross_attn_block.qkv_fn(src_seq, context)
-            k_img = v_img = None
-
-        x = attn(q, k, v)
-
-        if k_img is not None:
-            img_x = attn(q, k_img, v_img)
-            x = x + img_x
-
-        target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq)
-        k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d)
-        v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d)
-
-        q = rope_apply(q, src_grid_sizes, src_freqs, ref_lengths=src_ref_lengths, freqs_scaling=src_freqs_scaling)
-        k_target = rope_apply(
-            k_target,
-            target_grid_sizes,
-            target_freqs,
-            ref_lengths=target_ref_lengths,
-            freqs_scaling=target_freqs_scaling,
-        )
-
-        target_x = attn(q, k_target, v_target)
-
-        x = x + target_x
-        x = x.flatten(2)
-        x = cross_attn_block.o(x)
-        return x
-
-    def _cross_attention_ffn_forward(
-        self,
-        attn: Attention,
-        attn_block,
-        src_seq,
-        src_grid_sizes,
-        src_freqs,
-        target_seq,
-        target_seq_lens,
-        target_grid_sizes,
-        target_freqs,
-        context,
-        context_lens,
-        src_e,
-        src_ref_lengths=None,
-        src_freqs_scaling=None,
-        target_ref_lengths=None,
-        target_freqs_scaling=None,
-    ):
-        src_seq = src_seq + self._cross_attention_forward(
-            attn,
-            attn_block.cross_attn,
-            attn_block.norm3(src_seq),
-            src_grid_sizes=src_grid_sizes,
-            src_freqs=src_freqs,
-            target_seq=target_seq,
-            target_seq_lens=target_seq_lens,
-            target_grid_sizes=target_grid_sizes,
-            target_freqs=target_freqs,
-            context=context,
-            context_lens=context_lens,
-            src_ref_lengths=src_ref_lengths,
-            src_freqs_scaling=src_freqs_scaling,
-            target_ref_lengths=target_ref_lengths,
-            target_freqs_scaling=target_freqs_scaling,
-        )
-        y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2))
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-            src_seq = src_seq + y * src_e[5].squeeze(2)
-        return src_seq
-
-    def forward(
-        self,
-        hidden_states,
-        encoder_hidden_states,
-        attn: Attention,
-        vid_e,
-        vid_seq_lens,
-        vid_grid_sizes,
-        vid_freqs,
-        vid_context,
-        vid_context_lens,
-        vid_ref_lengths,
-        vid_freqs_scaling,
-        audio_e,
-        audio_seq_lens,
-        audio_grid_sizes,
-        audio_freqs,
-        audio_context,
-        audio_context_lens,
-        audio_ref_lengths,
-        audio_freqs_scaling,
-    ):
-        vid = hidden_states
-        audio = encoder_hidden_states
-        vid_block = self.vid_block
-        audio_block = self.audio_block
-
-        ## audio modulation
-        assert audio_e.dtype == torch.bfloat16
-        assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], (
-            f"{audio_e.shape}, {audio.shape}"
-        )
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-            audio_e = audio_block.modulation(audio_e).chunk(6, dim=2)
-        assert audio_e[0].dtype == torch.bfloat16
-
-        # audio self-attention
-        audio_y = audio_block.self_attn(
-            audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2),
-            audio_seq_lens,
-            audio_grid_sizes,
-            audio_freqs,
-            ref_lengths=audio_ref_lengths,
-            freqs_scaling=audio_freqs_scaling,
-        )
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-            audio = audio + audio_y * audio_e[2].squeeze(2)
-
-        ## video modulation
-        assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], (
-            f"{vid_e.shape}, {vid.shape}"
-        )
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-            vid_e = vid_block.modulation(vid_e).chunk(6, dim=2)
-
-        # video self-attention
-        vid_y = vid_block.self_attn(
-            vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2),
-            vid_seq_lens,
-            vid_grid_sizes,
-            vid_freqs,
-            ref_lengths=vid_ref_lengths,
-        )
-
-        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-            vid = vid + vid_y * vid_e[2].squeeze(2)
-
-        og_audio = audio
-
-        # audio cross-attention
-        audio = self._cross_attention_ffn_forward(
-            attn,
-            audio_block,
-            audio,
-            audio_grid_sizes,
-            audio_freqs,
-            vid,
-            vid_seq_lens,
-            vid_grid_sizes,
-            vid_freqs,
-            audio_context,
-            audio_context_lens,
-            audio_e,
-            src_ref_lengths=audio_ref_lengths,
-            src_freqs_scaling=audio_freqs_scaling,
-            target_ref_lengths=vid_ref_lengths,
-            target_freqs_scaling=None,
-        )
-
-        assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!"
-
-        # video cross-attention
-        vid = self._cross_attention_ffn_forward(
-            attn,
-            vid_block,
-            vid,
-            vid_grid_sizes,
-            vid_freqs,
-            og_audio,
-            audio_seq_lens,
-            audio_grid_sizes,
-            audio_freqs,
-            vid_context,
-            vid_context_lens,
-            vid_e,
-            src_ref_lengths=vid_ref_lengths,
-            src_freqs_scaling=None,
-            target_ref_lengths=audio_ref_lengths,
-            target_freqs_scaling=audio_freqs_scaling,
-        )
-
-        return vid, audio
-
-
-class FusionModel(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "fused_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-    )
-
-    _layerwise_offload_blocks_attrs = ["fused_blocks"]
-
-    packed_modules_mapping = {
-        "to_qkv": ["q", "k", "v"],
-    }
-
-    @staticmethod
-    def _is_fused_block(name: str, module) -> bool:
-        return "fused_blocks" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_fused_block]
-
-    def __init__(
-        self,
-        video_config=None,
-        audio_config=None,
-        *,
-        quant_config: "QuantizationConfig | None" = None,
-    ):
-        super().__init__()
-        has_video = True
-        has_audio = True
-        self.device = get_local_device()
-        if video_config is not None:
-            self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config)
-        else:
-            has_video = False
-            self.video_model = None
-            logger.warning("No video model is provided!")
-
-        if audio_config is not None:
-            self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config)
-        else:
-            has_audio = False
-            self.audio_model = None
-            logger.warning("No audio model is provided!")
-
-        if has_video and has_audio:
-            assert len(self.video_model.blocks) == len(self.audio_model.blocks)
-            self.num_blocks = len(self.video_model.blocks)
-
-            self.inject_cross_attention_kv_projections(quant_config=quant_config)
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.full_num_heads = self.video_model.num_heads
-        self.num_heads = self.full_num_heads // tp_size
-        self.head_dim = self.video_model.dim // self.video_model.num_heads
-        self.attn = Attention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=self.num_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-        )
-
-        if has_video and has_audio:
-            self.fused_blocks = nn.ModuleList(
-                [
-                    FusedBlock(
-                        self.video_model.blocks[i],
-                        self.audio_model.blocks[i],
-                        self.device,
-                    )
-                    for i in range(self.num_blocks)
-                ]
-            )
-
-    def load_state_dict(self, state_dict, strict=True, assign=False):
-        """Remap checkpoints where blocks are stored under
-        `video_model.blocks.N.*` / `audio_model.blocks.N.*` to the current
-        `fused_blocks.N.vid_block.*` / `fused_blocks.N.audio_block.*`.
-        """
-        needs_remap = any(re.match(r"^(video_model|audio_model)\.blocks\.\d+\.", k) for k in state_dict)
-        if needs_remap:
-            remapped = {}
-            for k, v in state_dict.items():
-                new_k = re.sub(r"^video_model\.blocks\.(\d+)\.", r"fused_blocks.\1.vid_block.", k)
-                new_k = re.sub(r"^audio_model\.blocks\.(\d+)\.", r"fused_blocks.\1.audio_block.", new_k)
-                remapped[new_k] = v
-            state_dict = remapped
-
-        self._detach_blocks_from_backbones()
-
-        return super().load_state_dict(state_dict, strict=strict, assign=assign)
-
-    def inject_cross_attention_kv_projections(
-        self,
-        quant_config: "QuantizationConfig | None" = None,
-    ):
-        def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear:
-            # K/V projections for the cross-modal (a2v / v2a) attention path.
-            # Inputs are full latents, so we quantize matched to self_attn.
-            return ColumnParallelLinear(
-                dim,
-                dim,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=prefix,
-            )
-
-        tp_size = get_tensor_model_parallel_world_size()
-        for i, vid_block in enumerate(self.video_model.blocks):
-            cross = vid_block.cross_attn
-            block_prefix = f"video_model.blocks.{i}.cross_attn"
-            cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion")
-            cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion")
-            cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
-            cross.norm_k_fusion = (
-                DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity()
-            )
-
-        for i, audio_block in enumerate(self.audio_model.blocks):
-            cross = audio_block.cross_attn
-            block_prefix = f"audio_model.blocks.{i}.cross_attn"
-            cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion")
-            cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion")
-            cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
-            cross.norm_k_fusion = (
-                DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity()
-            )
-
-    def _detach_blocks_from_backbones(self) -> None:
-        """Keep offloadable blocks owned only by a single place.
-
-        NOTE: This is a special workaround to support layerwise offloading.
-        The model registers the same Wan blocks under both the video/audio
-        backbones and `fused_blocks` which is a wrapper for unified blocks
-        walking through. However, layerwise offloading will only consider
-        `fused_blocks` as offloadable components and will materialize all
-        other modules onto device, including the same blocks owned by both
-        `fused_blocks` and `video_model` and `audio_model`.
-        """
-        video_blocks = list(self.video_model.blocks)
-        audio_blocks = list(self.audio_model.blocks)
-        self.video_model._modules.pop("blocks", None)
-        self.audio_model._modules.pop("blocks", None)
-        self.video_model.blocks = tuple(video_blocks)
-        self.audio_model.blocks = tuple(audio_blocks)
-
-    def merge_kwargs(self, vid_kwargs, audio_kwargs):
-        """
-        keys in each kwarg:
-        e
-        seq_lens
-        grid_sizes
-        freqs
-        context
-        context_lens
-        """
-        merged_kwargs = {}
-        for key in vid_kwargs:
-            merged_kwargs[f"vid_{key}"] = vid_kwargs[key]
-        for key in audio_kwargs:
-            merged_kwargs[f"audio_{key}"] = audio_kwargs[key]
-        return merged_kwargs
-
-    def forward(
-        self,
-        vid,
-        audio,
-        t,
-        vid_context,
-        audio_context,
-        vid_seq_len,
-        audio_seq_len,
-        ref_ip_lengths=None,
-        ref_audio_lengths=None,
-        slg_layer=False,
-        freqs_scaling=None,
-    ):
-        vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs(
-            x=vid, t=t, context=vid_context, seq_len=vid_seq_len, ref_lengths=ref_ip_lengths
-        )
-
-        audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs(
-            x=audio,
-            t=t,
-            context=audio_context,
-            seq_len=audio_seq_len,
-            ref_lengths=ref_audio_lengths,
-            freqs_scaling=freqs_scaling,
-        )
-
-        kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs)
-
-        for fused_block in self.fused_blocks:
-            vid, audio = fused_block(vid, audio, self.attn, **kwargs)
-
-        vid = self.video_model.post_transformer_block_out(vid, vid_kwargs["grid_sizes"], vid_e)
-        audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs["grid_sizes"], audio_e)
-
-        return vid, audio
-
-    def set_rope_params(self):
-        self.video_model.set_rope_params()
-        self.audio_model.set_rope_params()

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -1,608 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import logging
-import math
-import os
-from collections.abc import Iterable
-from typing import ClassVar
-
-import torch
-import torch.distributed
-from diffusers import FlowMatchEulerDiscreteScheduler
-from PIL import Image, ImageOps
-from torch import nn
-from torchvision.transforms import Compose, Normalize
-from tqdm import tqdm
-from vllm.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-
-try:
-    from dreamid_omni.utils.divisible_crop import DivisibleCrop
-    from dreamid_omni.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
-    from dreamid_omni.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
-    from dreamid_omni.utils.model_loading_utils import (
-        init_mmaudio_vae,
-        init_text_model,
-        init_wan_vae_2_2,
-    )
-    from dreamid_omni.utils.rearrange import Rearrange
-    from dreamid_omni.utils.resize import NaResize
-except ImportError:
-    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
-
-logger = logging.getLogger(__name__)
-
-
-def get_dreamid_omni_post_process_func(*args, **kwargs):
-    def post_process(output):
-        if isinstance(output, tuple) and len(output) == 2:
-            video, audio = output
-            return {
-                "video": video,
-                "audio": audio,
-                "audio_sample_rate": 16000,
-                "fps": 24,
-            }
-        return output
-
-    return post_process
-
-
-AUDIO_CONFIG = {
-    "patch_size": [1],
-    "model_type": "t2a",
-    "dim": 3072,
-    "ffn_dim": 14336,
-    "freq_dim": 256,
-    "num_heads": 24,
-    "num_layers": 30,
-    "in_dim": 20,
-    "out_dim": 20,
-    "text_len": 512,
-    "window_size": [-1, -1],
-    "qk_norm": True,
-    "cross_attn_norm": True,
-    "eps": 1e-6,
-    "temporal_rope_scaling_factor": 0.19676,
-}
-
-VIDEO_CONFIG = {
-    "patch_size": [1, 2, 2],
-    "model_type": "ti2v",
-    "dim": 3072,
-    "ffn_dim": 14336,
-    "freq_dim": 256,
-    "num_heads": 24,
-    "num_layers": 30,
-    "in_dim": 48,
-    "out_dim": 48,
-    "text_len": 512,
-    "window_size": [-1, -1],
-    "qk_norm": True,
-    "cross_attn_norm": True,
-    "eps": 1e-6,
-}
-
-
-class DreamIDOmniPipeline(
-    nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput, SupportsComponentDiscovery
-):
-    """DreamID-Omni pipeline for vLLM-Omni."""
-
-    _dit_modules: ClassVar[list[str]] = ["model"]
-    _encoder_modules: ClassVar[list[str]] = ["text_model"]
-    _vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"]
-
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-        self.parallel_config = od_config.parallel_config
-
-        self.device = get_local_device()
-        model = od_config.model
-        self.target_dtype = od_config.dtype
-
-        # Init Models
-        ## Load VAEs
-        vae_model_video = init_wan_vae_2_2(model, rank=self.device)
-        vae_model_video.model.requires_grad_(False).eval()
-        vae_model_video.model = vae_model_video.model.bfloat16()
-        self.vae_model_video = vae_model_video
-        self.vae = self.vae_model_video
-
-        vae_model_audio = init_mmaudio_vae(model, rank=self.device)
-        vae_model_audio.requires_grad_(False).eval()
-        self.vae_model_audio = vae_model_audio.bfloat16()
-
-        # Load T5 text model
-        self.text_model = init_text_model(model, rank=self.device)
-        self.text_encoder = self.text_model.model
-
-        # Fusion model — weights are loaded later via load_weights()
-        quant_config = getattr(self.od_config, "quantization_config", None)
-        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config)
-        self.transformer = self.model
-
-        fusion_path = self.od_config.tf_model_config.get("fusion", None)
-        assert fusion_path is not None, "fusion checkpoint path is None in transformer config"
-        fusion_subfolder = os.path.dirname(fusion_path) or None
-        fusion_filename = os.path.basename(fusion_path)
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=model,
-                subfolder=fusion_subfolder,
-                revision=None,
-                prefix="model.",
-                allow_patterns_overrides=[fusion_filename],
-            )
-        ]
-
-        # Fixed attributes, non-configurable
-        self.audio_latent_channel = AUDIO_CONFIG.get("in_dim")
-        self.video_latent_channel = VIDEO_CONFIG.get("in_dim")
-        self.video_latent_length = 31
-        self.audio_latent_length = 157
-        self.target_area = 960 * 960
-
-        # Fixed attributes, CFG scales
-        self.video_cfg_scale = 3.0
-        self.video_ref_cfg_scale = 1.5
-        self.audio_cfg_scale = 4.0
-        self.audio_ref_cfg_scale = 2.0
-
-        # Schedulers will be set in forward
-        self.scheduler_video = None
-        self.scheduler_audio = None
-
-    def load_image_latent_ref_ip_video(
-        self,
-        images,
-        audios,
-        video_frame_height_width,
-    ):
-        # Load size.
-        patch_size = self.model.video_model.patch_size
-        vae_stride = [4, 16, 16]
-
-        # Load image and video.
-        ref_vae_latents = {
-            "image": [],
-            "audio": [],
-        }
-        video_h = video_frame_height_width[0]
-        video_w = video_frame_height_width[1]
-        ref_audio_lengths = []
-        for img in images:
-            # Calculate the required size to keep aspect ratio and fill the rest with padding.
-            img_ratio = img.width / img.height
-            target_ratio = video_w / video_h
-
-            if img_ratio > target_ratio:  # Image is wider than target
-                new_width = video_w
-                new_height = int(new_width / img_ratio)
-            else:  # Image is taller than target
-                new_height = video_h
-                new_width = int(new_height * img_ratio)
-
-            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
-
-            # Create a new image with the target size and place the resized image in the center
-            delta_w = video_w - img.size[0]
-            delta_h = video_h - img.size[1]
-            padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
-            new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
-
-            # Transform to tensor and normalize.
-            image_transform = Compose(
-                [
-                    NaResize(
-                        resolution=math.sqrt(
-                            video_frame_height_width[0] * video_frame_height_width[1]
-                        ),  # 256*448, 480*832
-                        mode="area",
-                        downsample_only=True,
-                    ),
-                    DivisibleCrop((vae_stride[1] * patch_size[1], vae_stride[2] * patch_size[2])),
-                    Normalize(0.5, 0.5),
-                    Rearrange("t c h w -> c t h w"),
-                ]
-            )
-            new_img = image_transform([new_img])
-            new_img = new_img.transpose(0, 1)
-            new_img = new_img.to(self.device)
-            new_img = new_img.to(self.target_dtype)
-
-            img_vae_latent = self.vae_model_video.wrapped_encode(new_img[:, :, None]).to(self.target_dtype).squeeze(0)
-            ref_vae_latents["image"].append(img_vae_latent)
-
-        for audio_array in audios:
-            audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0)
-            # print(audio_tensor.shape)#torch.Size([1, 81280])
-            audio_vae_latent = self.vae_model_audio.wrapped_encode(audio_tensor)
-            # print(audio_vae_latent.shape)#torch.Size([1, 20, length])
-            audio_length = audio_vae_latent.shape[2]
-            ref_audio_lengths.append(audio_length)
-            audio_vae_latent = audio_vae_latent.squeeze(0).transpose(0, 1)
-            ref_vae_latents["audio"].append(audio_vae_latent)
-
-        ref_vae_latents["image"] = torch.cat(ref_vae_latents["image"], dim=1)
-
-        ref_vae_latents["audio"] = torch.cat(ref_vae_latents["audio"], dim=0)
-
-        return ref_vae_latents, ref_audio_lengths
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        pipeline_prefix = "model."
-
-        qkv_stacked_mapping = [
-            (".self_attn.to_qkv", ".self_attn.q", "q"),
-            (".self_attn.to_qkv", ".self_attn.k", "k"),
-            (".self_attn.to_qkv", ".self_attn.v", "v"),
-        ]
-
-        params_dict = dict(self.model.named_parameters())
-        loaded_local: set[str] = set()
-
-        tp_size = get_tensor_model_parallel_world_size()
-        tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
-
-        def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
-            """Slice a checkpoint tensor to this rank's TP shard when the
-            param is a TP-local shape but the checkpoint stores the full one.
-            """
-            if tp_size <= 1 or weight.shape == param.shape:
-                return weight
-            if weight.ndim == 1 and weight.numel() == param.numel() * tp_size:
-                return weight.chunk(tp_size, dim=0)[tp_rank]
-            if weight.ndim == 2:
-                if weight.shape[0] == param.shape[0] * tp_size:
-                    return weight.chunk(tp_size, dim=0)[tp_rank]
-                if weight.shape[1] == param.shape[1] * tp_size:
-                    return weight.chunk(tp_size, dim=1)[tp_rank]
-            return weight
-
-        for full_name, loaded_weight in weights:
-            if not full_name.startswith(pipeline_prefix):
-                continue
-            name = full_name[len(pipeline_prefix) :]
-
-            # Fused-QKV stacking: route checkpoint's separate
-            # - ...self_attn.{q,k,v}: weights into the corresponding
-            # - ...self_attn.to_qkv : via its sharded weight_loader.
-            stacked_hit = False
-            for fused_name, raw_name, shard_id in qkv_stacked_mapping:
-                if raw_name not in name:
-                    continue
-                stacked = name.replace(raw_name, fused_name)
-                if stacked not in params_dict:
-                    continue
-                param = params_dict[stacked]
-                param.weight_loader(param, loaded_weight, shard_id)
-                loaded_local.add(stacked)
-                stacked_hit = True
-                break
-            if stacked_hit:
-                continue
-
-            if name not in params_dict:
-                logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters")
-                continue
-            param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", None)
-            if weight_loader is not None:
-                # Parallel Linears handle TP sharding internally
-                weight_loader(param, loaded_weight)
-            else:
-                default_weight_loader(param, _maybe_shard_weight(loaded_weight, param))
-            loaded_local.add(name)
-
-        # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading`
-        # and before any forward pass that triggers layerwise-offload hooks.
-        self.model._detach_blocks_from_backbones()
-
-        return {pipeline_prefix + k for k in loaded_local}
-
-    def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
-        torch.manual_seed(4)
-
-        if solver_name == "unipc":
-            sample_scheduler = FlowUniPCMultistepScheduler(
-                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
-            )
-            sample_scheduler.set_timesteps(sampling_steps, device=device, shift=shift)
-            timesteps = sample_scheduler.timesteps
-
-        elif solver_name == "dpm++":
-            sample_scheduler = FlowDPMSolverMultistepScheduler(
-                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
-            )
-            sampling_sigmas = get_sampling_sigmas(sampling_steps, shift=shift)
-            timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=sampling_sigmas)
-
-        elif solver_name == "euler":
-            sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
-            timesteps, sampling_steps = retrieve_timesteps(
-                sample_scheduler,
-                sampling_steps,
-                device=device,
-            )
-
-        else:
-            raise NotImplementedError("Unsupported solver.")
-
-        return sample_scheduler, timesteps
-
-    def predict_noise(self, **kwargs):
-        pred_vid, pred_audio = self.model(**kwargs)
-        return (pred_vid[0], pred_audio[0])
-
-    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
-        vid_pos, audio_pos = predictions[0]
-        vid_neg, audio_neg = predictions[1]
-        vid_ip_neg, _ = predictions[2]
-        _, refaudio_neg = predictions[3]
-
-        pred_video = (
-            vid_neg
-            + true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg)
-            + true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg)
-        )
-        pred_audio = (
-            audio_neg
-            + true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg)
-            + true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg)
-        )
-        return (pred_video, pred_audio)
-
-    def diffuse(
-        self,
-        video_noise: torch.Tensor,
-        audio_noise: torch.Tensor,
-        latents_ref_image: torch.Tensor,
-        latents_ref_audio: torch.Tensor,
-        timesteps_video: torch.Tensor,
-        timesteps_audio: torch.Tensor,
-        text_embeddings_video_pos: torch.Tensor,
-        text_embeddings_video_neg: torch.Tensor,
-        text_embeddings_audio_pos: torch.Tensor,
-        text_embeddings_audio_neg: torch.Tensor,
-        max_seq_len_video: int,
-        max_seq_len_audio: int,
-        freqs_scaling_tensor: torch.Tensor,
-        ref_ip_num: int,
-        ref_audio_length: int,
-        ref_audio_lengths: list,
-        scheduler_video,
-        scheduler_audio,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Diffusion loop with CFG parallel support for DreamID-Omni.
-        """
-        for i, (t_v, t_a) in tqdm(enumerate(zip(timesteps_video, timesteps_audio))):
-            timestep_input = torch.full((1,), t_v, device=self.device)
-
-            model_input_video = torch.cat([video_noise[:, :-ref_ip_num], latents_ref_image], dim=1)
-            model_input_video_neg = torch.cat(
-                [video_noise[:, :-ref_ip_num], torch.zeros_like(latents_ref_image)], dim=1
-            )
-
-            model_input_audio = torch.cat([audio_noise[:-ref_audio_length, :], latents_ref_audio], dim=0)
-
-            model_input_audio_neg = torch.cat(
-                [audio_noise[:-ref_audio_length, :], torch.zeros_like(latents_ref_audio)], dim=0
-            )
-            ref_ip_lengths = [ref_ip_num]
-
-            common_args = {
-                "vid_seq_len": max_seq_len_video,
-                "audio_seq_len": max_seq_len_audio,
-                "freqs_scaling": freqs_scaling_tensor,
-                "ref_ip_lengths": [ref_ip_lengths],
-                "ref_audio_lengths": [ref_audio_lengths],
-            }
-            pos_args = {
-                **common_args,
-                "audio_context": [text_embeddings_audio_pos],
-                "vid_context": [text_embeddings_video_pos],
-            }
-            neg_args = {
-                **common_args,
-                "audio_context": [text_embeddings_audio_neg],
-                "vid_context": [text_embeddings_video_neg],
-            }
-
-            branches_kwargs = [
-                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args},
-                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args},
-                {"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args},
-                {"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args},
-            ]
-
-            pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg(
-                do_true_cfg=True,
-                true_cfg_scale={
-                    "video_cfg_scale": self.video_cfg_scale,
-                    "video_ref_cfg_scale": self.video_ref_cfg_scale,
-                    "audio_cfg_scale": self.audio_cfg_scale,
-                    "audio_ref_cfg_scale": self.audio_ref_cfg_scale,
-                },
-                branches_kwargs=branches_kwargs,
-            )
-            video_noise = scheduler_video.step(
-                pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False
-            )[0].squeeze(0)
-            audio_noise = scheduler_audio.step(
-                pred_audio_guided.unsqueeze(0), t_a, audio_noise.unsqueeze(0), return_dict=False
-            )[0].squeeze(0)
-        return video_noise, audio_noise
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        video_negative_prompt: str = "",
-        audio_negative_prompt: str = "",
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Encode the positive and negative prompts via ``self.text_model``.
-
-        This is the single text-encoder entrypoint so that the runner-level
-        prompt-embedding cache (see
-        ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
-        memoize results when the same prompts are submitted repeatedly (e.g.
-        GRPO rollouts that sample the same prompt with different seeds).
-
-        Returns:
-            (audio_pos, video_pos, video_neg, audio_neg) embeddings cast to
-            ``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the
-            same tensor, matching the original inline behavior.
-        """
-        text_embeddings = self.text_model(
-            [prompt, video_negative_prompt, audio_negative_prompt],
-            device=self.device,
-        )
-        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
-        # Index 0 is the positive embedding, shared between video and audio branches.
-        return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2]
-
-    def forward(
-        self,
-        request: DiffusionRequestBatch,
-        **kwargs,
-    ) -> DiffusionOutput:
-        """Main forward pass for DreamID-Omni pipeline for R2AV task."""
-        # Extract parameters from request
-        r_prompts = request.prompts[0]
-        if isinstance(r_prompts, str):
-            prompt = r_prompts
-            video_negative_prompt = ""
-            audio_negative_prompt = ""
-        else:
-            prompt = r_prompts.get("prompt", "")
-            video_negative_prompt = r_prompts.get("video_negative_prompt", "")
-            audio_negative_prompt = r_prompts.get("audio_negative_prompt", "")
-
-        multi_modal_data = r_prompts.get("multi_modal_data", {}) if not isinstance(r_prompts, str) else None
-        raw_images = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-        raw_audios = multi_modal_data.get("audio", None) if multi_modal_data is not None else None
-        if raw_images is None or raw_audios is None:
-            raise ValueError("This model requires image and audio to run.")
-        if not isinstance(raw_images, list):
-            raw_images = [raw_images]
-        if not isinstance(raw_audios, list):
-            raw_audios = [raw_audios]
-        if len(raw_images) != len(raw_audios):
-            raise ValueError("'image' and 'audio' lists must have the same length.")
-
-        height = request.sampling_params.height
-        width = request.sampling_params.width
-        num_inference_steps = request.sampling_params.num_inference_steps
-        shift = request.sampling_params.extra_args.get("shift", 5.0)
-        solver_name = request.sampling_params.extra_args.get("solver_name", "unipc")
-        seed = request.sampling_params.seed if request.sampling_params.seed is not None else 42
-
-        # 1. Prepare reference latents
-        ref_vae_latents, ref_audio_lengths = self.load_image_latent_ref_ip_video(
-            images=raw_images,
-            audios=raw_audios,
-            video_frame_height_width=(height, width),
-        )
-
-        latents_ref_image = ref_vae_latents["image"]
-        latents_ref_audio = ref_vae_latents["audio"]
-        ref_ip_num = latents_ref_image.shape[1]
-        ref_audio_length = latents_ref_audio.shape[0]
-
-        # 2. scheduler
-        scheduler_video, timesteps_video = self.get_scheduler_time_steps(
-            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
-        )
-        scheduler_audio, timesteps_audio = self.get_scheduler_time_steps(
-            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
-        )
-
-        # 3. text embedding
-        text_embeddings_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = (
-            self.encode_prompt(
-                prompt=prompt,
-                video_negative_prompt=video_negative_prompt,
-                audio_negative_prompt=audio_negative_prompt,
-            )
-        )
-
-        video_latent_h, video_latent_w = height // 16, width // 16
-
-        video_noise_len = self.video_latent_length + ref_ip_num
-        audio_noise_len = self.audio_latent_length + ref_audio_length
-        freqs_scaling_tensor = torch.tensor(
-            self.video_latent_length / self.audio_latent_length, device=self.device, dtype=self.target_dtype
-        )
-
-        video_noise = torch.randn(
-            (self.video_latent_channel, video_noise_len, video_latent_h, video_latent_w),
-            device=self.device,
-            dtype=self.target_dtype,
-            generator=torch.Generator(device=self.device).manual_seed(seed),
-        )
-        audio_noise = torch.randn(
-            (audio_noise_len, self.audio_latent_channel),
-            device=self.device,
-            dtype=self.target_dtype,
-            generator=torch.Generator(device=self.device).manual_seed(seed),
-        )
-
-        _patch_size_h, _patch_size_w = self.model.video_model.patch_size[1], self.model.video_model.patch_size[2]
-
-        max_seq_len_video = (
-            video_noise.shape[1] * video_noise.shape[2] * video_noise.shape[3] // (_patch_size_h * _patch_size_w)
-        )
-
-        max_seq_len_audio = audio_noise_len
-
-        with torch.amp.autocast("cuda", enabled=self.target_dtype != torch.float32, dtype=self.target_dtype):
-            video_noise, audio_noise = self.diffuse(
-                video_noise,
-                audio_noise,
-                latents_ref_image,
-                latents_ref_audio,
-                timesteps_video,
-                timesteps_audio,
-                text_embeddings_video_pos,
-                text_embeddings_video_neg,
-                text_embeddings_audio_pos,
-                text_embeddings_audio_neg,
-                max_seq_len_video,
-                max_seq_len_audio,
-                freqs_scaling_tensor,
-                ref_ip_num,
-                ref_audio_length,
-                ref_audio_lengths,
-                scheduler_video,
-                scheduler_audio,
-            )
-
-        video_noise_for_decode = video_noise[:, :-ref_ip_num]
-        audio_noise_for_decode = audio_noise[:-ref_audio_length, :]
-
-        audio_latents_for_vae = audio_noise_for_decode.unsqueeze(0).transpose(1, 2)
-        generated_audio = self.vae_model_audio.wrapped_decode(audio_latents_for_vae).squeeze().cpu().float().numpy()
-
-        video_latents_for_vae = video_noise_for_decode.unsqueeze(0)
-        generated_video = self.vae_model_video.wrapped_decode(video_latents_for_vae).squeeze(0).cpu().float().numpy()
-
-        return DiffusionOutput(output=(generated_video, generated_audio))

vllm_omni/diffusion/models/dreamid_omni/wan2_2.py

@@ -1,694 +0,0 @@
-# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
-import math
-
-import torch
-import torch.amp as amp
-import torch.nn as nn
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.modeling_utils import ModelMixin
-from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_reduce,
-)
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-from vllm_omni.diffusion.attention.layer import Attention
-
-try:
-    from dreamid_omni.modules.model import (
-        ChannelLastConv1d,
-        ConvMLP,
-        Head,
-        MLPProj,
-        ModulationAdd,
-        WanLayerNorm,
-        rope_apply,  # diff from  wan2.2
-        rope_params,
-        sinusoidal_embedding_1d,
-    )
-except ImportError:
-    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
-
-
-class DistributedRMSNorm(nn.Module):
-    """RMSNorm that computes global RMS across tensor parallel ranks.
-
-    Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm`
-    """
-
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        tp_size = get_tensor_model_parallel_world_size()
-        x_dtype = x.dtype
-        x_float = x.float()
-        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
-        local_count = x.shape[-1]
-
-        if tp_size > 1:
-            local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
-            local_count = local_count * tp_size
-
-        rms = torch.sqrt(local_sum_sq / local_count + self.eps)
-        return (x_float / rms * self.weight.float()).to(x_dtype)
-
-
-class WanSelfAttention(nn.Module):
-    """Optimized self-attention module using vLLM layers."""
-
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        window_size=(-1, -1),
-        qk_norm=True,
-        eps=1e-6,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        assert dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.head_dim = dim // num_heads
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
-        )
-        self.num_heads = self.to_qkv.num_heads
-        self.num_kv_heads = self.to_qkv.num_kv_heads
-        self.tp_inner_dim = self.num_heads * self.head_dim
-        self.tp_kv_dim = self.num_kv_heads * self.head_dim
-
-        self.o = RowParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o" if prefix else "o",
-        )
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity()
-        # Unified attention layer
-        self.attn = Attention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=self.num_kv_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-            prefix=prefix,
-        )
-
-    # query, key, value function
-    def qkv_fn(self, x):
-        b, s = x.shape[:2]
-        qkv, _ = self.to_qkv(x)
-        q_size = self.num_heads * self.head_dim
-        kv_size = self.num_kv_heads * self.head_dim
-        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
-
-        q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim)
-        k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim)
-        v = v.view(b, s, self.num_kv_heads, self.head_dim)
-        return q, k, v
-
-    def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L, C]
-            seq_lens(Tensor): Shape [B]
-            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
-            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
-            ref_lengths(Tensor): Shape [B]
-        """
-        q, k, v = self.qkv_fn(x)
-
-        x = self.attn(
-            rope_apply(q, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
-            rope_apply(k, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
-            v,
-        )
-
-        # output
-        x = x.flatten(2)
-        x = self.o(x)
-        return x
-
-
-class WanT2VCrossAttention(nn.Module):
-    """Text -> latent cross-attention."""
-
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        window_size=(-1, -1),
-        qk_norm=True,
-        eps=1e-6,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        assert dim % num_heads == 0
-        self.dim = dim
-        self.head_dim = dim // num_heads
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-
-        # Separate Q/K/V — Q's input differs from K/V's input (text context)
-        self.q = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.q" if prefix else "q",
-        )
-        self.k = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.k" if prefix else "k",
-        )
-        self.v = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.v" if prefix else "v",
-        )
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.num_heads = num_heads // tp_size
-        self.num_kv_heads = self.num_heads
-        self.tp_inner_dim = self.num_heads * self.head_dim
-
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-
-        self.o = RowParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o" if prefix else "o",
-        )
-
-        # Kept for interface parity with WanSelfAttention;
-        # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn.
-        # Disable kv quant to prevent affecting quality.
-        self.attn = Attention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=self.num_kv_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-            prefix=prefix,
-            disable_kv_quant=True,
-        )
-
-    def qkv_fn(self, x, context):
-        b = x.size(0)
-
-        q = self.q(x)
-        k = self.k(context)
-        v = self.v(context)
-        q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim)
-        k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim)
-        v = v.view(b, -1, self.num_kv_heads, self.head_dim)
-
-        return q, k, v
-
-    def forward(self, x, context, context_lens):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L1, C]
-            context(Tensor): Shape [B, L2, C]
-            context_lens(Tensor): Shape [B]
-        """
-        q, k, v = self.qkv_fn(x, context)
-
-        # compute attention
-        x = self.attn(q, k, v)
-
-        # output
-        x = x.flatten(2)
-        x = self.o(x)
-        return x
-
-
-class WanI2VCrossAttention(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-
-    def forward(self, *args, **kwargs):
-        raise RuntimeError("WanT2VCrossAttention is currently disabled and should not be used.")
-
-
-WAN_CROSSATTENTION_CLASSES = {
-    "t2v_cross_attn": WanT2VCrossAttention,
-    "i2v_cross_attn": WanI2VCrossAttention,
-}
-
-
-class WanAttentionBlock(nn.Module):
-    def __init__(
-        self,
-        cross_attn_type,
-        dim,
-        ffn_dim,
-        num_heads,
-        window_size=(-1, -1),
-        qk_norm=True,
-        cross_attn_norm=False,
-        eps=1e-6,
-        additional_emb_length=None,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.num_heads = num_heads
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-
-        # layers
-        self.norm1 = WanLayerNorm(dim, eps)
-        self.self_attn = WanSelfAttention(
-            dim,
-            num_heads,
-            window_size,
-            qk_norm,
-            eps,
-            quant_config=quant_config,
-            prefix=f"{prefix}.self_attn" if prefix else "self_attn",
-        )
-        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
-        if cross_attn_type == "i2v_cross_attn":
-            assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
-            self.cross_attn = WanI2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, additional_emb_length)
-        else:
-            assert additional_emb_length is None, "additional_emb_length should be None for t2v_cross_attn"
-            self.cross_attn = WanT2VCrossAttention(
-                dim,
-                num_heads,
-                (-1, -1),
-                qk_norm,
-                eps,
-                quant_config=quant_config,
-                prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
-            )
-        self.norm2 = WanLayerNorm(dim, eps)
-        # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints
-        self.ffn = nn.Sequential(
-            ColumnParallelLinear(
-                dim,
-                ffn_dim,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
-            ),
-            nn.GELU(approximate="tanh"),
-            RowParallelLinear(
-                ffn_dim,
-                dim,
-                bias=True,
-                input_is_parallel=True,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
-            ),
-        )
-
-        # modulation
-        self.modulation = ModulationAdd(dim, 6)
-
-    def forward(
-        self,
-        x,
-        e,
-        seq_lens,
-        grid_sizes,
-        freqs,
-        context,
-        context_lens,
-    ):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L, C]
-            e(Tensor): Shape [B, L1, 6, C]
-            seq_lens(Tensor): Shape [B], length of each sequence in batch
-            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
-            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
-        """
-        assert e.dtype == torch.bfloat16
-        assert len(e.shape) == 4 and e.size(2) == 6 and e.shape[1] == x.shape[1], f"{e.shape}, {x.shape}"
-        with amp.autocast("cuda", dtype=torch.bfloat16):
-            e = self.modulation(e).chunk(6, dim=2)
-        assert e[0].dtype == torch.bfloat16
-
-        # self-attention
-        y = self.self_attn(
-            self.norm1(x).bfloat16() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs
-        )
-        with amp.autocast("cuda", dtype=torch.bfloat16):
-            x = x + y * e[2].squeeze(2)
-
-        # cross-attention & ffn function
-        def cross_attn_ffn(x, context, context_lens, e):
-            x = x + self.cross_attn(self.norm3(x), context, context_lens)
-            y = self.ffn(self.norm2(x).bfloat16() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
-            with amp.autocast("cuda", dtype=torch.bfloat16):
-                x = x + y * e[5].squeeze(2)
-            return x
-
-        x = cross_attn_ffn(x, context, context_lens, e)
-        return x
-
-
-class WanModel(ModelMixin, ConfigMixin):
-    r"""
-    Wan diffusion backbone supporting both text-to-video and image-to-video, text-to-audio.
-    """
-
-    @register_to_config
-    def __init__(
-        self,
-        model_type="t2v",
-        patch_size=(1, 2, 2),
-        text_len=512,
-        in_dim=16,
-        dim=2048,
-        ffn_dim=8192,
-        freq_dim=256,
-        text_dim=4096,
-        additional_emb_dim=None,
-        additional_emb_length=None,
-        out_dim=16,
-        num_heads=16,
-        num_layers=32,
-        window_size=(-1, -1),
-        qk_norm=True,
-        cross_attn_norm=True,
-        temporal_rope_scaling_factor=1.0,
-        eps=1e-6,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        r"""
-        Initialize the diffusion model backbone.
-
-        Args:
-            model_type (`str`, *optional*, defaults to 't2v'):
-                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
-            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
-                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
-            text_len (`int`, *optional*, defaults to 512):
-                Fixed length for text embeddings
-            in_dim (`int`, *optional*, defaults to 16):
-                Input video channels (C_in)
-            dim (`int`, *optional*, defaults to 2048):
-                Hidden dimension of the transformer
-            ffn_dim (`int`, *optional*, defaults to 8192):
-                Intermediate dimension in feed-forward network
-            freq_dim (`int`, *optional*, defaults to 256):
-                Dimension for sinusoidal time embeddings
-            text_dim (`int`, *optional*, defaults to 4096):
-                Input dimension for text embeddings
-            out_dim (`int`, *optional*, defaults to 16):
-                Output video channels (C_out)
-            num_heads (`int`, *optional*, defaults to 16):
-                Number of attention heads
-            num_layers (`int`, *optional*, defaults to 32):
-                Number of transformer blocks
-            window_size (`tuple`, *optional*, defaults to (-1, -1)):
-                Window size for local attention (-1 indicates global attention)
-            qk_norm (`bool`, *optional*, defaults to True):
-                Enable query/key normalization
-            cross_attn_norm (`bool`, *optional*, defaults to False):
-                Enable cross-attention normalization
-            eps (`float`, *optional*, defaults to 1e-6):
-                Epsilon value for normalization layers
-        """
-
-        super().__init__()
-
-        assert model_type in [
-            "t2v",
-            "i2v",
-            "t2a",
-            "tt2a",
-            "ti2v",
-        ]  ## tt2a means text transcript + text description to audio (to support both TTS and T2A
-        self.model_type = model_type
-        is_audio_type = "a" in self.model_type
-        is_video_type = "v" in self.model_type
-        assert is_audio_type ^ is_video_type, "Either audio or video model should be specified"
-        if is_audio_type:
-            ## audio model
-            assert len(patch_size) == 1 and patch_size[0] == 1, (
-                "Audio model should only accept 1 dimensional input, and we dont do patchify"
-            )
-
-        self.patch_size = patch_size
-        self.text_len = text_len
-        self.in_dim = in_dim
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.freq_dim = freq_dim
-        self.text_dim = text_dim
-        self.out_dim = out_dim
-        self.num_heads = num_heads
-        self.num_layers = num_layers
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-        self.temporal_rope_scaling_factor = temporal_rope_scaling_factor
-        self.is_audio_type = is_audio_type
-        self.is_video_type = is_video_type
-        # embeddings
-        if is_audio_type:
-            ## hardcoded to MMAudio
-            self.patch_embedding = nn.Sequential(
-                ChannelLastConv1d(in_dim, dim, kernel_size=7, padding=3),
-                nn.SiLU(),
-                ConvMLP(dim, dim * 4, kernel_size=7, padding=3),
-            )
-        else:
-            self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
-
-        self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim))
-
-        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
-
-        # blocks
-        ## so i2v and tt2a share the same cross attention while t2v and t2a share the same cross attention
-        cross_attn_type = "t2v_cross_attn" if model_type in ["t2v", "t2a", "ti2v"] else "i2v_cross_attn"
-
-        if cross_attn_type == "t2v_cross_attn":
-            assert additional_emb_dim is None and additional_emb_length is None, (
-                "additional_emb_length should be None for t2v and t2a model"
-            )
-        else:
-            assert additional_emb_dim is not None and additional_emb_length is not None, (
-                "additional_emb_length should be specified for i2v and tt2a model"
-            )
-
-        self.blocks = nn.ModuleList(
-            [
-                WanAttentionBlock(
-                    cross_attn_type,
-                    dim,
-                    ffn_dim,
-                    num_heads,
-                    window_size,
-                    qk_norm,
-                    cross_attn_norm,
-                    eps,
-                    additional_emb_length,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}",
-                )
-                for layer_idx in range(num_layers)
-            ]
-        )
-
-        # head
-        self.head = Head(dim, out_dim, patch_size, eps)
-
-        self.set_rope_params()
-
-        if model_type in ["i2v", "tt2a"]:
-            self.img_emb = MLPProj(additional_emb_dim, dim)
-
-    def set_rope_params(self):
-        # buffers (don't use register_buffer otherwise dtype will be changed in to())
-        dim = self.dim
-        num_heads = self.num_heads
-        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
-        d = dim // num_heads
-
-        if self.is_audio_type:
-            ## to be determined
-            # self.freqs = rope_params(1024, d, freqs_scaling=temporal_rope_scaling_factor)
-            self.freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=self.temporal_rope_scaling_factor)
-        else:
-            self.freqs = torch.cat(
-                [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
-                dim=1,
-            )
-
-    def prepare_transformer_block_kwargs(
-        self,
-        x,
-        t,
-        context,
-        seq_len,
-        ref_lengths=None,
-        freqs_scaling=None,
-    ):
-        # params
-        ## need to change!
-        device = next(self.patch_embedding.parameters()).device
-
-        if self.is_audio_type and freqs_scaling is not None:
-            if isinstance(freqs_scaling, torch.Tensor):
-                scale_val = freqs_scaling.item()
-            else:
-                scale_val = freqs_scaling
-            d = self.dim // self.num_heads
-
-            current_freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=scale_val).to(device)
-        else:
-            current_freqs = self.freqs
-            if current_freqs.device != device:
-                current_freqs = current_freqs.to(device)
-
-        # embeddings
-        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]  ## x is list of [B L D] or [B C F H W]
-        if self.is_audio_type:
-            # [B, 1]
-            grid_sizes = torch.stack([torch.tensor(u.shape[1:2], dtype=torch.long) for u in x])
-        else:
-            # [B, 3]
-            grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
-            x = [u.flatten(2).transpose(1, 2) for u in x]  # [B C F H W] -> [B (F H W) C] -> [B L C]
-
-        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
-        assert seq_lens.max() <= seq_len, f"Sequence length {seq_lens.max()} exceeds maximum {seq_len}."
-        x = torch.cat(
-            [torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]
-        )  # single [B, L, C]
-
-        # time embeddings
-        if t.dim() == 1:
-            t = t.unsqueeze(1).expand(t.size(0), seq_len)
-        with amp.autocast("cuda", dtype=torch.bfloat16):
-            bt = t.size(0)
-            t = t.flatten()
-            e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).bfloat16())
-            e0 = self.time_projection(e).unflatten(2, (6, self.dim))  # [1, 26784, 6, 3072] - B, seq_len, 6, dim
-            assert e.dtype == torch.bfloat16 and e0.dtype == torch.bfloat16
-
-        # context
-        context_lens = None
-        context = self.text_embedding(
-            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
-        )
-
-        # arguments
-        kwargs = dict(
-            e=e0,
-            seq_lens=seq_lens,
-            grid_sizes=grid_sizes,
-            freqs=current_freqs,
-            context=context,
-            context_lens=context_lens,
-            ref_lengths=ref_lengths,
-            freqs_scaling=freqs_scaling,
-        )
-
-        return x, e, kwargs
-
-    def post_transformer_block_out(self, x, grid_sizes, e):
-        # head
-        x = self.head(x, e)
-        # unpatchify
-        if self.is_audio_type:
-            ## grid_sizes is [B 1] where 1 is L,
-            # converting grid_sizes from [B 1] -> [B]
-            grid_sizes = [gs[0] for gs in grid_sizes]
-            assert len(x) == len(grid_sizes)
-            x = [u[:gs] for u, gs in zip(x, grid_sizes)]
-        else:
-            ## grid_sizes is [B 3] where 3 is F H w
-            x = self.unpatchify(x, grid_sizes)
-
-        return [u.bfloat16() for u in x]
-
-    def forward(self, *args, **kwargs):
-        raise NotImplementedError("WanModel model does not support forward pass.")
-
-    def unpatchify(self, x, grid_sizes) -> list[torch.Tensor]:
-        r"""
-        Reconstruct video tensors from patch embeddings.
-
-        Args:
-            x (List[Tensor]):
-                List of patchified features, each with shape [L, C_out * prod(patch_size)]
-            grid_sizes (Tensor):
-                Original spatial-temporal grid dimensions before patching,
-                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
-
-        Returns:
-            List[Tensor]:
-                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
-        """
-
-        c = self.out_dim
-        out = []
-        for u, v in zip(x, grid_sizes.tolist()):
-            # v is [F H w] F * H * 80, 100, it was right padded by 20.
-            u = u[: math.prod(v)].view(*v, *self.patch_size, c)
-            u = torch.einsum("fhwpqrc->cfphqwr", u)
-            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
-            out.append(u)
-        # out is list of [C F H W]
-        return out

vllm_omni/diffusion/models/magi_human/__init__.py

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

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -402,6 +402,25 @@ def _apply_rotary_pos_emb(
     return q_embed, k_embed
 
 
+def _scaled_dot_product_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+) -> torch.Tensor:
+    """Run causal SDPA with the reference expanded-K/V GQA behavior."""
+    num_key_value_groups = query.shape[1] // key.shape[1]
+    if num_key_value_groups != 1:
+        key = key.repeat_interleave(num_key_value_groups, dim=1)
+        value = value.repeat_interleave(num_key_value_groups, dim=1)
+    return F.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        dropout_p=0.0,
+        is_causal=True,
+    )
+
+
 def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor:
     """Reorganize chunked [TTT...HHH...WWW] into interleaved [THWTHW...] layout."""
     freqs_t = freqs[0]
@@ -810,19 +829,7 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
         cos, sin = position_embeddings
         query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
-        # Expand GQA keys/values to the local query-head count so SDPA works
-        # uniformly across backends (mirrors the reference eager path).
-        num_key_value_groups = self.num_heads // self.num_kv_heads
-        key_states = key_states.repeat_interleave(num_key_value_groups, dim=1)
-        value_states = value_states.repeat_interleave(num_key_value_groups, dim=1)
-
-        attn_output = F.scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            dropout_p=0.0,
-            is_causal=True,
-        )
+        attn_output = _scaled_dot_product_attention(query_states, key_states, value_states)
         attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
         return attn_output
@@ -907,10 +914,7 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings)
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
         return hidden_states

vllm_omni/diffusion/registry.py

@@ -266,11 +266,6 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2",
         "Flux2Pipeline",
     ),
-    "DreamIDOmniPipeline": (
-        "dreamid_omni",
-        "pipeline_dreamid_omni",
-        "DreamIDOmniPipeline",
-    ),
     "SenseNovaU1Pipeline": (
         "sensenova_u1",
         "pipeline_sensenova_u1",
@@ -301,11 +296,6 @@ _DIFFUSION_MODELS = {
         "pipeline_sana_video_i2v",
         "SanaImageToVideoPipeline",
     ),
-    "MagiHumanPipeline": (
-        "magi_human",
-        "pipeline_magi_human",
-        "MagiHumanPipeline",
-    ),
     "OmniVoicePipeline": (
         "omnivoice",
         "pipeline_omnivoice",
@@ -598,9 +588,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_post_process_func",
     "SanaVideoPipeline": "get_sana_video_post_process_func",
     "SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func",
-    "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
-    "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_post_process_func",
@@ -645,7 +633,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_pre_process_func",
     "SanaImageToVideoPipeline": "get_sana_video_i2v_pre_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
-    "MagiHumanPipeline": "get_magi_human_pre_process_func",
     "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",

vllm_omni/diffusion/request.py

@@ -67,6 +67,11 @@ class OmniDiffusionRequest:
     # This is populated by a pipeline preprocessor before the request reaches
     # the scheduler; ``None`` keeps the default behavior for other pipelines.
     batch_compatibility_key: tuple[Any, ...] | None = None
+    # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward
+    # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric.
+    kv_recv_ms: float = 0.0
+    # Time spent waiting for initial admission by the diffusion scheduler.
+    scheduler_queue_wait_ms: float | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/sched/base_scheduler.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import time
 from abc import ABC, abstractmethod
 from collections import deque
 from dataclasses import fields
@@ -117,6 +118,7 @@ class BaseScheduler(ABC):
         if request_id in self._request_states:
             raise ValueError(f"request_id {request_id!r} is already active.")
         state = self._make_request_state(request_id, request)
+        state.queued_at = time.perf_counter()
         self._request_states[request_id] = state
         self._waiting.append(request_id)
         logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting))
@@ -178,6 +180,7 @@ class BaseScheduler(ABC):
             state.status = DiffusionRequestStatus.RUNNING
             self._running.append(request_id)
             if was_new_request:
+                state.req.scheduler_queue_wait_ms = max((time.perf_counter() - state.queued_at) * 1000.0, 0.0)
                 scheduled_new_reqs.append(
                     NewRequestData.from_state(
                         state,

vllm_omni/diffusion/sched/interface.py

@@ -157,6 +157,7 @@ class SchedulerRequestState:
     diffusion_kv_requests: tuple[DiffusionKVRequest, ...] = ()
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
     error: str | None = None
+    queued_at: float = 0.0
 
     def is_finished(self) -> bool:
         return DiffusionRequestStatus.is_finished(self.status)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Stage Diffusion Client for vLLM-Omni multi-stage runtime.
 
 Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the
@@ -30,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
 )
 from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -240,6 +244,14 @@ class StageDiffusionClient(StageClientBase):
 
             if msg_type == "result":
                 self._output_queue.put_nowait(msg["output"])
+            elif msg_type == "metrics":
+                self._output_queue.put_nowait(
+                    OmniRequestOutput(
+                        request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID,
+                        finished=True,
+                        metrics=msg.get("metrics") or {},
+                    )
+                )
             elif msg_type == "rpc_result":
                 self._rpc_results[msg["rpc_id"]] = msg["result"]
             elif msg_type == "error":
@@ -264,14 +276,14 @@ class StageDiffusionClient(StageClientBase):
                 # Route request errors as error outputs so the Orchestrator
                 # sees the request complete (instead of hanging forever).
                 if req_id is not None:
-                    self._output_queue.put_nowait(
-                        OmniRequestOutput.from_error(
-                            req_id,
-                            error_msg,
-                            status_code=status_code,
-                            error_type=error_type,
-                        )
+                    error_output = OmniRequestOutput.from_error(
+                        req_id,
+                        error_msg,
+                        status_code=status_code,
+                        error_type=error_type,
                     )
+                    error_output.metrics.update(msg.get("metrics") or {})
+                    self._output_queue.put_nowait(error_output)
 
     # Fields that are subprocess-local and cannot be serialized across
     # process boundaries.  They are recreated in the subprocess with

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Subprocess entry point for the diffusion engine.
 
 StageDiffusionProc runs DiffusionEngine in a child process,
@@ -36,6 +39,7 @@ from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
 from vllm_omni.engine.stage_init_utils import set_death_signal
 from vllm_omni.errors import client_error_metadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.metrics.utils import diffusion_exception_metrics
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -361,6 +365,16 @@ class StageDiffusionProc:
                     request_id,
                     str(e),
                 )
+                metrics = diffusion_exception_metrics(e)
+                if metrics:
+                    await response_socket.send(
+                        encoder.encode(
+                            {
+                                "type": "metrics",
+                                "metrics": metrics,
+                            }
+                        )
+                    )
             except Exception as e:
                 logger.exception("Diffusion request %s failed: %s", request_id, e)
                 status_code, error_type = client_error_metadata(e)
@@ -372,6 +386,7 @@ class StageDiffusionProc:
                             "error": str(e),
                             "status_code": status_code,
                             "error_type": error_type,
+                            "metrics": diffusion_exception_metrics(e),
                         }
                     )
                 )
@@ -426,9 +441,8 @@ class StageDiffusionProc:
 
                 elif msg_type == "abort":
                     for rid in msg.get("request_ids", []):
-                        task = tasks.pop(rid, None)
-                        if task:
-                            task.cancel()
+                        # Let the request task consume the terminal abort
+                        # output so it can publish the scheduler snapshot.
                         self._engine.abort(rid)
 
                 elif msg_type == "collective_rpc":

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -539,6 +539,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 target_device=self._target_device if use_prefetch else getattr(self.pipeline, "device", None),
             )
         kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000
+        req.kv_recv_ms = kv_recv_ms
         logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms)
 
         # Kick off the next request's prefetch (+ H2D) to overlap this forward.
@@ -596,6 +597,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         reqs: list[OmniDiffusionRequest],
         outputs: list[DiffusionOutput],
     ) -> BatchRunnerOutput:
+        for i in range(len(reqs)):
+            # Carry the runner-measured KV-recv timing onto the output so the
+            # engine's step_streaming can surface it as diffusion_kv_load_s.
+            outputs[i].kv_recv_ms = reqs[i].kv_recv_ms
         return BatchRunnerOutput.from_list(
             [
                 RunnerOutput(

vllm_omni/engine/async_omni_engine.py

@@ -107,6 +107,7 @@ class AsyncOmniEngine:
     _log_stats: bool = False
     _coordinator_runtime: Any = None
     _transfer_emitter: Any = None
+    _prom_metrics: Any = None
     _enable_orch_monitor: bool = False
 
     def __init__(
@@ -117,6 +118,7 @@ class AsyncOmniEngine:
         diffusion_batch_size: int = 1,
         single_stage_mode: bool = False,
         transfer_emitter: Any = None,
+        prom_metrics: Any = None,
         log_stats: bool = False,
         tokenizer: str | None = None,
         trust_remote_code: bool | None = None,
@@ -133,6 +135,7 @@ class AsyncOmniEngine:
         # Optional: when None, Orchestrator silently skips TX emit (existing
         # RX path still works via OrchestratorAggregator).
         self._transfer_emitter = transfer_emitter
+        self._prom_metrics = prom_metrics
         # Drives upstream EngineCore + scheduler stats production. When False
         # the engine skips SchedulerStats / IterationStats; the per-(stage,
         # replica) vllm:* wrap stays registered but reads zero. Respects the
@@ -398,6 +401,7 @@ class AsyncOmniEngine:
                 membership_controller=membership_controller,
                 running_counter=self._running_counter,
                 transfer_emitter=self._transfer_emitter,
+                prom_metrics=self._prom_metrics,
                 log_stats=self._log_stats,
                 enable_orch_monitor=self._enable_orch_monitor,
                 duplex_runtime_extension=duplex_runtime_extension,

vllm_omni/engine/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
 
 """MembershipController: distributed replica lifecycle management.
 
@@ -60,6 +60,7 @@ class MembershipController:
         self._watcher_task: asyncio.Task[None] | None = None
         self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None
         self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None
+        self._replica_removed_callback: Callable[[int, int], None] | None = None
 
         self._hub = OmniCoordClientForHub(coordinator_pub_address)
         factory = load_balancer_factory
@@ -118,8 +119,12 @@ class MembershipController:
             return
         effective_output_queue = output_queue if output_queue is not None else self._output_queue
         effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback
+        replica_id = pool.get_replica_id_by_addr(input_addr)
+        client = pool.get_client_by_addr(input_addr)
         affected = pool.invalidate_addr(input_addr)
         self._detach_replica(stage_id, input_addr)
+        if client is not None and replica_id is not None and self._replica_removed_callback is not None:
+            self._replica_removed_callback(stage_id, replica_id)
         if affected and effective_cleanup_callback is not None:
             await effective_cleanup_callback(affected)
         if affected and effective_output_queue is not None:
@@ -283,10 +288,12 @@ class MembershipController:
         *,
         output_queue: asyncio.Queue[EngineQueueMessage],
         cleanup_callback: Callable[[list[str]], Awaitable[None]],
+        replica_removed_callback: Callable[[int, int], None] | None = None,
     ) -> None:
         """Install shared cleanup sinks for watcher-driven unregister events."""
         self._output_queue = output_queue
         self._cleanup_callback = cleanup_callback
+        self._replica_removed_callback = replica_removed_callback
 
     def _pool_for_stage_id(self, stage_id: int) -> StagePool | None:
         if not (0 <= stage_id < len(self._stage_pools)):

vllm_omni/engine/orchestrator.py

@@ -57,8 +57,10 @@ from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_k
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError
 from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE
+from vllm_omni.metrics import definitions as metric_defs
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
+from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -360,6 +362,7 @@ class Orchestrator:
     # don't AttributeError when transfer / counter emit paths access them.
     _running_counter: OmniRequestCounter | None = None
     _transfer_emitter: Any = None
+    _prom_metrics: Any = None
     _stat_logger: OmniPrometheusStatLogger | None = None
     duplex_control_plane: DuplexControlPlanePort | None = None
 
@@ -375,6 +378,7 @@ class Orchestrator:
         membership_controller: MembershipController | None = None,
         running_counter: OmniRequestCounter | None = None,
         transfer_emitter: Any = None,
+        prom_metrics: Any = None,
         log_stats: bool = False,
         enable_orch_monitor: bool = False,
         duplex_runtime_extension: DuplexRuntimeExtension | None = None,
@@ -389,6 +393,8 @@ class Orchestrator:
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
         self.log_stats = log_stats
+        self._prom_metrics = prom_metrics
+        self._stage_replica_waiting: dict[tuple[int, int], int] = {}
         self._orch_monitor = create_orch_monitor(
             enabled=enable_orch_monitor,
             replica_sampler=self._sample_replica_metrics,
@@ -542,6 +548,7 @@ class Orchestrator:
             self._membership.install_unregister_handlers(
                 output_queue=self.output_async_queue,
                 cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True),
+                replica_removed_callback=self._remove_stage_replica_waiting,
             )
             membership_watcher = self._membership.start()
 
@@ -980,6 +987,16 @@ class Orchestrator:
                             if diffusion_output is None:
                                 continue
 
+                            output_metrics = getattr(diffusion_output, "metrics", None)
+                            if isinstance(output_metrics, dict):
+                                n_waiting = output_metrics.pop(metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY, None)
+                                if n_waiting is not None:
+                                    self._update_stage_replica_waiting(stage_id, replica_id, int(n_waiting))
+
+                            if diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID:
+                                idle = False
+                                continue
+
                             pool.record_output_timestamps([diffusion_output])
                             processed = [diffusion_output]
                         else:
@@ -989,6 +1006,18 @@ class Orchestrator:
 
                             await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs)
                             for eco in raw_outputs.outputs:
+                                # Emit kv_wait_s before _handle_kv_ready_raw_outputs'
+                                # async_chunk early-return so it lands in all modes.
+                                kv_params = getattr(eco, "kv_transfer_params", None)
+                                if (
+                                    self._prom_metrics is not None
+                                    and isinstance(kv_params, dict)
+                                    and (kv_wait_s := kv_params.get("kv_wait_s")) is not None
+                                ):
+                                    self._prom_metrics.observe_kv_wait(
+                                        kv_params.get("connector_type") or "unknown",
+                                        float(kv_wait_s),
+                                    )
                                 req_state = self.request_states.get(getattr(eco, "request_id", None))
                                 if req_state is None or not req_state.streaming.enabled:
                                     continue
@@ -1027,6 +1056,17 @@ class Orchestrator:
                                     iteration_stats,
                                     engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)],
                                 )
+                            _sched_stats = raw_outputs.scheduler_stats
+                            if (
+                                self._prom_metrics is not None
+                                and _sched_stats is not None
+                                and getattr(_sched_stats, "num_waiting_reqs", None) is not None
+                            ):
+                                self._update_stage_replica_waiting(
+                                    stage_id,
+                                    replica_id,
+                                    int(_sched_stats.num_waiting_reqs),
+                                )
                     except asyncio.CancelledError:
                         raise
                     except EngineDeadError as e:
@@ -1104,6 +1144,26 @@ class Orchestrator:
             close_duplex_sessions=True,
         )
 
+    def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None:
+        """Update one replica snapshot and expose the stage-wide total."""
+        self._stage_replica_waiting[(stage_id, replica_id)] = max(int(n_waiting), 0)
+        self._set_stage_waiting_total(stage_id)
+
+    def _remove_stage_replica_waiting(self, stage_id: int, replica_id: int) -> None:
+        """Drop a dead replica's snapshot and refresh the stage total."""
+        self._stage_replica_waiting.pop((stage_id, replica_id), None)
+        self._set_stage_waiting_total(stage_id)
+
+    def _set_stage_waiting_total(self, stage_id: int) -> None:
+        if self._prom_metrics is None:
+            return
+        total = sum(
+            n_waiting
+            for (snapshot_stage_id, _), n_waiting in self._stage_replica_waiting.items()
+            if snapshot_stage_id == stage_id
+        )
+        self._prom_metrics.set_stage_waiting_requests(stage_id, total)
+
     async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None:
         """Evict a dead stage replica and fail the requests stranded on it (#4285).
 
@@ -1122,6 +1182,7 @@ class Orchestrator:
             error,
         )
         pool.evict_replica(replica_id)
+        self._remove_stage_replica_waiting(stage_id, replica_id)
         stage_has_live = bool(pool.live_replica_ids())
         failed_ids: list[str] = []
         for req_id, req_state in list(self.request_states.items()):

vllm_omni/engine/stage_pool.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unified stage-local runtime abstraction for vLLM-Omni."""
 
 from __future__ import annotations
@@ -653,11 +656,6 @@ class StagePool:
         )
         image_pixels = self._count_image_pixels(request_outputs) if output_unit_type == "image" else 0
         num_inference_steps = coerce_positive_int_scalar(getattr(sampling_params, "num_inference_steps", None)) or 0
-        denoise_step_latency_ms = (
-            defs.compute_denoise_step_latency(stage_gen_time_ms, num_inference_steps)
-            if output_unit_type == "image"
-            else 0.0
-        )
         has_output_timestamps = bool(output_timestamps)
         first_ts = output_timestamps[0] if has_output_timestamps else now
         serving_time_to_first_output_ms = (
@@ -665,6 +663,11 @@ class StagePool:
             if non_empty_first_output_ts is not None
             else 0.0
         )
+        image_time_to_first_output_ms = (
+            max((non_empty_first_output_ts - submit_ts) * 1000.0, 0.0)
+            if non_empty_first_output_ts is not None and output_unit_type == "image"
+            else 0.0
+        )
         remaining_ms = max((now - first_ts) * 1000.0, 0.0)
         if output_unit_count > 1 and self._is_streaming_output_unit_type(output_unit_type):
             time_per_output_unit_ms = remaining_ms / float(output_unit_count - 1)
@@ -709,10 +712,11 @@ class StagePool:
             audio_sample_rate=audio_sample_rate,
             audio_duration_s=audio_duration_s,
             image_pixels=image_pixels,
-            denoise_step_latency_ms=denoise_step_latency_ms,
+            num_inference_steps=num_inference_steps,
             output_unit_type=output_unit_type,
             output_unit_count=output_unit_count,
             serving_time_to_first_output_ms=serving_time_to_first_output_ms,
+            image_time_to_first_output_ms=image_time_to_first_output_ms,
             time_per_output_unit_ms=time_per_output_unit_ms,
             inter_output_latency_ms=inter_output_latency_ms,
             inter_output_latencies_ms=inter_output_latencies_ms,

vllm_omni/entrypoints/async_omni.py

@@ -645,12 +645,12 @@ class AsyncOmni(EngineClient, OmniBase):
         except (asyncio.CancelledError, GeneratorExit):
             if input_stream_task is not None and not input_stream_task.done():
                 input_stream_task.cancel()
-            self._fire_failure_counter_if_alive(request_id)
+            self._record_request_failure_once(request_id, reason="client_disconnect")
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} aborted.")
             raise
         except Exception as e:
-            self._fire_failure_counter_if_alive(request_id)
+            self._record_request_failure_once(request_id, reason="stage_error")
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
             raise
@@ -1115,6 +1115,7 @@ class AsyncOmni(EngineClient, OmniBase):
         """
         await self.engine.abort_async(request_ids)
         for rid in request_ids:
+            self._record_request_failure_once(rid, reason="client_abort")
             state = self.request_states.pop(rid, None)
             input_stream_task = getattr(state, "input_stream_task", None)
             if input_stream_task is not None and not input_stream_task.done():

vllm_omni/entrypoints/client_request_state.py

@@ -17,6 +17,8 @@ class ClientRequestState:
         self.stage_id: int | None = None
         self.queue = queue if queue is not None else asyncio.Queue()
         self.metrics: OrchestratorAggregator | None = None
+        # Request-scoped idempotency guard for Prometheus failure counters.
+        self.failure_recorded = False
         # Wall-clock time at which the user's request arrived in the engine
         # entrypoint. Set in async_omni.generate() before the orchestrator
         # accepts the request. Used as the t0 anchor for audio_ttfp.

vllm_omni/entrypoints/omni.py

@@ -200,6 +200,8 @@ class Omni(OmniBase):
                     self._log_summary_and_cleanup(req_id)
         except Exception:
             if "active_reqs" in locals() and active_reqs:
+                for req_id in active_reqs:
+                    self._record_request_failure_once(req_id, reason="stage_error")
                 self.abort(list(active_reqs))
             raise
         finally:
@@ -210,6 +212,7 @@ class Omni(OmniBase):
         request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
         self.engine.abort(request_ids)
         for req_id in request_ids:
+            self._record_request_failure_once(req_id, reason="client_abort")
             self.request_states.pop(req_id, None)
         if self.log_stats:
             logger.info("[Omni] Aborted request(s) %s", ",".join(request_ids))

vllm_omni/entrypoints/omni_base.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 os
@@ -30,6 +33,11 @@ from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_
 from vllm_omni.metrics.prometheus import OmniPrometheusMetrics
 from vllm_omni.metrics.stats import OrchestratorAggregator
 from vllm_omni.metrics.transfer import OmniTransferMetrics
+from vllm_omni.metrics.utils import (
+    extract_queue_wait_s,
+    normalize_failure_reason,
+    observe_stage_workload_metrics,
+)
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.tracking_parser import TrackingNamespace
@@ -199,6 +207,7 @@ class OmniBase(PDDisaggregationMixin):
         # (which forwards it to the Orchestrator background thread for
         # TX-side emit; see Orchestrator._forward_to_next_stage).
         self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats)
+        self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats)
         st = time.time()
         self.engine = AsyncOmniEngine(
             model=model,
@@ -206,6 +215,7 @@ class OmniBase(PDDisaggregationMixin):
             stage_init_timeout=stage_init_timeout,
             diffusion_batch_size=diffusion_batch_size,
             transfer_emitter=self.transfer_metrics,
+            prom_metrics=self.prom_metrics,
             log_stats=log_stats,
             **kwargs,
         )
@@ -221,7 +231,6 @@ class OmniBase(PDDisaggregationMixin):
 
         self.request_states: dict[str, ClientRequestState] = {}
         self._consumed_metric_messages: dict[str, set[int]] = {}
-        self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats)
         self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats)
 
         self.default_sampling_params_list = self.engine.default_sampling_params_list
@@ -384,29 +393,26 @@ class OmniBase(PDDisaggregationMixin):
             raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}")
         return type(params)(**{**values, **constraints})
 
-    def _fire_failure_counter_if_alive(self, request_id: str) -> None:
-        """Fire the abort/exception bucket of requests_success_total.
-
-        Called from cancel / exception paths in async_omni.generate() BEFORE
-        _abort_internal_requests pops request_states — that method resolves
-        the internal id by dict lookup, so popping first would no-op it. We
-        keep this counter fire separate from _log_summary_and_cleanup (which
-        pops) so the abort path can still find the state to clean up.
-        """
+    def _record_request_failure_once(self, request_id: str, reason: str) -> None:
         req_state = self.request_states.get(request_id)
         prom = getattr(self, "prom_metrics", None)
-        if req_state is None or req_state.metrics is None or prom is None:
+        metrics = getattr(req_state, "metrics", None)
+        if metrics is None or prom is None:
+            return
+        if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False):
             return
-        if str(request_id) not in req_state.metrics.e2e_done:
-            prom.request_failed()
 
-    def _log_summary_and_cleanup(self, request_id: str) -> None:
+        req_state.failure_recorded = True
+        prom.request_failed()
+        prom.inc_requests_failed(normalize_failure_reason(reason))
+
+    def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") -> None:
         req_state = self.request_states.get(request_id)
         try:
             if req_state is None or req_state.metrics is None:
                 return
             if str(request_id) not in req_state.metrics.e2e_done:
-                self.prom_metrics.request_failed()
+                self._record_request_failure_once(request_id, reason)
             if self.log_stats:
                 # Emit per-request orchestrator timing (including e2e_total_ms)
                 # before dropping request state.
@@ -623,6 +629,30 @@ class OmniBase(PDDisaggregationMixin):
                 metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs)
                 metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
                 consumed.add(msg_id)
+                if peak_memory_mb > 0:
+                    self.prom_metrics.set_peak_memory(stage_id, peak_memory_mb)
+                self.prom_metrics.observe_stage_gen_time(
+                    stage_id,
+                    stage_meta.stage_type,
+                    _m.stage_gen_time_ms / 1000.0,
+                )
+                observe_stage_workload_metrics(
+                    self.prom_metrics,
+                    stage_type=stage_meta.stage_type,
+                    stage_metrics=_m,
+                )
+                # Only meaningful for diffusion stages (text stages have no exec key).
+                _diff_m = _m.diffusion_metrics or {}
+                stage_in_queue_s = _diff_m.get("scheduler_queue_wait_s")
+                if stage_in_queue_s is not None:
+                    self.prom_metrics.observe_stage_in_queue(stage_id, float(stage_in_queue_s))
+                if _m.output_unit_type == "image":
+                    if result.replica_id is not None:
+                        self.mod_metrics.observe_image_ttfp(
+                            str(stage_id),
+                            str(result.replica_id),
+                            _m.image_time_to_first_output_ms / 1000.0,
+                        )
 
         if not stage_meta.final_output:
             return None
@@ -656,6 +686,10 @@ class OmniBase(PDDisaggregationMixin):
                     _gen_tok += int(evt.num_tokens_out)
                 self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok)
 
+                queue_wait_s = extract_queue_wait_s(_m.pipeline_timings if _m is not None else None)
+                if queue_wait_s is not None:
+                    self.prom_metrics.observe_queue_wait(queue_wait_s)
+
                 # Modality observe inside the same finalize guard so it fires
                 # once per request and inherits the try/except isolation.
                 observe_modality_at_finalize(

vllm_omni/experimental/fullduplex/minicpmo45/adapter.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 base64
@@ -97,10 +100,14 @@ class MiniCPMO45NativeDuplexServingAdapter:
         cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config)
 
         ref_audio = config.ref_audio
-        if ref_audio is None and isinstance(extra_body.get("ref_audio"), str):
-            ref_audio = extra_body.pop("ref_audio")
-        if ref_audio is None and isinstance(extra_body.get("tts_ref_audio"), str):
-            ref_audio = extra_body.pop("tts_ref_audio")
+        extra_ref_audio = extra_body.get("ref_audio")
+        if ref_audio is None and isinstance(extra_ref_audio, str):
+            extra_body.pop("ref_audio")
+            ref_audio = extra_ref_audio
+        extra_tts_ref_audio = extra_body.get("tts_ref_audio")
+        if ref_audio is None and isinstance(extra_tts_ref_audio, str):
+            extra_body.pop("tts_ref_audio")
+            ref_audio = extra_tts_ref_audio
 
         if ref_audio is None:
             if any(str(modality).lower() == "audio" for modality in config.modalities):
@@ -161,7 +168,11 @@ class MiniCPMO45NativeDuplexServingAdapter:
         stop_token_ids = cls._native_stage0_stop_token_ids(model_config)
         if stop_token_ids:
             stage0_params["stop_token_ids"] = stop_token_ids
-        runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params}
+        # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's
+        # min_new_token=50: that floor is for whole-utterance chat TTS, while a
+        # duplex chunk is 26 codec samples, so check_stop would never release
+        # the Talker and Thinker would stall on the next model turn.
+        runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}}
         scheduler_token_id = cls._native_scheduler_token_id(model_config)
         if scheduler_token_id is not None:
             runtime_config["duplex_scheduler_token_id"] = scheduler_token_id
@@ -236,6 +247,8 @@ class MiniCPMO45NativeDuplexServingAdapter:
             if token_id is not None:
                 return token_id
         eos_id = getattr(tokenizer, "eos_token_id", None)
+        if eos_id is None:
+            return None
         try:
             return int(eos_id)
         except (TypeError, ValueError):
@@ -265,10 +278,13 @@ class MiniCPMO45NativeDuplexServingAdapter:
             value = convert(token)
             if isinstance(value, list):
                 value = value[0] if len(value) == 1 else None
-        try:
-            token_id = int(value)
-        except (TypeError, ValueError):
+        if value is None:
             token_id = -1
+        else:
+            try:
+                token_id = int(value)
+            except (TypeError, ValueError):
+                token_id = -1
         unk_token_id = getattr(tokenizer, "unk_token_id", None)
         if token_id >= 0 and token_id != unk_token_id:
             return token_id

vllm_omni/metrics/definitions.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Single source of truth for vLLM-Omni Prometheus + bench CLI metric naming.
 
 Consumed by:
@@ -61,6 +64,7 @@ IMAGE_COUNT = "image_count"
 IMAGE_GENERATION = "image_generation"
 IMAGE_GENERATION_TIME_MS = f"{IMAGE_GENERATION}_time_ms"
 IMAGE_PIXELS = "image_pixels"
+DIFFUSION_SCHEDULER_WAITING_KEY = "scheduler_num_waiting_reqs"
 TOTAL_IMAGES = "total_images"
 IMAGE_THROUGHPUT = "image_throughput"
 AVERAGE_PIXELS_PER_IMAGE = "average_pixels_per_image"
@@ -91,6 +95,7 @@ POSTPROCESS_TIME_MS = f"{POSTPROCESS_TIME}_ms"
 POSTPROCESS_TIMES_MS = f"{POSTPROCESS_TIME}s_ms"
 OUTPUT_UNIT_COUNT = "output_unit_count"
 SERVING_TIME_TO_FIRST_OUTPUT_MS = "serving_time_to_first_output_ms"
+IMAGE_TIME_TO_FIRST_OUTPUT_MS = "image_time_to_first_output_ms"
 SERVING_TIME_TO_FIRST_OUTPUTS_MS = "serving_time_to_first_outputs_ms"
 TIME_PER_OUTPUT_UNIT_MS = "time_per_output_unit_ms"
 TIME_PER_OUTPUT_UNITS_MS = "time_per_output_units_ms"
@@ -152,6 +157,30 @@ TRANSFER_RX_S = METRIC_PREFIX + "transfer_rx_s"
 TRANSFER_IN_FLIGHT_S = METRIC_PREFIX + "transfer_in_flight_s"
 
 
+# ============================================================================
+# Image / diffusion service-level families
+# ============================================================================
+STAGE_GEN_TIME_S = METRIC_PREFIX + "stage_gen_time_s"
+REQUEST_QUEUE_WAIT_S = METRIC_PREFIX + "request_queue_wait_s"
+STAGE_WAITING_REQUESTS = METRIC_PREFIX + "stage_waiting_requests"
+NUM_INFERENCE_STEPS = METRIC_PREFIX + "num_inference_steps"
+IMAGE_COUNT_METRIC = METRIC_PREFIX + IMAGE_COUNT
+IMAGE_PIXELS_METRIC = METRIC_PREFIX + IMAGE_PIXELS
+PEAK_MEMORY_MB = METRIC_PREFIX + "peak_memory_mb"
+REQUESTS_FAILED = METRIC_PREFIX + "requests_failed"
+KV_WAIT_S = METRIC_PREFIX + "kv_wait_s"
+DIFFUSION_FORWARD_S = METRIC_PREFIX + "diffusion_forward_s"
+DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s"
+DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s"
+DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s"
+DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s"
+VAE_DECODE_S = METRIC_PREFIX + "vae_decode_s"
+DENOISE_STEP_LATENCY_S = METRIC_PREFIX + DENOISE_STEP_LATENCY + "_s"
+DIFFUSION_KV_LOAD_S = METRIC_PREFIX + "diffusion_kv_load_s"
+IMAGE_TTFP_S = METRIC_PREFIX + "image_ttfp_s"
+STAGE_IN_QUEUE_S = METRIC_PREFIX + "stage_in_queue_s"
+
+
 # ============================================================================
 # Label sets
 # ============================================================================
@@ -163,6 +192,11 @@ SUCCESS_LABELS = ("model_name", "finished_reason")
 # ``stage`` + ``replica``.
 STAGE_LABELS = ("model_name", "stage", "replica")
 
+STAGE_GEN_TIME_LABELS = ("model_name", "stage", "stage_type")
+DIFFUSION_LABELS = ("model_name", "stage")
+FAILED_LABELS = ("model_name", "reason")
+KV_WAIT_LABELS = ("model_name", "connector_type")
+
 # Audio continuity Counter carries an extra ``threshold_ms`` label so multiple
 # threshold buckets can be tracked simultaneously. The ``_ms`` suffix names a
 # numeric threshold *value* in ms, not a time-bearing metric.
@@ -311,16 +345,6 @@ def compute_audio_frames(
     return pcm_nbytes // frame_width
 
 
-def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int) -> float:
-    """Mean denoise step latency = image stage generation time / step count.
-
-    The returned value uses the same time unit as ``stage_gen_time``.
-    """
-    if num_inference_steps <= 0:
-        return 0.0
-    return stage_gen_time / float(num_inference_steps)
-
-
 # ============================================================================
 # Audio sample-rate resolution
 # ============================================================================

vllm_omni/metrics/modality.py

@@ -27,7 +27,7 @@ from typing import Any
 from prometheus_client import Counter, Histogram
 
 from vllm_omni.metrics import definitions as defs
-from vllm_omni.metrics.utils import count_audio_frames, extract_mm_output
+from vllm_omni.metrics.utils import observe_audio_finalize, observe_diffusion_finalize
 
 _stage_labels = list(defs.STAGE_LABELS)
 
@@ -105,6 +105,40 @@ _diffusion_postprocess_family = Histogram(
     labelnames=_stage_labels,
     buckets=defs.SECONDS_FAST_BUCKETS,
 )
+_vae_decode_family = Histogram(
+    defs.VAE_DECODE_S,
+    "VAE decode latency in seconds (latents -> pixels/audio/video).",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_diffusion_forward_family = Histogram(
+    defs.DIFFUSION_FORWARD_S,
+    "Diffusion forward-only latency in seconds (denoise loop; excludes "
+    "preprocess / postprocess / VAE decode / KV load). Absent when the "
+    "pipeline profiler is off.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_diffusion_kv_load_family = Histogram(
+    defs.DIFFUSION_KV_LOAD_S,
+    "Diffusion KV-recv latency in seconds (AR→diffusion KV fetch; absent when "
+    "the stage has no upstream KV to receive).",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+_image_ttfp_family = Histogram(
+    defs.IMAGE_TTFP_S,
+    "Image time-to-first-output in seconds (stage submit → first image materialized; non-streaming single-image). "
+    "Stage-level, not e2e — excludes API queue and inter-stage transfer before the image stage.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_denoise_step_latency_family = Histogram(
+    defs.DENOISE_STEP_LATENCY_S,
+    "Mean per-step denoise forward latency in seconds (forward_time / num_inference_steps).",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
 
 
 class OmniModalityMetrics:
@@ -206,6 +240,31 @@ class OmniModalityMetrics:
             replica=replica,
         ).observe(seconds)
 
+    def observe_vae_decode(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats or seconds < 0:
+            return
+        _vae_decode_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds)
+
+    def observe_diffusion_forward(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats or seconds < 0:
+            return
+        _diffusion_forward_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds)
+
+    def observe_diffusion_kv_load(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats or seconds < 0:
+            return
+        _diffusion_kv_load_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds)
+
+    def observe_image_ttfp(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats or seconds < 0:
+            return
+        _image_ttfp_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds)
+
+    def observe_denoise_step_latency(self, stage: str, replica: str, seconds: float) -> None:
+        if not self._log_stats or seconds <= 0:
+            return
+        _denoise_step_latency_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds)
+
 
 def observe_modality_at_finalize(
     mod_metrics: OmniModalityMetrics,
@@ -216,61 +275,34 @@ def observe_modality_at_finalize(
     stage_metrics: Any,
     engine_outputs: Any,
 ) -> None:
-    """Route audio-path observations for a finalized request.
+    """Route per-modality observations for a finalized request.
 
     Used by ``omni_base._process_single_result`` inside the e2e_done finalize
-    guard so it fires once per request. Skips text path (covered by upstream
-    ``vllm:*{stage="thinker", ...}``) and any case where required inputs are
-    missing — caller should not need to pre-validate.
+    guard so it fires once per request. Text path falls through — covered by
+    upstream ``vllm:*{stage="thinker", ...}``. Caller should not need to
+    pre-validate; missing inputs are silently skipped.
 
     audio_ttfp is intentionally NOT observed here; it's emitted by the
     streaming hook at first-packet time, not at finalize.
     """
-    if replica_id is None or stage_metrics is None or output_type is None:
+    if replica_id is None or stage_metrics is None:
         return
 
-    stage_label = str(stage_id)
-    replica_label = str(replica_id)
-    gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
-    mm_out = extract_mm_output(engine_outputs)
+    observe_diffusion_finalize(
+        mod_metrics,
+        stage_id=stage_id,
+        replica_id=replica_id,
+        stage_metrics=stage_metrics,
+    )
 
     if output_type == "audio":
-        sample_rate = defs.resolve_audio_sample_rate(mm_out)
-        n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
-        if n_frames == 0:
-            n_frames = count_audio_frames(mm_out)
-        mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
-        duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
-        if duration_s > 0:
-            mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
-            mod_metrics.observe_audio_rtf(
-                stage_label,
-                replica_label,
-                defs.compute_audio_rtf(gen_time_s, duration_s),
-            )
-        else:
-            mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
-
-    dm = getattr(stage_metrics, "diffusion_metrics", None)
-    if dm:
-        _key_map = {
-            "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec,
-            "preprocess_time_s": mod_metrics.observe_diffusion_preprocess,
-            "postprocess_time_s": mod_metrics.observe_diffusion_postprocess,
-        }
-        for key, observe_fn in _key_map.items():
-            val = dm.get(key)
-            if val is not None:
-                observe_fn(stage_label, replica_label, float(val))
-
-        exec_time = dm.get("diffusion_engine_exec_time_s")
-        num_steps = dm.get("num_inference_steps")
-        if exec_time is not None and num_steps and num_steps > 0:
-            mod_metrics.observe_diffusion_exec_per_step(
-                stage_label,
-                replica_label,
-                float(exec_time) / int(num_steps),
-            )
+        observe_audio_finalize(
+            mod_metrics,
+            stage_id=stage_id,
+            replica_id=replica_id,
+            stage_metrics=stage_metrics,
+            engine_outputs=engine_outputs,
+        )
 
 
 def observe_audio_first_packet(

vllm_omni/metrics/prometheus.py

@@ -38,6 +38,66 @@ _generation_tokens_family = Counter(
     labelnames=_labelnames,
 )
 
+_stage_gen_time_family = Histogram(
+    defs.STAGE_GEN_TIME_S,
+    "Per-stage generation time in seconds for diffusion / image stages.",
+    labelnames=list(defs.STAGE_GEN_TIME_LABELS),
+    buckets=defs.SECONDS_BUCKETS,
+)
+_request_queue_wait_family = Histogram(
+    defs.REQUEST_QUEUE_WAIT_S,
+    "Per-request queue wait between submit arrival and stage submit, in seconds.",
+    labelnames=_labelnames,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_stage_waiting_requests_family = Gauge(
+    defs.STAGE_WAITING_REQUESTS,
+    "Sum of the latest scheduler waiting snapshots across a stage's replicas; "
+    "diffusion snapshots update when outputs arrive.",
+    labelnames=list(defs.DIFFUSION_LABELS),
+)
+_num_inference_steps_family = Histogram(
+    defs.NUM_INFERENCE_STEPS,
+    "Diffusion step count distribution per request.",
+    labelnames=_labelnames,
+)
+_image_count_family = Counter(
+    defs.IMAGE_COUNT_METRIC,
+    "Cumulative image count; per-model rate. Throughput recovered via rate().",
+    labelnames=_labelnames,
+)
+_image_pixels_family = Histogram(
+    defs.IMAGE_PIXELS_METRIC,
+    "Image pixel count distribution per request.",
+    labelnames=_labelnames,
+)
+_peak_memory_family = Gauge(
+    defs.PEAK_MEMORY_MB,
+    "Peak GPU memory in MB observed during a stage's generation.",
+    labelnames=list(defs.DIFFUSION_LABELS),
+)
+_stage_in_queue_family = Histogram(
+    defs.STAGE_IN_QUEUE_S,
+    "Time in seconds from diffusion scheduler enqueue to initial execution admission. "
+    "Zero is the healthy low-load baseline and is counted.",
+    labelnames=list(defs.DIFFUSION_LABELS),
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+
+_requests_failed_family = Counter(
+    defs.REQUESTS_FAILED,
+    "Total requests by failure reason. Pairs with requests_success_total "
+    "(abort bucket) to give the full failure surface. Reasons are bounded to "
+    "client_abort, client_disconnect, stage_error, and unknown.",
+    labelnames=list(defs.FAILED_LABELS),
+)
+_kv_wait_s_family = Histogram(
+    defs.KV_WAIT_S,
+    "KV cache transfer wait time per request, sliced by connector backend.",
+    labelnames=list(defs.KV_WAIT_LABELS),
+    buckets=defs.SECONDS_BUCKETS,
+)
+
 
 class OmniPrometheusMetrics:
     """Label-bound wrapper around the raw Prometheus metrics.
@@ -54,6 +114,10 @@ class OmniPrometheusMetrics:
         self._e2e_latency = _e2e_latency_family.labels(model_name=model_name)
         self._prompt_tokens = _prompt_tokens_family.labels(model_name=model_name)
         self._generation_tokens = _generation_tokens_family.labels(model_name=model_name)
+        self._request_queue_wait = _request_queue_wait_family.labels(model_name=model_name)
+        self._num_inference_steps = _num_inference_steps_family.labels(model_name=model_name)
+        self._image_count = _image_count_family.labels(model_name=model_name)
+        self._image_pixels = _image_pixels_family.labels(model_name=model_name)
 
     def set_running(self, n: int) -> None:
         if not self._log_stats:
@@ -89,13 +153,80 @@ class OmniPrometheusMetrics:
     def request_failed(self) -> None:
         if not self._log_stats:
             return
-        # Pipeline-level "fail" maps to the upstream FinishReason.ABORT bucket;
-        # a single counter family now covers both normal stops and aborts.
         _completion_family.labels(
             model_name=self._model_name,
             finished_reason="abort",
         ).inc()
 
+    def observe_stage_gen_time(self, stage: int, stage_type: str, gen_time_s: float) -> None:
+        if not self._log_stats:
+            return
+        _stage_gen_time_family.labels(
+            model_name=self._model_name,
+            stage=str(stage),
+            stage_type=stage_type,
+        ).observe(max(gen_time_s, 0.0))
+
+    def observe_stage_in_queue(self, stage: int, in_queue_s: float) -> None:
+        if not self._log_stats:
+            return
+        _stage_in_queue_family.labels(
+            model_name=self._model_name,
+            stage=str(stage),
+        ).observe(max(in_queue_s, 0.0))
+
+    def observe_queue_wait(self, queue_wait_s: float) -> None:
+        if not self._log_stats:
+            return
+        self._request_queue_wait.observe(max(queue_wait_s, 0.0))
+
+    def set_stage_waiting_requests(self, stage: int, n_waiting: int) -> None:
+        if not self._log_stats:
+            return
+        _stage_waiting_requests_family.labels(
+            model_name=self._model_name,
+            stage=str(stage),
+        ).set(max(n_waiting, 0))
+
+    def observe_num_inference_steps(self, n_steps: int) -> None:
+        if not self._log_stats or n_steps <= 0:
+            return
+        self._num_inference_steps.observe(n_steps)
+
+    def inc_image_count(self, n_images: int = 1) -> None:
+        if not self._log_stats or n_images <= 0:
+            return
+        self._image_count.inc(n_images)
+
+    def observe_image_pixels(self, n_pixels: int) -> None:
+        if not self._log_stats or n_pixels <= 0:
+            return
+        self._image_pixels.observe(n_pixels)
+
+    def set_peak_memory(self, stage: int, peak_memory_mb: float) -> None:
+        if not self._log_stats:
+            return
+        _peak_memory_family.labels(
+            model_name=self._model_name,
+            stage=str(stage),
+        ).set(max(peak_memory_mb, 0.0))
+
+    def inc_requests_failed(self, reason: str) -> None:
+        if not self._log_stats:
+            return
+        _requests_failed_family.labels(
+            model_name=self._model_name,
+            reason=reason or "unknown",
+        ).inc()
+
+    def observe_kv_wait(self, connector_type: str, kv_wait_s: float) -> None:
+        if not self._log_stats or kv_wait_s < 0:
+            return
+        _kv_wait_s_family.labels(
+            model_name=self._model_name,
+            connector_type=connector_type or "unknown",
+        ).observe(kv_wait_s)
+
 
 class OmniRequestCounter:
     """Running-request counter written by the orchestrator thread, read by the client thread."""

vllm_omni/metrics/stats.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 time
@@ -49,11 +52,13 @@ class StageRequestStats:
     audio_sample_rate: int = 0
     audio_duration_s: float = 0.0
     image_pixels: int = 0
+    num_inference_steps: int = 0
     denoise_step_latency_ms: float = 0.0
     pipeline_timings: dict[str, float] | None = None
     output_unit_type: str | None = None
     output_unit_count: int = 0
     serving_time_to_first_output_ms: float = 0.0
+    image_time_to_first_output_ms: float = 0.0
     time_per_output_unit_ms: float = 0.0
     inter_output_latency_ms: float = 0.0
     inter_output_latencies_ms: list[float] | None = None
@@ -480,6 +485,7 @@ class OrchestratorAggregator:
                 "output_unit_type": evt.output_unit_type,
                 defs.OUTPUT_UNIT_COUNT: int(evt.output_unit_count),
                 defs.SERVING_TIME_TO_FIRST_OUTPUT_MS: float(evt.serving_time_to_first_output_ms),
+                defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS: float(evt.image_time_to_first_output_ms),
                 defs.TIME_PER_OUTPUT_UNIT_MS: float(evt.time_per_output_unit_ms),
                 defs.INTER_OUTPUT_LATENCY_MS: float(evt.inter_output_latency_ms),
                 defs.INTER_OUTPUT_LATENCIES_MS: list(evt.inter_output_latencies_ms or []),
@@ -513,6 +519,11 @@ class OrchestratorAggregator:
         if current_first_output_ms <= 0 < first_output_ms:
             current[defs.SERVING_TIME_TO_FIRST_OUTPUT_MS] = first_output_ms
 
+        image_first_output_ms = float(evt.image_time_to_first_output_ms)
+        current_image_first_output_ms = float(current.get(defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS, 0.0))
+        if current_image_first_output_ms <= 0 < image_first_output_ms:
+            current[defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS] = image_first_output_ms
+
         if evt.output_unit_type:
             current["output_unit_type"] = evt.output_unit_type
         if evt.final_output_type:
@@ -578,11 +589,20 @@ class OrchestratorAggregator:
         stats.request_id = req_id
         if final_output_type is not None:
             stats.final_output_type = final_output_type
+        diffusion_metrics_by_request = getattr(self, "diffusion_metrics", None)
         stats.diffusion_metrics = (
-            {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
-            if req_id in self.diffusion_metrics
+            {k: float(v) for k, v in diffusion_metrics_by_request.pop(req_id, {}).items()}
+            if diffusion_metrics_by_request is not None and req_id in diffusion_metrics_by_request
             else None
         )
+        forward_time_s = (stats.diffusion_metrics or {}).get("forward_time_s")
+        num_inference_steps = int(getattr(stats, "num_inference_steps", 0) or 0)
+        if num_inference_steps > 0 and getattr(stats, "output_unit_type", None) == "image":
+            if forward_time_s is not None:
+                denoise_time_ms = float(forward_time_s) * 1000.0
+            else:
+                denoise_time_ms = float(getattr(stats, "stage_gen_time_ms", 0.0) or 0.0)
+            stats.denoise_step_latency_ms = denoise_time_ms / num_inference_steps
         return stats
 
     def on_stage_metrics(
@@ -628,10 +648,24 @@ class OrchestratorAggregator:
             _postproc_ms = (time.perf_counter() - _t0) * 1000.0
             self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms)
 
+    _MS_TO_S: dict[str, str] = {
+        "preprocess_time_ms": "preprocess_time_s",
+        "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s",
+        "postprocess_time_ms": "postprocess_time_s",
+        "vae_decode_time_ms": "vae_decode_time_s",
+        "forward_time_ms": "forward_time_s",
+        "scheduler_queue_wait_ms": "scheduler_queue_wait_s",
+        "kv_recv_time_ms": "kv_recv_time_s",
+    }
+
     def accumulate_diffusion_metrics(self, stage_type: str, req_id: Any, engine_outputs: Any) -> None:
         """Accumulate diffusion metrics for a request.
 
-        Handles extraction and accumulation of diffusion stage metrics.
+        Engine emits ``*_ms`` timings; the accumulator converts them to ``_s``
+        keys via ``_MS_TO_S`` so downstream observers read a uniform
+        seconds-bearing dict. Per-chunk timing keys are summed; non-timing
+        keys (e.g. ``image_num`` / ``resolution`` from ``format_diffusion_outputs``)
+        preserve the existing ``+=`` semantics.
 
         Args:
             req_id: Request ID
@@ -643,20 +677,16 @@ class OrchestratorAggregator:
         diffusion_metrics: dict = getattr(engine_output, "metrics", {})
         if isinstance(diffusion_metrics, list):
             diffusion_metrics = diffusion_metrics[0]
-        if diffusion_metrics:
-            _MS_TO_S = {
-                "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s",
-                "preprocess_time_ms": "preprocess_time_s",
-                "postprocess_time_ms": "postprocess_time_s",
-                "diffusion_engine_total_time_ms": "diffusion_engine_total_time_s",
-            }
-            for key, value in diffusion_metrics.items():
-                if value is None:
-                    continue
-                if key in _MS_TO_S:
-                    self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0
-                else:
-                    self.diffusion_metrics[req_id][key] = value
+        if not diffusion_metrics:
+            return
+        bucket = self.diffusion_metrics[req_id]
+        for key, value in diffusion_metrics.items():
+            if value is None:
+                continue
+            if key in self._MS_TO_S:
+                bucket[self._MS_TO_S[key]] += float(value) / 1000.0
+            else:
+                bucket[key] += float(value)
 
     def on_forward(
         self,

vllm_omni/metrics/utils.py

@@ -1,9 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import fields
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from prettytable import PrettyTable
 
+from vllm_omni.metrics import definitions as metric_defs
+
+if TYPE_CHECKING:
+    from vllm_omni.metrics.modality import OmniModalityMetrics
+    from vllm_omni.metrics.prometheus import OmniPrometheusMetrics
+    from vllm_omni.metrics.stats import StageRequestStats
+
+FAILURE_REASONS = frozenset({"client_abort", "client_disconnect", "stage_error", "unknown"})
+DIFFUSION_METRICS_ONLY_REQUEST_ID = "__vllm_omni_diffusion_metrics__"
+
+
+def extract_queue_wait_s(pipeline_timings: Mapping[str, float] | None) -> float | None:
+    if pipeline_timings is None or "queue_wait_ms" not in pipeline_timings:
+        return None
+    return float(pipeline_timings["queue_wait_ms"] or 0.0) / 1000.0
+
+
+def observe_stage_workload_metrics(
+    prom_metrics: OmniPrometheusMetrics,
+    *,
+    stage_type: str,
+    stage_metrics: StageRequestStats,
+) -> None:
+    if stage_type == "diffusion":
+        prom_metrics.observe_num_inference_steps(stage_metrics.num_inference_steps)
+
+    if stage_metrics.output_unit_type == "image":
+        prom_metrics.observe_image_pixels(stage_metrics.image_pixels)
+        prom_metrics.inc_image_count(stage_metrics.output_unit_count)
+
+
+def observe_audio_finalize(
+    mod_metrics: OmniModalityMetrics,
+    *,
+    stage_id: int,
+    replica_id: int,
+    stage_metrics: Any,
+    engine_outputs: Any,
+) -> None:
+    stage_label = str(stage_id)
+    replica_label = str(replica_id)
+    gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
+    mm_out = extract_mm_output(engine_outputs)
+
+    sample_rate = metric_defs.resolve_audio_sample_rate(mm_out)
+    n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
+    if n_frames == 0:
+        n_frames = count_audio_frames(mm_out)
+    mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
+    duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
+    if duration_s > 0:
+        mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
+        mod_metrics.observe_audio_rtf(
+            stage_label,
+            replica_label,
+            metric_defs.compute_audio_rtf(gen_time_s, duration_s),
+        )
+    else:
+        mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
+
+
+def observe_diffusion_finalize(
+    mod_metrics: OmniModalityMetrics,
+    *,
+    stage_id: int,
+    replica_id: int,
+    stage_metrics: Any,
+) -> None:
+    diff_metrics = getattr(stage_metrics, "diffusion_metrics", None)
+    if not diff_metrics:
+        return
+
+    stage_label = str(stage_id)
+    replica_label = str(replica_id)
+
+    exec_s = diff_metrics.get("diffusion_engine_exec_time_s")
+    if exec_s is not None:
+        mod_metrics.observe_diffusion_exec(stage_label, replica_label, float(exec_s))
+        num_steps = int(
+            getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0
+        )
+        if num_steps > 0:
+            mod_metrics.observe_diffusion_exec_per_step(stage_label, replica_label, float(exec_s) / num_steps)
+
+    pre_s = diff_metrics.get("preprocess_time_s")
+    if pre_s is not None:
+        mod_metrics.observe_diffusion_preprocess(stage_label, replica_label, float(pre_s))
+
+    post_s = diff_metrics.get("postprocess_time_s")
+    if post_s is not None:
+        mod_metrics.observe_diffusion_postprocess(stage_label, replica_label, float(post_s))
+
+    vae_s = diff_metrics.get("vae_decode_time_s")
+    if vae_s is not None:
+        mod_metrics.observe_vae_decode(stage_label, replica_label, float(vae_s))
+
+    forward_s = diff_metrics.get("forward_time_s")
+    if forward_s is not None:
+        mod_metrics.observe_diffusion_forward(stage_label, replica_label, float(forward_s))
+        num_steps = int(
+            getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0
+        )
+        if num_steps > 0 and getattr(stage_metrics, "output_unit_type", None) == "image":
+            mod_metrics.observe_denoise_step_latency(stage_label, replica_label, float(forward_s) / num_steps)
+
+    kv_load_s = diff_metrics.get("kv_recv_time_s")
+    if kv_load_s is not None:
+        mod_metrics.observe_diffusion_kv_load(stage_label, replica_label, float(kv_load_s))
+
+
+def normalize_failure_reason(reason: str | None) -> str:
+    """Map request failure details to the bounded metrics taxonomy."""
+    return reason if reason in FAILURE_REASONS else "unknown"
+
+
+def diffusion_scheduler_waiting_metrics(n_waiting: int) -> dict[str, int]:
+    """Build the diffusion scheduler snapshot consumed by the orchestrator."""
+    return {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: max(int(n_waiting), 0)}
+
+
+def diffusion_exception_metrics(exc: BaseException) -> dict[str, Any]:
+    """Return metrics attached to a terminal diffusion error."""
+    metrics = getattr(exc, "diffusion_metrics", None)
+    return dict(metrics) if isinstance(metrics, dict) else {}
+
+
+def sum_diffusion_stage_durations_ms(output: Any, suffix: str) -> float | None:
+    """Sum matching diffusion stage durations in milliseconds."""
+    durations = getattr(output, "stage_durations", None)
+    if not isinstance(durations, dict):
+        return None
+
+    total = 0.0
+    found = False
+    for key, value in durations.items():
+        if not key.endswith(suffix):
+            continue
+        try:
+            total += float(value)
+            found = True
+        except (TypeError, ValueError):
+            continue
+    return total * 1000.0 if found else None
+
+
+def extract_diffusion_vae_decode_ms(output: Any) -> float | None:
+    return sum_diffusion_stage_durations_ms(output, ".vae.decode")
+
+
+def extract_diffusion_denoise_ms(output: Any) -> float | None:
+    """Extract denoise-loop time (transformer and per-step scheduler)."""
+    return sum_diffusion_stage_durations_ms(output, ".diffuse")
+
 
 def coerce_positive_int_scalar(value: object) -> int | None:
     """Coerce a value to a positive int without importing tensor libs.

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py

@@ -148,9 +148,9 @@ class ConditionalCFM(BASECFM):
             # ``.contiguous().data_ptr()`` could free the temp -> dangling ptr).
             io_dtype = getattr(self.estimator, "io_dtype", x.dtype)
             [estimator, stream], trt_engine = self.estimator.acquire_estimator()
-            # NOTE need to synchronize when switching stream
-            torch.cuda.current_stream().synchronize()
-            with stream:
+            caller_stream = torch.cuda.current_stream(x.device)
+            stream.wait_stream(caller_stream)
+            with torch.cuda.stream(stream):
                 x_e = x.to(io_dtype).contiguous()
                 mask_e = mask.to(io_dtype).contiguous()
                 mu_e = mu.to(io_dtype).contiguous()
@@ -176,8 +176,13 @@ class ConditionalCFM(BASECFM):
                 for i, j in enumerate(data_ptrs):
                     estimator.set_tensor_address(trt_engine.get_tensor_name(i), j)
                 # run trt engine
-                assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True
-                torch.cuda.current_stream().synchronize()
+                assert estimator.execute_async_v3(stream.cuda_stream) is True
+                for tensor in (x_e, mask_e, mu_e, t_e, spks_e, cond_e, out_e):
+                    if tensor.is_cuda:
+                        tensor.record_stream(stream)
+            caller_stream.wait_stream(stream)
+            if out_e.is_cuda:
+                out_e.record_stream(caller_stream)
             self.estimator.release_estimator(estimator, stream)
             return out_e.to(x.dtype)
 

vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py

@@ -121,7 +121,7 @@ class TrtContextWrapper:
         for _ in range(trt_concurrent):
             ctx = engine.create_execution_context()
             assert ctx is not None, "failed to create TRT execution context (out of memory?)"
-            stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device)))
+            stream = torch.cuda.Stream(torch.device(device))
             self._pool.put([ctx, stream])
 
     def acquire_estimator(self):

vllm_omni/model_extras/magi_human.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-# MagiHuman is a text -> video+audio model. Its model-specific knobs are read
-# from sampling_params.extra_args in pipeline_magi_human.py; declaring them here
-# routes request `extra_body` fields into OmniDiffusionSamplingParams.extra_args
-# so the model can be driven through the standard task example.
-MAGI_HUMAN_EXTRA_BODY_PARAMS = frozenset(
-    {
-        "seconds",
-        "audio_path",
-        "image_path",
-        "sr_height",
-        "sr_width",
-        "sr_num_inference_steps",
-    }
-)
-MAGI_HUMAN_EXTRA_OUTPUT_PARAMS = frozenset()

vllm_omni/model_extras/registry.py

@@ -36,10 +36,6 @@ from vllm_omni.model_extras.ltx2 import (
     ltx_preserves_reference_image_size,
     ltx_transformer_config_subfolder,
 )
-from vllm_omni.model_extras.magi_human import (
-    MAGI_HUMAN_EXTRA_BODY_PARAMS,
-    MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
-)
 from vllm_omni.model_extras.ming_flash_omni import (
     MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
     MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
@@ -182,10 +178,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
     },
-    "MagiHumanPipeline": {
-        "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
-        "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
-    },
     "HeliosPipeline": {
         "extra_body_params": HELIOS_EXTRA_BODY_PARAMS,
         "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,

vllm_omni/platforms/npu/models/minimax_h3.py

@@ -19,6 +19,7 @@ from vllm_omni.platforms.npu.layers.rotary_embedding import (
 logger = init_logger(__name__)
 
 _ROPE_PATCHED = False
+_SDPA_PATCHED = False
 _SWIGLU_PATCHED = False
 
 
@@ -48,6 +49,50 @@ def apply_minimax_h3_qwen3vl_patch() -> None:
     logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder")
 
 
+def _scaled_dot_product_attention_npu(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+) -> torch.Tensor:
+    """Run causal SDPA with compressed K/V heads through NPU native GQA."""
+    num_heads = query.shape[1]
+    num_key_value_heads = key.shape[1]
+    num_value_heads = value.shape[1]
+    if num_key_value_heads != num_value_heads:
+        raise ValueError(
+            "GQA requires key and value to have the same number of heads, "
+            f"got k_heads={num_key_value_heads} and v_heads={num_value_heads}."
+        )
+    if num_key_value_heads == 0:
+        raise ValueError("GQA requires at least one KV head.")
+    if num_heads % num_key_value_heads != 0:
+        raise ValueError(
+            "GQA requires query heads to be a multiple of KV heads, "
+            f"got q_heads={num_heads} and kv_heads={num_key_value_heads}."
+        )
+    return F.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        dropout_p=0.0,
+        is_causal=True,
+        enable_gqa=num_heads != num_key_value_heads,
+    )
+
+
+def apply_minimax_h3_qwen3vl_sdpa_patch() -> None:
+    """Route MiniMax H3 Qwen3-VL text attention to NPU native GQA."""
+    global _SDPA_PATCHED
+    if _SDPA_PATCHED:
+        return
+
+    from vllm_omni.diffusion.models.minimax_h3 import encoder
+
+    encoder._scaled_dot_product_attention = _scaled_dot_product_attention_npu
+    _SDPA_PATCHED = True
+    logger.debug("Applied NPU SDPA patch for MiniMax H3 Qwen3-VL text encoder")
+
+
 def npu_swiglu_from_packed(gate_up: torch.Tensor) -> torch.Tensor:
     """Apply fused SwiGLU to a packed gate/up projection tensor."""
     return torch_npu.npu_swiglu(gate_up, dim=-1)

vllm_omni/platforms/npu/platform.py

@@ -107,16 +107,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         from vllm_omni.platforms.npu.models.minimax_h3 import (
             apply_minimax_h3_qwen3vl_patch,
+            apply_minimax_h3_qwen3vl_sdpa_patch,
             apply_minimax_h3_qwen3vl_swiglu_patch,
         )
 
-        # Both patches import the MiniMax encoder package, whose __init__ loads
+        # These patches import the MiniMax encoder package, whose __init__ loads
         # pipeline_minimax_h3 → diffusion.data. Doing that during platform
         # construction races vllm_omni/__init__.py (patch before config) and
         # closes a cycle through pipeline_registry → PI0_PIPELINE →
         # DiffusionOutput. Apply them only after the platform exists, before
         # the diffusion pipeline is loaded.
         apply_minimax_h3_qwen3vl_patch()
+        apply_minimax_h3_qwen3vl_sdpa_patch()
         apply_minimax_h3_qwen3vl_swiglu_patch()
         set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1)
         set_mc2_mask(vllm_config, device)