[Performance] Add per-component DLO offload policy

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

.buildkite/cuda/test-nightly.yml

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

.buildkite/cuda/test-weekly.yml

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

collect_env.py

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

docs/cli/README.md

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

docs/cli/bench/serve.md

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

docs/configuration/README.md

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

docs/configuration/environment_variables.md

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

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

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -29,6 +29,7 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
 | **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
 | **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
 | **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
+| **Auxiliary component policy** | ✅ Encoders/VAEs may remain resident. | ✅ Encoders/VAEs may remain resident. |
 
 See [Parallelism compatibility](#parallelism-compatibility) and
 [Request and loading constraints](#request-and-loading-constraints) for the
@@ -169,18 +170,42 @@ This mode means:
   copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
+### Component allocator-cache retention
+
+MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged
+encoder, DiT, and VAE component boundaries so allocations released by one
+component can be reused by the next. The DiT prefetch path and its two shared
+device buffers are unchanged.
+
+After a component is offloaded, cached-but-unallocated memory is retained only
+while it is at most 25% of device capacity and at least 5% of device capacity
+remains physically free. Crossing either bound releases the allocator cache.
+Missing allocator telemetry also releases it conservatively. Component or
+staging failure forces a release, and an out-of-memory allocation gets one
+retry after release. This is a component-local policy rather than a global
+`empty_cache` override, so the executor's unconditional shutdown cleanup is
+preserved.
+
+Retaining allocator blocks trades higher reserved and peak device memory for
+lower lifecycle latency. The benefit is device- and workload-dependent, and a
+smaller or externally loaded device may cross a bound and return to the
+ordinary per-component cache-release behavior. Current performance validation
+covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer
+topologies and platforms receive the same bounded policy but are not claimed
+performance targets; other DLO models remain unchanged unless they explicitly
+adopt it.
+
 ### Final-layout Host Weight Runtime consumer
 
-> **Status:** PR2 implementation. The representation-independent identity,
-> producer, and restoration contracts landed in
-> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
-> wires the final-layout BF16 path into eligible no-AllGather DLO.
+> **Status:** PR2 landed in
+> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
+> registered direct H2D plus generic HWR publication/source-digest
+> optimizations without changing AllGather behavior.
 
 Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
 only to no-AllGather DLO. Enable it with
 `--host-weight-runtime-mode preferred` (or `required`) and
-`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
-and direct asynchronous H2D remain a separate PR3 transport change.
+`--host-weight-runtime-root <node-local-root>`.
 
 The modes express operator fallback policy, not different artifact formats:
 
@@ -229,6 +254,44 @@ outcomes are authoritative:
 Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
 not selected.
 
+Local source identity and publication avoid repeated full-payload passes while
+preserving the same semantic contract:
+
+- Canonical files without a trusted immutable Hub blob identity are hashed in
+  parallel. Their content IDs are cached per storage domain under the HWR root
+  and reused only when the path, inode, size, timestamps, symlink target, and
+  cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
+  cache-I/O failure falls back to hashing the canonical file directly.
+- Producers that emit payload bytes in canonical storage-key order may declare
+  `ordered=True`. The filesystem writer then computes file and tensor SHA256
+  values during the write and overlaps payload `fsync` with later producer
+  work. Producers without that guarantee retain unordered writes and use a
+  parallel readback checksum fallback.
+- Manifest and `READY` publication still wait for every payload `fsync` and
+  checksum to complete. These optimizations do not weaken artifact validation
+  or change lease ownership.
+
+#### Registered mmap transport
+
+After DLO takes an HWR lease, the no-AllGather backend may register the lease's
+complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
+Registration is transport state: it does not change artifact identity, the
+store, tensor ownership, or H2D payload. On success, each tensor view copies
+directly into the existing rotating HBM block buffers and the two private host
+staging slots are not allocated.
+
+`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
+over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
+memory policy, unsupported platform/capability, over-budget mapping, or fully
+rolled-back registration error selects the existing two-slot staging path. A
+partial registration that cannot be rolled back aborts startup because closing
+the lease would unmap memory still owned by the platform.
+
+Direct checkpoint mmap remains unchanged and continues to use staging. It may
+require loader-owned per-block transforms, while the HWR artifact already
+contains final runtime bytes. DLO AllGather never receives an HWR final-layout
+lease and therefore never enters this registration path.
+
 #### Pre-service transaction boundary
 
 The startup transaction does not end at restore commit. A lease-backed model is
@@ -292,6 +355,7 @@ loader owns
   -> backend takes once before asynchronous work
   -> backend drains pending H2D work
   -> backend releases hooks, staging, and model references
+  -> backend unregisters every mapped range
   -> backend closes lease
 ```
 
@@ -301,13 +365,15 @@ it. Once the backend takes ownership, only backend abort or teardown may drain
 work and close it.
 
 The implementation keeps the final-layout lease through backend setup and
-initial prefetch. The backend uses the existing two bounded rank-local staging
-slots for the immutable HWR tensors, and its transactional `enable()` cleanup
-removes partial hooks and closes the lease before reporting a setup failure.
-Preferred mode then uses the runner's fresh canonical retry; required mode
-propagates the failure.
+initial prefetch. Successful registration copies directly from immutable HWR
+views; otherwise the backend uses the existing two bounded rank-local staging
+slots. Transactional cleanup drains device work, removes source references,
+unregisters mappings, and only then closes the lease. An unregistration failure
+retains both registration and lease for retry/process teardown. Preferred mode
+then uses the runner's fresh canonical retry; required mode propagates the
+failure.
 
-#### PR2 promotion gates
+#### PR2 and PR3 promotion gates
 
 - Warm hit performs zero ordinary DiT materialization and zero producer calls.
 - Shared warm finalization changes neither restored bytes nor backing pointers.
@@ -319,6 +385,10 @@ propagates the failure.
 - Disabled mode and AllGather emit zero final-layout HWR interaction.
 - The lease carrier rejects duplicate take and serialization.
 - Backend setup or prefetch failure drains asynchronous work before lease close.
+- Registration is all-or-nothing; rollback/unregistration failure never closes
+  a mapping still owned by the platform.
+- Registered HWR transport bypasses host staging while preserving output and
+  H2D payload; unsupported registration retains the bounded staging path.
 - Compatible checkpoint mmap remains an unchanged benchmark and control path.
 - Prewarm uses one matching producer cohort per TP/SP topology and storage
   domain.
@@ -350,6 +420,22 @@ propagates the failure.
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
+### Auxiliary component policy
+
+`dlo_offload_components` maps discovered encoder/VAE paths to booleans. A
+component set to `False` is materialized on the target device once and is not
+entered through the model's manual offload lifecycle. A `True` value permits
+the existing `OffloadPlan.on_demand_component_paths` and encoder block-hook
+behavior; it does not add support where the model declares none. The optional
+`default` key handles unmatched auxiliary components, and an empty map
+preserves the current policy.
+
+The policy is resolved after component discovery and before DLO mutates DiT
+storage. Unknown names, including DiT paths, are rejected. DLO continues to
+own DiT streaming regardless of this map, which avoids an invalid state where
+the loader has skipped ordinary DiT materialization but the backend is asked
+to keep that DiT resident.
+
 ## Request and loading constraints
 
 AllGather DP multi-concurrency requires:
@@ -361,6 +447,15 @@ AllGather DP multi-concurrency requires:
 The no-AllGather path does not impose these DLO-specific synchronized-wave
 requirements.
 
+MiniMax-H3 currently builds one global encoder TP group and broadcasts the
+resulting conditioning from global rank 0. Therefore an encoder TP group that
+spans multiple DP replicas is not valid for independent-prompt AllGather
+waves. Exact-recipe DP2/TP2 and DP2/SP2 experiments with encoder TP4 completed
+the collective schedule, but two distinct prompts with the same seed produced
+byte-identical video and audio in every measured wave. Enabling or disabling
+encoder offload produced paired-identical results, so this is a pre-existing
+model topology limitation rather than component-policy behavior.
+
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
 ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
@@ -369,11 +464,12 @@ reconstructs those tensors with their recorded layouts. Other online methods
 must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
-The Host Weight Runtime representation and publication contracts are merged;
-see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
-[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
-no-AllGather consumer above is the PR2 implementation; registration and direct
-asynchronous H2D remain a separate transport follow-up.
+The Host Weight Runtime representation, publication, and no-AllGather consumer
+contracts are merged; see
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
+[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
+direct H2D is an optional transport layer over that merged lease contract.
 
 ## Validation coverage
 
@@ -386,10 +482,16 @@ Current source-level validation includes:
 - ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
   sharding of finalized weights and scales;
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
+- ordered publication hashing, overlapped durability, unordered parallel
+  checksum fallback, and node-local source-digest reuse/invalidation;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;
+- read-only registration preflight, direct source-to-device copies, safe
+  fallback, partial rollback, retryable unregistration, and lease ordering;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
+- bounded component allocator-cache retention, conservative/forced release,
+  OOM retry, and immutable encoder non-block staging;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
   tests.
 
@@ -418,6 +520,9 @@ AllGather, because DP replicas did not share checkpoint-backed runtime
 weights. This is a functional and memory smoke test, not a production-quality
 performance or output-quality benchmark.
 
+These B300 DP waves intentionally reused the same prompt and seed, so they do
+not establish independent-prompt isolation across DP replicas.
+
 ### Host-memory measurement
 
 A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the

docs/design/module/diffusion/offloader.md

@@ -18,7 +18,7 @@ validation_paths:
   - tests/diffusion/offloader/**
 upstream_refs:
   - torch.nn.Module.to
-last_reviewed: 2026-07-16
+last_reviewed: 2026-08-25
 ---
 
 # Diffusion offloader
@@ -48,6 +48,11 @@ and correctness required by the selected execution mode.
 **Rule:** Retained host and device copies MUST respect configured limits and
 provide deterministic teardown.
 
+Allocator-cache retention MUST remain local to an explicit component owner,
+declare both cache and physical-free-memory bounds, release conservatively
+when telemetry is unavailable, and force release on failure or memory
+pressure. It MUST NOT replace unconditional executor shutdown cleanup.
+
 ## Safe-change guide
 
 Test enabled and disabled paths, repeated execution, asynchronous transfer,

docs/design/module/host_weight_runtime.md

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

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -36,7 +36,8 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4 \
-  --dlo-no-use-allgather
+  --dlo-no-use-allgather \
+  --dlo-offload-components '{"text_encoder": false}'
 
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
@@ -51,6 +52,7 @@ omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=True,
     dlo_use_allgather=True,
+    dlo_offload_components={"text_encoder": False},
 )
 ```
 
@@ -63,8 +65,10 @@ omni = Omni(
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `--dlo-offload-components JSON` | Permit or disable offload for discovered encoder/VAE components | `{}` |
 | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
 | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
+| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
 
 ## Host-weight loading
 
@@ -178,12 +182,48 @@ returning to `required`.
 
 On a cold start, the canonical loader remains authoritative and publishes a
 validated final-layout artifact for later workers. A warm start restores the
-DiT final tensors without ordinary DiT materialization, then DLO streams them
-through the same two bounded host staging slots. The artifact identity includes
-the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
-layouts do not alias one another. Publication failure remains separate from a
-`preferred` serving startup; the next `required` startup provides the explicit
-artifact-availability check.
+DiT final tensors without ordinary DiT materialization. DLO then attempts to
+register the immutable mapping for direct H2D; if registration is unavailable,
+it streams through the same two bounded host staging slots. The artifact
+identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
+and distinct SP layouts do not alias one another. Publication failure remains
+separate from a `preferred` serving startup; the next `required` startup
+provides the explicit artifact-availability check.
+
+For local canonical checkpoints, the first eligible worker may hash source
+shards to establish immutable identity. HWR caches those digests in the same
+node-local storage domain and validates file metadata before reuse, so later
+workers normally avoid repeating that read. Cold BF16 publication also hashes
+ordered payloads as they are written and overlaps payload durability work with
+later shards; this changes startup work only, not artifact contents or runtime
+H2D behavior.
+
+#### Registered direct H2D
+
+Registration is attempted automatically only for an eligible warm HWR hit when
+`pin_cpu_memory` is enabled. A successful path is:
+
+```text
+shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
+```
+
+It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
+the two host staging slots. It does not reduce the H2D payload or make GPU
+kernels access host memory directly.
+
+CUDA requires read-only host-registration capability for the immutable HWR
+mapping. Unsupported capability, a positive registration limit smaller than
+the complete page-aligned mapping, or a safely rolled-back registration error
+falls back to bounded staging. Programmatic configurations can set
+`pin_cpu_memory=False` to disable registration explicitly. A successful
+registration locks the complete mapped range in host memory for that worker's
+lifetime, so use
+`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
+required.
+
+Shutdown drains H2D work, releases hook/source references, unregisters every
+range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
+retain their existing transfer paths.
 
 When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
 not resolve HWR sources or construct its store.
@@ -207,6 +247,25 @@ class MyPipeline(nn.Module):
 When no plan exists, discovery falls back to
 `_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
 
+## Auxiliary component policy
+
+`dlo_offload_components` follows the same component-map style as diffusion
+quantization. Keys are discovered encoder or VAE paths, values are booleans,
+and `default` is the optional fallback:
+
+```python
+dlo_offload_components = {
+    "text_encoder": False,
+    "default": True,
+}
+```
+
+`True` permits the model's existing DLO behavior for that component. `False`
+keeps the whole component resident, including its child blocks and non-block
+modules. An empty map preserves current behavior. Unknown keys are rejected;
+the map does not control DiT blocks because enabling DLO gives the backend
+ownership of DiT streaming.
+
 ## Data-parallel concurrency
 
 With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
@@ -231,6 +290,13 @@ must enter each collective.
   are validated.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
+- Component policy applies only to discovered auxiliary encoders and VAEs;
+  it cannot disable DLO for a DiT.
+- MiniMax-H3's encoder TP group is global rather than DP-replica-local. An
+  encoder TP size that spans multiple DP replicas (for example TP4 with
+  DP2/TP2 or DP2/SP2 on four ranks) is not supported for independent-prompt
+  AllGather waves: the current pipeline broadcasts rank 0's conditioning to
+  both replicas. This limitation is independent of the component policy.
 - DP concurrency requires an explicit, identical inference-step count.
 
 Sharing quantized or otherwise unvalidated transformed layouts through a

docs/user_guide/examples/online_serving/text_to_audio.md

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

examples/offline_inference/minimax_h3/dlo_lifecycle.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
 """Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
 
 The DLO modes accept any positive data-parallel size.  With AllGather enabled,
@@ -85,12 +85,20 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--height", type=int, default=768)
     parser.add_argument("--batch-wait-ms", type=float, default=500.0)
     parser.add_argument("--init-timeout", type=float, default=1800.0)
+    parser.add_argument(
+        "--dlo-offload-components",
+        type=json.loads,
+        default=None,
+        help="Component policy JSON, e.g. '{\"text_encoder\": false}'",
+    )
     parser.add_argument("--output", type=Path)
     return parser.parse_args()
 
 
 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
     is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
+    if args.dlo_offload_components and not is_dlo:
+        raise ValueError("--dlo-offload-components requires a DLO mode")
     common: dict[str, Any] = {
         "model": args.model,
         "trust_remote_code": True,
@@ -114,6 +122,7 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
             enable_distributed_layerwise_offload=True,
             dlo_use_allgather=args.mode == "dlo",
             dlo_resident_layers=0,
+            dlo_offload_components=({} if args.dlo_offload_components is None else args.dlo_offload_components),
         )
     else:
         common.update(

examples/online_serving/minicpmo/realtime_duplex_demo.py

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

examples/online_serving/stable_audio/README.md

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

examples/online_serving/stable_audio/curl_examples.sh

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

examples/online_serving/stable_audio/stable_audio_client.py

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

recipes/StabilityAI/Stable-Audio-Open.md

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

tests/benchmarks/metrics/test_metrics.py

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

tests/benchmarks/test_omniinteract.py

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

tests/benchmarks/test_serve_cli.py

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

tests/config/test_environment_variables.py

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

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

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

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

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

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

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

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

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

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

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

tests/diffusion/batching/test_diffusion_batching.py

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

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlanResult,
     TensorBinding,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.diffusion.registry import initialize_model
@@ -173,6 +174,15 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
         str(canonical_root / "model.safetensors"),
     )
     store_root = tmp_path / "hwr-store"
+    hash_calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal hash_calls
+        hash_calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
 
     def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
         loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
@@ -205,6 +215,8 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     assert warm_plan is not None
     assert warm_plan.lease_carrier is not None
     warm_plan.lease_carrier.close()
+    assert hash_calls == 1
+    assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
 
 
 def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -6,6 +6,8 @@ from __future__ import annotations
 
 import dataclasses
 import gc
+import json
+import threading
 from collections.abc import Sequence
 from pathlib import Path
 from types import SimpleNamespace
@@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import (
     FinalLayoutRequest,
     FinalLayoutTensorRestorer,
     ImplementationIdentity,
+    NodeSourceDigestCache,
     PreparedWeightSource,
     WeightSourceKind,
     build_final_layout_identity,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.model_loader.host_weights.contracts import (
     FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
     FinalLayoutArtifactSpec,
@@ -380,6 +384,7 @@ def _identity(
     *,
     request: FinalLayoutRequest | None = None,
     policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
+    source_digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutIdentityContext:
     transformer = model.get_submodule("transformer")
     return build_final_layout_identity(
@@ -388,6 +393,7 @@ def _identity(
         prepared_sources=(source,),
         request=request or _request(),
         policy=policy,
+        source_digest_cache=source_digest_cache,
     )
 
 
@@ -564,6 +570,131 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path
     assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
 
 
+def test_node_source_digest_cache_reuses_recovers_and_invalidates(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache_root = tmp_path / "store"
+    calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal calls
+        calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
+    first = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    second = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+
+    assert second.identity == first.identity
+    assert calls == 1
+
+    cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
+    cache_entry.write_text("not-json", encoding="utf-8")
+    recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert recovered.identity == first.identity
+    assert calls == 2
+
+    cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
+    cache_document["content_id"] = f"sha256:{'0' * 64}"
+    cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
+    checksum_recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert checksum_recovered.identity == first.identity
+    assert calls == 3
+
+    source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
+    changed = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert changed.identity != first.identity
+    assert calls == 4
+
+
+def test_source_identity_hashes_checkpoint_shards_in_parallel(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    second_file = source.resolved_root / "model-00002.safetensors"
+    second_file.write_bytes(b"second-checkpoint-shard")
+    source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
+    barrier = threading.Barrier(2)
+    thread_ids: set[int] = set()
+    thread_ids_lock = threading.Lock()
+    original_sha256 = source_identity_module._sha256_file
+
+    def synchronized_sha256(path: Path, state: object) -> str:
+        with thread_ids_lock:
+            thread_ids.add(threading.get_ident())
+        barrier.wait(timeout=5)
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
+    _identity(model, source)
+
+    assert len(thread_ids) == 2
+
+
+def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+
+    def unavailable_lock(*_args: object, **_kwargs: object) -> object:
+        raise source_identity_module.FileLockTimeoutError("busy")
+
+    monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+    assert not tuple(cache.entries.glob("*.json"))
+
+
+def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    original_mkdir = Path.mkdir
+
+    def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
+        if "source-digests-v1" in path.parts:
+            raise OSError("injected cache directory failure")
+        original_mkdir(path, *args, **kwargs)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+
+
 def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
     model = _TinyPipeline()
     source_root = tmp_path / "local-symlink-source"

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _OffloadAbort(BaseException):
+    pass
+
+
+def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+    class VisionStack(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.patch_embed = torch.nn.Linear(2, 2)
+            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+    class TextStack(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed_tokens = torch.nn.Embedding(4, 2)
+            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+    encoder = object.__new__(encoder_module.MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    encoder.device_target = torch.device("cpu")
+    encoder.vision = VisionStack()
+    encoder.text_model = TextStack()
+    hook = mocker.Mock()
+    hook.pin_memory = False
+    encoder._omni_layerwise_hooks = [hook]
+    encoder._omni_layerwise_enabled = True
+    cache = mocker.Mock()
+    encoder.set_omni_component_cache(cache)
+    stager = mocker.Mock()
+    stager_cls = mocker.Mock(return_value=stager)
+    monkeypatch.setattr(encoder_module, "PinnedModuleStager", stager_cls)
+
+    encoder.load_to_device()
+    encoder.offload_to_cpu()
+
+    stager_cls.assert_called_once_with(
+        [encoder.vision.patch_embed, encoder.text_model.embed_tokens],
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=cache,
+    )
+    stager.load.assert_called_once_with()
+    stager.offload.assert_called_once_with()
+    hook.offload_layer.assert_called_once_with()
+
+
+def test_manual_component_failure_forces_retained_cache_release(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    pipeline._dlo_component_cache = mocker.Mock()
+    component = mocker.Mock()
+    component.offload_to_cpu.side_effect = _OffloadAbort("offload failed")
+
+    with pytest.raises(RuntimeError, match="component failed"):
+        with pipeline._component_on_device(component):
+            raise RuntimeError("component failed")
+
+    component.load_to_device.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+    pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    pipeline._dlo_component_cache = mocker.Mock()
+    component = mocker.Mock()
+    component.offload_to_cpu.side_effect = [None, _OffloadAbort("offload failed")]
+
+    with pipeline._component_on_device(component):
+        pass
+    with pytest.raises(_OffloadAbort, match="offload failed"):
+        with pipeline._component_on_device(component):
+            pass
+
+    assert component.load_to_device.call_count == 2
+    assert component.offload_to_cpu.call_count == 2
+    pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_resident_component_skips_manual_dlo_lifecycle(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    component = mocker.Mock()
+    component._omni_dlo_offload_enabled = False
+
+    with pipeline._component_on_device(component):
+        pass
+
+    component.load_to_device.assert_not_called()
+    component.offload_to_cpu.assert_not_called()

tests/diffusion/offloader/test_component_policy.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
+    DistributedLayerwiseOffloadBackend,
+)
+from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _od_config(**overrides):
+    values = {
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": True,
+        "dlo_use_allgather": False,
+        "dlo_resident_layers": 0,
+        "dlo_offload_components": {},
+        "pin_cpu_memory": False,
+        "parallel_config": SimpleNamespace(
+            use_hsdp=False,
+            data_parallel_size=1,
+            sequence_parallel_size=1,
+        ),
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _backend(policy: dict[str, bool]) -> DistributedLayerwiseOffloadBackend:
+    backend = object.__new__(DistributedLayerwiseOffloadBackend)
+    backend.config = OffloadConfig(
+        strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+        pin_cpu_memory=False,
+        dlo_use_allgather=False,
+        dlo_offload_components=policy,
+    )
+    backend.device = torch.device("cpu")
+    return backend
+
+
+def _modules() -> PipelineModules:
+    return PipelineModules(
+        dits=[],
+        dit_names=[],
+        encoders=[nn.Module()],
+        encoder_names=["text_encoder"],
+        vaes=[nn.Module(), nn.Module()],
+        vae_names=["video_vae", "audio_vae"],
+    )
+
+
+def test_component_policy_resolves_exact_then_default():
+    backend = _backend({"text_encoder": False, "default": True})
+
+    assert not backend._component_offload_enabled("text_encoder")
+    assert backend._component_offload_enabled("video_vae")
+
+
+@pytest.mark.parametrize(
+    ("value", "error", "match"),
+    [
+        ([], TypeError, "must be a mapping"),
+        (["text_encoder"], TypeError, "must be a mapping"),
+        (False, TypeError, "must be a mapping"),
+        ({1: False}, TypeError, "keys must be non-empty strings"),
+        ({"text_encoder": 0}, TypeError, "must be a boolean"),
+    ],
+)
+def test_component_policy_rejects_invalid_shapes(value, error, match):
+    with pytest.raises(error, match=match):
+        OffloadConfig.from_od_config(_od_config(dlo_offload_components=value))
+
+
+def test_component_policy_accepts_none_as_default():
+    config = OffloadConfig.from_od_config(_od_config(dlo_offload_components=None))
+
+    assert config.dlo_offload_components == {}
+
+
+def test_component_policy_requires_dlo():
+    with pytest.raises(ValueError, match="requires distributed layerwise offload"):
+        OffloadConfig.from_od_config(
+            _od_config(
+                enable_distributed_layerwise_offload=False,
+                dlo_offload_components={"text_encoder": False},
+            )
+        )
+
+
+def test_component_policy_rejects_unknown_or_dit_names():
+    backend = _backend({"transformer": False})
+
+    with pytest.raises(ValueError, match="unknown auxiliary components.*transformer"):
+        backend._validate_component_offload_policy(_modules())
+
+
+@pytest.mark.parametrize(
+    ("policy", "encoder_stage", "video_stage", "audio_stage", "try_encoder"),
+    [
+        ({}, True, True, True, True),
+        ({"text_encoder": False}, False, True, True, False),
+        ({"default": False, "video_vae": True}, False, True, False, False),
+    ],
+)
+def test_prepare_auxiliary_components_applies_policy(
+    mocker,
+    policy,
+    encoder_stage,
+    video_stage,
+    audio_stage,
+    try_encoder,
+):
+    backend = _backend(policy)
+    backend._try_layerwise_offload_encoder = mocker.Mock()
+    backend._register_on_demand_hook = mocker.Mock()
+    modules = _modules()
+    plan = OffloadPlan(
+        encoder_block_attrs={"text_encoder": ("layers",)},
+        on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
+    )
+
+    backend._prepare_auxiliary_components(modules, plan)
+
+    assert backend._try_layerwise_offload_encoder.called is try_encoder
+    assert modules.encoders[0]._omni_dlo_offload_enabled is encoder_stage
+    assert modules.vaes[0]._omni_dlo_offload_enabled is video_stage
+    assert modules.vaes[1]._omni_dlo_offload_enabled is audio_stage
+    assert backend._register_on_demand_hook.call_args_list == [
+        mocker.call(modules.encoders[0], "text_encoder", stage_on_demand=encoder_stage),
+        mocker.call(modules.vaes[0], "video_vae", stage_on_demand=video_stage),
+        mocker.call(modules.vaes[1], "audio_vae", stage_on_demand=audio_stage),
+    ]

tests/diffusion/offloader/test_cuda_host_registration.py

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

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -5,6 +5,7 @@
 
 import gc
 import json
+import weakref
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -17,6 +18,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
 from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     TensorBinding,
@@ -33,12 +35,14 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadHook,
     PinnedResidentLayerGroup,
 )
+from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
+from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -703,11 +707,15 @@ class _HWRPipeline(nn.Module):
 
 
 class _FakeHostWeightLease:
-    def __init__(self, resolution_id: str):
+    def __init__(self, resolution_id: str, events: list[str] | None = None):
         self.closed = False
         self.provenance = SimpleNamespace(resolution_id=resolution_id)
+        self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
+        self._events = events
 
     def close(self):
+        if self._events is not None:
+            self._events.append("lease")
         self.closed = True
 
 
@@ -757,6 +765,203 @@ def _hwr_backend():
     return _HWRPipeline(), backend, carrier, lease
 
 
+def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
+    current_block = nn.Linear(2, 2, bias=False)
+    next_block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    next_block.weight.data.copy_(expected)
+    hook = DistributedLayerwiseOffloadHook(
+        next_block=next_block,
+        device=torch.device("cpu"),
+        dp_group=None,
+        dp_size=1,
+        rank=0,
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    hook.initialize_hook(current_block)
+    hook.registered_mmap = True
+    hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging")  # type: ignore[method-assign]
+
+    hook.prefetch_layer(slot=0, non_blocking=True)
+
+    assert torch.equal(next_block.weight, expected)
+
+
+def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
+    block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    block.weight.data.copy_(expected)
+    group = PinnedResidentLayerGroup(
+        [block],
+        device=torch.device("cpu"),
+        copy_stream=DummyStream(),
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    group.registered_mmap = True
+    group._cpu_staging_buffers.clear()
+
+    group.load()
+
+    assert torch.equal(block.weight, expected)
+
+
+def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    plan, _, lease = _fake_hwr_plan()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+            dlo_host_registration_limit_gib=1.5,
+        ),
+        torch.device("cuda"),
+        host_weight_plan=plan,
+    )
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    calls: list[tuple[object, int | None]] = []
+
+    class Registration:
+        total_bytes = 4096
+        region_count = 1
+
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            return ()
+
+    def register(regions, *, device, max_bytes):
+        assert device == torch.device("cuda")
+        calls.append((regions, max_bytes))
+        return Registration()
+
+    monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+    pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
+
+    assert backend._try_register_hwr_mmap((pinned_source,))  # type: ignore[arg-type]
+    assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend._release_registered_mmap()
+
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+
+
+def test_hwr_registration_budget_validation_is_transport_scoped():
+    assert (
+        validate_dlo_host_registration_options(
+            limit_gib=1.5,
+            enable_dlo=True,
+            use_allgather=False,
+            hwr_mode="preferred",
+        )
+        == 1.5
+    )
+    invalid = (
+        dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
+    )
+    for options in invalid:
+        with pytest.raises(ValueError):
+            validate_dlo_host_registration_options(**options)
+
+
+def test_unregistration_precedes_lease_close_and_retries_failures(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    events: list[str] = []
+    responses = iter([("busy",), ()])
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return next(responses)
+
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retry", events)
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = Registration()
+    backend._using_rank_local_mmap = True
+    backend._using_registered_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    assert not lease.closed
+    assert backend._host_registration is not None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend.disable()
+
+    assert lease.closed
+    assert backend._host_registration is None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+    assert events == ["unregister", "unregister", "lease"]
+
+
+def test_failed_unregistration_retains_lease_after_backend_is_dropped(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    events: list[str] = []
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return ("busy",)
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retained", events)
+    registration = Registration()
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = registration
+    backend._using_rank_local_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    lease_ref = weakref.ref(lease)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    del backend
+    del lease
+    gc.collect()
+
+    retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
+    assert lease_ref() is retained_lease
+    assert not retained_lease.closed
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    retained_lease.close()
+
+
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 

tests/diffusion/offloader/test_module_residency.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+from torch import nn
+
+import vllm_omni.diffusion.offloader.module_residency as residency_module
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class _DummyEvent:
+    def __init__(self) -> None:
+        self.record_count = 0
+
+    def record(self, _stream) -> None:
+        self.record_count += 1
+
+
+@contextmanager
+def _dummy_stream(_stream):
+    yield
+
+
+@pytest.fixture
+def patched_runtime(monkeypatch, mocker):
+    empty_cache = mocker.Mock()
+    synchronize = mocker.Mock()
+    monkeypatch.setattr(residency_module.current_omni_platform, "Stream", _DummyStream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "Event", _DummyEvent)
+    monkeypatch.setattr(residency_module.current_omni_platform, "current_stream", _DummyStream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "stream", _dummy_stream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "empty_cache", empty_cache)
+    monkeypatch.setattr(residency_module.current_omni_platform, "synchronize", synchronize)
+    return empty_cache, synchronize
+
+
+def _aliased_modules() -> tuple[nn.Module, nn.Module]:
+    parameter_storage = torch.arange(12, dtype=torch.float32)
+    buffer_storage = torch.arange(8, dtype=torch.int64)
+    left = nn.Module()
+    right = nn.Module()
+    left.weight = nn.Parameter(parameter_storage[:8].view(2, 4))
+    right.weight = nn.Parameter(parameter_storage[4:].view(2, 4))
+    left.register_buffer("state", buffer_storage[:6])
+    right.register_buffer("state", buffer_storage[2:])
+    return left, right
+
+
+def test_module_group_preserves_parameter_and_buffer_storage_aliases(patched_runtime, mocker):
+    left, right = _aliased_modules()
+    expected_left_weight = left.weight.detach().clone()
+    expected_right_weight = right.weight.detach().clone()
+    expected_left_state = left.state.clone()
+    expected_right_state = right.state.clone()
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        [left, right],
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+
+    master_weight_ptr = left.weight.untyped_storage().data_ptr()
+    master_state_ptr = left.state.untyped_storage().data_ptr()
+    assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert right.state.untyped_storage().data_ptr() == master_state_ptr
+    assert (left.weight.storage_offset(), right.weight.storage_offset()) == (0, 4)
+    assert (left.state.storage_offset(), right.state.storage_offset()) == (0, 2)
+
+    stager.load()
+
+    assert left.weight.untyped_storage().data_ptr() != master_weight_ptr
+    assert right.weight.untyped_storage().data_ptr() == left.weight.untyped_storage().data_ptr()
+    assert right.state.untyped_storage().data_ptr() == left.state.untyped_storage().data_ptr()
+    left.weight.data[1, 0] = -1
+    left.state[2] = -1
+    assert right.weight[0, 0].item() == -1
+    assert right.state[0].item() == -1
+
+    stager.offload()
+
+    assert left.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert left.state.untyped_storage().data_ptr() == master_state_ptr
+    assert right.state.untyped_storage().data_ptr() == master_state_ptr
+    assert torch.equal(left.weight, expected_left_weight)
+    assert torch.equal(right.weight, expected_right_weight)
+    assert torch.equal(left.state, expected_left_state)
+    assert torch.equal(right.state, expected_right_state)
+    retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_staging_transitions_are_idempotent_and_reuse_one_event(patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    ready_event = stager._ready_event
+
+    stager.load()
+    first_storage = module.weight.untyped_storage().data_ptr()
+    stager.load()
+    assert module.weight.untyped_storage().data_ptr() == first_storage
+    assert stager._ready_event is ready_event
+    assert ready_event.record_count == 1
+
+    stager.offload()
+    stager.offload()
+    assert not stager.loaded
+    retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_failed_load_restores_master_and_forces_cache_release(monkeypatch, patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    expected = module.weight.detach().clone()
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    master_ptr = module.weight.untyped_storage().data_ptr()
+
+    def fail_after_rebind() -> None:
+        device_storages = [torch.empty_like(group.master) for group in stager._groups]
+        stager._bind(device_storages)
+        module.weight.data.zero_()
+        raise RuntimeError("copy failed")
+
+    monkeypatch.setattr(stager, "_load_once", fail_after_rebind)
+    with pytest.raises(RuntimeError, match="copy failed"):
+        stager.load()
+
+    assert not stager.loaded
+    assert module.weight.untyped_storage().data_ptr() == master_ptr
+    assert torch.equal(module.weight, expected)
+    retention.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_out_of_memory_releases_cache_and_retries_once(monkeypatch, patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    load_once = stager._load_once
+    attempts = 0
+
+    def fail_once() -> None:
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise torch.OutOfMemoryError("memory pressure")
+        load_once()
+
+    monkeypatch.setattr(stager, "_load_once", fail_once)
+    stager.load()
+
+    assert stager.loaded
+    assert attempts == 2
+    retention.release_if_needed.assert_called_once_with(force=True)
+
+
+@pytest.mark.parametrize(
+    ("reserved", "allocated", "free", "expected_release"),
+    [
+        (20, 10, 80, False),
+        (40, 10, 80, True),
+        (20, 10, 4, True),
+    ],
+)
+def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
+    monkeypatch,
+    patched_runtime,
+    reserved,
+    allocated,
+    free,
+    expected_release,
+):
+    empty_cache, _ = patched_runtime
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved)
+    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated)
+    monkeypatch.setattr(
+        residency_module.current_omni_platform,
+        "get_device_memory",
+        lambda _device: (free, 100),
+    )
+    retention = BoundedAllocatorCache(torch.device("cpu"))
+
+    released = retention.release_if_needed()
+
+    assert released is expected_release
+    assert empty_cache.call_count == int(expected_release)
+
+
+def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime):
+    empty_cache, _ = patched_runtime
+
+    def unavailable(_device):
+        raise NotImplementedError
+
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", unavailable)
+    retention = BoundedAllocatorCache(torch.device("cpu"))
+
+    assert retention.release_if_needed()
+    empty_cache.assert_called_once_with()

tests/diffusion/quantization/test_wan_autoround_mxfp4.py

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

tests/diffusion/test_worker_wrapper_base.py

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

tests/e2e/features/fullduplex/test_client.py

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

tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py

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

tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py

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

tests/e2e/online_serving/test_qwen3_omni_expansion.py

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

tests/e2e/online_serving/test_qwen_image_expansion.py

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

tests/e2e/online_serving/test_stable_audio_online_expansion.py

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

tests/engine/test_async_omni_engine_stage_init.py

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

tests/engine/test_orchestrator_error_handling.py

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

tests/engine/test_single_stage_mode.py

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

tests/entrypoints/openai_api/test_serving_speech.py

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

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -332,7 +332,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
     ]
 
 
-def test_serve_cli_forwards_distributed_offload_residency():
+def test_serve_cli_forwards_distributed_offload_options():
     """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
@@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-offload-components",
+            '{"text_encoder": false, "default": true}',
         ]
     )
 
@@ -357,9 +359,14 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
     assert args.dlo_resident_layers == 20
+    assert args.dlo_offload_components == {"text_encoder": False, "default": True}
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_resident_layers"] == 20
+    assert engine_args["dlo_offload_components"] == {
+        "text_encoder": False,
+        "default": True,
+    }
 
 
 def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
@@ -378,6 +385,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "preferred",
             "--host-weight-runtime-root",
             "/var/cache/vllm-omni/hwr",
+            "--dlo-host-registration-limit-gib",
+            "80",
         ]
     )
 
@@ -387,8 +396,10 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
 
     assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
     assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
     assert engine_args["host_weight_runtime_mode"] == "preferred"
     assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert engine_args["dlo_host_registration_limit_gib"] == 80
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

tests/host_weight_runtime/test_filesystem_store.py

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

vllm_omni/benchmarks/data_modules/omniinteract_dataset.py

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

vllm_omni/benchmarks/metrics/metrics.py

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

vllm_omni/benchmarks/omniinteract.py

@@ -0,0 +1,933 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniInteract session execution and artifacts for serving benchmarks."""
+
+from __future__ import annotations
+
+import asyncio
+import binascii
+import contextlib
+import copy
+import fcntl
+import hashlib
+import json
+import logging
+import math
+import os
+import subprocess
+import tempfile
+import time
+from collections import Counter
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, cast
+from urllib.parse import urljoin, urlsplit
+
+import pybase64 as base64
+
+from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+    DEFAULT_MAX_VIDEO_DURATION_S,
+    OmniInteractCase,
+    OmniInteractPreparedInput,
+    case_manifest,
+)
+from vllm_omni.experimental.fullduplex.client import (
+    PCM16_BYTES_PER_SAMPLE,
+    PCM16_SAMPLE_RATE,
+    RealtimeDuplexClient,
+    RealtimeEventCollector,
+    build_realtime_url,
+    chunk_period_ms,
+    has_residual_model_unit,
+    reference_audio_data_url,
+    summarize_session_request_metrics,
+    write_pcm16_wav,
+)
+
+OUTPUT_SAMPLE_RATE = 24_000
+SUCCESS_ARTIFACTS = (".done", "output.wav", "wav_transcript.json", "events.json", "result.json")
+BATCH_ARTIFACTS = ("batch_summary.json", "official_eval_manifest.jsonl")
+ARTIFACT_LOCK_FILE = ".omniinteract.lock"
+_INPUT_CHUNK_MS = 200
+VIDEO_FPS = 1.0
+_COMPLETION_SETTLE_S = 2.0
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniInteractBenchmarkConfig:
+    model: str
+    base_url: str = "http://127.0.0.1:8000"
+    endpoint: str = "/v1/realtime"
+    output_root: Path = Path("omniinteract-output")
+    timeout_s: float = 900.0
+    media_timeout_s: float = 600.0
+    max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
+    ref_audio: str | None = None
+    require_response: bool = False
+    extra_headers: dict[str, str] | None = None
+    extra_body: dict[str, object] | None = None
+
+
+@dataclass
+class OmniInteractCaseResult:
+    subset: str
+    video: str
+    output_dir: str
+    success: bool = False
+    error: str = ""
+    session_id: str = ""
+    latency_s: float = 0.0
+    input_audio_chunks: int = 0
+    input_video_frames: int = 0
+    pacing_mean_lag_s: float = 0.0
+    pacing_max_lag_s: float = 0.0
+    responses: int = 0
+    audio_bytes: int = 0
+    audio_clipped_bytes: int = 0
+    transcript: str = ""
+    eligible_for_official_eval: bool = False
+    official_eval_ineligible_reasons: list[str] = field(default_factory=list)
+    artifact_warnings: list[str] = field(default_factory=list)
+    output_tokens: int = 0
+    duplex_request_metrics: list[dict[str, object]] = field(default_factory=list)
+    duplex_session_metrics: dict[str, object] = field(default_factory=dict)
+    _artifact_context: _DeferredArtifactContext | None = field(default=None, repr=False, compare=False)
+
+    def as_dict(self) -> dict[str, Any]:
+        return copy.deepcopy({key: value for key, value in vars(self).items() if not key.startswith("_")})
+
+
+@dataclass(frozen=True)
+class _DeferredArtifactContext:
+    horizon_bytes: int
+    spans: tuple[tuple[int, bytes], ...]
+    chunks: list[dict[str, object]]
+    events: list[dict[str, object]]
+
+
+def benchmark_summary(results: list[OmniInteractCaseResult]) -> dict[str, Any]:
+    succeeded = sum(result.success for result in results)
+    eligible = sum(result.success and result.eligible_for_official_eval for result in results)
+    return {
+        "total": len(results),
+        "success": succeeded,
+        "failed": len(results) - succeeded,
+        "eligible_for_official_eval": eligible,
+        "successful_but_ineligible": succeeded - eligible,
+        "audio_clipped_bytes": sum(result.audio_clipped_bytes for result in results),
+        "results": [result.as_dict() for result in results],
+    }
+
+
+def _run_media_command(
+    command: list[str],
+    *,
+    timeout_s: float,
+    text: bool = False,
+) -> subprocess.CompletedProcess[Any]:
+    try:
+        result = subprocess.run(
+            command,
+            capture_output=True,
+            check=False,
+            text=text,
+            timeout=timeout_s,
+        )
+    except subprocess.TimeoutExpired as exc:
+        raise TimeoutError(f"Media command timed out after {timeout_s:g}s: {command[0]}") from exc
+    except FileNotFoundError as exc:
+        raise RuntimeError(f"Required media command is unavailable: {command[0]}") from exc
+    if result.returncode:
+        error = result.stderr if text else result.stderr.decode("utf-8", "ignore")
+        raise RuntimeError(f"{command[0]} failed: {error.strip()}")
+    return result
+
+
+def prepare_media(
+    video: Path,
+    fps: float,
+    *,
+    timeout_s: float,
+    max_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S,
+) -> tuple[float, bytes, list[str | None]]:
+    probe = [
+        "ffprobe",
+        *"-v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1".split(),
+        str(video),
+    ]
+    duration = float(_run_media_command(probe, text=True, timeout_s=timeout_s).stdout.strip())
+    if not math.isfinite(duration) or duration <= 0 or duration > max_duration_s:
+        raise ValueError(f"Invalid video duration for {video}: {duration!r}")
+    target_bytes = math.ceil(duration) * PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+    audio = [
+        "ffmpeg",
+        *"-loglevel error -i".split(),
+        str(video),
+        "-t",
+        str(duration),
+        *f"-vn -f s16le -ac 1 -ar {PCM16_SAMPLE_RATE} -fs".split(),
+        str(target_bytes),
+        "pipe:1",
+    ]
+    pcm = _run_media_command(audio, timeout_s=timeout_s).stdout
+    pcm = (pcm + bytes(max(0, target_bytes - len(pcm))))[:target_bytes]
+
+    frame_count = math.ceil(duration * fps)
+    frames: list[str | None] = [None] * frame_count
+    with tempfile.TemporaryDirectory(prefix="vllm-omni-frames-") as temp_dir:
+        output_pattern = Path(temp_dir) / "frame-%06d.jpg"
+        frame_filter = f"select=gte(t\\,(selected_n+0.5)/{fps}),scale=640:640:force_original_aspect_ratio=decrease"
+        frame_command = [
+            "ffmpeg",
+            *"-loglevel error -i".split(),
+            str(video),
+            "-an",
+            "-vf",
+            frame_filter,
+            *f"-frames:v {frame_count} -vsync vfr -q:v 5".split(),
+            str(output_pattern),
+        ]
+        _run_media_command(frame_command, timeout_s=timeout_s)
+        for index, frame_path in enumerate(sorted(Path(temp_dir).glob("frame-*.jpg"))[:frame_count]):
+            frames[index] = base64.b64encode(frame_path.read_bytes()).decode("ascii")
+    return duration, pcm, frames
+
+
+@dataclass(frozen=True)
+class _AudioSegment:
+    event_index: int
+    response_id: str
+    start_s: float
+    end_s: float
+    pcm16: bytes
+
+
+class _Playback:
+    def __init__(self) -> None:
+        self.cursor = 0
+        self.end_s = 0.0
+        self.segments: list[_AudioSegment] = []
+        self.completed: set[str] = set()
+        self.completion_acked: set[str] = set()
+        self.warnings: list[str] = []
+        self._total_samples: dict[str, int] = {}
+        self._response_end_s: dict[str, float] = {}
+
+    def _warn_once(self, warning: str) -> None:
+        if warning not in self.warnings:
+            self.warnings.append(warning)
+
+    def ingest(self, events: RealtimeEventCollector) -> None:
+        while self.cursor < len(events.events):
+            index, self.cursor = self.cursor, self.cursor + 1
+            event = events.events[index]
+            response_id = events.response_id(event)
+            if event.get("type") == "response.done":
+                if response_id:
+                    self.completed.add(response_id)
+                continue
+            if event.get("type") != "response.audio.delta":
+                continue
+            encoded = event.get("delta") or event.get("audio")
+            if not response_id:
+                raise ValueError("response audio has no response_id")
+            if not isinstance(encoded, str):
+                raise ValueError("response audio payload is missing")
+            if event.get("format") is None:
+                self._warn_once("response.audio.delta omitted format; assumed pcm16")
+            elif event.get("format") != "pcm16":
+                raise ValueError("OmniInteract output must be pcm16")
+            rate = event.get("sample_rate_hz")
+            if rate is None:
+                rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE
+                self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}")
+            if rate != OUTPUT_SAMPLE_RATE:
+                raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio")
+            try:
+                raw = base64.b64decode(encoded, validate=True)
+            except (ValueError, binascii.Error) as exc:
+                raise ValueError("response audio is not valid base64") from exc
+            if not raw or len(raw) % PCM16_BYTES_PER_SAMPLE:
+                raise ValueError("response audio is empty or not PCM16 aligned")
+            start = max(events.event_received_at_s[index], self.end_s)
+            samples = len(raw) // PCM16_BYTES_PER_SAMPLE
+            segment = _AudioSegment(index, response_id, start, start + samples / rate, raw)
+            self.segments.append(segment)
+            self.end_s = segment.end_s
+            self._total_samples[response_id] = self._total_samples.get(response_id, 0) + samples
+            self._response_end_s[response_id] = segment.end_s
+
+    async def acknowledge(self, client: RealtimeDuplexClient, now: float | None = None) -> None:
+        events = client.events
+        self.ingest(events)
+        now = time.monotonic() if now is None else now
+        for response_id in self.completed - self.completion_acked:
+            total_samples = self._total_samples.get(response_id, 0)
+            if not total_samples or now < self._response_end_s[response_id]:
+                continue
+            played_ms = total_samples * 1000 // OUTPUT_SAMPLE_RATE
+            await client.send_playback_ack(response_id, played_ms)
+            self.completion_acked.add(response_id)
+
+
+async def stream_inputs(
+    client: RealtimeDuplexClient,
+    pcm: bytes,
+    frames: Sequence[str | None],
+    playback: _Playback,
+) -> tuple[int, int, float, float]:
+    bytes_per_second = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+    chunk_bytes = bytes_per_second * _INPUT_CHUNK_MS // 1000
+    started_at = time.monotonic()
+    frame_cursor = sent_frames = 0
+    lags: list[float] = []
+    for chunks, offset in enumerate(range(0, len(pcm), chunk_bytes), start=1):
+        end = min(offset + chunk_bytes, len(pcm))
+        end_ms = end * 1000 // bytes_per_second
+        ready: list[str] = []
+        while frame_cursor < len(frames) and end_ms >= (frame_cursor + 0.5) * 1000 / VIDEO_FPS:
+            if frames[frame_cursor]:
+                ready.append(frames[frame_cursor] or "")
+            frame_cursor += 1
+        sent_frames += len(ready)
+        event: dict[str, object] = {
+            "type": "input_audio_buffer.append",
+            "audio": base64.b64encode(pcm[offset:end]).decode("ascii"),
+            "input_audio_format": "pcm16",
+            "sample_rate_hz": PCM16_SAMPLE_RATE,
+            "duration_ms": (end - offset) * 1000 // bytes_per_second,
+            "audio_end_ms": end_ms,
+        }
+        if ready:
+            event["video_frames"] = ready
+        lags.append(max(0.0, time.monotonic() - started_at - offset / bytes_per_second))
+        await client.send(event)
+        await playback.acknowledge(client)
+        await asyncio.sleep(max(0.0, started_at + end_ms / 1000 - time.monotonic()))
+    return chunks if pcm else 0, sent_frames, sum(lags) / len(lags) if lags else 0.0, max(lags, default=0.0)
+
+
+def _response_status(event: dict[str, object]) -> str | None:
+    response = event.get("response")
+    if isinstance(response, dict) and response.get("status"):
+        return str(response["status"])
+    status = event.get("status")
+    return str(status) if status else None
+
+
+def response_ledger(
+    collector: RealtimeEventCollector,
+    *,
+    end: int | None = None,
+) -> tuple[set[str], set[str]]:
+    created: set[str] = set()
+    done: set[str] = set()
+    for event in collector.events[:end]:
+        event_type = event.get("type")
+        if event_type not in {"response.created", "response.done"}:
+            continue
+        response_id = collector.response_id(event)
+        if not response_id:
+            raise ValueError(f"{event_type} has no response_id")
+        if event_type == "response.created":
+            if response_id in created:
+                raise ValueError(f"duplicate response.created for {response_id}")
+            created.add(response_id)
+            continue
+        if response_id not in created:
+            raise ValueError(f"response.done without response.created for {response_id}")
+        if response_id in done:
+            raise ValueError(f"duplicate response.done for {response_id}")
+        response = event.get("response") if isinstance(event.get("response"), dict) else {}
+        details = event.get("status_details") if isinstance(event.get("status_details"), dict) else {}
+        nested = response.get("status_details") if isinstance(response.get("status_details"), dict) else {}
+        if "failed" in {event.get("status"), response.get("status"), details.get("type"), nested.get("type")}:
+            raise ValueError(f"response.done reports failure for {response_id}")
+        done.add(response_id)
+    return created, done
+
+
+def _is_model_listen_decision(event: dict[str, object]) -> bool:
+    response = event.get("response")
+    metadata = response.get("metadata") if isinstance(response, dict) else event.get("metadata")
+    return isinstance(metadata, dict) and metadata.get("buffering") is not True and metadata.get("model_listen") is True
+
+
+def _has_post_commit_decision(
+    collector: RealtimeEventCollector,
+    events: list[dict[str, object]],
+    *,
+    prior_response_ids: set[str],
+) -> bool:
+    pending_prior = set(prior_response_ids)
+    for event in events:
+        event_type = event.get("type")
+        if event_type == "response.done":
+            response_id = collector.response_id(event)
+            if response_id in pending_prior:
+                pending_prior.remove(response_id)
+                continue
+            if not pending_prior and _response_status(event) != "cancelled":
+                return True
+        elif event_type == "response.listen" and not pending_prior and _is_model_listen_decision(event):
+            return True
+    return False
+
+
+def _raise_if_session_terminated(
+    collector: RealtimeEventCollector,
+    from_index: int,
+    *,
+    explicit_close_from: int | None = None,
+) -> None:
+    errors = collector.errors()
+    if errors:
+        raise RuntimeError(str(errors[-1]))
+    for index, event in enumerate(collector.events[from_index:], start=from_index):
+        event_type = event.get("type")
+        if event_type not in {"session.expired", "session.closed"}:
+            continue
+        reason = event.get("reason")
+        nested = event.get("event")
+        if reason is None and isinstance(nested, dict):
+            reason = nested.get("reason")
+        expected = (
+            event_type == "session.closed"
+            and explicit_close_from is not None
+            and index >= explicit_close_from
+            and reason is None
+        )
+        if expected:
+            continue
+        detail = f": {reason}" if reason else ""
+        prefix = "Unexpected " if explicit_close_from is not None else ""
+        raise RuntimeError(f"{prefix}{event_type}{detail}")
+
+
+def _ensure_final_commit_tail(pcm: bytes, events: list[dict[str, object]]) -> bytes:
+    """Reserve one PCM sample so an exact model unit is flushed by commit.
+
+    A complete unit is otherwise emitted before commit, so its asynchronous
+    decision cannot be correlated with the accepted final input. Keeping an
+    almost-full unit buffered makes the server's final residual flush carry
+    that correlation; the server pads the missing sample back to a full unit.
+    """
+
+    period_ms = chunk_period_ms(events)
+    if len(pcm) >= PCM16_BYTES_PER_SAMPLE and not has_residual_model_unit(pcm, chunk_period_ms=period_ms):
+        return pcm[:-PCM16_BYTES_PER_SAMPLE]
+    return pcm
+
+
+async def wait_for_session_completion(
+    client: RealtimeDuplexClient,
+    playback: _Playback,
+    *,
+    commit_from: int,
+    session_from: int | None = None,
+    timeout_s: float,
+    settle_s: float,
+) -> int:
+    deadline = time.monotonic() + timeout_s
+    committed_index: int | None = None
+    prior_response_ids: set[str] = set()
+    last_event_count = len(client.events.events)
+    stable_since = time.monotonic()
+    while time.monotonic() < deadline:
+        client.raise_if_reader_stopped()
+        _raise_if_session_terminated(client.events, commit_from if session_from is None else session_from)
+        await playback.acknowledge(client)
+        if len(client.events.events) != last_event_count:
+            last_event_count = len(client.events.events)
+            stable_since = time.monotonic()
+        if committed_index is None:
+            committed_index = next(
+                (
+                    index
+                    for index in range(commit_from, len(client.events.events))
+                    if client.events.events[index].get("type") == "input_audio_buffer.committed"
+                ),
+                None,
+            )
+            if committed_index is not None:
+                committed_event = client.events.events[committed_index]
+                committed = committed_event.get("event")
+                if isinstance(committed, dict) and committed.get("overlap_deferred") is True:
+                    prior_created, prior_done = response_ledger(client.events, end=committed_index)
+                    prior_response_ids = prior_created - prior_done
+                stable_since = time.monotonic()
+        if committed_index is not None:
+            created, done = response_ledger(client.events)
+            post_commit = client.events.events[committed_index + 1 :]
+            decision = _has_post_commit_decision(
+                client.events,
+                post_commit,
+                prior_response_ids=prior_response_ids,
+            )
+            if created == done and decision and time.monotonic() - stable_since >= settle_s:
+                return committed_index
+        await asyncio.sleep(0.05)
+    missing: set[str] = set()
+    with contextlib.suppress(ValueError):
+        created, done = response_ledger(client.events)
+        missing = created - done
+    raise TimeoutError(
+        "Timed out waiting for committed input and stable responses"
+        + (f"; unfinished response_ids={sorted(missing)}" if missing else "")
+    )
+
+
+def _output_dir(root: Path, case: OmniInteractCase) -> Path:
+    relative = case.video_rel.replace("\\", "/")
+    stem = Path(relative).with_suffix("").as_posix().replace("/", "__")
+    digest = hashlib.sha256(f"{case.subset}/{relative}".encode()).hexdigest()[:8]
+    return root / case.subset / f"{stem}--{digest}"
+
+
+@contextlib.contextmanager
+def omniinteract_output_lock(root: Path) -> Iterator[None]:
+    root.mkdir(parents=True, exist_ok=True)
+    with (root / ARTIFACT_LOCK_FILE).open("a+b") as lock_file:
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+def _atomic_replace(path: Path, writer: Callable[[Path], None]) -> None:
+    descriptor, temporary_name = tempfile.mkstemp(
+        dir=path.parent,
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+    )
+    os.close(descriptor)
+    temporary = Path(temporary_name)
+    try:
+        writer(temporary)
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def _atomic_write_text(path: Path, value: str) -> None:
+    _atomic_replace(path, lambda temporary: temporary.write_text(value, encoding="utf-8"))
+
+
+def _atomic_write_json(path: Path, value: object) -> None:
+    _atomic_write_text(path, json.dumps(value, ensure_ascii=False, indent=2) + "\n")
+
+
+def _atomic_write_wav(path: Path, pcm: bytes | bytearray, rate: int) -> None:
+    _atomic_replace(
+        path,
+        lambda temporary: write_pcm16_wav(
+            temporary,
+            cast(bytes, pcm),  # wave.writeframes accepts bytearray.
+            sample_rate_hz=rate,
+        ),
+    )
+
+
+def _clear_artifacts(directory: Path, names: Sequence[str]) -> None:
+    directory.mkdir(parents=True, exist_ok=True)
+    for name in names:
+        (directory / name).unlink(missing_ok=True)
+
+
+def _artifact_summary(case: OmniInteractCase, result: OmniInteractCaseResult) -> dict[str, Any]:
+    return {
+        **result.as_dict(),
+        "annotation": str(case.annotation_path),
+        "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
+    }
+
+
+def _publish_success_artifacts(
+    directory: Path,
+    result: OmniInteractCaseResult,
+    context: _DeferredArtifactContext,
+    summary: dict[str, Any],
+) -> None:
+    pcm = bytearray(context.horizon_bytes)
+    for offset, chunk in context.spans:
+        pcm[offset : offset + len(chunk)] = chunk
+    try:
+        _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
+        _atomic_write_wav(directory / "output.wav", pcm, OUTPUT_SAMPLE_RATE)
+        _atomic_write_json(
+            directory / "wav_transcript.json",
+            {
+                "text": result.transcript,
+                "chunks": context.chunks,
+                "timestamp_semantics": "serialized playback queue time relative to input streaming start",
+            },
+        )
+        _atomic_write_json(directory / "events.json", context.events)
+        _atomic_write_json(directory / "result.json", result.as_dict())
+        _atomic_write_json(directory / ".done", summary)
+    except Exception:
+        _clear_artifacts(directory, SUCCESS_ARTIFACTS)
+        result.success = result.eligible_for_official_eval = False
+        if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
+            result.official_eval_ineligible_reasons.append("artifact_write_failed")
+        raise
+
+
+def _collect_output(
+    collector: RealtimeEventCollector,
+    playback: _Playback,
+    *,
+    stream_start: float,
+    video_duration_s: float,
+    require_response: bool,
+    retain_events: bool,
+    result: OmniInteractCaseResult,
+) -> _DeferredArtifactContext:
+    playback.ingest(collector)
+    created, done = response_ledger(collector)
+    if created != done:
+        raise ValueError(f"unfinished response_ids: {sorted(created - done)}")
+    done_events = [event for event in collector.events if event.get("type") == "response.done"]
+    cancelled = any(_response_status(event) == "cancelled" for event in done_events)
+    terminal_responses = {
+        response_id
+        for event in done_events
+        if _response_status(event) != "cancelled" and (response_id := collector.response_id(event))
+    }
+    horizon_s = math.ceil(video_duration_s)
+    horizon_bytes = horizon_s * OUTPUT_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+    spans: list[tuple[int, bytes]] = []
+    clipped_bytes = 0
+    response_times: dict[str, list[float]] = {}
+    responses_with_audio: set[str] = set()
+    for segment in playback.segments:
+        response_id = segment.response_id
+        if response_id not in created:
+            raise ValueError("response audio has no matching response.created")
+        start_s = max(0.0, segment.start_s - stream_start)
+        end_s = max(start_s, segment.end_s - stream_start)
+        offset = round(start_s * OUTPUT_SAMPLE_RATE) * PCM16_BYTES_PER_SAMPLE
+        writable = max(0, min(len(segment.pcm16), horizon_bytes - offset))
+        clipped_bytes += len(segment.pcm16) - writable
+        if writable:
+            clipped = segment.pcm16[:writable]
+            spans.append((offset, clipped))
+            if any(clipped):
+                responses_with_audio.add(response_id)
+        timing = response_times.setdefault(response_id, [start_s, end_s])
+        timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s)
+
+    text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"}
+    text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types}
+    if None in text_ids or not text_ids <= created:
+        raise ValueError("response transcript has no matching response.created")
+    chunks: list[dict[str, object]] = []
+    texts: list[str] = []
+    for response_id in collector.response_ids:
+        text = collector.response_text(response_id).strip()
+        if not text:
+            continue
+        response_timing = response_times.get(response_id)
+        if response_timing is None or response_timing[0] >= horizon_s:
+            continue
+        texts.append(text)
+        chunks.append(
+            {
+                "response_id": response_id,
+                "text": text,
+                "timestamp": [round(response_timing[0], 6), round(min(response_timing[1], horizon_s), 6)],
+            }
+        )
+    complete_outputs = terminal_responses & responses_with_audio & {str(chunk["response_id"]) for chunk in chunks}
+    if require_response and not complete_outputs:
+        raise ValueError("OmniInteract E2E requires a response with audio and transcript")
+    result.audio_bytes, result.audio_clipped_bytes = (
+        sum(len(segment.pcm16) for segment in playback.segments),
+        clipped_bytes,
+    )
+    result.transcript, result.responses, result.success = " ".join(texts).strip(), len(created), True
+    result.artifact_warnings = list(playback.warnings)
+    result.official_eval_ineligible_reasons = []
+    if clipped_bytes:
+        result.official_eval_ineligible_reasons.append("audio_clipped")
+    if cancelled:
+        result.official_eval_ineligible_reasons.append("cancelled_response")
+    result.eligible_for_official_eval = not result.official_eval_ineligible_reasons
+    audio_sizes = {segment.event_index: len(segment.pcm16) for segment in playback.segments}
+    events = (
+        [
+            {
+                **{key: value for key, value in event.items() if key not in {"delta", "audio"}},
+                "audio_bytes": audio_sizes[index],
+            }
+            if event.get("type") == "response.audio.delta"
+            else dict(event)
+            for index, event in enumerate(collector.events)
+        ]
+        if retain_events
+        else []
+    )
+    return _DeferredArtifactContext(horizon_bytes, tuple(spans), chunks, events)
+
+
+def prepare_success_artifacts(
+    root: Path,
+    case: OmniInteractCase,
+    collector: RealtimeEventCollector,
+    *,
+    playback: _Playback | None = None,
+    stream_start: float,
+    video_duration_s: float,
+    require_response: bool,
+    result: OmniInteractCaseResult,
+    capture_artifacts: bool = True,
+) -> dict[str, Any]:
+    directory = _output_dir(root, case)
+    context = _collect_output(
+        collector,
+        playback or _Playback(),
+        stream_start=stream_start,
+        video_duration_s=video_duration_s,
+        require_response=require_response,
+        retain_events=capture_artifacts,
+        result=result,
+    )
+    result.output_dir = str(directory.resolve())
+    summary = _artifact_summary(case, result)
+    if capture_artifacts:
+        result._artifact_context = context
+    return summary
+
+
+def write_failure_artifacts(root: Path, case: OmniInteractCase, result: OmniInteractCaseResult) -> None:
+    directory = _output_dir(root, case)
+    result.output_dir = str(directory.resolve())
+    result.success = result.eligible_for_official_eval = False
+    _clear_artifacts(directory, SUCCESS_ARTIFACTS)
+    _atomic_write_json(directory / ".failed.json", _artifact_summary(case, result))
+
+
+def publish_deferred_case_artifacts(
+    root: Path,
+    case: OmniInteractCase,
+    result: OmniInteractCaseResult,
+) -> None:
+    context = result._artifact_context
+    result._artifact_context = None
+    if not result.success:
+        write_failure_artifacts(root, case, result)
+        return
+    if context is None:
+        raise RuntimeError("OmniInteract benchmark output lost its deferred artifact context")
+    directory = _output_dir(root, case)
+    _publish_success_artifacts(directory, result, context, _artifact_summary(case, result))
+
+
+def clear_case_artifacts(root: Path, case: OmniInteractCase) -> None:
+    directory = _output_dir(root, case)
+    _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
+
+
+def clear_batch_artifacts(root: Path) -> None:
+    _clear_artifacts(root, BATCH_ARTIFACTS)
+
+
+def write_batch_artifacts(
+    root: Path,
+    cases: list[OmniInteractCase],
+    results: list[OmniInteractCaseResult],
+) -> None:
+    root.mkdir(parents=True, exist_ok=True)
+    rows = [
+        case_manifest(case, _output_dir(root, case))
+        for case, result in zip(cases, results, strict=True)
+        if result.success and result.eligible_for_official_eval
+    ]
+    ineligible = [result for result in results if result.success and not result.eligible_for_official_eval]
+    if ineligible:
+        reasons = Counter(
+            reason for result in ineligible for reason in (result.official_eval_ineligible_reasons or ["unspecified"])
+        )
+        logger.warning(
+            "%d successful OmniInteract cases were excluded from official evaluation: %s",
+            len(ineligible),
+            ", ".join(f"{reason}={count}" for reason, count in sorted(reasons.items())),
+        )
+    _atomic_write_json(root / "batch_summary.json", benchmark_summary(results))
+    _atomic_write_text(
+        root / "official_eval_manifest.jsonl",
+        "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
+    )
+
+
+def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str:
+    endpoint = (
+        config.endpoint
+        if urlsplit(config.endpoint).scheme
+        else urljoin(config.base_url.rstrip("/") + "/", config.endpoint.lstrip("/"))
+    )
+    return build_realtime_url(
+        endpoint,
+        config.model,
+        autostart=False,
+        native_duplex=True,
+        session_id=session_id,
+    )
+
+
+def _populate_response_metrics(
+    result: OmniInteractCaseResult,
+    collector: RealtimeEventCollector,
+    *,
+    stream_start: float,
+) -> None:
+    measurement_origin = {
+        "ttft": "response.created client receive to first non-empty text delta",
+        "ttfp": "response.created client receive to first audio packet",
+        "rtf": "response.created client receive to last audio packet divided by emitted audio duration",
+    }
+    request_metrics: list[dict[str, object]] = []
+    output_tokens = 0
+    for request_index, response_id in enumerate(collector.response_ids):
+        timing = collector.timing_summary(
+            after_s=stream_start,
+            input_committed_at_s=None,
+            response_id=response_id,
+            measurement_origin=measurement_origin,
+        )
+        raw_metric = timing.get("request_metrics")
+        stage0 = timing.get("stage0_tokens")
+        metric = {
+            "session_id": result.session_id,
+            "request_index": request_index,
+            "response_id": response_id,
+            **(raw_metric if isinstance(raw_metric, dict) else {}),
+        }
+        if isinstance(stage0, dict):
+            metric["stage0_tokens"] = dict(stage0)
+            output_tokens += int(stage0.get("output_token_count") or 0)
+        if isinstance(raw_metric, dict) or isinstance(stage0, dict):
+            request_metrics.append(metric)
+    result.output_tokens = output_tokens
+    result.duplex_request_metrics = request_metrics
+    result.duplex_session_metrics = summarize_session_request_metrics(
+        request_metrics,
+        session_id=result.session_id,
+    )
+
+
+async def run_omniinteract_case(
+    case: OmniInteractCase,
+    config: OmniInteractBenchmarkConfig,
+    *,
+    request_index: int | str,
+    capture_artifacts: bool = True,
+    prepared_input: OmniInteractPreparedInput | None = None,
+) -> OmniInteractCaseResult:
+    if not config.ref_audio:
+        raise ValueError("ref_audio is required for MiniCPM-o native-duplex audio output")
+    for name, value in (
+        ("timeout_s", config.timeout_s),
+        ("media_timeout_s", config.media_timeout_s),
+        ("max_video_duration_s", config.max_video_duration_s),
+    ):
+        if not math.isfinite(value) or value <= 0:
+            raise ValueError(f"{name} must be finite and positive")
+    session_id = f"omniinteract:{case.subset}:{request_index}:{time.monotonic_ns()}"
+    result = OmniInteractCaseResult(
+        case.subset,
+        str(case.video_path),
+        str(_output_dir(config.output_root, case).resolve()),
+        session_id=session_id,
+    )
+    started_at = time.monotonic()
+    try:
+        if prepared_input is None:
+            reference_audio = reference_audio_data_url(config.ref_audio)
+            duration, pcm, frames = await asyncio.to_thread(
+                prepare_media,
+                case.video_path,
+                VIDEO_FPS,
+                timeout_s=config.media_timeout_s,
+                max_duration_s=config.max_video_duration_s,
+            )
+        else:
+            reference_audio = prepared_input.ref_audio_data_url
+            duration = prepared_input.duration_s
+            pcm = prepared_input.pcm16
+            frames = prepared_input.video_frames
+        if not any(frames):
+            raise ValueError(f"No video frames were decoded from {case.video_path}")
+        async with RealtimeDuplexClient(
+            _websocket_url(config, session_id), additional_headers=config.extra_headers
+        ) as client:
+            session_from = len(client.events.events)
+            await client.configure(
+                config.model,
+                ref_audio=reference_audio,
+                session_id=session_id,
+                extra_body=config.extra_body,
+                idle_timeout_s=config.timeout_s,
+                timeout_s=min(config.timeout_s, 20.0),
+            )
+            pcm = _ensure_final_commit_tail(pcm, client.events.events)
+            playback = _Playback()
+            stream_start = time.monotonic()
+            try:
+                input_duration_s = len(pcm) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+                upload_timeout_s = config.timeout_s + input_duration_s
+                try:
+                    chunks, frame_count, mean_lag, max_lag = await asyncio.wait_for(
+                        stream_inputs(client, pcm, frames, playback), timeout=upload_timeout_s
+                    )
+                except asyncio.TimeoutError as exc:
+                    raise TimeoutError(f"Realtime upload timed out after {upload_timeout_s:g}s") from exc
+                commit_from = len(client.events.events)
+                await client.commit()
+                await wait_for_session_completion(
+                    client,
+                    playback,
+                    commit_from=commit_from,
+                    session_from=session_from,
+                    timeout_s=config.timeout_s,
+                    settle_s=_COMPLETION_SETTLE_S,
+                )
+                await playback.acknowledge(client, playback.end_s)
+                _raise_if_session_terminated(client.events, session_from)
+                close_from = len(client.events.events)
+                await client.close_session(timeout_s=min(config.timeout_s, 20.0))
+            except Exception:
+                with contextlib.suppress(Exception):
+                    await client.close_session(timeout_s=min(config.timeout_s, 20.0))
+                raise
+            errors = client.events.errors()
+            if errors:
+                raise RuntimeError(str(errors[-1]))
+            _raise_if_session_terminated(client.events, session_from, explicit_close_from=close_from)
+            result.latency_s = time.monotonic() - started_at
+            result.input_audio_chunks, result.input_video_frames = chunks, frame_count
+            result.pacing_mean_lag_s, result.pacing_max_lag_s = mean_lag, max_lag
+            _populate_response_metrics(result, client.events, stream_start=stream_start)
+            await asyncio.to_thread(
+                prepare_success_artifacts,
+                config.output_root,
+                case,
+                client.events,
+                playback=playback,
+                stream_start=stream_start,
+                video_duration_s=duration,
+                require_response=config.require_response,
+                result=result,
+                capture_artifacts=capture_artifacts,
+            )
+    except Exception as exc:
+        result.success = result.eligible_for_official_eval = False
+        if "case_failed" not in result.official_eval_ineligible_reasons:
+            result.official_eval_ineligible_reasons.append("case_failed")
+        result.error = str(exc)
+        result.latency_s = max(0.0, time.monotonic() - started_at)
+    return result

vllm_omni/benchmarks/patch/patch.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import contextlib
 import io
@@ -16,7 +19,6 @@ from dataclasses import dataclass, replace
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Literal
-from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 import aiohttp
 import numpy as np
@@ -31,6 +33,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (
     RequestFuncOutput,
     StreamedResponseHandler,
     _get_chat_content,
+    _get_headers,
     _update_headers_common,
     _update_payload_common,
     _validate_api_url,
@@ -46,6 +49,13 @@ from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
     daily_omni_local_videos_dir,
     resolve_daily_omni_local_root,
 )
+from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+    DEFAULT_OMNIINTERACT_REPO,
+    OmniInteractDataset,
+    OmniInteractPreparedInput,
+    OmniInteractSampleRequest,
+    OmniInteractSessionOptions,
+)
 from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
 from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
     SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT,
@@ -56,8 +66,27 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
 )
 from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
 from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
+from vllm_omni.benchmarks.omniinteract import (
+    VIDEO_FPS,
+    OmniInteractBenchmarkConfig,
+    OmniInteractCaseResult,
+    clear_batch_artifacts,
+    clear_case_artifacts,
+    prepare_media,
+    publish_deferred_case_artifacts,
+    run_omniinteract_case,
+    write_failure_artifacts,
+)
+from vllm_omni.benchmarks.omniinteract import (
+    benchmark_summary as omniinteract_benchmark_summary,
+)
+from vllm_omni.benchmarks.omniinteract import (
+    write_batch_artifacts as write_omniinteract_batch_artifacts,
+)
 from vllm_omni.experimental.fullduplex.client import (
     RealtimeDuplexClient,
+    build_realtime_url,
+    reference_audio_data_url,
     summarize_session_request_metrics,
     wait_for,
 )
@@ -221,6 +250,91 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu
     rfi.extra_body = base
 
 
+def _attach_omniinteract_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None:
+    if not isinstance(sample, OmniInteractSampleRequest):
+        return
+    setattr(rfi, "omniinteract_case", sample.omniinteract_case)
+    setattr(rfi, "omniinteract_options", sample.omniinteract_options)
+    setattr(rfi, "omniinteract_prepared_input", sample.omniinteract_prepared_input)
+
+
+def _append_error(existing: str, message: str) -> str:
+    return f"{existing}\n{message}" if existing else message
+
+
+def _finalize_omniinteract_batch(
+    input_requests: list[SampleRequest],
+    outputs: list[RequestFuncOutput],
+) -> dict[str, object] | None:
+    rows = [
+        (sample, output)
+        for sample, output in zip(input_requests, outputs, strict=True)
+        if isinstance(sample, OmniInteractSampleRequest)
+    ]
+    if not rows:
+        return None
+    options = rows[0][0].omniinteract_options
+    if not isinstance(options, OmniInteractSessionOptions):
+        raise RuntimeError("OmniInteract benchmark output lost its artifact options")
+    cases, results = [], []
+    artifact_errors: list[str] = []
+    for sample, output in rows:
+        case = sample.omniinteract_case
+        result = getattr(output, "omniinteract_case_result", None)
+        if case is None or not isinstance(result, OmniInteractCaseResult):
+            raise RuntimeError("OmniInteract benchmark output lost its dataset identity")
+        cases.append(case)
+        results.append(result)
+        try:
+            publish_deferred_case_artifacts(options.output_root, case, result)
+        except Exception as exc:  # noqa: BLE001 - artifact backends raise heterogeneous errors
+            result.success = result.eligible_for_official_eval = False
+            if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
+                result.official_eval_ineligible_reasons.append("artifact_write_failed")
+            artifact_error = f"Artifact publication failed: {exc}"
+            result.error = _append_error(result.error, artifact_error)
+            output.error = _append_error(output.error, artifact_error)
+            message = f"{case.subset}/{case.video_rel}: {artifact_error}"
+            artifact_errors.append(message)
+            logger.exception(message)
+            try:
+                write_failure_artifacts(options.output_root, case, result)
+            except Exception as recovery_exc:  # noqa: BLE001 - preserve benchmark metrics after I/O failure
+                message = f"{case.subset}/{case.video_rel}: failure artifact publication failed: {recovery_exc}"
+                artifact_errors.append(message)
+                logger.exception(message)
+    try:
+        write_omniinteract_batch_artifacts(options.output_root, cases, results)
+    except Exception as exc:  # noqa: BLE001 - preserve benchmark metrics after I/O failure
+        message = f"Batch artifact publication failed: {exc}"
+        artifact_errors.append(message)
+        logger.exception(message)
+        try:
+            clear_batch_artifacts(options.output_root)
+        except Exception as cleanup_exc:  # noqa: BLE001 - best-effort removal of partial batch artifacts
+            message = f"Partial batch artifact cleanup failed: {cleanup_exc}"
+            artifact_errors.append(message)
+            logger.exception(message)
+    summary = omniinteract_benchmark_summary(results)
+    compact_summary = {key: value for key, value in summary.items() if key != "results"}
+    compact_summary["artifacts_complete"] = not artifact_errors
+    if artifact_errors:
+        compact_summary["artifact_errors"] = artifact_errors
+    return compact_summary
+
+
+def _prepare_omniinteract_batch(input_requests: list[SampleRequest]) -> None:
+    roots: set[Path] = set()
+    for sample in input_requests:
+        if not isinstance(sample, OmniInteractSampleRequest):
+            continue
+        root = sample.omniinteract_options.output_root.resolve()
+        roots.add(root)
+        clear_case_artifacts(sample.omniinteract_options.output_root, sample.omniinteract_case)
+    for root in roots:
+        clear_batch_artifacts(root)
+
+
 def _daily_omni_repo_from_args(args) -> str | None:
     """Resolve HuggingFace repo id for Daily-Omni from CLI args.
 
@@ -250,6 +364,7 @@ def get_samples(args, tokenizer):
         "ttsd",
         "sound-effect",
     )
+    is_omniinteract = args.dataset_name == "omniinteract"
 
     # Check if we need to handle omni-related backends/datasets
     is_omni_backend = args.backend in [
@@ -258,12 +373,68 @@ def get_samples(args, tokenizer):
         "openai-realtime-duplex",
         "daily-omni",
     ]
-    is_omni_dataset = is_daily_omni or is_seed_tts or args.dataset_name == "random-mm"
+    is_omni_dataset = is_daily_omni or is_seed_tts or is_omniinteract or args.dataset_name == "random-mm"
 
     if not is_omni_backend and not is_omni_dataset:
         # Not an omni-related request, delegate to original implementation
         return get_samples_old(args, tokenizer)
 
+    if is_omniinteract:
+        dataset_path = getattr(args, "dataset_path", None)
+        data_root = None
+        dataset_repo = DEFAULT_OMNIINTERACT_REPO
+        if dataset_path:
+            candidate = Path(dataset_path).expanduser()
+            if candidate.exists() or candidate.is_absolute():
+                data_root = str(candidate)
+            else:
+                dataset_repo = str(dataset_path)
+        dataset = OmniInteractDataset(
+            data_root=data_root,
+            dataset_repo=dataset_repo,
+            subsets=tuple(getattr(args, "omniinteract_subsets")),
+            random_seed=args.seed,
+            disable_shuffle=getattr(args, "disable_shuffle", False),
+        )
+        options = OmniInteractSessionOptions(
+            output_root=Path(getattr(args, "omniinteract_output_dir")),
+            timeout_s=float(getattr(args, "omniinteract_timeout_s")),
+            media_timeout_s=float(getattr(args, "omniinteract_media_timeout_s")),
+            ref_audio=str(getattr(args, "omniinteract_ref_audio")),
+            require_response=bool(getattr(args, "omniinteract_require_response")),
+            max_video_duration_s=float(getattr(args, "omniinteract_max_video_duration_s")),
+        )
+        requests = dataset.sample(
+            tokenizer,
+            args.num_prompts,
+            request_id_prefix=args.request_id_prefix,
+            options=options,
+        )
+        if not requests:
+            raise ValueError("No OmniInteract sessions were selected")
+        encoded_ref_audio = reference_audio_data_url(options.ref_audio)
+        assert encoded_ref_audio is not None
+        for request in requests:
+            case = request.omniinteract_case
+            assert isinstance(request, OmniInteractSampleRequest) and case is not None
+            duration, pcm, frames = prepare_media(
+                case.video_path,
+                VIDEO_FPS,
+                timeout_s=options.media_timeout_s,
+                max_duration_s=options.max_video_duration_s,
+            )
+            if not any(frames):
+                raise ValueError(f"No video frames were decoded from {case.video_path}")
+            request.omniinteract_prepared_input = OmniInteractPreparedInput(
+                duration_s=duration,
+                pcm16=pcm,
+                video_frames=tuple(frames),
+                ref_audio_data_url=encoded_ref_audio,
+            )
+        # Replace OmniInteract's ``0`` (all) with the measured request count.
+        args.num_prompts = len(requests)
+        return requests
+
     # Handle Daily-Omni dataset
     if is_daily_omni:
         # Support:
@@ -476,6 +647,7 @@ class MixRequestFuncOutput(RequestFuncOutput):
     image_pixels: int = 0
     denoise_step_latency_ms: float = 0.0
     text_latency: float = 0.0
+    tpot_measured: bool = True
     #: Worst-case streaming-audio underrun (wall-clock seconds the player
     #: would have been starved). Populated by the audio-speech backend; ``0.0``
     #: for backends that do not run continuity analysis.
@@ -1313,11 +1485,115 @@ async def async_request_openai_audio_speech(
 
 
 def _realtime_websocket_url(api_url: str) -> str:
-    parts = urlsplit(api_url)
-    scheme = "wss" if parts.scheme == "https" else "ws"
-    query = dict(parse_qsl(parts.query, keep_blank_values=True))
-    query.setdefault("duplex", "1")
-    return urlunsplit((scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
+    return build_realtime_url(api_url, None, native_duplex=None)
+
+
+def _nonnegative_number(value: object) -> bool:
+    return isinstance(value, int | float) and not isinstance(value, bool) and np.isfinite(value) and value >= 0
+
+
+async def _async_request_omniinteract(
+    request_func_input: RequestFuncInput,
+    *,
+    pbar: tqdm | None,
+) -> MixRequestFuncOutput:
+    case = getattr(request_func_input, "omniinteract_case", None)
+    options = getattr(request_func_input, "omniinteract_options", None)
+    output = MixRequestFuncOutput()
+    output.prompt_len = request_func_input.prompt_len
+    output.start_time = time.perf_counter()
+    try:
+        if case is None or not isinstance(options, OmniInteractSessionOptions):
+            raise ValueError("OmniInteract RequestFuncInput is missing its dataset session layout")
+        prepared_input = getattr(request_func_input, "omniinteract_prepared_input", None)
+        if not isinstance(prepared_input, OmniInteractPreparedInput):
+            raise ValueError("OmniInteract media must be prepared before benchmark request timing")
+        headers = _get_headers()
+        _update_headers_common(headers, request_func_input)
+        config = OmniInteractBenchmarkConfig(
+            endpoint=request_func_input.api_url,
+            model=request_func_input.model_name or request_func_input.model,
+            output_root=options.output_root,
+            timeout_s=options.timeout_s,
+            media_timeout_s=options.media_timeout_s,
+            max_video_duration_s=options.max_video_duration_s,
+            ref_audio=options.ref_audio,
+            require_response=options.require_response,
+            extra_headers=headers or None,
+            extra_body=dict(request_func_input.extra_body or {}) or None,
+        )
+        case_result = await run_omniinteract_case(
+            case,
+            config,
+            request_index=request_func_input.request_id or uuid.uuid4().hex,
+            capture_artifacts=request_func_input.request_id is not None,
+            prepared_input=prepared_input,
+        )
+        output.latency = case_result.latency_s
+        output.generated_text = case_result.transcript
+        output.output_tokens = case_result.output_tokens
+        output.audio_duration = case_result.audio_bytes / (24_000 * 2)
+        output.audio_frames = case_result.audio_bytes // 2
+        session_metrics = case_result.duplex_session_metrics
+        output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0
+        output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0
+        output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0)
+        stages: list[tuple[int, dict[str, object]]] = []
+        for request_metric in case_result.duplex_request_metrics:
+            stage0 = request_metric.get("stage0_tokens")
+            if not isinstance(stage0, dict):
+                continue
+            token_count = stage0.get("output_token_count")
+            if isinstance(token_count, int) and not isinstance(token_count, bool) and token_count >= 0:
+                stages.append((token_count, stage0))
+        timed = [(count - 1, stage) for count, stage in stages if count > 1]
+        complete = (
+            len(stages) == len(case_result.duplex_request_metrics)
+            and sum(count for count, _ in stages) == output.output_tokens
+        )
+        exact = complete and all(
+            isinstance(values := stage.get("itls_ms"), list)
+            and len(values) == intervals
+            and all(_nonnegative_number(value) for value in values)
+            for intervals, stage in timed
+        )
+        output.itl = [float(value) / 1000.0 for _, stage in timed for value in stage["itls_ms"]] if exact else []
+        if output.itl:
+            output.text_latency = output.ttft + sum(output.itl)
+        elif complete and timed and all(_nonnegative_number(stage.get("tpot_ms")) for _, stage in timed):
+            intervals = sum(count for count, _ in timed)
+            weighted_tpot = sum(float(stage["tpot_ms"]) * count for count, stage in timed) / intervals / 1000.0
+            output.text_latency = output.ttft + weighted_tpot * (output.output_tokens - 1)
+        else:
+            output.text_latency = output.ttft
+            output.tpot_measured = False
+            if output.output_tokens > 1 or (output.output_tokens == 0 and output.generated_text):
+                logger.warning(
+                    "OmniInteract session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
+                    case_result.session_id,
+                )
+        output.duplex_request_metrics = case_result.duplex_request_metrics
+        output.duplex_session_metrics = session_metrics
+        output.success = case_result.success
+        output.error = case_result.error
+        setattr(output, "omniinteract_case", case)
+        setattr(output, "omniinteract_case_result", case_result)
+    except Exception:
+        output.success = False
+        output.error = traceback.format_exc()
+        logger.error("OmniInteract Realtime request failed: %s", output.error)
+        if case is not None:
+            case_result = OmniInteractCaseResult(
+                subset=case.subset,
+                video=str(case.video_path),
+                output_dir="",
+                error=output.error,
+            )
+            setattr(output, "omniinteract_case", case)
+            setattr(output, "omniinteract_case_result", case_result)
+    if pbar:
+        pbar.update(1)
+    return output
 
 
 async def async_request_openai_realtime_duplex(
@@ -1325,8 +1601,9 @@ async def async_request_openai_realtime_duplex(
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
 ) -> MixRequestFuncOutput:
-    """Run one or more Seed-TTS turns through one explicit Realtime TTS session."""
     del session
+    if getattr(request_func_input, "omniinteract_case", None) is not None:
+        return await _async_request_omniinteract(request_func_input, pbar=pbar)
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
     output.start_time = time.perf_counter()
@@ -1387,8 +1664,10 @@ async def async_request_openai_realtime_duplex(
                 )
                 await client.send({"type": "response.create"})
                 await wait_for(
-                    lambda: client.events.count("response.done") > done_before
-                    or len(client.events.errors()) > errors_before,
+                    lambda: (
+                        client.events.count("response.done") > done_before
+                        or len(client.events.errors()) > errors_before
+                    ),
                     timeout_s=180.0,
                     label=f"Seed-TTS Realtime TTS turn {request_index} response.done",
                 )
@@ -1433,19 +1712,7 @@ async def async_request_openai_realtime_duplex(
                         channels=1,
                     )
                 )
-                turn_transcripts.append(
-                    "".join(
-                        str(event.get("delta") or "")
-                        for event in client.events.events
-                        if client.events.response_id(event) == response_id
-                        and event.get("type")
-                        in {
-                            "response.audio_transcript.delta",
-                            "response.output_text.delta",
-                            "response.text.delta",
-                        }
-                    )
-                )
+                turn_transcripts.append(client.events.response_text(response_id))
                 await client.acknowledge_playback()
             request_finished_at = time.perf_counter()
             session_metrics = summarize_session_request_metrics(
@@ -1520,6 +1787,7 @@ from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, ge
 from vllm_omni.benchmarks.metrics.metrics import (
     MultiModalsBenchmarkMetrics,
     calculate_metrics,
+    has_metric_samples,
 )
 
 # ruff: noqa: E402
@@ -1626,6 +1894,7 @@ async def benchmark(
     )
     _attach_daily_omni_to_request_func_input(input_requests[0], test_input)
     _attach_seed_tts_to_request_func_input(input_requests[0], test_input)
+    _attach_omniinteract_to_request_func_input(input_requests[0], test_input)
 
     if ready_check_timeout_sec > 0:
         test_output = await wait_for_endpoint(
@@ -1741,6 +2010,7 @@ async def benchmark(
         print(f"Probe request rate: {probe_request_rate} req/s")
         probe_task = asyncio.create_task(probe_loop())
 
+    _prepare_omniinteract_batch(input_requests)
     benchmark_start_time = time.perf_counter()
     tasks: list[asyncio.Task] = []
 
@@ -1801,6 +2071,7 @@ async def benchmark(
         )
         _attach_daily_omni_to_request_func_input(request, request_func_input)
         _attach_seed_tts_to_request_func_input(request, request_func_input)
+        _attach_omniinteract_to_request_func_input(request, request_func_input)
         tasks.append(
             asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar))
         )
@@ -1815,6 +2086,8 @@ async def benchmark(
 
     benchmark_duration = time.perf_counter() - benchmark_start_time
 
+    omniinteract_summary = _finalize_omniinteract_batch(input_requests, outputs)
+
     if task_type == TaskType.GENERATION:
         metrics, actual_output_lens = calculate_metrics(
             input_requests=input_requests,
@@ -1839,6 +2112,13 @@ async def benchmark(
         actual_output_lens = 0
 
     if isinstance(metrics, MultiModalsBenchmarkMetrics):
+
+        def measured_ttft(output: RequestFuncOutput) -> float | None:
+            session_metrics = getattr(output, "duplex_session_metrics", None)
+            if isinstance(session_metrics, dict) and session_metrics.get("mean_ttft_ms") is None:
+                return None
+            return output.ttft
+
         result = {
             "duration": benchmark_duration,
             "completed": metrics.completed,
@@ -1859,7 +2139,7 @@ async def benchmark(
             "input_lens": [output.prompt_len for output in outputs],
             "start_times": [output.start_time for output in outputs],
             "output_lens": actual_output_lens,
-            "ttfts": [output.ttft for output in outputs],
+            "ttfts": [measured_ttft(output) for output in outputs],
             "itls": [output.itl for output in outputs],
             "generated_texts": [output.generated_text for output in outputs],
             "errors": [output.error for output in outputs],
@@ -1893,6 +2173,8 @@ async def benchmark(
     ]
     if duplex_session_metrics:
         result["duplex_session_metrics"] = duplex_session_metrics
+    if omniinteract_summary is not None:
+        result["omniinteract"] = omniinteract_summary
 
     from vllm_omni.benchmarks.data_modules.daily_omni_eval import (
         compute_daily_omni_accuracy_metrics,
@@ -1949,6 +2231,8 @@ async def benchmark(
         # metric.
         if metric_attribute_name not in result_percentile_metrics:
             return
+        if not has_metric_samples(metrics, metric_attribute_name):
+            return
         # No text tokens generated (e.g. pure TTS speech endpoint): per-token
         # latency metrics (ttft/tpot/itl) are undefined, so skip them.
         is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio"))

vllm_omni/benchmarks/serve.py

@@ -1,6 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
 import argparse
 import asyncio
+import contextlib
 import os
+from pathlib import Path
 from typing import Any
 
 from vllm.benchmarks.serve import main_async
@@ -9,6 +16,7 @@ from vllm.benchmarks.serve import main_async
 # This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
 # Must be imported before any vllm.benchmarks module usage
 import vllm_omni.benchmarks.patch.patch  # noqa: F401
+from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock
 from vllm_omni.benchmarks.patch.patch import (
     maybe_enable_stage_metrics,
     set_print_stage,
@@ -28,4 +36,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
         getattr(args, "extra_body", None),
         enabled=should_request_stage_metrics(args),
     )
-    return asyncio.run(main_async(args))
+    lock = (
+        omniinteract_output_lock(Path(args.omniinteract_output_dir))
+        if getattr(args, "dataset_name", None) == "omniinteract"
+        else contextlib.nullcontext()
+    )
+    with lock:
+        return asyncio.run(main_async(args))

vllm_omni/config/environment_variable_inventory.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reviewed inventory of environment variables used by vLLM-Omni.
+
+This module classifies environment-variable names; it does not read them or
+make every classified name part of the public configuration contract. In
+particular, model-specific variables remain transitional until their recorded
+disposition is implemented.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import Enum
+from types import MappingProxyType
+
+
+class EnvironmentVariableCategory(str, Enum):
+    """Ownership and documentation boundary for an environment variable."""
+
+    PUBLIC_OMNI = "public_omni"
+    INHERITED_VLLM = "inherited_vllm"
+    PLATFORM_EXTERNAL = "platform_external"
+    MODEL_SPECIFIC = "model_specific"
+    BENCHMARK_TRANSITIONAL = "benchmark_transitional"
+    INTERNAL = "internal"
+
+
+class ModelEnvironmentVariableDisposition(str, Enum):
+    """Required destination for a model-specific environment variable."""
+
+    PROMOTE = "promote"
+    REQUEST_SCOPE = "request_scope"
+    EXTERNAL = "external"
+    INTERNALIZE = "internalize"
+    DEPRECATE_REMOVE = "deprecate_remove"
+
+
+@dataclass(frozen=True)
+class EnvironmentVariableClassification:
+    """Classification metadata for one statically identifiable name."""
+
+    name: str
+    category: EnvironmentVariableCategory
+    model_disposition: ModelEnvironmentVariableDisposition | None = None
+    redact_value: bool = False
+
+    @property
+    def is_public_omni(self) -> bool:
+        return self.category is EnvironmentVariableCategory.PUBLIC_OMNI
+
+
+# New entries must use the VLLM_OMNI_ prefix. The regression test explicitly
+# grandfathers the older public names below.
+_PUBLIC_OMNI = (
+    "DIFFUSION_ATTENTION_BACKEND",
+    "DIFFUSION_CACHE_ADAPTER",
+    "DIFFUSION_CACHE_BACKEND",
+    "OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
+    "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
+    "OMNI_DIFFUSION_SESSION_STATE_MANAGER",
+    "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
+    "SPEAKER_MAX_UPLOADED",
+    "SPEAKER_SAMPLES_DIR",
+    "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
+    "VLLM_OMNI_ORCH_MONITOR_PATH",
+    "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
+    "VLLM_OMNI_SERVER_STORAGE__FILE_TTL",
+    "VLLM_OMNI_SERVER_STORAGE__PATH",
+    "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
+    "VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP",
+    "VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
+    "VLLM_OMNI_STORAGE_PATH",
+    "VLLM_OMNI_USE_QUACK_FP8",
+    "VLLM_OMNI_VIDEO_SYNC_TIMEOUT",
+    "VLLM_VIDEO_ASYNC_CHUNK",
+    "VLLM_VIDEO_AUDIO_DELTA_MODE",
+)
+
+_INHERITED_VLLM = (
+    "CUDA_HOME",
+    "CUDA_VISIBLE_DEVICES",
+    "LOCAL_RANK",
+    "NO_COLOR",
+    "VLLM_ALLOW_LONG_MAX_MODEL_LEN",
+    "VLLM_BATCH_INVARIANT",
+    "VLLM_CACHE_ROOT",
+    "VLLM_DISABLE_LOG_LOGO",
+    "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE",
+    "VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
+    "VLLM_LOGGING_COLOR",
+    "VLLM_MOONCAKE_BOOTSTRAP_PORT",
+    "VLLM_PLUGINS",
+    "VLLM_ROCM_USE_AITER",
+    "VLLM_ROCM_USE_AITER_RMSNORM",
+    "VLLM_TUNED_CONFIG_FOLDER",
+    "VLLM_USE_MODELSCOPE",
+    "VLLM_USE_OINK_OPS",
+    "VLLM_WORKER_MULTIPROC_METHOD",
+    "VLLM_XPU_USE_SAMPLER_KERNEL",
+)
+
+_PLATFORM_EXTERNAL = (
+    "ASCEND_LAUNCH_BLOCKING",
+    "ASCEND_RT_VISIBLE_DEVICES",
+    "FLASHINFER_DISABLE_VERSION_CHECK",
+    "HF_HOME",
+    "HF_HUB_DISABLE_XET",
+    "HF_TOKEN",
+    "HIP_VISIBLE_DEVICES",
+    "HUGGINGFACE_HUB_TOKEN",
+    "MASTER_ADDR",
+    "MASTER_PORT",
+    "MINDIE_SD_FA_TYPE",
+    "MORI_RDMA_DEVICES",
+    "MUSA_VISIBLE_DEVICES",
+    "MV2_COMM_WORLD_LOCAL_RANK",
+    "NCCL_ASYNC_ERROR_HANDLING",
+    "OMPI_COMM_WORLD_LOCAL_RANK",
+    "OPENAI_API_KEY",
+    "PYTHONPATH",
+    "QUACK_CACHE_DIR",
+    "RANK",
+    "RANK_ID",
+    "RAY_RAYLET_PID",
+    "RDMA_DEVICE_NAME",
+    "SAGE_ATTN_TENSOR_LAYOUT",
+    "VLLM_LOCAL_RANK",
+    "WORLD_SIZE",
+    "XDG_CACHE_HOME",
+)
+
+_MODEL_PROMOTE = (
+    "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH",
+    "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
+    "COSMOS3_SOUND_TOKENIZER_PATH",
+    "COSYVOICE3_ESTIMATOR_ONNX",
+    "COSYVOICE3_TRT",
+    "COSYVOICE3_TRT_CACHE",
+    "DYNIN_CONFIG_PATH",
+    "DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY",
+    "DYNIN_MAGVIT_REMOTE_CODE_REPO_ID",
+    "DYNIN_MAGVIT_REMOTE_CODE_REVISION",
+    "DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY",
+    "DYNIN_REMOTE_CODE_REPO_ID",
+    "DYNIN_REMOTE_CODE_REVISION",
+    "HIGGS_AUDIO_TOKENIZER_PATH",
+    "INTERNVLA_A1_COSMOS_DECODER_PATH",
+    "INTERNVLA_A1_COSMOS_DIR",
+    "INTERNVLA_A1_COSMOS_ENCODER_PATH",
+    "INTERNVLA_A1_PROCESSOR_DIR",
+    "MIMO_AUDIO_TOKENIZER_CUDA_GRAPH",
+    "MIMO_AUDIO_TOKENIZER_DEVICE",
+    "MIMO_AUDIO_TOKENIZER_PATH",
+    "MING_CFM_CUDAGRAPH",
+    "MINICPMO_TOKEN2WAV_TRT",
+    "MINICPMO_TOKEN2WAV_TRT_CACHE",
+    "MINICPMO_TOKEN2WAV_TRT_DTYPE",
+    "NEMOTRON_VOICECHAT_LLM_PATH",
+    "OMNIVOICE_CUDA_GRAPH",
+    "OMNIVOICE_TF32",
+    "VLLM_GEPARD_CHUNK_FRAMES",
+    "VLLM_GEPARD_FIRST_CHUNK_FRAMES",
+    "VLLM_GEPARD_LOOKBACK_FRAMES",
+    "VLLM_OMNI_FISH_KVCACHE_ATTN",
+    "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
+)
+
+_MODEL_REQUEST_SCOPE = (
+    "DOTS_TTS_DIT_NUM_STEPS",
+    "GR00T_NOISE_SEED",
+    "STEP_AUDIO2_DEFAULT_PROMPT_WAV",
+    "VLLM_GEPARD_END_FADE_MS",
+    "VLLM_GEPARD_END_SILENCE_MS",
+    "VLLM_GEPARD_GREEDY",
+)
+
+# No audited model variable is currently owned by a supported third-party
+# contract. Keep the disposition explicit so future reviews cannot silently
+# fold an externally owned name into "promote".
+_MODEL_EXTERNAL: tuple[str, ...] = ()
+
+_MODEL_INTERNALIZE = (
+    "DOTS_TTS_BETA_TRACE",
+    "DYNIN_S2U_VENDOR_ROOT",
+    "DZ_PHASE_TIMING",
+    "MING_TTS_STAGE1_FINAL_LOG",
+    "MOSS_TTS_DEBUG_STOP",
+    "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE",
+    "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS",
+    "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY",
+    "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS",
+    "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_FILE",
+    "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY",
+)
+
+_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",
+    "XCODEC1_PATH",
+    "model_stage",
+)
+
+_BENCHMARK_TRANSITIONAL = (
+    "DAILY_OMNI_EXTRACT_MODE",
+    "DAILY_OMNI_SAVE_EVAL_ITEMS",
+    "MAX_NUM_FRAMES",
+    "SEED_TTS_EVAL_DEVICE",
+    "SEED_TTS_HF_WHISPER_MODEL",
+    "SEED_TTS_SIM_DEVICE",
+    "SEED_TTS_SIM_EVAL",
+    "SEED_TTS_UTMOS_EVAL",
+    "SEED_TTS_UTMOS_HF_REPO",
+    "SEED_TTS_UTMOS_JIT_FILE",
+    "SEED_TTS_WAVLM_MAX_SECONDS",
+    "SEED_TTS_WAVLM_MIN_SAMPLES",
+    "SEED_TTS_WAVLM_MODEL",
+    "SEED_TTS_WER_EVAL",
+    "SEED_TTS_WER_SAVE_AUDIO_DIR",
+    "SEED_TTS_WER_SAVE_ITEMS",
+    "VLLM_DAILY_OMNI_MEDIA_REPO",
+    "VLLM_OMNI_BENCH_AUDIO_CHANNELS",
+    "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S",
+    "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
+)
+
+_INTERNAL = (
+    "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT",
+    "VLLM_OMNI_REPLICA_ID",
+)
+
+_REDACTED = frozenset(
+    {
+        "HF_TOKEN",
+        "HUGGINGFACE_HUB_TOKEN",
+        "OPENAI_API_KEY",
+    }
+)
+
+
+def _build_inventory() -> Mapping[str, EnvironmentVariableClassification]:
+    inventory: dict[str, EnvironmentVariableClassification] = {}
+
+    def add(
+        names: tuple[str, ...],
+        category: EnvironmentVariableCategory,
+        disposition: ModelEnvironmentVariableDisposition | None = None,
+    ) -> None:
+        for name in names:
+            if name in inventory:
+                raise ValueError(f"Environment variable {name!r} is classified more than once")
+            inventory[name] = EnvironmentVariableClassification(
+                name=name,
+                category=category,
+                model_disposition=disposition,
+                redact_value=name in _REDACTED,
+            )
+
+    add(_PUBLIC_OMNI, EnvironmentVariableCategory.PUBLIC_OMNI)
+    add(_INHERITED_VLLM, EnvironmentVariableCategory.INHERITED_VLLM)
+    add(_PLATFORM_EXTERNAL, EnvironmentVariableCategory.PLATFORM_EXTERNAL)
+    add(_MODEL_PROMOTE, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.PROMOTE)
+    add(
+        _MODEL_REQUEST_SCOPE,
+        EnvironmentVariableCategory.MODEL_SPECIFIC,
+        ModelEnvironmentVariableDisposition.REQUEST_SCOPE,
+    )
+    add(_MODEL_EXTERNAL, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.EXTERNAL)
+    add(
+        _MODEL_INTERNALIZE,
+        EnvironmentVariableCategory.MODEL_SPECIFIC,
+        ModelEnvironmentVariableDisposition.INTERNALIZE,
+    )
+    add(
+        _MODEL_DEPRECATE_REMOVE,
+        EnvironmentVariableCategory.MODEL_SPECIFIC,
+        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE,
+    )
+    add(_BENCHMARK_TRANSITIONAL, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL)
+    add(_INTERNAL, EnvironmentVariableCategory.INTERNAL)
+    return MappingProxyType(dict(sorted(inventory.items())))
+
+
+# All reviewed environment-variable names, indexed by exact name. This is
+# classification metadata, not the executable value registry used by
+# ``vllm.envs``.
+ENVIRONMENT_VARIABLE_INVENTORY: Mapping[str, EnvironmentVariableClassification] = _build_inventory()

vllm_omni/config/omni_config.py

@@ -697,8 +697,10 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_offload_components: dict[str, bool] = Field(default_factory=dict)
     host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -756,6 +758,7 @@ class _DiffusionConfigProjection:
             TransformerConfig,
             build_attention_config,
             parse_kv_cache_skip_selector,
+            validate_dlo_host_registration_options,
             validate_host_weight_runtime_options,
         )
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
@@ -838,6 +841,12 @@ class _DiffusionConfigProjection:
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
         )
+        self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=self.dlo_host_registration_limit_gib,
+            enable_dlo=self.enable_distributed_layerwise_offload,
+            use_allgather=self.dlo_use_allgather,
+            hwr_mode=self.host_weight_runtime_mode,
+        )
 
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(

vllm_omni/config/stage_config.py

@@ -389,8 +389,10 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_offload_components: dict[str, bool] | None = None
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_limit_gib: float | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/deploy/wan2_2_ti2v.yaml

@@ -1,6 +1,7 @@
 # Wan2.2 TI2V-5B deploy.
 #
 # Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
+# Request-level batching uses scheduler max_num_seqs, not diffusion_batch_size.
 async_chunk: false
 trust_remote_code: true
 distributed_executor_backend: mp
@@ -8,7 +9,7 @@ distributed_executor_backend: mp
 stages:
   - stage_id: 0
     devices: "0"
-    max_num_seqs: 1
+    max_num_seqs: 4  # in-flight request batch width (CLI: --max-num-seqs)
     tensor_parallel_size: 1
     enforce_eager: true
     # Keep VAE tiling on by default for higher resolutions / longer clips.

vllm_omni/diffusion/data.py

@@ -113,6 +113,22 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
         raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
 
 
+def validate_dlo_host_registration_options(
+    *,
+    limit_gib: object,
+    enable_dlo: bool,
+    use_allgather: bool,
+    hwr_mode: object,
+) -> float:
+    """Validate the optional transport budget without probing CUDA or HWR."""
+    value = float(limit_gib)
+    if not math.isfinite(value) or value < 0:
+        raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
+    if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
+        raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime")
+    return value
+
+
 def parse_kv_cache_skip_selector(
     selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -779,11 +795,17 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Auxiliary encoder/VAE policy. False keeps the whole component resident;
+    # DiT offload remains owned by DLO.
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
     # Final-layout Host Weight Runtime policy. The loader only activates this
     # for eligible no-AllGather DLO; all other configurations preserve their
     # existing loader/storage path.
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    # Optional per-worker ceiling for registering final-layout HWR mappings.
+    # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
+    dlo_host_registration_limit_gib: float = 0.0
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
@@ -1164,6 +1186,12 @@ class OmniDiffusionConfig:
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
         )
+        self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=self.dlo_host_registration_limit_gib,
+            enable_dlo=self.enable_distributed_layerwise_offload,
+            use_allgather=self.dlo_use_allgather,
+            hwr_mode=self.host_weight_runtime_mode,
+        )
 
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(

vllm_omni/diffusion/model_loader/host_weight_loader.py

@@ -27,7 +27,10 @@ logger = init_logger(__name__)
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
-    from vllm_omni.diffusion.model_loader.host_weights.source_identity import PreparedWeightSource
+    from vllm_omni.diffusion.model_loader.host_weights.source_identity import (
+        NodeSourceDigestCache,
+        PreparedWeightSource,
+    )
 
 
 class _HWRCommitError(RuntimeError):
@@ -215,6 +218,7 @@ class HWRLoaderMixin:
         *,
         load_format: str,
         sources: Sequence[object],
+        source_digest_cache: NodeSourceDigestCache | None = None,
     ) -> FinalLayoutIdentityContext:
         from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
 
@@ -293,6 +297,7 @@ class HWRLoaderMixin:
             prepared_sources=prepared_sources,
             request=request,
             policy=FINAL_LAYOUT_BF16_POLICY,
+            source_digest_cache=source_digest_cache,
         )
 
     def _resolve_hwr(
@@ -306,7 +311,10 @@ class HWRLoaderMixin:
         sources: Sequence[object],
     ) -> dict[str, object] | None:
         """Resolve an eligible no-AllGather final-layout HWR transaction."""
-        from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutTensorRestorer
+        from vllm_omni.diffusion.model_loader.host_weights import (
+            FinalLayoutTensorRestorer,
+            NodeSourceDigestCache,
+        )
         from vllm_omni.host_weight_runtime import (
             HostWeightLeaseCarrier,
             HostWeightRuntime,
@@ -347,7 +355,6 @@ class HWRLoaderMixin:
                 raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
             logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
             return None
-        context = self._build_hwr_context(model, modules, load_format=load_format, sources=sources)
         root_value = getattr(self.od_config, "host_weight_runtime_root", None)
         if not isinstance(root_value, str) or not root_value.strip():
             raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
@@ -362,6 +369,17 @@ class HWRLoaderMixin:
                 ),
             )
         )
+        source_digest_cache = NodeSourceDigestCache(
+            root,
+            timeout_seconds=runtime.config.wait.coordination_timeout_seconds,
+        )
+        context = self._build_hwr_context(
+            model,
+            modules,
+            load_format=load_format,
+            sources=sources,
+            source_digest_cache=source_digest_cache,
+        )
         resolution = runtime.resolve(context.identity)
         state: dict[str, object] = {
             "mode": mode,

vllm_omni/diffusion/model_loader/host_weights/__init__.py

@@ -19,7 +19,7 @@ from .producers import (
     FinalLayoutBF16Producer,
 )
 from .restorer import FinalLayoutTensorRestorer
-from .source_identity import PreparedWeightSource, WeightSourceKind
+from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind
 
 __all__ = [
     "FINAL_LAYOUT_BF16_POLICY",
@@ -32,6 +32,7 @@ __all__ = [
     "FinalLayoutRequest",
     "FinalLayoutTensorRestorer",
     "ImplementationIdentity",
+    "NodeSourceDigestCache",
     "PreparedWeightSource",
     "WeightSourceKind",
     "build_final_layout_identity",

vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py

@@ -27,6 +27,7 @@ from .contracts import (
 )
 from .source_identity import (
     FinalLayoutSourceContext,
+    NodeSourceDigestCache,
     PreparedWeightSource,
     resolve_final_layout_source_identity,
 )
@@ -78,6 +79,7 @@ def build_final_layout_identity(
     prepared_sources: Sequence[PreparedWeightSource],
     request: FinalLayoutRequest,
     policy: FinalLayoutTensorPolicy,
+    source_digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutIdentityContext:
     """Build exact source, representation, layout, and implementation identity."""
     if not isinstance(request, FinalLayoutRequest):
@@ -99,6 +101,7 @@ def build_final_layout_identity(
         prepared_sources,
         model_id=request.model_id,
         target_names=target_names,
+        digest_cache=source_digest_cache,
     )
     dit_names = tuple(name for name, _ in dit_modules)
     model_contracts = tuple(binding.contract for binding in bindings)

vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py

@@ -286,7 +286,7 @@ class FinalLayoutBF16Producer:
                 for record in shard
             )
             file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
-            with writer.open_tensor_file(file_name, specs) as output:
+            with writer.open_tensor_file(file_name, specs, ordered=True) as output:
                 for record in shard:
                     output.write_tensor(record.name, record.tensor.detach())
 

vllm_omni/diffusion/model_loader/host_weights/source_identity.py

@@ -5,9 +5,13 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import os
 import stat
+import time
+import uuid
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -15,6 +19,7 @@ from pathlib import Path
 import regex as re
 
 from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError
 from vllm_omni.host_weight_runtime.identity import canonical_json
 
 from .contracts import (
@@ -101,6 +106,140 @@ class _FileSnapshot:
         )
 
 
+@dataclass(frozen=True)
+class _FileState:
+    size: int
+    device: int
+    inode: int
+    mtime_ns: int
+    ctime_ns: int
+    symlink_target: str | None
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "size": self.size,
+            "device": self.device,
+            "inode": self.inode,
+            "mtime_ns": self.mtime_ns,
+            "ctime_ns": self.ctime_ns,
+            "symlink_target": self.symlink_target,
+        }
+
+
+def _observe(path: Path) -> _FileState:
+    try:
+        current = path.stat()
+    except OSError as exc:
+        raise ValueError(f"cannot stat canonical weight file {path}") from exc
+    if not stat.S_ISREG(current.st_mode):
+        raise ValueError(f"canonical weight source is not a regular file: {path}")
+    return _FileState(
+        size=current.st_size,
+        device=current.st_dev,
+        inode=current.st_ino,
+        mtime_ns=current.st_mtime_ns,
+        ctime_ns=current.st_ctime_ns,
+        symlink_target=os.readlink(path) if path.is_symlink() else None,
+    )
+
+
+def _sha256_file(path: Path, expected: _FileState) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        before = os.fstat(handle.fileno())
+        while chunk := handle.read(_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+        after = os.fstat(handle.fileno())
+    observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns)
+    completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns)
+    expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns)
+    if observed != expected_stat or completed != expected_stat or _observe(path) != expected:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.SOURCE_CHANGED,
+            f"canonical weight file changed while hashing: {path}",
+        )
+    return digest.hexdigest()
+
+
+class NodeSourceDigestCache:
+    """Share content-backed local-source fingerprints across node workers."""
+
+    def __init__(self, root: Path, *, timeout_seconds: float) -> None:
+        self.entries = root / "source-digests-v1" / "entries"
+        self.locks = root / "source-digests-v1" / "locks"
+        self.timeout_seconds = timeout_seconds
+        self._available = True
+        try:
+            self.entries.mkdir(parents=True, exist_ok=True, mode=0o700)
+            self.locks.mkdir(parents=True, exist_ok=True, mode=0o700)
+        except OSError:
+            self._available = False
+
+    @staticmethod
+    def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None:
+        try:
+            cached = json.loads(entry.read_bytes())
+            if not isinstance(cached, dict):
+                return None
+            cached_state = cached["state"]
+            content_id = cached["content_id"]
+            record_sha256 = cached["record_sha256"]
+        except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError):
+            return None
+        record = {
+            "content_id": content_id,
+            "schema_version": cached.get("schema_version"),
+            "state": cached_state,
+        }
+        if (
+            record["schema_version"] != 1
+            or cached_state != state.to_dict()
+            or not isinstance(content_id, str)
+            or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None
+            or not isinstance(record_sha256, str)
+            or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256
+            or _observe(path) != state
+        ):
+            return None
+        return content_id
+
+    def content_id(self, path: Path, state: _FileState) -> str:
+        if not self._available:
+            return f"sha256:{_sha256_file(path, state)}"
+        key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest()
+        entry = self.entries / f"{key}.json"
+        content_id: str | None = None
+        try:
+            with FileLock(
+                self.locks / f"{key}.lock",
+                exclusive=True,
+                deadline=time.monotonic() + self.timeout_seconds,
+            ):
+                if (cached := self._cached_content_id(entry, path, state)) is not None:
+                    return cached
+                content_id = f"sha256:{_sha256_file(path, state)}"
+                record: dict[str, object] = {
+                    "content_id": content_id,
+                    "schema_version": 1,
+                    "state": state.to_dict(),
+                }
+                record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest()
+                temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+                try:
+                    temporary.write_bytes(canonical_json(record))
+                    os.chmod(temporary, 0o600)
+                    os.replace(temporary, entry)
+                finally:
+                    temporary.unlink(missing_ok=True)
+                return content_id
+        except (FileLockTimeoutError, OSError):
+            # Digest sharing is an optimization. Coordination or cache I/O
+            # failures must not make a readable canonical source unavailable.
+            if content_id is not None and _observe(path) == state:
+                return content_id
+            return f"sha256:{_sha256_file(path, state)}"
+
+
 @dataclass(frozen=True)
 class _SourceSnapshot:
     semantic: CanonicalJson
@@ -129,14 +268,6 @@ class FinalLayoutSourceContext:
             )
 
 
-def _sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        while chunk := handle.read(_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
 def _hf_blob_content_id(
     path: Path,
     *,
@@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str:
     return value
 
 
-def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
+def _snapshot_source(
+    source: PreparedWeightSource,
+    digest_cache: NodeSourceDigestCache | None,
+) -> _SourceSnapshot:
     root = source.resolved_root.expanduser().resolve()
     if not root.is_dir():
         raise ValueError(f"canonical source root is not a directory: {root}")
@@ -210,8 +344,7 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
         if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
             immutable_revision = requested.lower()
 
-    files: list[_FileSnapshot] = []
-    for path in sorted(source.weight_files, key=lambda item: str(item)):
+    def snapshot_file(path: Path) -> _FileSnapshot:
         candidate = path.expanduser()
         if not candidate.is_absolute():
             candidate = root / candidate
@@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
             relative_name = candidate.relative_to(root).as_posix()
         except ValueError as exc:
             raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
-        try:
-            current = candidate.stat()
-        except OSError as exc:
-            raise ValueError(f"cannot stat canonical weight file {candidate}") from exc
-        if not stat.S_ISREG(current.st_mode):
-            raise ValueError(f"canonical weight source is not a regular file: {candidate}")
-        symlink_target = os.readlink(candidate) if candidate.is_symlink() else None
+        current = _observe(candidate)
         content_id = _hf_blob_content_id(
             candidate,
             source=source,
             source_root=root,
         )
         if content_id is None:
-            content_id = f"sha256:{_sha256_file(candidate)}"
-        files.append(
-            _FileSnapshot(
-                path=candidate,
-                relative_name=relative_name,
-                size=current.st_size,
-                device=current.st_dev,
-                inode=current.st_ino,
-                mtime_ns=current.st_mtime_ns,
-                ctime_ns=current.st_ctime_ns,
-                symlink_target=symlink_target,
-                content_id=content_id,
+            content_id = (
+                digest_cache.content_id(candidate, current)
+                if digest_cache is not None
+                else f"sha256:{_sha256_file(candidate, current)}"
             )
+        return _FileSnapshot(
+            path=candidate,
+            relative_name=relative_name,
+            size=current.size,
+            device=current.device,
+            inode=current.inode,
+            mtime_ns=current.mtime_ns,
+            ctime_ns=current.ctime_ns,
+            symlink_target=current.symlink_target,
+            content_id=content_id,
         )
 
+    paths = sorted(source.weight_files, key=lambda item: str(item))
+    with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor:
+        files = list(executor.map(snapshot_file, paths))
+
     file_semantics = [file.semantic_dict() for file in files]
     file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
     revision = immutable_revision or f"content-{file_fingerprint}"
@@ -306,11 +439,12 @@ def resolve_final_layout_source_identity(
     *,
     model_id: str,
     target_names: frozenset[str],
+    digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutSourceContext:
     """Resolve exact source identity for files covering the owned tensors."""
     bindings = _resolve_target_sources(prepared_sources, target_names)
     selected_sources = frozenset(bindings.values())
-    snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources}
+    snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources}
     snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
     source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
     target_bindings = [
@@ -345,6 +479,7 @@ def resolve_final_layout_source_identity(
 
 __all__ = [
     "FinalLayoutSourceContext",
+    "NodeSourceDigestCache",
     "PreparedWeightSource",
     "WeightSourceKind",
     "resolve_final_layout_source_identity",

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 Qwen3-VL layer-50 text/vision encoder.
 
 The encoder is reimplemented on top of vLLM-style tensor-parallel building
@@ -42,6 +43,10 @@ from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMetho
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
 
 MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
@@ -1211,12 +1216,18 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         if not self.is_loaded:
             return
         if getattr(self, "_omni_layerwise_enabled", False):
-            for name, child in self.vision.named_children():
-                if name != "blocks":
-                    child.to(self.device_target)
-            for name, child in self.text_model.named_children():
-                if name != "layers":
-                    child.to(self.device_target)
+            stager = getattr(self, "_omni_non_block_stager", None)
+            if stager is None:
+                hooks = getattr(self, "_omni_layerwise_hooks", ())
+                pin_memory = bool(getattr(hooks[0], "pin_memory", True)) if hooks else True
+                stager = PinnedModuleStager(
+                    self._omni_non_block_modules(),
+                    self.device_target,
+                    pin_memory=pin_memory,
+                    cache_retention=getattr(self, "_omni_component_cache", None),
+                )
+                self._omni_non_block_stager = stager
+            stager.load()
             return
         self.vision.to(self.device_target)
         self.text_model.to(self.device_target)
@@ -1227,18 +1238,39 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         if getattr(self, "_omni_layerwise_enabled", False):
             for hook in self._omni_layerwise_hooks:
                 hook.offload_layer()
-            for name, child in self.vision.named_children():
-                if name != "blocks":
+            stager = getattr(self, "_omni_non_block_stager", None)
+            if stager is not None:
+                stager.offload()
+            else:
+                # The initial DLO setup reaches this branch before the first
+                # load. These modules were constructed on CPU, so this is not
+                # a device-to-host rematerialization.
+                for child in self._omni_non_block_modules():
                     child.to("cpu")
-            for name, child in self.text_model.named_children():
-                if name != "layers":
-                    child.to("cpu")
-            torch.accelerator.empty_cache()
+                self._release_omni_component_cache()
             return
         self.vision.to("cpu")
         self.text_model.to("cpu")
         torch.accelerator.empty_cache()
 
+    def _omni_non_block_modules(self) -> list[nn.Module]:
+        modules = [child for name, child in self.vision.named_children() if name != "blocks"]
+        modules.extend(child for name, child in self.text_model.named_children() if name != "layers")
+        return modules
+
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        stager = getattr(self, "_omni_non_block_stager", None)
+        if stager is not None:
+            stager.set_cache_retention(cache)
+
+    def _release_omni_component_cache(self) -> None:
+        cache = getattr(self, "_omni_component_cache", None)
+        if cache is None:
+            torch.accelerator.empty_cache()
+        else:
+            cache.release_if_needed()
+
     def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None:
         """Stream the TP-local Qwen vision/text blocks for low-HBM serving.
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -46,6 +46,7 @@ from vllm_omni.diffusion.models.interface import (
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.offloader import (
+    BoundedAllocatorCache,
     OffloadPlan,
     apply_sequential_offload,
     remove_sequential_offload,
@@ -952,6 +953,12 @@ class MiniMaxH3Pipeline(
         # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
         self.vae = self.video_vae
 
+        self._dlo_component_cache = None
+        if getattr(od_config, "enable_distributed_layerwise_offload", False):
+            self._dlo_component_cache = BoundedAllocatorCache(self.device)
+            for component in (self.text_encoder, self.video_vae, self.audio_vae):
+                component.set_omni_component_cache(self._dlo_component_cache)
+
         self._quality_policy = MiniMaxH3QualityPolicy(od_config)
         self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
 
@@ -1371,11 +1378,8 @@ class MiniMaxH3Pipeline(
         ):
             # Layerwise DiT offload already provides the low-residency encoder
             # phase used by the checkpoint reference.
-            self.text_encoder.load_to_device()
-            try:
+            with self._component_on_device(self.text_encoder):
                 return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
-            finally:
-                self.text_encoder.offload_to_cpu()
 
         # Keep both Qwen and DiT resident across requests. Moving either model
         # here makes encoder latency include a tens-of-gigabytes PCIe transfer,
@@ -1432,14 +1436,36 @@ class MiniMaxH3Pipeline(
             with sequential_offload_component(component):
                 yield
             return
-        staged = self._uses_manual_component_offload()
-        if staged:
-            component.load_to_device()
+        staged = self._uses_manual_component_offload() and getattr(component, "_omni_dlo_offload_enabled", True)
         try:
-            yield
-        finally:
             if staged:
-                component.offload_to_cpu()
+                component.load_to_device()
+            yield
+        except BaseException:
+            if staged:
+                try:
+                    component.offload_to_cpu()
+                except BaseException:
+                    logger.exception("Failed to release %s after component failure", component.__class__.__name__)
+                cache = getattr(self, "_dlo_component_cache", None)
+                if cache is not None:
+                    try:
+                        cache.release_if_needed(force=True)
+                    except BaseException:
+                        logger.exception("Failed to release retained allocator cache after component failure")
+            raise
+        else:
+            if staged:
+                try:
+                    component.offload_to_cpu()
+                except BaseException:
+                    cache = getattr(self, "_dlo_component_cache", None)
+                    if cache is not None:
+                        try:
+                            cache.release_if_needed(force=True)
+                        except BaseException:
+                            logger.exception("Failed to release retained allocator cache after offload failure")
+                    raise
 
     def _encode_visual_conditions(
         self,

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 remote-code VAE adapters and exact latent contracts."""
 
 from __future__ import annotations
@@ -21,7 +22,10 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     DistributedVaeMixin,
 )
 from vllm_omni.diffusion.distributed.parallel_state import get_world_group
-from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
 
 from .packed_tokens import minimax_h3_patchify_video_latent
 
@@ -148,12 +152,21 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         else:
             self.remote.to(self._device_target)
 
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        if self._stager is not None:
+            self._stager.set_cache_retention(cache)
+
     def offload_to_cpu(self) -> None:
         if self._stager is not None:
             self._stager.offload()
         else:
             self.remote.to("cpu")
-            torch.accelerator.empty_cache()
+            cache = getattr(self, "_omni_component_cache", None)
+            if cache is None:
+                torch.accelerator.empty_cache()
+            else:
+                cache.release_if_needed()
 
     def set_parallel_size(
         self,
@@ -411,12 +424,21 @@ class MiniMaxH3AudioVAE(nn.Module):
         else:
             self.remote.to(self._device_target)
 
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        if self._stager is not None:
+            self._stager.set_cache_retention(cache)
+
     def offload_to_cpu(self) -> None:
         if self._stager is not None:
             self._stager.offload()
         else:
             self.remote.to("cpu")
-            torch.accelerator.empty_cache()
+            cache = getattr(self, "_omni_component_cache", None)
+            if cache is None:
+                torch.accelerator.empty_cache()
+            else:
+                cache.release_if_needed()
 
     @torch.inference_mode()
     def encode_waveform(

vllm_omni/diffusion/offloader/__init__.py

@@ -17,7 +17,7 @@ from .distributed_layerwise_backend import (
     remove_distributed_block_hook,
 )
 from .layerwise_backend import LayerWiseOffloadBackend
-from .module_residency import PinnedModuleStager
+from .module_residency import BoundedAllocatorCache, PinnedModuleStager
 from .offload_plan import OffloadPlan, get_offload_plan
 from .sequential_backend import (
     ModelLevelOffloadBackend,
@@ -46,6 +46,7 @@ __all__ = [
     "DistributedLayerwiseOffloadBackend",
     "DistributedLayerwiseOffloadHook",
     "ModelLevelOffloadBackend",
+    "BoundedAllocatorCache",
     "PinnedModuleStager",
     "apply_sequential_offload",
     "remove_sequential_offload",

vllm_omni/diffusion/offloader/base.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol, runtime_checkable
 
@@ -10,7 +11,7 @@ import torch
 from torch import nn
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options
 
 logger = init_logger(__name__)
 
@@ -52,6 +53,12 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Per auxiliary component: True permits model-declared DLO offload;
+    # False keeps the whole component resident. "default" is the fallback.
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
+    # Optional per-worker ceiling for registering an HWR mmap. Zero means no
+    # additional ceiling; pin_cpu_memory controls whether registration is tried.
+    dlo_host_registration_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -122,6 +129,31 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        raw_component_policy = getattr(od_config, "dlo_offload_components", {})
+        if raw_component_policy is None:
+            raw_component_policy = {}
+        if not isinstance(raw_component_policy, Mapping):
+            raise TypeError(
+                "dlo_offload_components must be a mapping of component names to booleans, "
+                f"got {type(raw_component_policy).__name__}"
+            )
+        dlo_offload_components: dict[str, bool] = {}
+        for component, enabled in raw_component_policy.items():
+            if not isinstance(component, str) or not component:
+                raise TypeError("dlo_offload_components keys must be non-empty strings")
+            if not isinstance(enabled, bool):
+                raise TypeError(
+                    f"dlo_offload_components[{component!r}] must be a boolean, got {type(enabled).__name__}"
+                )
+            dlo_offload_components[component] = enabled
+        if dlo_offload_components and not enable_distributed_layerwise_offload:
+            raise ValueError("dlo_offload_components requires distributed layerwise offload")
+        dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
+            enable_dlo=enable_distributed_layerwise_offload,
+            use_allgather=dlo_use_allgather,
+            hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"),
+        )
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -157,6 +189,8 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_offload_components=dlo_offload_components,
+            dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA implementation of read-only host-mapping registration."""
+
+from __future__ import annotations
+
+import ctypes
+import mmap
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+from .host_registration import (
+    HostRegistrationBudgetError,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+)
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
+
+
+class _CudaRuntime(Protocol):
+    def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ...
+
+    def cudaHostUnregister(self, address: int) -> int: ...
+
+    def cudaGetErrorString(self, error: int) -> str | bytes: ...
+
+    def cudaGetLastError(self) -> int: ...
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _mapped_regions(regions: Sequence[MappedHostRegion]) -> tuple[_AddressRange, ...]:
+    """Resolve page spans without merging unrelated artifact files."""
+    by_file: dict[str, list[tuple[int, int]]] = defaultdict(list)
+    for region in regions:
+        if not region.read_only:
+            raise HostRegistrationError(
+                f"CUDA direct H2D requires an immutable mapping, but {region.file_name!r} is writable"
+            )
+        by_file[region.file_name].append((region.address, region.length))
+    return tuple(span for ranges in by_file.values() for span in _coalesce_ranges(ranges))
+
+
+def _error_message(runtime: _CudaRuntime, error: int) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+def _consume_last_cuda_error(runtime: _CudaRuntime, expected_error: int) -> None:
+    """Clear one handled CUDA Runtime error before returning to PyTorch."""
+    try:
+        get_last_error = getattr(runtime, "cudaGetLastError", None)
+        if get_last_error is None:
+            get_last_error = ctypes.CDLL(None).cudaGetLastError
+            get_last_error.argtypes = []
+            get_last_error.restype = ctypes.c_int
+        pending_error = int(get_last_error())
+    except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc:
+        raise HostRegistrationError("cannot clear CUDA's pending error after host-registration failure") from exc
+
+    if pending_error not in (0, expected_error):
+        raise HostRegistrationError(
+            f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
+        )
+
+
+def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
+    error_code = int(error)
+    message = _error_message(runtime, error_code)
+    try:
+        _consume_last_cuda_error(runtime, error_code)
+    except HostRegistrationError as exc:
+        raise HostRegistrationError(f"{message}; {exc}") from exc
+    return message
+
+
+def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
+    """Query support required by immutable artifact mappings."""
+    try:
+        get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None)
+        if get_attribute is None:
+            get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute
+            get_attribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
+            get_attribute.restype = ctypes.c_int
+    except (AttributeError, OSError) as exc:
+        raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc
+
+    supported = ctypes.c_int()
+    try:
+        error = get_attribute(
+            ctypes.byref(supported),
+            _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED,
+            torch.accelerator.current_device_index(),
+        )
+    except Exception as exc:
+        raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc
+    if int(error) != 0:
+        raise HostRegistrationError(
+            f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}"
+        )
+    return bool(supported.value)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for immutable file-backed CPU regions."""
+
+    def __init__(self, runtime: _CudaRuntime, regions: tuple[_AddressRange, ...]) -> None:
+        self._runtime = runtime
+        self._regions = regions
+
+    @classmethod
+    def create(
+        cls,
+        regions: Sequence[MappedHostRegion],
+        *,
+        max_bytes: int | None,
+    ) -> CudaHostRegistration:
+        mapped = _mapped_regions(regions)
+        total_bytes = sum(region.size for region in mapped)
+        if max_bytes is not None and total_bytes > max_bytes:
+            raise HostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not mapped:
+            raise HostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise HostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+        if not _supports_read_only_host_registration(runtime):
+            raise HostRegistrationError(
+                "CUDA device does not support read-only host registration required by immutable HWR mappings"
+            )
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in mapped:
+                error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
+                if int(error) != 0:
+                    raise HostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            rollback_failed: list[_AddressRange] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
+                        )
+                        rollback_failed.append(region)
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+                    rollback_failed.append(region)
+            if rollback_errors:
+                active_registration = cls(runtime, tuple(reversed(rollback_failed)))
+                raise HostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}",
+                    active_registration=active_registration,
+                ) from exc
+            if isinstance(exc, HostRegistrationError):
+                raise
+            raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, mapped)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> tuple[str, ...]:
+        """Unregister every range and retain failures for a later retry."""
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        return tuple(errors)

vllm_omni/diffusion/offloader/distributed_layerwise_backend.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
 """Distributed Layerwise Offload backend with double-buffered H2D.
 
 This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
@@ -18,6 +18,8 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import threading
+import time
 from itertools import chain
 from typing import Any
 
@@ -35,7 +37,13 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
-from .module_collector import ModuleDiscovery
+from .host_registration import (
+    HostRegistration,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+    register_host_mappings,
+)
+from .module_collector import ModuleDiscovery, PipelineModules
 from .offload_plan import (
     OffloadPlan,
     get_offload_plan,
@@ -51,6 +59,38 @@ from .tensor_utils import (
 
 logger = init_logger(__name__)
 
+# A backend normally owns both objects below. This process-lifetime safety
+# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup
+# failure unwinds startup and the backend itself becomes unreachable. A clean
+# retry removes the pair before closing the lease.
+_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = []
+_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock()
+
+
+def _retain_active_hwr_registration(
+    registration: HostRegistration,
+    lease: HostWeightLease,
+) -> None:
+    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
+        if not any(
+            candidate is registration and candidate_lease is lease
+            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
+        ):
+            _ACTIVE_HWR_REGISTRATIONS.append((registration, lease))
+
+
+def _forget_active_hwr_registration(
+    registration: HostRegistration,
+    lease: HostWeightLease,
+) -> None:
+    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
+        _ACTIVE_HWR_REGISTRATIONS[:] = [
+            (candidate, candidate_lease)
+            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
+            if candidate is not registration or candidate_lease is not lease
+        ]
+
+
 # Threshold (in MB) for deciding whether a non-block DiT submodule should
 # use layerwise offload (streaming hooks) or be moved to GPU as a resident
 # module.  Submodules larger than this are offloaded to save HBM; smaller
@@ -94,6 +134,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -445,6 +486,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -457,18 +518,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -500,6 +550,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
         return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     @torch.compiler.disable
     def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
         """Prepare next block's weights into the shared device buffer for *slot*.
@@ -515,17 +591,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-            with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
-                    gw = gpu_weights[dtype]
-                    async_copy = non_blocking and cpu_shard.is_pinned()
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                evt.record(self.copy_stream)
-            if self.rank_local_mmap:
-                # The CPU slot may be overwritten only after this H2D copy has
-                # finished.  The shared event protects reuse by another hook.
-                self.cpu_staging_events[slot] = evt
+            if self.rank_local_mmap and self.registered_mmap:
+                with current_omni_platform.stream(self.copy_stream):
+                    self._copy_mmap_sources_to_device(
+                        self.cpu_sources,
+                        self.metadata,
+                        gpu_weights,
+                        non_blocking=non_blocking,
+                    )
+                    evt.record(self.copy_stream)
+            else:
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -735,6 +821,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -805,6 +892,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -902,8 +997,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._host_weight_lease: HostWeightLease | None = None
+        self._host_registration: HostRegistration | None = None
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
 
     def load_resident_layers(self) -> None:
@@ -916,6 +1013,136 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if self._resident_layer_group is not None:
             self._resident_layer_group.offload()
 
+    def _rank_local_source_tensors(
+        self,
+        hooks: list[DistributedLayerwiseOffloadHook],
+    ) -> tuple[torch.Tensor, ...]:
+        """Return the exact CPU sources that direct H2D would consume."""
+        tensors: list[torch.Tensor] = []
+        seen: set[int] = set()
+
+        def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None:
+            for sources in cpu_sources.values():
+                for source in sources:
+                    tensor = source["tensor"]
+                    if id(tensor) not in seen:
+                        seen.add(id(tensor))
+                        tensors.append(tensor)
+
+        for hook in hooks:
+            collect(hook.cpu_sources)
+        if self._resident_layer_group is not None:
+            for state in self._resident_layer_group._states:
+                collect(state["cpu_sources"])
+        return tuple(tensors)
+
+    def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool:
+        """Register the complete final-layout lease under pinned-memory policy."""
+        lease = self._host_weight_lease
+        if lease is None:
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
+            return False
+        if not lease.mapped_regions or not source_tensors:
+            logger.warning("HWR mmap registration found no mapped sources; using bounded host staging")
+            return False
+
+        limit_gib = self.config.dlo_host_registration_limit_gib
+        max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
+        started = time.perf_counter()
+        try:
+            registration = register_host_mappings(
+                lease.mapped_regions,
+                device=self.device,
+                max_bytes=max_bytes,
+            )
+            try:
+                unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()]
+            except Exception as exc:
+                errors = registration.close()
+                if errors:
+                    self._host_registration = registration
+                    _retain_active_hwr_registration(registration, lease)
+                    raise HostRegistrationCleanupError(
+                        "CUDA registration succeeded but pinned-source verification failed, "
+                        f"and rollback failed: {errors[:3]}"
+                    ) from exc
+                raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc
+            if unpinned:
+                errors = registration.close()
+                if errors:
+                    self._host_registration = registration
+                    _retain_active_hwr_registration(registration, lease)
+                    raise HostRegistrationCleanupError(
+                        "CUDA registration succeeded but PyTorch rejected mapped sources, "
+                        f"and rollback failed: {errors[:3]}"
+                    )
+                raise HostRegistrationError(
+                    "CUDA registration succeeded but PyTorch did not recognize "
+                    f"{len(unpinned)} mapped source(s) as pinned"
+                )
+        except HostRegistrationCleanupError as exc:
+            # Falling back could close a lease while the platform still owns
+            # one of its mappings. Fail startup and retain ownership for retry.
+            active_registration = exc.active_registration
+            if active_registration is not None:
+                self._host_registration = active_registration
+                _retain_active_hwr_registration(active_registration, lease)
+            logger.exception("HWR mmap registration rollback failed")
+            raise
+        except HostRegistrationError as exc:
+            logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc)
+            return False
+
+        self._host_registration = registration
+        _retain_active_hwr_registration(registration, lease)
+        logger.info(
+            "Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None:
+        """Select registered direct H2D or bounded staging once per backend."""
+        plan = self.host_weight_plan
+        if (
+            (not hooks and self._resident_layer_group is None)
+            or not self._using_rank_local_mmap
+            or plan is None
+            or plan.backing_kind != "host_weight_runtime"
+        ):
+            return
+
+        source_tensors = self._rank_local_source_tensors(hooks)
+        self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors)
+        for hook in hooks:
+            hook.registered_mmap = self._using_registered_mmap
+        if self._resident_layer_group is not None:
+            self._resident_layer_group.registered_mmap = self._using_registered_mmap
+            if self._using_registered_mmap:
+                self._resident_layer_group._cpu_staging_buffers.clear()
+
+    def _release_registered_mmap(self) -> None:
+        """Release every platform registration before closing the HWR lease."""
+        registration = self._host_registration
+        if registration is None:
+            return
+        errors = registration.close()
+        if errors:
+            lease = self._host_weight_lease
+            if lease is not None and not lease.closed:
+                _retain_active_hwr_registration(registration, lease)
+            logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3])
+            raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)")
+        lease = self._host_weight_lease
+        if lease is not None:
+            _forget_active_hwr_registration(registration, lease)
+        self._host_registration = None
+        logger.info("Unregistered HWR mmap ranges")
+
     def _load_weights_via_mmap(
         self,
         pipeline: nn.Module,
@@ -1173,6 +1400,39 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         module.to(self.device)
         logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)
 
+    def _component_offload_enabled(self, name: str) -> bool:
+        policy = self.config.dlo_offload_components
+        return policy.get(name, policy.get("default", True))
+
+    def _validate_component_offload_policy(self, modules: PipelineModules) -> None:
+        known_components = {"default", *modules.encoder_names, *modules.vae_names}
+        unknown = sorted(self.config.dlo_offload_components.keys() - known_components)
+        if unknown:
+            raise ValueError(
+                f"dlo_offload_components contains unknown auxiliary components: {unknown}. DLO always owns DiT offload."
+            )
+
+    def _prepare_auxiliary_components(self, modules: PipelineModules, plan: OffloadPlan | None) -> None:
+        for encoder, name in zip(modules.encoders, modules.encoder_names, strict=True):
+            offload = self._component_offload_enabled(name)
+            encoder._omni_dlo_offload_enabled = offload
+            if offload:
+                self._try_layerwise_offload_encoder(encoder, name, plan)
+            self._register_on_demand_hook(
+                encoder,
+                name,
+                stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
+            )
+
+        for vae, name in zip(modules.vaes, modules.vae_names, strict=True):
+            offload = self._component_offload_enabled(name)
+            vae._omni_dlo_offload_enabled = offload
+            self._register_on_demand_hook(
+                vae,
+                name,
+                stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
+            )
+
     def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
         """Stream plan-declared encoder blocks on each rank without AllGather."""
         if plan is None or name not in plan.encoder_block_attrs:
@@ -1441,6 +1701,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Retrieve optional declarative OffloadPlan from the pipeline.
         # When present, replaces heuristic block discovery.
         plan = get_offload_plan(pipeline)
+        self._validate_component_offload_policy(modules)
 
         if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
             logger.warning(
@@ -1465,11 +1726,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     raise RuntimeError("DLO received a closed Host Weight Runtime lease")
                 # The final-layout restorer has already rebound the model to
                 # immutable host tensors.  Treat those tensors as mmap-like
-                # sources for the bounded two-slot rank-local staging path;
-                # do not reinitialize or rematerialize the restored model.
+                # sources; transport setup below selects registered direct
+                # H2D or the bounded two-slot staging fallback.
                 self._using_rank_local_mmap = True
                 logger.info(
-                    "DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging",
+                    "DLO consuming final-layout Host Weight Runtime lease %s",
                     self._host_weight_lease.provenance.resolution_id,
                 )
             elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
@@ -1483,7 +1744,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             if self._using_rank_local_mmap:
                 logger.info(
                     "DLO rank-local host storage enabled: source pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "node-shared; transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1501,21 +1762,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 )
             logger.info("DLO is using host tensors materialized by the ordinary loader")
 
-        # Keep VAE/encoders on CPU; move to GPU on-demand via hooks.
-        # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
-        # during the DiT forward pass.  They are only needed briefly for
-        # text-encoding (before DiT) and VAE-decode (after DiT).
-        for enc, enc_name in zip(modules.encoders, modules.encoder_names):
-            self._try_layerwise_offload_encoder(enc, enc_name, plan)
-            self._register_on_demand_hook(
-                enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths
-            )
-        for vae, vae_name in zip(modules.vaes, modules.vae_names):
-            self._register_on_demand_hook(
-                vae,
-                "vae",
-                stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths),
-            )
+        # Apply the user policy to discovered auxiliary components. The default
+        # keeps the model-declared staged lifecycle; False leaves a component
+        # resident for repeated requests.
+        self._prepare_auxiliary_components(modules, plan)
 
         # Move resident modules to GPU (small modules needed every forward)
         for name, module in zip(modules.resident_names, modules.resident_modules):
@@ -1653,6 +1903,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             )
             pipeline._dlo_residency_controller = self
 
+        all_hooks = [hook for group in self._all_hook_groups for hook in group]
+        self._configure_hwr_transfer(all_hooks)
+
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
             if self._using_mmap and not self.enabled:
@@ -1662,17 +1915,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers
         # sized to the max block across ALL module groups (gen_layers +
         # language_model).  Groups execute sequentially, so 2 buffers suffice.
-        all_hooks: list[DistributedLayerwiseOffloadHook] = []
-        for group in self._all_hook_groups:
-            all_hooks.extend(group)
-
         unified_buffers = self._allocate_shared_buffers(all_hooks)
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1753,6 +2002,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def _release_mmap_handles(self) -> None:
         """Release source handles and the transport-owned HWR lease."""
+        if self._host_registration is not None:
+            raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active")
         self._mmap_transforms_by_tensor_id.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
@@ -1781,6 +2032,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def disable(self) -> None:
         has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
+        has_registration = self._host_registration is not None
         has_carrier = (
             self.host_weight_plan is not None
             and self.host_weight_plan.lease_carrier is not None
@@ -1796,12 +2048,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             not self.enabled
             and not hasattr(self, "_mmap_file_cache")
             and not has_open_lease
+            and not has_registration
             and not has_carrier
             and not has_partial_hooks
         ):
             return
 
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap or has_registration:
             current_omni_platform.synchronize()
 
         for blocks in self._blocks:
@@ -1818,9 +2071,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
+        self._release_registered_mmap()
         self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/offloader/host_registration.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Platform-neutral lifecycle for registering existing host mappings."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol
+
+import torch
+
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+
+class HostRegistrationError(RuntimeError):
+    """Raised when existing host ranges cannot be registered safely."""
+
+
+class HostRegistrationBudgetError(HostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class HostRegistrationCleanupError(HostRegistrationError):
+    """Raised when registered ranges cannot be released safely."""
+
+    def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None:
+        super().__init__(message)
+        self.active_registration = active_registration
+
+
+class HostRegistration(Protocol):
+    """Transport-owned registration for one process-local host mapping."""
+
+    @property
+    def total_bytes(self) -> int: ...
+
+    @property
+    def region_count(self) -> int: ...
+
+    def close(self) -> tuple[str, ...]: ...
+
+
+def register_host_mappings(
+    regions: Sequence[MappedHostRegion],
+    *,
+    device: torch.device,
+    max_bytes: int | None,
+) -> HostRegistration:
+    """Register immutable ranges with the active platform implementation."""
+    if device.type == "cuda":
+        from .cuda_host_registration import CudaHostRegistration
+
+        return CudaHostRegistration.create(regions, max_bytes=max_bytes)
+    raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}")
+
+
+__all__ = [
+    "HostRegistration",
+    "HostRegistrationBudgetError",
+    "HostRegistrationCleanupError",
+    "HostRegistrationError",
+    "register_host_mappings",
+]

vllm_omni/diffusion/offloader/module_residency.py

@@ -1,89 +1,271 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""On-demand module staging backed by an immutable pinned CPU snapshot."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""On-demand module staging backed by immutable pinned CPU storage."""
 
 from __future__ import annotations
 
+from collections.abc import Iterable
+from dataclasses import dataclass
 from itertools import chain
 from typing import Any
 
 import torch
 from torch import nn
+from vllm.logger import init_logger
 
 from vllm_omni.platforms import current_omni_platform
 
-from .tensor_utils import set_tensor_storage
+from .tensor_utils import is_dtensor, set_tensor_storage
+
+logger = init_logger(__name__)
 
 
-class PinnedModuleStager:
-    """Stage an immutable module without copying device weights back to CPU.
+class BoundedAllocatorCache:
+    """Retain reusable allocator blocks without monopolizing device memory.
 
-    ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter
-    after each forward. In inference the weights are immutable, so retain one
-    pinned CPU master instead. ``load`` materializes fresh device tensors from
-    that master; ``offload`` only rebinds the Parameters/Buffers to it and
-    releases the device allocations.
+    Component offload normally calls ``empty_cache`` after every stage. That
+    makes the next stage return to the device allocator even though PyTorch's
+    cached blocks are immediately reusable. This policy keeps the cache while
+    both of these bounds hold:
+
+    * cached-but-unallocated memory is at most 25% of device capacity; and
+    * at least 5% of device capacity is physically free.
+
+    Missing memory telemetry is handled conservatively by releasing the cache.
+    Failure paths can force release; normal executor shutdown keeps its own
+    unconditional device-cache cleanup because this policy is not global.
     """
 
     def __init__(
         self,
-        module: nn.Module,
+        device: torch.device,
+        *,
+        max_cached_fraction: float = 0.25,
+        min_free_fraction: float = 0.05,
+    ) -> None:
+        if not 0.0 <= max_cached_fraction <= 1.0:
+            raise ValueError(f"max_cached_fraction must be in [0, 1], got {max_cached_fraction}")
+        if not 0.0 <= min_free_fraction <= 1.0:
+            raise ValueError(f"min_free_fraction must be in [0, 1], got {min_free_fraction}")
+        self.device = device
+        self.max_cached_fraction = max_cached_fraction
+        self.min_free_fraction = min_free_fraction
+
+    def _should_release(self) -> bool:
+        reserved = int(torch.accelerator.memory_reserved(self.device))
+        allocated = int(torch.accelerator.memory_allocated(self.device))
+        free, total = current_omni_platform.get_device_memory(self.device)
+        cached = max(0, reserved - allocated)
+        return cached > int(total * self.max_cached_fraction) or free < int(total * self.min_free_fraction)
+
+    def release_if_needed(self, *, force: bool = False) -> bool:
+        """Release cached blocks when a bound is crossed or release is forced."""
+        if not force:
+            try:
+                if not self._should_release():
+                    return False
+            except Exception as exc:
+                # Preserve the pre-retention behavior on platforms that do not
+                # expose allocator telemetry through torch.accelerator.
+                logger.debug("Allocator cache telemetry unavailable; releasing cache: %s", exc)
+        current_omni_platform.empty_cache()
+        return True
+
+
+@dataclass(frozen=True)
+class _TensorBinding:
+    target: torch.Tensor
+    dtype: torch.dtype
+    shape: tuple[int, ...]
+    stride: tuple[int, ...]
+    storage_offset: int
+
+
+@dataclass
+class _StorageGroup:
+    master: torch.Tensor
+    bindings: list[_TensorBinding]
+
+
+class PinnedModuleStager:
+    """Stage immutable module groups without copying device weights to CPU.
+
+    ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter
+    after each forward. In inference the weights are immutable, so retain one
+    pinned CPU master instead. ``load`` materializes device storage from that
+    master; ``offload`` only rebinds Parameters and buffers to the master.
+
+    A module iterable is treated as one staging group. It uses one copy stream
+    and one reusable completion event. Tensors sharing storage keep their
+    shapes, strides, offsets, dtypes, and aliases across every transition.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module | Iterable[nn.Module],
         device: torch.device,
         *,
         pin_memory: bool = True,
         copy_stream: Any | None = None,
+        cache_retention: BoundedAllocatorCache | None = None,
     ) -> None:
-        self.device = device
-        self.copy_stream = copy_stream or current_omni_platform.Stream()
-        self.loaded = False
-        self._entries: list[tuple[torch.Tensor, torch.Tensor]] = []
-        self._device_tensors: list[torch.Tensor] = []
+        modules = (module,) if isinstance(module, nn.Module) else tuple(module)
+        if not modules or not all(isinstance(item, nn.Module) for item in modules):
+            raise ValueError("PinnedModuleStager requires at least one nn.Module")
 
-        # parameters()/buffers() remove duplicate objects by default, which
-        # preserves tied weights and buffer aliases when their storage moves.
-        for target in chain(module.parameters(), module.buffers()):
-            master = target.detach()
-            if master.device.type != "cpu":
-                master = master.to("cpu")
+        self.device = device
+        self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream()
+        self._ready_event = current_omni_platform.Event()
+        self.cache_retention = cache_retention
+        self.loaded = False
+        self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
+        self._device_storages: list[torch.Tensor] = []
+        self._restore_masters()
+
+    @staticmethod
+    def _local_tensor(target: torch.Tensor) -> torch.Tensor:
+        return target.to_local() if is_dtensor(target) else target
+
+    @classmethod
+    def _snapshot_groups(
+        cls,
+        modules: tuple[nn.Module, ...],
+        *,
+        pin_memory: bool,
+    ) -> list[_StorageGroup]:
+        targets: list[torch.Tensor] = []
+        seen_targets: set[int] = set()
+        for target in chain.from_iterable(chain(item.parameters(), item.buffers()) for item in modules):
+            if id(target) not in seen_targets:
+                seen_targets.add(id(target))
+                targets.append(target)
+
+        grouped: dict[tuple[Any, ...], tuple[torch.Tensor, list[_TensorBinding]]] = {}
+        for target in targets:
+            local = cls._local_tensor(target)
+            if local.is_meta:
+                raise ValueError("PinnedModuleStager cannot snapshot a meta tensor")
+            storage = local.untyped_storage()
+            if storage.nbytes() == 0:
+                storage_key: tuple[Any, ...] = ("empty", id(target))
+            else:
+                storage_key = (
+                    local.device.type,
+                    local.device.index,
+                    storage.data_ptr(),
+                    storage.nbytes(),
+                )
+            binding = _TensorBinding(
+                target=target,
+                dtype=local.dtype,
+                shape=tuple(local.shape),
+                stride=tuple(local.stride()),
+                storage_offset=local.storage_offset(),
+            )
+            if storage_key not in grouped:
+                grouped[storage_key] = (local, [])
+            grouped[storage_key][1].append(binding)
+
+        groups: list[_StorageGroup] = []
+        for source, bindings in grouped.values():
+            storage = source.untyped_storage()
+            storage_view = torch.empty(0, dtype=torch.uint8, device=source.device).set_(
+                storage,
+                0,
+                (storage.nbytes(),),
+                (1,),
+            )
+            master = storage_view.detach() if storage_view.device.type == "cpu" else storage_view.to("cpu")
             if pin_memory and not master.is_pinned():
                 master = master.pin_memory()
-            set_tensor_storage(target, master)
-            self._entries.append((target, master))
+            groups.append(_StorageGroup(master=master, bindings=bindings))
+        return groups
+
+    @staticmethod
+    def _view(backing: torch.Tensor, binding: _TensorBinding) -> torch.Tensor:
+        return torch.empty(0, dtype=binding.dtype, device=backing.device).set_(
+            backing.untyped_storage(),
+            binding.storage_offset,
+            binding.shape,
+            binding.stride,
+        )
+
+    def _bind(self, storages: list[torch.Tensor]) -> None:
+        for storage, group in zip(storages, self._groups):
+            for binding in group.bindings:
+                set_tensor_storage(binding.target, self._view(storage, binding))
+
+    def _restore_masters(self) -> None:
+        self._bind([group.master for group in self._groups])
+
+    def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None:
+        self.cache_retention = cache_retention
+
+    def _release_cache(self, *, force: bool = False) -> None:
+        if self.cache_retention is None:
+            current_omni_platform.empty_cache()
+        else:
+            self.cache_retention.release_if_needed(force=force)
+
+    def _load_once(self) -> None:
+        device_storages = [torch.empty_like(group.master, device=self.device) for group in self._groups]
+        compute_stream = current_omni_platform.current_stream()
+        self.copy_stream.wait_stream(compute_stream)
+        with current_omni_platform.stream(self.copy_stream):
+            for device_storage, group in zip(device_storages, self._groups):
+                device_storage.copy_(group.master, non_blocking=group.master.is_pinned())
+            self._ready_event.record(self.copy_stream)
+
+        self._bind(device_storages)
+        compute_stream.wait_event(self._ready_event)
+        self._device_storages = device_storages
+        self.loaded = True
+
+    def _cleanup_failed_load(self) -> None:
+        try:
+            current_omni_platform.synchronize()
+        except Exception:
+            logger.debug("Device synchronization failed while cleaning up module staging", exc_info=True)
+        self._restore_masters()
+        self._device_storages.clear()
+        self.loaded = False
+        self._release_cache(force=True)
 
     def load(self) -> None:
         if self.loaded:
             return
-
-        # Allocate on the compute stream so the caching allocator can reuse
-        # memory released by the preceding encoder/DiT stage.
-        device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries]
-        compute_stream = current_omni_platform.current_stream()
-        self.copy_stream.wait_stream(compute_stream)
-        ready = current_omni_platform.Event()
-        with current_omni_platform.stream(self.copy_stream):
-            for device_tensor, (_, master) in zip(device_tensors, self._entries):
-                device_tensor.copy_(master, non_blocking=master.is_pinned())
-            ready.record(self.copy_stream)
-
-        for device_tensor, (target, _) in zip(device_tensors, self._entries):
-            set_tensor_storage(target, device_tensor)
-        compute_stream.wait_event(ready)
-        self._device_tensors = device_tensors
-        self.loaded = True
+        try:
+            self._load_once()
+        except torch.OutOfMemoryError:
+            # A retained cache is normally reusable by this process, but an
+            # explicit flush gives external memory pressure one bounded retry.
+            self._cleanup_failed_load()
+            try:
+                self._load_once()
+            except BaseException:
+                self._cleanup_failed_load()
+                raise
+        except BaseException:
+            self._cleanup_failed_load()
+            raise
 
     def offload(self) -> None:
         if not self.loaded:
             return
 
         # The module has completed on the compute stream. Synchronize once at
-        # the stage boundary, then release HBM by rebinding to the unchanged
-        # host masters. No device-to-host transfer is necessary.
-        current_omni_platform.synchronize()
-        for target, master in self._entries:
-            set_tensor_storage(target, master)
-        self._device_tensors.clear()
+        # the stage boundary, then discard device storage without any D2H copy.
+        try:
+            current_omni_platform.synchronize()
+            self._restore_masters()
+        except BaseException:
+            self._device_storages.clear()
+            self.loaded = False
+            self._release_cache(force=True)
+            raise
+        self._device_storages.clear()
         self.loaded = False
-        current_omni_platform.empty_cache()
+        self._release_cache()
 
 
-__all__ = ["PinnedModuleStager"]
+__all__ = ["BoundedAllocatorCache", "PinnedModuleStager"]

vllm_omni/diffusion/worker/diffusion_worker.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
 
 """
 Diffusion Worker for vLLM-Omni.
@@ -897,11 +897,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-            if mgr is not None:
-                mgr.shutdown_prefetch()
-        destroy_distributed_env()
+        try:
+            if self.model_runner is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                try:
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -569,8 +569,10 @@ class OrchestratorArgs:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1059,8 +1059,12 @@ class AsyncOmniEngine:
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
+            "dlo_offload_components": (
+                {} if kwargs.get("dlo_offload_components") is None else kwargs["dlo_offload_components"]
+            ),
             "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
             "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
+            "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
                 "regional"

vllm_omni/engine/orchestrator.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Orchestrator for vLLM-Omni multi-stage runtime.
 
@@ -363,8 +366,8 @@ class Orchestrator:
     def __init__(
         self,
         request_async_queue: janus.AsyncQueue[EngineQueueMessage],
-        output_async_queue: janus.AsyncQueue[dict[str, Any]],
-        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
+        output_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
         stage_pools: list[StagePool],
         *,
         async_chunk: bool = False,
@@ -1173,6 +1176,35 @@ class Orchestrator:
             abort=True,
         )
 
+    async def _fail_request_client_error(
+        self,
+        req_id: str,
+        stage_id: int,
+        error: str,
+        *,
+        close_duplex_sessions: bool = False,
+    ) -> None:
+        """Fail one request with a non-fatal 400 (bad input, engine survives).
+
+        The non-fatal counterpart of `_fail_request_dead_stage`: emits a
+        client-error ErrorMessage (default `fatal=False`) so the engine keeps
+        serving, then releases the request's state across every stage pool.
+        """
+        await self.output_async_queue.put(
+            ErrorMessage(
+                error=error,
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                error_type=DEFAULT_CLIENT_ERROR_TYPE,
+                request_id=req_id,
+                stage_id=stage_id,
+            )
+        )
+        await self._cleanup_request_ids(
+            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
+            abort=True,
+            close_duplex_sessions=close_duplex_sessions,
+        )
+
     async def _dispatch_or_fail_request(
         self,
         dispatch: Callable[[], Awaitable[Any]],
@@ -1194,7 +1226,9 @@ class Orchestrator:
         StageUnavailableError (no live replica / evicted slot) or
         EngineDeadError (replica died mid-submit before the poll loop evicted
         it). Both mean "this request cannot be placed", not "the server is
-        broken": fail the request and keep serving (#4285). Unrelated errors
+        broken": fail the request and keep serving (#4285). An OverflowError
+        from encoding an out-of-range request value is likewise failed as a
+        client error (400) rather than killing the loop. Other unrelated errors
         propagate. Returns False when the request was failed.
 
         ``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
@@ -1237,6 +1271,24 @@ class Orchestrator:
             if dead_replica is not None and dead_replica in pool.live_replica_ids():
                 await self._handle_dead_replica(stage_id, dead_replica, e)
             return False
+        except OverflowError as e:
+            # Catch overflow errors in the request payload to ensure that msgpack
+            # encoding can't kill the engine; instead, we just fail this request.
+            # The frontend should generally validate this as well, but this is needed
+            # in case we get values like seed > 2**64-1 from sampling params.
+            logger.warning(
+                "[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s",
+                operation,
+                req_id,
+                stage_id,
+                e,
+            )
+            await self._fail_request_client_error(
+                req_id,
+                stage_id,
+                f"Request contains a value that cannot be serialized: {e}",
+            )
+            return False
 
     # ---- Shared helpers ----
 
@@ -2292,21 +2344,11 @@ class Orchestrator:
                 "and none were provided; failing the request",
                 request_id,
             )
-            await self.output_async_queue.put(
-                ErrorMessage(
-                    error=(
-                        "async_chunk requires prompt_token_ids on the stage-0 request; "
-                        "an embeds-only prompt cannot prewarm downstream stages"
-                    ),
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    error_type=DEFAULT_CLIENT_ERROR_TYPE,
-                    request_id=request_id,
-                    stage_id=0,
-                )
-            )
-            await self._cleanup_request_ids(
-                [request_id, *self._cfg_tracker.cleanup_parent(request_id)],
-                abort=True,
+            await self._fail_request_client_error(
+                request_id,
+                0,
+                "async_chunk requires prompt_token_ids on the stage-0 request; "
+                "an embeds-only prompt cannot prewarm downstream stages",
                 close_duplex_sessions=True,
             )
             return False

vllm_omni/engine/stage_engine_startup.py

@@ -1558,7 +1558,6 @@ def launch_diffusion_stage_replica(
     from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 
     od_config = build_diffusion_config(model, stage_config, metadata)
-    od_config.max_num_seqs = batch_size
     parallel_config = getattr(od_config, "parallel_config", None)
     world_size = getattr(parallel_config, "world_size", 1)
     try:

vllm_omni/engine/stage_init_utils.py

@@ -1506,15 +1506,14 @@ def initialize_diffusion_stage(
         stage_cfg: Stage configuration.
         metadata: Extracted stage metadata.
         stage_init_timeout: Timeout in seconds for stage initialization handshake
-        batch_size: Maximum number of requests to batch together in the
-            diffusion engine.  Passed through to ``StageDiffusionClient``
-            and ultimately to ``AsyncOmni``.
+        batch_size: Client-side request batch width. Does not set scheduler
+            ``max_num_seqs``; pass ``--max-num-seqs`` or stage YAML for that.
+            Forwarded to ``StageDiffusionClient``.
         use_inline: If True, uses the inline diffusion client instead of subprocess.
     """
     from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
     od_config = build_diffusion_config(model, stage_cfg, metadata)
-    od_config.max_num_seqs = batch_size
     return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
 
 

vllm_omni/entrypoints/cli/benchmark/cli_args.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """vLLM-Omni extensions for the ``vllm bench serve`` CLI.
 
 Core functions:
@@ -16,6 +19,63 @@ by ``add_omni_args``.
 """
 
 import argparse
+import math
+from pathlib import Path
+
+_DEFAULT_OMNIINTERACT_NUM_PROMPTS = 3
+
+
+def _positive_finite_float(value: str) -> float:
+    parsed = float(value)
+    if not math.isfinite(parsed) or parsed <= 0:
+        raise argparse.ArgumentTypeError(f"must be a finite positive number, got {value!r}")
+    return parsed
+
+
+def _existing_file(value: str) -> str:
+    path = Path(value).expanduser()
+    if not path.is_file():
+        raise argparse.ArgumentTypeError(f"file does not exist: {value!r}")
+    return str(path)
+
+
+def add_omniinteract_cli_args(parser: argparse.ArgumentParser) -> None:
+    from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+        DEFAULT_MAX_VIDEO_DURATION_S,
+        OMNIINTERACT_SUBSETS,
+    )
+
+    group = parser.add_argument_group("OmniInteract Benchmark Options")
+    group.add_argument(
+        "--omniinteract-subsets", nargs="+", choices=OMNIINTERACT_SUBSETS, default=list(OMNIINTERACT_SUBSETS)
+    )
+    group.add_argument(
+        "--omniinteract-timeout-s", type=_positive_finite_float, default=900.0, help="Complete session timeout."
+    )
+    group.add_argument(
+        "--omniinteract-media-timeout-s",
+        type=_positive_finite_float,
+        default=600.0,
+        help="Per-command media timeout.",
+    )
+    group.add_argument(
+        "--omniinteract-max-video-duration-s",
+        type=_positive_finite_float,
+        default=DEFAULT_MAX_VIDEO_DURATION_S,
+        help="Reject media longer than this safety limit before decoding.",
+    )
+    group.add_argument(
+        "--omniinteract-ref-audio", type=_existing_file, help="Reference WAV for native-duplex audio output."
+    )
+    group.add_argument(
+        "--omniinteract-require-response", action="store_true", help="Fail LISTEN-only functional E2E cases."
+    )
+    group.add_argument(
+        "--omniinteract-output-dir",
+        type=Path,
+        default=Path("omniinteract-output"),
+        help="Directory for case and evaluator artifacts.",
+    )
 
 
 def add_multi_stage_cli_args(parser: argparse.ArgumentParser) -> None:
@@ -188,6 +248,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
 
 _OMNI_BENCH_DATASET_CHOICES = (
     "daily-omni",
+    "omniinteract",
     "seed-tts",
     "seed-tts-text",
     "seed-tts-design",
@@ -218,6 +279,11 @@ def extend_omni_choices(parser: argparse.ArgumentParser) -> None:
 def update_omni_help(parser: argparse.ArgumentParser) -> None:
     """Update upstream argument help text to describe Omni-specific behavior."""
     for action in parser._actions:
+        if action.dest == "num_prompts":
+            action.help = (
+                f"{action.help} OmniInteract uses {_DEFAULT_OMNIINTERACT_NUM_PROMPTS} when this option is omitted; "
+                "0 selects all available cases."
+            )
         if action.dest == "percentile_metrics":
             action.help = (
                 "Comma-separated list of selected metrics to report percentiles. "
@@ -254,6 +320,7 @@ def update_omni_help(parser: argparse.ArgumentParser) -> None:
 def add_omni_args(parser: argparse.ArgumentParser) -> None:
     """Register all vLLM-Omni serving benchmark arguments."""
     add_daily_omni_cli_args(parser)
+    add_omniinteract_cli_args(parser)
     add_seed_tts_cli_args(parser)
     add_multi_stage_cli_args(parser)
     add_diffusion_cli_args(parser)
@@ -261,6 +328,28 @@ def add_omni_args(parser: argparse.ArgumentParser) -> None:
 
 def preprocess_serve_args(args: argparse.Namespace) -> None:
     """Apply serving benchmark CLI transformations after parsing."""
+    if getattr(args, "dataset_name", None) == "omniinteract":
+        if getattr(args, "backend", None) != "openai-realtime-duplex":
+            raise ValueError("OmniInteract requires --backend openai-realtime-duplex")
+        if getattr(args, "endpoint", None) != "/v1/realtime":
+            raise ValueError("OmniInteract requires --endpoint /v1/realtime")
+        if not getattr(args, "omniinteract_ref_audio", None):
+            raise ValueError("OmniInteract requires --omniinteract-ref-audio")
+        if getattr(args, "ignore_eos", False):
+            raise ValueError("OmniInteract does not support --ignore-eos")
+        if getattr(args, "profile", False):
+            raise ValueError("OmniInteract does not support --profile")
+        if getattr(args, "skip_tokenizer_init", False):
+            raise ValueError("OmniInteract does not support --skip-tokenizer-init")
+        if float(getattr(args, "probe_request_rate", 0.0) or 0.0) > 0:
+            raise ValueError("OmniInteract does not support --probe-request-rate")
+        if "num_prompts" not in getattr(args, "explicit_keys", ()):
+            args.num_prompts = _DEFAULT_OMNIINTERACT_NUM_PROMPTS
+        max_concurrency = getattr(args, "max_concurrency", None)
+        if max_concurrency is None:
+            args.max_concurrency = 1
+        elif max_concurrency <= 0:
+            raise ValueError("OmniInteract requires --max-concurrency to be positive")
     extra_body = dict(getattr(args, "extra_body", None) or {})
     bot_task = getattr(args, "bot_task", None)
     if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None:

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Online serving benchmark subcommand for vLLM-Omni.
 
 ``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark
@@ -24,7 +27,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
     """The `serve` subcommand for vllm bench."""
 
     name = "serve"
-    help = "Benchmark the online serving throughput. Supports Daily-Omni and Seed-TTS datasets."
+    help = "Benchmark online serving. Supports Daily-Omni, OmniInteract, and Seed-TTS datasets."
 
     @classmethod
     def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

vllm_omni/entrypoints/cli/serve.py

@@ -760,6 +760,28 @@ class OmniServeCommand(CLISubcommand):
                 "the same persistent path for population and serving."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-host-registration-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. "
+                "Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration "
+                "under the existing pinned-memory policy and fall back to bounded staging when unavailable."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-offload-components",
+            type=json.loads,
+            default=None,
+            metavar="JSON",
+            help=(
+                "Per auxiliary encoder/VAE DLO policy as a JSON object, for example "
+                "'{\"text_encoder\": false}'. False keeps the whole component "
+                "resident; true permits the model's existing offload behavior. "
+                "DiT offload is not controlled by this map."
+            ),
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -1,9 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import math
 from typing import Any, Literal
 
 import numpy as np
 from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 _MAX_EMBEDDING_DIM = 8192
 
 SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
@@ -142,18 +148,21 @@ class OpenAICreateSpeechRequest(BaseModel):
     )
     max_new_tokens: int | None = Field(
         default=None,
+        ge=1,
+        le=_INT64_MAX,
         description="Maximum tokens to generate",
     )
     seed: int | None = Field(
         default=None,
-        ge=0,
-        le=2**63 - 1,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description="Random seed for reproducible generation. When set, ensures "
         "deterministic output for the same input text and seed value.",
     )
     initial_codec_chunk_frames: int | None = Field(
         default=None,
         ge=0,
+        le=_INT64_MAX,
         description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
     )
     non_streaming_mode: bool | None = Field(
@@ -365,10 +374,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
     )
     num_inference_steps: int | None = Field(
         default=None,
+        ge=1,
+        le=_INT64_MAX,
         description="Number of inference steps",
     )
     seed: int | None = Field(
         default=None,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description="Random seed for reproducibility",
     )
 
@@ -413,8 +426,8 @@ class SpeechBatchItem(BaseModel):
     ref_audio: str | None = None
     ref_text: str | None = None
     x_vector_only_mode: bool | None = None
-    max_new_tokens: int | None = None
-    initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
+    initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
     non_streaming_mode: bool | None = None
 
 
@@ -433,8 +446,8 @@ class BatchSpeechRequest(BaseModel):
     ref_audio: str | None = None
     ref_text: str | None = None
     x_vector_only_mode: bool | None = None
-    max_new_tokens: int | None = None
-    initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
+    initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX)
     non_streaming_mode: bool | None = None
 
 
@@ -517,10 +530,11 @@ class StreamingSpeechSessionConfig(BaseModel):
     instructions: str | None = None
     response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
-    max_new_tokens: int | None = Field(default=None, ge=1)
+    max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX)
     initial_codec_chunk_frames: int | None = Field(
         default=None,
         ge=0,
+        le=_INT64_MAX,
         description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.",
     )
     non_streaming_mode: bool | None = Field(

vllm_omni/entrypoints/openai/protocol/images.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
 """
 OpenAI-compatible protocol definitions for image generation.
 
@@ -21,6 +21,10 @@ from pydantic import BaseModel, Field, field_validator
 
 from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 _IMAGE_FILE_METADATA = {
     "jpg": ("jpg", "image/jpeg"),
     "jpeg": ("jpeg", "image/jpeg"),
@@ -144,7 +148,7 @@ class ImageGenerationRequest(BaseModel):
         default=None,
         description="Optional model-specific parameters passed directly to the model's extra_args.",
     )
-    seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
     generator_device: str | None = Field(
         default=None,
         description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",

vllm_omni/entrypoints/openai/protocol/videos.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
 """
 OpenAI-compatible protocol definitions for video generation.
 
@@ -20,6 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_vali
 from vllm_omni.entrypoints.openai.image_api_utils import parse_size
 from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS
 
+# Bound int request fields to avoid overflow issues.
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
+
 
 class VideoGenerationStatus(str, Enum):
     QUEUED = "queued"
@@ -48,10 +52,10 @@ def file_extension(media_type: str):
 class VideoParams(BaseModel):
     """Optional block for video-specific parameters."""
 
-    width: int | None = Field(default=None, ge=1, description="Video width in pixels")
-    height: int | None = Field(default=None, ge=1, description="Video height in pixels")
-    num_frames: int | None = Field(default=None, ge=1, description="Number of frames")
-    fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
+    width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
+    height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
+    num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames")
+    fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
 
     @property
     def size(self) -> str | None:
@@ -140,10 +144,10 @@ class VideoGenerationRequest(BaseModel):
     user: str | None = Field(default=None, description="User identifier for tracking")
 
     # Video-specific fields (top-level for OpenAI-style compatibility)
-    width: int | None = Field(default=None, ge=1, description="Video width in pixels")
-    height: int | None = Field(default=None, ge=1, description="Video height in pixels")
-    fps: int | None = Field(default=None, ge=1, description="Frames per second for output video")
-    num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate")
+    width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels")
+    height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels")
+    fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video")
+    num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate")
     aspect_ratio: str | None = Field(
         default=None,
         description=(
@@ -151,7 +155,12 @@ class VideoGenerationRequest(BaseModel):
             "FL2VA follows the input image; Ref2VA defaults to 16:9."
         ),
     )
-    short_edge: int | None = Field(default=None, ge=1, description="MiniMax H3 output short edge in pixels")
+    short_edge: int | None = Field(
+        default=None,
+        ge=1,
+        le=_INT64_MAX,
+        description="MiniMax H3 output short edge in pixels",
+    )
     num_outputs_per_prompt: int = Field(
         default=1,
         ge=1,
@@ -218,7 +227,7 @@ class VideoGenerationRequest(BaseModel):
         le=20.0,
         description="True CFG scale (model-specific parameter, may be ignored if not supported)",
     )
-    seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility")
     generate_sound: bool = Field(
         default=False,
         description="Request model-generated audio for video models that support sound generation.",
@@ -237,6 +246,7 @@ class VideoGenerationRequest(BaseModel):
     frame_interpolation_exp: int = Field(
         default=1,
         ge=1,
+        le=_INT64_MAX,
         description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.",
     )
     frame_interpolation_scale: float = Field(

vllm_omni/experimental/fullduplex/client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos."""
 
 from __future__ import annotations
@@ -30,6 +33,13 @@ def _rounded_ms(value: float) -> float:
     return round(float(value), 3)
 
 
+def _finite_number(value: object, *, nonnegative: bool = False) -> float | None:
+    if isinstance(value, bool) or not isinstance(value, int | float):
+        return None
+    parsed = float(value)
+    return parsed if math.isfinite(parsed) and (not nonnegative or parsed >= 0) else None
+
+
 def _interval_summary(values: list[float]) -> dict[str, float | int]:
     clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
     if not clean:
@@ -56,11 +66,7 @@ def summarize_session_request_metrics(
     """Average client-observed metrics across turns that emitted audio."""
 
     def mean(metric: str, *, digits: int = 3) -> float | None:
-        values = [
-            float(request[metric])
-            for request in request_metrics
-            if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric]))
-        ]
+        values = [value for request in request_metrics if (value := _finite_number(request.get(metric))) is not None]
         return round(sum(values) / len(values), digits) if values else None
 
     return {
@@ -93,24 +99,57 @@ def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None:
 
 def build_realtime_url(
     url: str,
-    model: str,
+    model: str | None,
     *,
     autostart: bool | None = None,
+    native_duplex: bool | None = True,
     session_id: str | None = None,
 ) -> str:
     """Add the explicit native-duplex query parameters to a Realtime URL."""
     parts = urlsplit(url)
+    if parts.scheme in {"http", "https"}:
+        parts = parts._replace(scheme="ws" if parts.scheme == "http" else "wss")
+    if parts.scheme not in {"ws", "wss"} or not parts.netloc:
+        raise ValueError(f"Unsupported Realtime URL: {url!r}")
     query = dict(parse_qsl(parts.query, keep_blank_values=True))
-    query.setdefault("duplex", "1")
-    query.setdefault("model", model)
-    query.setdefault("minicpmo45_native_duplex", "1")
+    query["duplex"] = "1"
+    if model:
+        query["model"] = model
+    if native_duplex is not None:
+        query["minicpmo45_native_duplex"] = "1" if native_duplex else "0"
     if autostart is not None:
-        query.setdefault("autostart", "1" if autostart else "0")
+        query["autostart"] = "1" if autostart else "0"
     if session_id:
-        query.setdefault("session_id", session_id)
+        query["session_id"] = session_id
     return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
 
 
+def reference_audio_data_url(path: str | None) -> str | None:
+    """Encode a local reference WAV for a Realtime session update."""
+    if path is None:
+        return None
+    audio = Path(path).expanduser().resolve()
+    if not audio.is_file():
+        raise FileNotFoundError(f"Reference audio does not exist: {audio}")
+    return "data:audio/wav;base64," + base64.b64encode(audio.read_bytes()).decode("ascii")
+
+
+def chunk_period_ms(events: list[dict[str, object]], *, default: int = 1000) -> int:
+    """Read the negotiated native-duplex model-unit duration."""
+    for event in reversed(events):
+        session = event.get("session")
+        capabilities = session.get("capabilities") if isinstance(session, dict) else None
+        value = capabilities.get("chunk_period_ms") if isinstance(capabilities, dict) else None
+        if isinstance(value, int) and not isinstance(value, bool) and value > 0:
+            return value
+    return default
+
+
+def has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
+    unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
+    return bool(unit_bytes and len(pcm16) % unit_bytes)
+
+
 def read_pcm16_wav(path: Path) -> bytes:
     """Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
     with wave.open(str(path), "rb") as wav_file:
@@ -202,6 +241,20 @@ class RealtimeEventCollector:
             chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ())
         )
 
+    def response_text(self, response_id: str) -> str:
+        """Join all text/transcript deltas for one response identity."""
+        return "".join(
+            str(event.get("delta") or "")
+            for event in self.events
+            if self.response_id(event) == response_id
+            and event.get("type")
+            in {
+                "response.audio_transcript.delta",
+                "response.output_text.delta",
+                "response.text.delta",
+            }
+        )
+
     def errors(self) -> list[dict[str, object]]:
         return [event for event in self.events if event.get("type") == "error"]
 
@@ -282,15 +335,16 @@ class RealtimeEventCollector:
         if stage0_metrics is not None:
             raw_itls = stage0_metrics.get("vllm_itls_ms")
             itls = (
-                [float(value) for value in raw_itls if isinstance(value, int | float)]
+                [parsed for value in raw_itls if (parsed := _finite_number(value, nonnegative=True)) is not None]
                 if isinstance(raw_itls, list)
                 else []
             )
             result["stage0_tokens"] = {
                 "source": "engine_stage_metrics",
-                "output_token_count": int(stage0_metrics.get("num_tokens_out") or 0),
-                "ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0),
-                "tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0),
+                "output_token_count": int(_finite_number(stage0_metrics.get("num_tokens_out"), nonnegative=True) or 0),
+                "ttft_ms": _finite_number(stage0_metrics.get("vllm_ttft_ms"), nonnegative=True) or 0.0,
+                "tpot_ms": _finite_number(stage0_metrics.get("vllm_tpot_ms"), nonnegative=True),
+                "itls_ms": itls,
                 "inter_token_interval_ms": _interval_summary(itls),
             }
 
@@ -368,15 +422,26 @@ class RealtimeEventCollector:
 class RealtimeDuplexClient:
     """Small async client used by the user demo and reusable smoke probes."""
 
-    def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
+    def __init__(
+        self,
+        url: str,
+        *,
+        max_size: int = 64 * 1024 * 1024,
+        additional_headers: dict[str, str] | None = None,
+    ) -> None:
         self.url = url
         self.max_size = max_size
+        self.additional_headers = additional_headers
         self.events = RealtimeEventCollector()
         self._ws: Any = None
         self._reader_task: asyncio.Task[None] | None = None
 
     async def __aenter__(self) -> RealtimeDuplexClient:
-        self._ws = await websockets.connect(self.url, max_size=self.max_size)
+        self._ws = await websockets.connect(
+            self.url,
+            max_size=self.max_size,
+            additional_headers=self.additional_headers,
+        )
         self._reader_task = asyncio.create_task(self._read_events())
         return self
 
@@ -405,6 +470,19 @@ class RealtimeDuplexClient:
     async def send(self, event: dict[str, object]) -> None:
         await self._ws.send(json.dumps(event))
 
+    def raise_if_reader_stopped(self) -> None:
+        """Fail a caller waiting on events after the WebSocket reader exits."""
+
+        task = self._reader_task
+        if task is None or not task.done():
+            return
+        if task.cancelled():
+            raise ConnectionError("Realtime WebSocket reader was cancelled")
+        error = task.exception()
+        if error is not None:
+            raise ConnectionError("Realtime WebSocket reader failed") from error
+        raise ConnectionError("Realtime WebSocket closed before the requested event arrived")
+
     async def configure(
         self,
         model: str,
@@ -418,6 +496,7 @@ class RealtimeDuplexClient:
         temperature: float | None = None,
         extra_body: dict[str, object] | None = None,
         session_id: str | None = None,
+        idle_timeout_s: float | None = None,
         timeout_s: float = 20.0,
     ) -> None:
         session_extra_body = dict(extra_body or {})
@@ -448,14 +527,23 @@ class RealtimeDuplexClient:
         if instructions is not None:
             session["instructions"] = instructions
         if initial_user_text is not None:
-            extra_body = session["extra_body"]
-            assert isinstance(extra_body, dict)
-            extra_body["duplex_initial_user_text"] = initial_user_text
+            session_extra = session["extra_body"]
+            assert isinstance(session_extra, dict)
+            session_extra["duplex_initial_user_text"] = initial_user_text
         if session_id:
             session["session_id"] = session_id
+        if idle_timeout_s is not None:
+            session["idle_timeout_s"] = idle_timeout_s
         await self.send({"type": "session.update", "session": session})
+
+        def session_created() -> bool:
+            if self.events.count("session.created") > 0:
+                return True
+            self.raise_if_reader_stopped()
+            return False
+
         await wait_for(
-            lambda: self.events.count("session.created") > 0,
+            session_created,
             timeout_s=timeout_s,
             label="session.created",
         )
@@ -498,20 +586,31 @@ class RealtimeDuplexClient:
             if not pcm16:
                 continue
             played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
-            await self.send(
-                {
-                    "type": "playback.ack",
-                    "response_id": response_id,
-                    "item_id": f"item_{response_id}",
-                    "played_ms": played_ms,
-                    "committed_ms": played_ms,
-                }
-            )
+            await self.send_playback_ack(response_id, played_ms)
+
+    async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
+        await self.send(
+            {
+                "type": "playback.ack",
+                "response_id": response_id,
+                "item_id": f"item_{response_id}",
+                "played_ms": played_ms,
+                "committed_ms": played_ms,
+            }
+        )
 
     async def close_session(self, *, timeout_s: float = 20.0) -> None:
+        from_index = len(self.events.events)
         await self.send({"type": "session.close"})
+
+        def session_closed() -> bool:
+            if any(event.get("type") == "session.closed" for event in self.events.events[from_index:]):
+                return True
+            self.raise_if_reader_stopped()
+            return False
+
         await wait_for(
-            lambda: self.events.count("session.closed") > 0,
+            session_closed,
             timeout_s=timeout_s,
             label="session.closed",
         )

vllm_omni/host_weight_runtime/filesystem/writer.py

@@ -12,17 +12,19 @@ import os
 import struct
 import uuid
 from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
 from types import TracebackType
 
 import torch
-from safetensors import safe_open
 
 from ..identity import WeightArtifactIdentity, canonical_json
 from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
 from ..protocols import ProductionMetadata, TensorWriteSpec
 
 _HASH_CHUNK_BYTES = 8 * 1024**2
+_MAX_HASH_WORKERS = 8
 _MANIFEST_FILE = "manifest.json"
 _READY_FILE = "READY.json"
 
@@ -79,21 +81,47 @@ def _fsync_directory(path: Path) -> None:
         os.close(fd)
 
 
-def _file_sha256(path: Path) -> str:
-    digest = hashlib.sha256()
+def _fsync_payload(fd: int) -> None:
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+@dataclass(frozen=True)
+class _PayloadDigests:
+    file: str
+    tensors: dict[str, str]
+
+
+def _scan_payload(path: Path, specs: tuple[TensorWriteSpec, ...]) -> _PayloadDigests:
+    """Calculate file and tensor digests in one ordered read."""
+    file_digest = hashlib.sha256()
+    tensor_digests: dict[str, str] = {}
     with path.open("rb") as handle:
-        while chunk := handle.read(_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _tensor_sha256(tensor: torch.Tensor) -> str:
-    digest = hashlib.sha256()
-    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
-    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
-        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
-        digest.update(memoryview(chunk.numpy()))
-    return digest.hexdigest()
+        prefix = handle.read(8)
+        if len(prefix) != 8:
+            raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
+        header_size = struct.unpack("<Q", prefix)[0]
+        header = handle.read(header_size)
+        if len(header) != header_size:
+            raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
+        file_digest.update(prefix)
+        file_digest.update(header)
+        for spec in specs:
+            digest = hashlib.sha256()
+            remaining = spec.nbytes
+            while remaining:
+                chunk = handle.read(min(remaining, _HASH_CHUNK_BYTES))
+                if not chunk:
+                    raise ArtifactWriterError(f"truncated tensor {spec.name!r} in {path.name!r}")
+                file_digest.update(chunk)
+                digest.update(chunk)
+                remaining -= len(chunk)
+            tensor_digests[spec.name] = digest.hexdigest()
+        if handle.read(1):
+            raise ArtifactWriterError(f"unexpected trailing bytes in {path.name!r}")
+    return _PayloadDigests(file_digest.hexdigest(), tensor_digests)
 
 
 class _Coverage:
@@ -130,6 +158,8 @@ class StreamingSafeTensorFile:
         owner: FilesystemArtifactWriter,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool,
     ) -> None:
         self._owner = owner
         self.relative_name = _safe_relative_name(relative_name)
@@ -144,6 +174,7 @@ class StreamingSafeTensorFile:
         self.specs = tuple(sorted(specs, key=lambda spec: spec.name))
         self._specs = {spec.name: spec for spec in self.specs}
         self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
+        self._ordered = ordered
         self._closed = False
 
         offsets: dict[str, int] = {}
@@ -167,6 +198,7 @@ class StreamingSafeTensorFile:
         total_size = self._payload_offset + cursor
         owner._reserve_file(total_size)
         self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600)
+        prefix = struct.pack("<Q", len(encoded_header))
         try:
             if hasattr(os, "posix_fallocate"):
                 try:
@@ -177,13 +209,16 @@ class StreamingSafeTensorFile:
                     os.ftruncate(self._fd, total_size)
             else:  # pragma: no cover - Linux production path has posix_fallocate
                 os.ftruncate(self._fd, total_size)
-            _pwrite_all(self._fd, struct.pack("<Q", len(encoded_header)), 0)
+            _pwrite_all(self._fd, prefix, 0)
             _pwrite_all(self._fd, encoded_header, 8)
         except BaseException:
             os.close(self._fd)
             self._path.unlink(missing_ok=True)
             owner._forget_reserved_file(total_size)
             raise
+        self._file_digest = hashlib.sha256(prefix + encoded_header) if ordered else None
+        self._tensor_digests = {spec.name: hashlib.sha256() for spec in self.specs} if ordered else {}
+        self._write_cursor = self._payload_offset
         owner._opened(self)
 
     @staticmethod
@@ -226,9 +261,19 @@ class StreamingSafeTensorFile:
             raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc
 
     def _write(self, name: str, byte_offset: int, data: memoryview) -> None:
+        file_offset = self._payload_offset + self._offsets[name] + byte_offset
+        if self._ordered and file_offset != self._write_cursor:
+            raise ArtifactWriterError(
+                f"ordered tensor payload expected offset {self._write_cursor}, got {file_offset} for {name!r}"
+            )
         coverage = self._coverage[name]
         coverage.add(byte_offset, byte_offset + len(data))
-        _pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset)
+        _pwrite_all(self._fd, data, file_offset)
+        if self._ordered:
+            assert self._file_digest is not None
+            self._file_digest.update(data)
+            self._tensor_digests[name].update(data)
+            self._write_cursor += len(data)
 
     def close(self) -> None:
         if self._closed:
@@ -238,11 +283,19 @@ class StreamingSafeTensorFile:
             self.abort()
             raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
         self._closed = True
+        digests = None
+        if self._ordered:
+            assert self._file_digest is not None
+            digests = _PayloadDigests(
+                self._file_digest.hexdigest(),
+                {name: digest.hexdigest() for name, digest in self._tensor_digests.items()},
+            )
         try:
-            os.fsync(self._fd)
-        finally:
+            os.fchmod(self._fd, 0o444)
+            self._owner._closed(self, digests, self._fd)
+        except BaseException:
             os.close(self._fd)
-        self._owner._closed(self)
+            raise
 
     def abort(self) -> None:
         if self._closed:
@@ -281,20 +334,25 @@ class FilesystemArtifactWriter:
         self._capacity_check = capacity_check
         self._active: set[StreamingSafeTensorFile] = set()
         self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
+        self._digests: dict[str, _PayloadDigests] = {}
         self._reserved_bytes = 0
         self._finalized = False
+        self._fsync_executor: ThreadPoolExecutor | None = None
+        self._fsync_futures: list[Future[None]] = []
 
     def open_tensor_file(
         self,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool = False,
     ) -> StreamingSafeTensorFile:
         if self._finalized:
             raise ArtifactWriterError("cannot add payloads after artifact finalization")
         relative_name = _safe_relative_name(relative_name)
         if relative_name in self._files or any(item.relative_name == relative_name for item in self._active):
             raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once")
-        return StreamingSafeTensorFile(self, relative_name, specs)
+        return StreamingSafeTensorFile(self, relative_name, specs, ordered=ordered)
 
     def _reserve_file(self, size: int) -> None:
         self._capacity_check(self._reserved_bytes + size, size)
@@ -306,9 +364,14 @@ class FilesystemArtifactWriter:
     def _opened(self, writer: StreamingSafeTensorFile) -> None:
         self._active.add(writer)
 
-    def _closed(self, writer: StreamingSafeTensorFile) -> None:
+    def _closed(self, writer: StreamingSafeTensorFile, digests: _PayloadDigests | None, fd: int) -> None:
         self._active.remove(writer)
         self._files[writer.relative_name] = writer.specs
+        if digests is not None:
+            self._digests[writer.relative_name] = digests
+        if self._fsync_executor is None:
+            self._fsync_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hwr-fsync")
+        self._fsync_futures.append(self._fsync_executor.submit(_fsync_payload, fd))
 
     def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None:
         self._active.discard(writer)
@@ -318,6 +381,19 @@ class FilesystemArtifactWriter:
         for writer in tuple(self._active):
             with contextlib.suppress(Exception):
                 writer.abort()
+        with contextlib.suppress(Exception):
+            self._wait_for_payload_fsync()
+
+    def _wait_for_payload_fsync(self) -> None:
+        if self._fsync_executor is None:
+            return
+        executor = self._fsync_executor
+        futures = tuple(self._fsync_futures)
+        self._fsync_executor = None
+        self._fsync_futures.clear()
+        executor.shutdown(wait=True)
+        for future in futures:
+            future.result()
 
     def finalize(
         self,
@@ -335,46 +411,50 @@ class FilesystemArtifactWriter:
         if metadata.restorer_schema != identity.producer.restorer_schema:
             raise ArtifactWriterError("production metadata does not match the identity's restorer schema")
 
+        self._wait_for_payload_fsync()
+
         tensor_entries: list[TensorManifestEntry] = []
         file_entries: list[FileManifestEntry] = []
         seen_tensors: set[str] = set()
-        for file_name in sorted(self._files):
+        file_names = sorted(self._files)
+        unhashed = [file_name for file_name in file_names if file_name not in self._digests]
+        if unhashed:
+            with ThreadPoolExecutor(max_workers=min(_MAX_HASH_WORKERS, len(unhashed))) as executor:
+                scans = executor.map(
+                    lambda file_name: _scan_payload(self.staging_dir / file_name, self._files[file_name]),
+                    unhashed,
+                )
+                self._digests.update(zip(unhashed, scans, strict=True))
+
+        for file_name in file_names:
             path = self.staging_dir / file_name
-            os.chmod(path, 0o444)
-            with path.open("rb") as handle:
-                os.fsync(handle.fileno())
+            specs = self._files[file_name]
+            digests = self._digests[file_name]
             file_entries.append(
                 FileManifestEntry(
                     file_name=file_name,
                     size=path.stat().st_size,
-                    sha256=_file_sha256(path),
+                    sha256=digests.file,
                 )
             )
-            expected = {spec.name: spec for spec in self._files[file_name]}
-            with safe_open(str(path), framework="pt", device="cpu") as handle:
-                if set(handle.keys()) != set(expected):
-                    raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}")
-                for name in sorted(expected):
-                    if name in seen_tensors:
-                        raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files")
-                    seen_tensors.add(name)
-                    spec = expected[name]
-                    tensor = handle.get_tensor(name)
-                    if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype:
-                        raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata")
-                    entry = TensorManifestEntry(
-                        name=name,
+            for spec in specs:
+                if spec.name in seen_tensors:
+                    raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files")
+                seen_tensors.add(spec.name)
+                tensor_entries.append(
+                    TensorManifestEntry(
+                        name=spec.name,
                         kind=spec.kind,
                         role=spec.role,
                         file_name=file_name,
-                        storage_key=name,
-                        shape=tuple(tensor.shape),
-                        dtype=str(tensor.dtype),
-                        stride=tuple(tensor.stride()),
-                        nbytes=tensor.numel() * tensor.element_size(),
-                        sha256=_tensor_sha256(tensor),
+                        storage_key=spec.name,
+                        shape=spec.shape,
+                        dtype=str(spec.dtype),
+                        stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()),
+                        nbytes=spec.nbytes,
+                        sha256=digests.tensors[spec.name],
                     )
-                    tensor_entries.append(entry)
+                )
 
         ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
         content_digest = hashlib.sha256()

vllm_omni/host_weight_runtime/protocols.py

@@ -88,6 +88,8 @@ class ArtifactWriter(Protocol):
         self,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool = False,
     ) -> TensorFileWriter: ...
 
 

vllm_omni/quantization/factory.py

@@ -470,10 +470,14 @@ def resolve_quant_config_from_disk(
         )
         return build_quant_config(qc_method, **qc_kwargs)
 
-    # AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of
-    # is_checkpoint_*_serialized; rebuild so the offline path is selected.
+    # AutoRound MXFP checkpoints use data_type="mx_fp" instead of
+    # is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
+    # selected according to the checkpoint's bit width.
     if qc_kwargs.get("data_type") == "mx_fp":
-        logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.")
+        logger.info(
+            "config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP%d.",
+            qc_kwargs.get("bits", getattr(quant_config, "weight_bits", 0)),
+        )
         return build_quant_config(qc_method, **qc_kwargs)
 
     if (