Add normalized online FP8 cache for DLO

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

.buildkite/cuda/test-weekly.yml

@@ -30,14 +30,6 @@ steps:
       - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
     mirror_hardwares: h100_4
 
-  - label: "Reliability Test - voxcpm2"
-    timeout_in_minutes: 180
-    depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
-    commands:
-      - pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
-    mirror_hardwares: l4_1
-
   - group: "Reliability Test - Invalid parameters Test"
     depends_on: upload-weekly-pipeline
     if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"

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

@@ -69,6 +69,12 @@ steps:
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
           HF_TOKEN: "${HF_TOKEN}"
         commands:
+          # Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
+          # auto-picks npu:0, which the omni server under test already fills with its
+          # KV cache, so Whisper OOMs on `.to(device)` and every clip reports
+          # "ASR / WER failed".
+          - export SEED_TTS_WER_EVAL=1
+          - export SEED_TTS_EVAL_DEVICE=npu:1
           - |
             set +e
             pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model

.buildkite/npu/test-npu-ready.yml

@@ -1,4 +1,15 @@
 steps:
+  - label: "Platforms Test"
+    key: ready-platforms-test
+    depends_on: upload-ready-pipeline
+    mirror_hardwares: a3_npu_2
+    commands:
+      - >-
+        pytest -s -v
+        tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
+        -m 'core_model and npu and A3 and omni'
+        --run-level 'core_model'
+
   - group: ":card_index_dividers: Omni Model Test"
     key: ready-omni-group
     depends_on: upload-ready-pipeline

docs/.nav.yml

@@ -43,6 +43,7 @@ nav:
     - Overview: features/README.md
     - Runtime and Stage Execution:
       - Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
+      - Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
       - Sleep Mode: features/sleep_mode.md
     - Quantization:
       - Overview: user_guide/quantization/overview.md

docs/configuration/README.md

@@ -14,10 +14,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
-## Multi-Stage Recipes
-
-- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
-
 ## Optimization Features
 
 - **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models

docs/configuration/pd_disaggregation.md

@@ -1,171 +0,0 @@
-# Prefill-Decode (PD) Disaggregation
-
-PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
-stages so prompt processing and token generation can run on different workers.
-
-This is documented as a stage-config recipe instead of a bundled YAML because the
-deployment-specific values usually change per environment:
-
-- GPU placement
-- `tensor_parallel_size`
-- connector backend and connector ports
-- connector IPs or bootstrap addresses
-
-Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
-and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
-changes below.
-
-## Requirements
-
-- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
-- A KV connector supported by vLLM, such as `MooncakeConnector`
-- Matching `tensor_parallel_size` on the prefill and decode thinker stages
-
-## 1. Split the thinker into prefill and decode stages
-
-Replace the original thinker stage with two stages:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_prefill_only: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_producer"
-        kv_rank: 0
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25201
-    final_output: false
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: llm
-    is_decode_only: true
-    runtime:
-      devices: "1"
-    engine_args:
-      max_num_seqs: 64
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_consumer"
-        kv_rank: 1
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25202
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-```
-
-Notes:
-
-- `is_prefill_only: true` marks the thinker stage that only saves KV.
-- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
-- `kv_transfer_config` is required on both stages.
-- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
-  prefill side only processes the prompt and exports KV.
-
-## 2. Shift the downstream stages by one index
-
-After inserting the extra thinker stage, renumber the remaining stages:
-
-```yaml
-  - stage_id: 2
-    runtime:
-      devices: "2"
-    engine_input_source: [1]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
-
-  - stage_id: 3
-    runtime:
-      devices: "2"
-    engine_args:
-      max_num_seqs: 1
-    engine_input_source: [2]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
-```
-
-Compared with the default Qwen3-Omni config:
-
-- the talker becomes stage `2` instead of stage `1`
-- the code2wav stage becomes stage `3` instead of stage `2`
-- the talker now reads from decode stage `1`
-
-## 3. Add runtime edges for the four-stage pipeline
-
-```yaml
-runtime:
-  enabled: true
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2
-    - from: 2
-      to: 3
-```
-
-## 4. Launch with your custom config
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --deploy-config /path/to/qwen3_omni_pd.yaml
-```
-
-## Operational Notes
-
-- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
-  pair. Keep prefill and decode at the same `tensor_parallel_size`.
-- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
-  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
-- Choose connector ports and addresses that match your deployment. The values
-  shown above are examples only.

docs/configuration/stage_configs.md

@@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
 |-------|------|----------|---------|-------------|
 | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
 | `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
-| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
+| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
+| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
 | `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
 | `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
-| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
-| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
+| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
+| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
 | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
 | `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
 | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |

docs/contributing/ci/test_execution_guide.md

@@ -184,9 +184,6 @@ Failed jobs: 1/2
     # Reliability: HunyuanImage DiT (H100 × 4)
     pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
 
-    # Reliability: VoxCPM2 (L4)
-    pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
-
     ```
 
     The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).

docs/contributing/ci/test_writing_guide.md

@@ -562,7 +562,6 @@ L5 level testing focuses on the performance of model services under ***long-runn
     -   `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
     -   `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
     -   `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
-    -   `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
     -   `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
     -   `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
     -   `README.md` — Short local run commands for this directory.
@@ -662,13 +661,12 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
 | `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
 | `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
 | `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
-| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
 | `README.md` | Minimal run / collect examples. |
 
 ##### Parametrization and markers
 
 - Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
-- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
 - **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
 
 ##### CI trigger
@@ -680,13 +678,11 @@ Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model sui
 | Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
 | Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
 | Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
-| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
 
 ```bash
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
 ```
 
 ##### Local commands
@@ -696,7 +692,6 @@ pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```
 
 ##### Adding a new model suite
@@ -709,4 +704,4 @@ pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 </details>
 
 -   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
+    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -10,6 +10,30 @@ For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
 and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
 
+## Feature compatibility
+
+Host-storage optimization and runtime compatibility are separate decisions.
+When direct checkpoint mmap is unavailable, DLO can still use tensors produced
+by the ordinary loader. "Compatibility path" below means that fallback is
+implemented but has less end-to-end coverage than the primary path.
+
+Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
+
+| Feature | DLO + AllGather | DLO without AllGather |
+|---|---|---|
+| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
+| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
+| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
+| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
+| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **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`. |
+
+See [Parallelism compatibility](#parallelism-compatibility) and
+[Request and loading constraints](#request-and-loading-constraints) for the
+detailed contracts and validation boundaries.
+
 ## Status
 
 DLO is implemented for multi-device diffusion execution. The default
@@ -81,6 +105,31 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
+### Phase-I normalized online-FP8 cache
+
+The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
+adds an explicit, opt-in normalized cache for the narrower case of one TP1
+transformer with per-tensor online FP8 and DLO AllGather:
+
+```text
+cache miss:  BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
+cache hit:   FP8 cache   -> direct mmap plan -> DLO host shard -> AllGather
+```
+
+The cache stores the serialized runtime representation, including finalized FP8
+weights, generated scales, and persistent transformer buffers. Its manifest
+records the schema, component, quantization contract, tensor count, source
+files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
+invalid artifact is a cache miss; cache publication is best-effort and falls
+back to the ordinary loader on failure.
+
+Phase I keeps the cache binding deliberately simple: one source per transformer
+component, with existing per-tensor transforms applied before publication. The
+cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
+and it normalizes the transformer only. Cache sharing and compatibility for
+other components, transformed TP coordinates, and broader DP/SP layouts remain
+in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -181,13 +230,16 @@ requirements.
 
 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. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
+ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
+ordinary loader finalizes their runtime weights and scales; DLO then shards and
+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.
 
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
+The general cache-compatibility contract for additional components, transformed
+TP coordinates, and cross-DP/SP reuse remains in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 ## Validation coverage
 
@@ -197,6 +249,8 @@ Current source-level validation includes:
 - HSDP + DLO without AllGather acceptance at configuration level;
 - loader preflight fallback for TP, HSDP, online quantization, unknown custom
   loaders, missing keys, and shape/dtype mismatches;
+- 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;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;

docs/design/feature/quantization.md

@@ -190,12 +190,15 @@ configured parallelism strategy:
 - `WeightsMapper` transformations must be applied before matching ignored-layer
   lists or component prefixes.
 
-Distributed layerwise offload has an additional boundary: its sharded mmap
-AllGather path is not compatible with online quantization because the path
-expects already materialized checkpoint weights. Use the rank-local
-`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
-combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
-for the full compatibility matrix.
+Distributed layerwise offload has an additional boundary: direct checkpoint
+mmap cannot prepare online-quantized tensors because it bypasses the ordinary
+loader that creates their runtime weights and scales. Per-tensor online FP8
+linears can still use DLO AllGather through the ordinary loader: each rank first
+finalizes the FP8 tensors, then DLO shards those runtime tensors for
+reconstruction during prefetch. Other online methods remain limited to
+`--dlo-no-use-allgather` until their packing and scale layouts are validated.
+See the [DLO design document](offloader/distributed_layerwise_offload.md) for
+the full compatibility matrix.
 
 ## Adding a quantization backend
 

docs/design/module/vllm_omni_config.md

@@ -18,8 +18,8 @@ validation_paths:
   - tests/config/**
 upstream_refs:
   - vllm.config
-last_reviewed: 2026-08-07
-last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+last_reviewed: 2026-08-11
+last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
 ---
 
 # vLLM-Omni configuration
@@ -39,11 +39,66 @@ The PR review proposal to capture environment-derived values in configuration
 objects at initialization is an open design question, not a current
 invariant.
 
+## Current structured-stage reuse behavior
+
+The current RFC #4021 implementation uses upstream vLLM config classes as the
+schema source for AR and generation stages:
+
+- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
+- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
+- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
+- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
+
+This inheritance preserves the existing load, cache, scheduler, and parallel
+concern boundaries while reducing duplicate downstream declarations. It does
+not, by itself, make every inherited field an effective vLLM-Omni runtime
+option. Construction and runtime consumption are separate surfaces:
+
+| Surface | Current behavior |
+| --- | --- |
+| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
+| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
+| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
+| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
+
+The generated projection maps include known upstream naming differences such
+as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
+`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
+upstream concern config and `EngineArgs` therefore enters the AR/generation
+projection without requiring a second downstream allowlist. Explicit inherited
+fields that have no projection raise an error rather than being silently
+dropped.
+
+Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
+cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
+`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
+fields are terminal internals. These exclusions prevent one input from being
+constructed or projected through two config concerns. Effective-engine-argument
+tests cover both the dynamic mapping and this exclusion boundary.
+
+`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
+`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
+inputs and complete upstream value objects. Mapping inputs are validated and
+materialized as their concrete upstream types when the structured config is
+constructed, so downstream consumers see one resolved representation.
+Prebuilt upstream objects retain their type across the structured projection
+boundary.
+
+Quantization remains on its existing Omni-owned transport contract. This
+reuse step does not adopt upstream `QuantizationConfigArgs`, because the
+current engine-specific quantization split needs to be resolved separately.
+
+Diffusion keeps its existing engine-field ownership and projection surface.
+Shared Python inheritance must not be treated as evidence that an LLM-only
+parallel input is effective for a diffusion stage.
+
 ## Safe-change guide while deferred
 
 Changes should trace every affected producer to its runtime consumer and test
 the structured, legacy, CLI, and deployment paths that are actually supported.
-Do not infer a stable contract from this placeholder page.
+For inherited vLLM fields, tests should separately cover constructor behavior
+and effective `EngineArgs` projection. Do not infer a stable contract from
+this placeholder page.
 
 ## Promotion gate
 

docs/features/README.md

@@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow.
 | Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
 | Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
 
+The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
+describes the experimental Qwen3-Omni runtime topology and its current
+integration boundaries. It is not yet a supported launch recipe.
+
 Some runtime designs are deliberately not promoted as standalone User Guide
 features yet:
 
-- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
+- Generic, OmniConnector-based
+  [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
+  yet a standalone User Guide feature. Its
   [OmniConnector implementations](../design/index.md#communication) describe
-  topology and transport contracts. Practical deployment configuration remains
-  under [Pipeline and deploy configurations](../configuration/stage_configs.md).
+  transport contracts, while practical connector configuration remains under
+  [Pipeline and deploy configurations](../configuration/stage_configs.md).
 - [Async Chunk](../design/feature/async_chunk.md) and
   [Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
   are model- and pipeline-dependent. Use the selected model's deploy

docs/features/pd_disaggregation.md

@@ -0,0 +1,140 @@
+# Prefill-Decode (PD) Disaggregation (experimental)
+
+!!! warning "Experimental design reference"
+    PD disaggregation is only partially integrated on the current `main` branch.
+    The runtime contains PD detection and routing scaffolding, but there is no
+    supported deploy overlay or end-to-end validated launch recipe yet. Treat this
+    page as a description of the current design, not as production guidance.
+
+Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
+stages: a prefill worker that processes the prompt and produces KV cache, and a
+decode worker that imports that cache and generates tokens. The Talker and
+Code2Wav stages remain downstream; PD does not disaggregate the entire
+Qwen3-Omni pipeline.
+
+## Current design
+
+### Runtime topology
+
+```mermaid
+flowchart LR
+    client["Client request"] --> orchestrator["Omni Orchestrator"]
+    orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
+    prefill -->|"KV blocks"| connector[("vLLM KV connector")]
+    prefill -->|"KV routing metadata"| orchestrator
+    orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
+    connector -->|"remote KV load"| decode
+    decode -->|"generated text"| client
+```
+
+The orchestrator owns the request lifecycle, while the vLLM KV connector owns
+the bulk KV transfer. The original prompt is submitted to both Thinker stages:
+the decode worker still needs the request tokens and multimodal features to
+construct its request, while the imported KV cache prevents it from repeating
+the completed prefill computation.
+
+For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
+
+| Stage | Role | Output or handoff |
+| --- | --- | --- |
+| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
+| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
+| 2 | Talker | Converts combined Thinker conditioning into codec codes |
+| 3 | Code2Wav | Converts codec codes into audio |
+
+Text remains a final output of the decode stage, and audio remains a final
+output of Code2Wav.
+
+### Request lifecycle
+
+1. At startup, PD detection looks for one stage marked `is_prefill_only` and
+   one downstream stage marked `is_decode_only`. The decode stage must list the
+   prefill stage as an input source.
+2. The orchestrator clones the Thinker sampling parameters for prefill, forces
+   `max_tokens=1`, clears stop conditions, and sets producer-side
+   `kv_transfer_params`. This makes prefill finish by length after exporting KV.
+3. When prefill finishes, the orchestrator captures its KV routing metadata and
+   Qwen multimodal output under the original request ID.
+4. The orchestrator submits the original prompt to decode with consumer-side
+   `kv_transfer_params`. Those parameters identify the transfer, prefill engine,
+   bootstrap endpoint, and, when required by the connector, remote request.
+5. The decode engine imports the KV cache and continues generation. The normal
+   Omni stage path then carries Thinker conditioning through Talker and
+   Code2Wav.
+
+### Control and data paths
+
+PD uses two different connector layers. They solve different handoffs and are
+not interchangeable.
+
+| Path | Payload | Owner |
+| --- | --- | --- |
+| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
+| PD data path | Thinker attention KV blocks | vLLM KV connector |
+| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
+
+### Detection and validation
+
+The current runtime accepts one logical prefill/decode pair per pipeline. Each
+stage may use tensor parallelism, but both sides of the pair must satisfy these
+invariants:
+
+- both stages define `kv_transfer_config`
+- prefill uses `kv_producer` or `kv_both`
+- decode uses `kv_consumer` or `kv_both`
+- both stages use the same KV connector
+- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
+- both stages use the same `tensor_parallel_size`
+
+These checks happen before request processing so an inconsistent PD topology
+fails during engine initialization.
+
+## What is implemented today
+
+| Area | Current state |
+| --- | --- |
+| Pair detection and topology validation | Implemented in the shared PD helper |
+| Prefill sampling preparation | Implemented for sync and async entrypoints |
+| Prefill-to-decode routing | Implemented in the current Orchestrator |
+| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
+| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
+| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
+| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
+| Performance | No non-regression or throughput/latency result is established for the current path |
+
+Because these pieces are still converging, the previous legacy `stage_args` YAML
+and `vllm serve` command have been removed. The serve CLI now accepts the
+deploy-based `stages:` schema, so retaining the legacy recipe would imply a
+launch path that current `main` rejects.
+
+## Configuration contract under development
+
+The intended deploy integration starts from the
+[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
+and expands its frozen three-stage pipeline into the four stages shown above.
+A complete integration must:
+
+- split the Thinker definition without duplicating the Talker or Code2Wav model
+  contracts
+- remap downstream stage IDs, input sources, and stage connectors
+- inject producer and consumer `kv_transfer_config` values into the two Thinker
+  stages
+- allocate separate device sets for prefill and decode
+- preserve both final outputs: decode text and Code2Wav audio
+- provide connector-specific routing metadata without hard-coding one connector
+  contract
+
+The required number of GPUs depends on tensor-parallel sizes and whether Talker
+and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
+for decode, and one for the two downstream audio stages.
+
+## Implementation map
+
+- PD detection, validation, and sampling-parameter preparation:
+  [`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
+- startup detection and connector metadata:
+  [`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
+- request-scoped prefill-to-decode routing:
+  [`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
+- Qwen3 downstream conditioning helpers:
+  [Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)

docs/models/supported_models.md

@@ -53,6 +53,7 @@ th {
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -77,6 +78,7 @@ th {
 | `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
 | `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
 | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
+| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/serving/speech_api.md

@@ -141,7 +141,20 @@ The response shape depends on the streaming parameters:
 
 **Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
 complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
-`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
+`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
+non-streaming responses from non-diffusion speech servers report usage through
+response headers:
+
+| Header | Description |
+| --- | --- |
+| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
+| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
+| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
+| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
+| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
+
+Diffusion-mode speech servers route through a separate response path and do not
+emit these headers.
 
 **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
 WAV) as they are decoded.

docs/serving/videos_api.md

@@ -167,6 +167,14 @@ transfer options such as `control_path`, `control_guidance`,
 `num_conditional_frames`, `show_control_condition`, and `show_input`; see the
 Cosmos3 recipe for complete examples.
 
+HTTP redirects for `image_reference.image_url` follow vLLM's
+`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
+`1` (the default) to allow redirects or `0` to reject them. A redirect target
+can differ from the original host, and vLLM-Omni does not yet fully implement
+upstream vLLM's media URL allowlist protection. Deployments should therefore
+treat remote media URLs as untrusted and choose this setting as part of their
+URL access policy.
+
 ### Speech-to-Video
 
 For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --usp 4
+
+# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
+vllm serve /path/to/model --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload \
+  --dlo-fp8-cache-dir /path/to/node-local/fp8-cache
 ```
 
 ```python
@@ -63,6 +69,7 @@ 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-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
 
 ## Host-weight loading
 
@@ -104,6 +111,24 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+## Normalized online-FP8 cache
+
+`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
+TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
+transformer once and publishes finalized FP8 weights and scales as an atomic,
+serialized safetensors artifact. On a valid cache hit, the loader selects that
+serialized source before DLO planning, so the transformer can be loaded through
+the direct checkpoint-mmap path without re-materializing the full BF16 source.
+
+The cache is keyed to the source files through the manifest's schema, component,
+quantization contract, and source fingerprint. Incomplete, incompatible, or
+stale artifacts are treated as cache misses. Cache creation is best-effort: a
+publication failure logs a warning and keeps the ordinary online-FP8 path.
+
+Phase I normalizes the transformer only. Other pipeline components keep their
+ordinary loading and quantization behavior. The cache is not selected when DLO
+is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -138,16 +163,18 @@ must enter each collective.
   but it provides no shared-mmap host-memory benefit or guarantee.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
-- Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
+- Per-tensor online FP8 linears use the ordinary loader and can run with either
+  DLO transfer path. With AllGather, every rank temporarily materializes the
+  complete FP8 model in host memory before DLO retains only its shard. Other
+  online quantization methods require no-AllGather until their runtime layouts
+  are validated.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
+cache remains a follow-up design in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

docs/user_guide/diffusion/parallelism/hsdp.md

@@ -152,7 +152,8 @@ In `DiffusionParallelConfig`:
 **Constraints:**
 
 - `hsdp_replicate_size × hsdp_shard_size == world_size`
-- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
+  Pipeline Parallelism, or Expert Parallelism.
 
 ---
 

docs/user_guide/fault_injection_reliability_matrix.md

@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
 
 ## Source of Conclusions
 
-The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
+The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
 
 ### Example: `SIGTERM` fault injection log (Qwen3-Omni)
 

docs/user_guide/quantization/fp8.md

@@ -179,6 +179,24 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
 vllm serve <your-model> --omni --quantization fp8
 ```
 
+### Distributed layerwise offload
+
+Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
+
+```bash
+vllm serve <your-model> \
+  --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload
+```
+
+Direct checkpoint mmap is not used for this combination. Each rank runs the
+ordinary loader to create the final FP8 weights and scales before DLO shards
+them, so startup temporarily materializes the complete FP8 model in host memory
+on every rank. Runtime DLO host residency is sharded normally. Use
+`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
+does not fit the workload.
+
 ## Parameters
 
 | Parameter | Type | Default | Description |

examples/online_serving/pi0/README.md

@@ -0,0 +1,104 @@
+# π0 (Pi-Zero) VLA — OpenPI realtime serving
+
+[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
+model from Physical Intelligence: multi-camera images + a language instruction +
+robot proprioceptive state → a continuous action chunk via flow-matching denoising
+(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
+websocket protocol at `/v1/realtime/robot/openpi`.
+
+## Install extras
+
+The core `pip install -e .` does not include the OpenPI client used here:
+
+- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
+  repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
+download it (`MODEL=lerobot/pi0_base`) or point at a local copy
+(`MODEL=/path/to/pi0_base`).
+
+## Run the server
+
+```bash
+vllm serve lerobot/pi0_base --omni --port 8000 \
+    --served-model-name pi0 \
+    --deploy-config vllm_omni/deploy/pi0.yaml \
+    --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
+metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
+
+### Commonly adjusted `model_config` keys
+
+These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
+
+| Key | Default | Meaning |
+|---|---|---|
+| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
+| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
+| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
+| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
+| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
+| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
+| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
+| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
+| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
+
+The `policy_server_config` block below them is the OpenPI handshake metadata
+advertised to the client; keep its `action_horizon` / `action_dim` /
+`image_resolution` in sync with the `model_config` values above. The e2e test
+`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
+live server and asserts the advertised metadata matches these `pi0.yaml` values.
+
+## Run the client
+
+```bash
+python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
+    --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
+cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
+proprioceptive state to drive a robot.
+
+### Observation format
+
+The client sends a flat dict per inference:
+
+```python
+{
+    "observation.images.base_0_rgb":       np.uint8[H, W, 3],
+    "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
+    "observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
+    "state":   np.float32[state_dim],   # zero-padded to max_state_dim=32 server-side
+    "prompt":  "pick up the red block",
+    "session_id": "<uuid>",             # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
+(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
+version-stable across transformers releases (the version-stability checks in
+`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
+lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
+
+## Limitations
+
+- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
+  fully supported (state passes through, actions are returned in the model's space).
+  Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
+  stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
+  not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
+  added before its actions are in real-world units.

examples/online_serving/pi0/openpi_client.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
+
+Connects to a running π0 server (see README.md), sends a single robot
+observation (3 cameras + state + language instruction), and prints the returned
+continuous action chunk ``[action_horizon, action_dim]``.
+
+The server speaks the OpenPI websocket protocol:
+    connect -> server sends msgpack(policy_server_config metadata)
+    infer   -> client sends msgpack(obs), server replies msgpack(actions ndarray)
+
+Usage:
+    python openpi_client.py --host 127.0.0.1 --port 8000 \
+        --prompt "pick up the red block"
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+import uuid
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    # Avoid shadowing the installed ``openpi_client`` package with this dir.
+    example_dir = str(Path(__file__).resolve().parent)
+    removed = sys.path and sys.path[0] == example_dir
+    if removed:
+        sys.path.pop(0)
+    try:
+        from openpi_client.websocket_client_policy import WebsocketClientPolicy
+    finally:
+        if removed:
+            sys.path.insert(0, example_dir)
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
+
+logger = logging.getLogger("pi0_openpi_client")
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+
+# π0 / pi0_base camera identities (must match the server's image_feature_keys).
+CAMERA_KEYS = (
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+)
+STATE_DIM = 32
+
+
+def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
+    """Build a single π0 observation with blank cameras + zero state.
+
+    Replace the zeros with real camera frames (HWC uint8) and proprioceptive
+    state to drive an actual robot.
+    """
+    obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
+    obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _ws_uri(host: str, port: int, path: str) -> str:
+    """Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
+    ``host`` (it has no separate path arg)."""
+    if host.startswith(("ws://", "wss://")):
+        return host
+    return f"ws://{host}:{port}{path}"
+
+
+def _as_action_array(response: Any) -> np.ndarray:
+    if isinstance(response, dict) and response.get("type") == "error":
+        raise RuntimeError(f"Inference failed: {response.get('message', response)}")
+    return np.asarray(response, dtype=np.float32)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="π0 OpenPI client")
+    parser.add_argument("--host", default=DEFAULT_HOST)
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
+    parser.add_argument("--path", default=DEFAULT_PATH)
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+    parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
+    metadata = dict(policy.get_server_metadata())
+    logger.info("Server metadata: %s", metadata)
+
+    session_id = str(uuid.uuid4())
+    for step in range(args.num_steps):
+        obs = _make_dummy_obs(args.prompt, session_id)
+        actions = _as_action_array(policy.infer(obs))
+        logger.info(
+            "[step %d] actions shape=%s mean=%.4f std=%.4f",
+            step,
+            actions.shape,
+            float(actions.mean()),
+            float(actions.std()),
+        )
+        if not np.isfinite(actions).all():
+            raise RuntimeError("Server returned non-finite actions.")
+    logger.info("Done — π0 OpenPI round-trip OK.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mkdocs.yml

@@ -79,6 +79,7 @@ plugins:
         user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
         user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
         user_guide/quantization/index.md: user_guide/quantization/overview.md
+        configuration/pd_disaggregation.md: features/pd_disaggregation.md
         design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
         design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
         design/feature/cpu_offloading.md: design/feature/offloader/README.md

recipes/MammothModa2/MammothModa2.md

@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
   --model ./MammothModa2-Preview \
-  --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
+  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
   --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
   --height 1024 \
   --width 1024 \
@@ -132,6 +132,44 @@ ls -lh mammoth_t2i.png
 python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
 ```
 
+### 1x AMD MI300X, MammothModa2 Preview
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Offline Commands
+
+The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+    --model bytedance-research/MammothModa2-Preview \
+    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+    --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+    --height 1024 \
+    --width 1024 \
+    --seed 42 \
+    --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+    --enable-diffusion-pipeline-profiler \
+    --log-stats \
+    --output mammoth_t2i.png
+```
+
+#### Verification
+
+The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
+
+The output was a valid 1024 by 1024 RGB PNG.
+
 ## MammothModa2-Dev unified inference
 
 MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses

recipes/MiniMaxAI/MiniMax-H3.md

@@ -416,9 +416,10 @@ For example, keep the first main block's attention projections in BF16 with:
 ```
 
 The structured option replaces `--quantization fp8`. Online FP8 can be used
-with H3 layerwise offload and with distributed layerwise offload's full-weight
-per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
-supported for runtime-created FP8 weights.
+with H3 layerwise offload and with both DLO transfer paths. The default
+AllGather path uses the ordinary loader to finalize FP8 weights and scales
+before sharding them across ranks. `--dlo-no-use-allgather` instead retains
+complete rank-local tensors and avoids the synchronized request-wave contract.
 
 ## AMD ROCm (gfx942 / gfx950)
 
@@ -870,8 +871,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
   H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
   pure Ulysses.
-- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
-  use `--dlo-no-use-allgather` when combining online FP8 with DLO.
+- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
+  in host memory on every rank during startup before retaining only each rank's
+  shard. Size startup host memory for that transient peak.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone

recipes/MiniMaxAI/MiniMax-Music3.md

@@ -0,0 +1,173 @@
+# MiniMax Music 3
+
+> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out
+
+## Summary
+
+- Vendor: MiniMax
+- Model: `MiniMaxAI/MiniMax-Music3`
+- Task: Text-to-music generation (vocals and instrumental together)
+- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: @linyueqian
+
+## When to use this recipe
+
+Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six
+minutes. This is not a TTS model with a musical voice: there is no reference
+speaker, no `voice`, and no `temperature`. What you control is the lyrics, the
+caption, the seed, and the length.
+
+## References
+
+- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3)
+- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/)
+- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU),
+  `vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two)
+
+## Architecture
+
+| Component | Spec |
+|---|---|
+| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k |
+| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 |
+| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder |
+| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 |
+| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling |
+| Solver | Euler, 30 steps, acoustic guidance scale 1.7 |
+| Window | 200 frames per acoustic chunk, 100-frame hop |
+| Frame rate | 25 frames per second |
+| Context length | 10,240 tokens |
+| Output | 32 kHz stereo WAV |
+
+Generation is two-stage. The backbone predicts one audio frame per decode step,
+each frame eight RVQ codebooks deep. Every 200 frames are handed to the
+flow-matching transformer, which solves a latent that the decoder turns into
+waveform.
+
+**Guidance doubles the rows.** Every request decodes twice: one row sees the
+real prompt, the other sees the same prompt with the caption-and-lyrics span
+replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole
+song, so a request costs twice the KV its length suggests. Size the pool for
+rows, not requests.
+
+## Hardware Support
+
+## GPU
+
+### 1x H200 141GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- CUDA: 12.9
+- vLLM version: 0.27.0
+- vLLM-Omni version or commit: branch `feat/minimax-music3`
+
+#### Command
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+    --host 0.0.0.0 --port 8000 \
+    --trust-remote-code --omni
+```
+
+The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered
+automatically from the checkpoint's `model_type=minimax_music3`. It colocates
+both stages on device 0.
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "MiniMaxAI/MiniMax-Music3",
+    "input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me",
+    "instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.",
+    "seed": 1,
+    "max_new_tokens": 750
+  }' \
+  --output song.wav
+```
+
+Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames
+at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a
+target: the model ends the song itself when it emits the audio-end token, and a
+response shorter than the cap is the model finishing rather than a truncation.
+
+#### Notes
+
+- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone
+  and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights,
+  which sit outside any KV budget.
+- **Key flags:** none required. Guidance, the solver step count and the window
+  geometry are fixed by the checkpoint and are not request parameters.
+- **Known limitations:**
+  - The external API is non-streaming; `stream: true` is rejected.
+  - `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a
+    maximum-length song (9,000 frames) cannot both fit. Long captions shorten
+    the maximum song. This is a property of the checkpoint.
+  - The acoustic stage runs in float32. bfloat16 is accepted but measurably
+    degrades the solver.
+
+### 2x H200 141GB
+
+Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device:
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+    --host 0.0.0.0 --port 8000 \
+    --trust-remote-code --omni \
+    --deploy-config minimax_music3_2gpu.yaml
+```
+
+Only the placement differs; both layouts run the acoustic stage in float32.
+
+## Writing the request
+
+Two fields carry the request, and both are required.
+
+- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`,
+  `[Outro]`) steer the arrangement and are part of the prompt contract.
+  **Put a tag on its own line.** Normalization keeps only the tags on a line
+  that starts with one and drops the rest of that line, so lyrics written next
+  to a tag are silently lost:
+
+  ```text
+  "[Verse]\nWalking down the street"   ->   [start] [verse] Walking down the street
+  "[Verse] Walking down the street"    ->   [start] [verse]
+  ```
+
+- `instructions` is the **caption** describing genre, instrumentation, tempo,
+  mood and production. It is the strongest control available: vague captions
+  give generic arrangements. Long structured captions work well, since the
+  model was trained on descriptions covering global attributes, emotional
+  progression, vocal detail and arrangement.
+
+A request is deterministic in its seed: the same lyrics, caption, seed and
+length return the same audio. Changing only the seed gives another take of the
+same song.
+
+## Request parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `input` | string | required | Lyrics, non-empty. Tags on their own lines |
+| `instructions` | string | required | Caption, non-empty |
+| `seed` | int | `0` | Non-negative. Fixes the output for a given request |
+| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 |
+| `response_format` | string | `"wav"` | Output container |
+
+## Parameters the model rejects
+
+The request is refused rather than silently ignored, so a mistake is visible
+immediately.
+
+| Parameter | Reason |
+|---|---|
+| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw |
+| `voice` | There is no speaker to select; the vocal comes from the caption |
+| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag |
+| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" |
+| `stream: true` | The external API is non-streaming |

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

@@ -71,6 +71,35 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
 two-GPU profile. The removed fused two-stage implementation is not retained as
 a fallback because it would duplicate state machines and correctness paths.
 
+### Graph execution
+
+Graph boundaries follow each stage's state model:
+
+| Stage | Graph mode on Ascend | Eager boundary |
+| --- | --- | --- |
+| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
+| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
+| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
+
+Stage 2 keeps `enforce_eager: true` for the generation runner because its
+Python request metadata and per-request cache dictionaries are not valid
+outer-graph inputs. The backend captures only the deterministic CFM DiT
+estimator, with an exact graph key for each tensor shape. Timestep embedding
+stays eager because the upstream implementation creates a CPU frequency tensor;
+HiFT stays eager because it generates random phase. The graph cache stores up
+to 32 entries by default, after which unseen shapes run eagerly.
+
+An ACL graph capture failure is fatal to that Stage-2 process because older
+torch-npu releases can leave allocator and RNG capture state invalid. Restart
+the service after a capture failure. To run without capture, set
+`code2wav_enable_npu_graph: false` under the Stage-2
+`platforms.npu.stages[].additional_config` block before startup. Tune the cache
+limit there with `code2wav_max_npu_graphs`. Graph mode also requires
+`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
+
+The inner NPUGraph is independent of the outer runner setting, so do not use a
+global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
+
 ## GPU
 
 ### 1 x GPU (default — single command)
@@ -211,9 +240,9 @@ speech output (TTS)"** checkbox on / off.
   model processes still share one CUDA device.
 - `--trust-remote-code` is required — the HF repo ships a custom
   `MiniCPMO` config / model class.
-- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
-  eager because its request-owned Flow/HiFT caches and variable chunk/cache
-  shapes are not yet exposed through a static exact-shape graph wrapper.
+- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
+  request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
+  Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
 - All default stages use `max_num_seqs: 4` to reduce cross-process GPU
   contention. Talker AR
   state and Code2Wav caches are request-owned; Code2Wav batches only
@@ -224,8 +253,8 @@ speech output (TTS)"** checkbox on / off.
 - `StageRequestStats.batch_size` is a request-scoped placeholder, not the
   scheduler's execution batch.
 - Single-GPU co-location trades throughput for hardware density: Stage 0/1
-  CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
-  CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
+  CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
+  contexts. Use the 8x4090 config or a custom multi-GPU mapping for
   throughput-sensitive serving.
 
 ### 8 x RTX 4090 24GB (consumer-GPU layout)

recipes/Qwen/Qwen3-TTS.md

@@ -252,3 +252,53 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     - stage_id: 1
       gpu_memory_utilization: 0.15
   ```
+
+### 1x AMD MI300X, 1.7B checkpoints
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
+- transformers: 5.15.0
+
+#### Command
+
+The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
+
+Set this environment variable before you start Qwen3-TTS:
+
+```bash
+export MIOPEN_FIND_MODE=FAST
+```
+
+#### Verification
+
+Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
+
+| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
+| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
+| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
+| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
+| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
+| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
+| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
+| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
+
+All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
+
+#### Notes
+
+- The current ROCm config keeps Code2Wav in eager mode.
+- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
+- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
+- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
+- Each setup was tested with one request, so the timing is not a performance benchmark.
+- GPU memory was sampled once per second with `rocm-smi`.

recipes/README.md

@@ -49,6 +49,7 @@ recipes/
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
+| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
 | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
 | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
@@ -73,9 +74,9 @@ recipes/
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
-| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
-| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
-| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
+| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
+| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
+| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |

recipes/SenseNova/SenseNova-U1.md

@@ -178,6 +178,48 @@ with the expected 1536×2720 resolution.
 - The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
   for fused QKV and gate/up projections with TP support.
 
+### 1x AMD MI300X 192GB
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Command
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+    --model SenseNova/SenseNova-U1-8B-MoT \
+    --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
+    --width 1536 \
+    --height 2720 \
+    --seed 42 \
+    --num-inference-steps 50 \
+    --cfg-scale 4.0 \
+    --extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+    --enable-diffusion-pipeline-profiler \
+    --log-stats \
+    --output sensenova_u1_mi300x.png
+```
+
+#### Verification
+
+The command completed and wrote a valid 1536 by 2720 RGB PNG.
+
+#### Notes
+
+- Generation with 50 inference steps took 34.011 seconds.
+- Model loading used 32.774 GiB and took 13.987 seconds.
+- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
+- The highest one second whole device memory sample was 37.93 GiB.
+
 ## Online Serving
 
 SenseNova-U1 supports all four modalities via the OpenAI-compatible

recipes/StabilityAI/Stable-Audio-Open.md

@@ -164,7 +164,50 @@ PY
   `torchaudio==2.11.0+cu130`.
 - The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
   warnings observed during validation did not prevent successful generation.
-- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
-  ROCm, XPU, and NPU setups are not covered here.
+- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
 - Long generations, higher inference-step counts, and non-WAV response formats
   were not benchmarked in this recipe.
+
+#### 1x AMD MI300X 192GB
+
+##### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+##### Commands
+
+```bash
+python3 examples/offline_inference/text_to_audio/text_to_audio.py \
+    --model stabilityai/stable-audio-open-1.0 \
+    --prompt "A gentle piano melody with soft room ambience" \
+    --negative-prompt "Low quality, distorted, noisy" \
+    --seed 42 \
+    --guidance-scale 7.0 \
+    --audio-length 10.0 \
+    --num-inference-steps 50 \
+    --cache-backend tea_cache \
+    --enable-diffusion-pipeline-profiler \
+    --output stable_audio_10s.wav
+```
+
+##### Verification
+
+The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
+
+##### Notes
+
+- TeaCache ran with `rel_l1_thresh=0.2`.
+- Model loading used 2.7891 GiB and took 3.706 seconds.
+- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
+- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
+- The highest one second whole device memory sample was 19.61 GiB.
+- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
+- The full process took 384 seconds, including startup and compilation.

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -189,8 +189,10 @@ wait
 
 - `--dlo-use-allgather` (default) requires all concurrent requests to have the same
   `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
-- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
-  use `--dlo-no-use-allgather` or disable online quantization
+- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader
+  instead of direct checkpoint mmap, then shard the finalized FP8 weights and
+  scales. Size host memory for a transient complete FP8 model on every rank at
+  startup
 - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
   `--dlo-no-use-allgather` retains the standard TP loader and should be treated
   as an experimental compatibility path

recipes/inclusionAI/Ming-omni-tts.md

@@ -170,6 +170,35 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
 `--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
 for the dialect case.
 
+**Offline smoke result on ROCm 7.2:**
+
+A later offline smoke run used the same MI300X model with the following software:
+
+- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
+- Kernel: Linux 6.8.0-134-generic, x86_64
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- ROCm or HIP: 7.2.53211
+- vLLM: 0.27.0+rocm723
+- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+
+```bash
+python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --model inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --case basic \
+    --text "你好，这是 AMD MI300X ROCm 配方验证。" \
+    --output-dir ./ming_tts_output \
+    --output-name ming_omni_tts_mi300x.wav \
+    --enforce-eager \
+    --log-stats \
+    --metadata-json ./ming_tts_manifest.json
+```
+
+The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
+
+The run used the official ROCm image built from `docker/Dockerfile.rocm`.
+
 ## Notes
 
 - The official ROCm image includes the platform dependencies.

recipes/k2-fsa/OmniVoice.md

@@ -0,0 +1,98 @@
+# OmniVoice
+
+> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
+
+## Summary
+
+- Vendor: k2-fsa
+- Model: `k2-fsa/OmniVoice`
+- Task: Multilingual text to speech
+- Mode: Offline inference for all modes and online serving for automatic voice
+- Maintainer: Community
+
+## References
+
+- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
+- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
+- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
+- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
+
+## Requirements
+
+Download the checkpoint before the first run:
+
+```bash
+huggingface-cli download k2-fsa/OmniVoice
+```
+
+Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
+
+The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
+
+## ROCm
+
+### 1x AMD MI300X
+
+#### Environment
+
+- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
+- Kernel: Linux 6.8.0-134-generic, x86_64
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- ROCm or HIP: 7.2.53211
+- vLLM: 0.27.0+rocm723
+- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- transformers: 5.15.0
+
+The run used the official ROCm image built from `docker/Dockerfile.rocm`.
+
+#### Automatic voice
+
+Run the offline example from the repository root:
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+    --model k2-fsa/OmniVoice \
+    --deploy-config vllm_omni/deploy/omnivoice.yaml \
+    --text "Hello, this is OmniVoice running on one AMD MI300X." \
+    --seed 42 \
+    --output omnivoice_mi300x.wav
+```
+
+The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
+
+#### Reference voice cloning
+
+The offline interface also accepts a reference WAV and its transcript:
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+    --model k2-fsa/OmniVoice \
+    --text "Hello, this is a cloned voice." \
+    --ref-audio ref.wav \
+    --ref-text "This is the reference transcription." \
+    --output cloned.wav
+```
+
+#### Voice design
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+    --model k2-fsa/OmniVoice \
+    --text "Hello, this is a designed voice." \
+    --instruct "female, low pitch, British accent" \
+    --output designed.wav
+```
+
+## Online serving
+
+Online serving currently supports automatic voice:
+
+```bash
+vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
+```
+
+```bash
+cd examples/online_serving/text_to_speech/omnivoice
+python speech_client.py --text "Hello, how are you?"
+```

recipes/lerobot/Pi0.md

@@ -0,0 +1,108 @@
+# Pi0 (π0, Pi-Zero) VLA
+
+> Vision-Language-Action policy serving over the OpenPI realtime websocket
+
+## Summary
+
+- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
+- Model: `lerobot/pi0_base`
+- Task: Vision-Language-Action — multi-camera images + a language instruction +
+  robot proprioceptive state → a continuous action chunk
+  `[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
+  **not** emit text tokens.
+- Mode: Online serving via the OpenPI realtime robot API
+  (`/v1/realtime/robot/openpi` websocket)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
+client connects over the OpenPI websocket and, per control step, sends an
+observation (camera frames + state + instruction) and receives an action chunk.
+
+A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
+calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
+ignored.
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
+server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
+
+## Hardware Support
+
+### 1x H200 141GB (Online serving)
+
+#### Environment
+
+- OS: Ubuntu 22.04+
+- Python: 3.12+
+- Driver / runtime: NVIDIA CUDA environment
+- vLLM / vLLM-Omni: use the commit you are deploying from
+- OpenPI client extras (for the example client): `openpi-client`
+  (`pip install -e packages/openpi-client` from the
+  [openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
+  `msgpack`, `msgpack-numpy`
+
+#### Command
+
+```bash
+vllm serve lerobot/pi0_base \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --served-model-name pi0 \
+  --deploy-config vllm_omni/deploy/pi0.yaml \
+  --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
+handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
+32) that enables the OpenPI realtime serving layer.
+
+#### Verification
+
+Run the example client (sends a synthetic observation and prints the returned
+action chunk):
+
+```bash
+python examples/online_serving/pi0/openpi_client.py \
+  --host 127.0.0.1 --port 8000 \
+  --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
+is a flat dict per inference:
+
+```python
+{
+    "observation.images.base_0_rgb":        np.uint8[H, W, 3],
+    "observation.images.left_wrist_0_rgb":  np.uint8[H, W, 3],
+    "observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
+    "state":   np.float32[state_dim],   # zero-padded to max_state_dim=32 server-side
+    "prompt":  "pick up the red block",
+    "session_id": "<uuid>",             # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
+reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
+
+## References
+
+- Model blog: <https://www.physicalintelligence.company/blog/pi0>
+- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
+- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
+- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
+- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
+- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
+  OpenPI websocket e2e in
+  [`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)

tests/benchmarks/test_daily_omni_local_source.py

@@ -9,13 +9,17 @@ from pathlib import Path
 import pytest
 
 from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
+    DailyOmniDataset,
     daily_omni_local_qa_json,
     daily_omni_local_videos_dir,
+    ensure_daily_omni_hub_root,
     resolve_daily_omni_local_root,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
 
+_HUB_ID = "liarliar/Daily-Omni"
+
 
 def _write_qa_json(root: Path) -> Path:
     qa = root / "qa.json"
@@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path:
     return qa
 
 
+def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
+    """Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
+
+    Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
+    only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
+    """
+    constants = pytest.importorskip("huggingface_hub").constants
+
+    cache = tmp_path / "hub"
+    repo_dir = cache / "datasets--liarliar--Daily-Omni"
+    snapshot = repo_dir / "snapshots" / "deadbeef"
+    snapshot.mkdir(parents=True)
+    if with_qa:
+        _write_qa_json(snapshot)
+    (repo_dir / "refs").mkdir()
+    (repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
+
+    monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
+    monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+    return snapshot
+
+
 def test_resolve_local_root_returns_none_for_hub_id() -> None:
     assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
     assert resolve_daily_omni_local_root("") is None
@@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
     root = resolve_daily_omni_local_root(str(tmp_path))
     assert daily_omni_local_qa_json(root) is None
     assert daily_omni_local_videos_dir(root) is None
+
+
+def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+
+    root = ensure_daily_omni_hub_root(_HUB_ID)
+
+    assert root == snapshot.resolve()
+    assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
+
+
+def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """The regression this file exists for: a Hub id must not require ``load_dataset``."""
+    snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+    # Any use of `datasets` here means we went back to the path that cannot read the hub cache.
+    monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+    dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
+
+    assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
+    assert [row["video_id"] for row in dataset.data] == ["vid0"]
+
+
+def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    _stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
+    monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
+
+    message = str(excinfo.value)
+    assert "someone-else/Daily-Omni" in message
+    assert "snapshot_download failed" in message
+    assert "--daily-omni-qa-json" in message

tests/config/test_config_import_cycle.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Guard the NPU pytest circular import through DiffusionOutput."""
+
+from __future__ import annotations
+
+import ast
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _top_level_import_modules(path: Path) -> set[str]:
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    modules: set[str] = set()
+    for node in tree.body:
+        if isinstance(node, ast.ImportFrom) and node.module:
+            modules.add(node.module)
+        elif isinstance(node, ast.Import):
+            modules.update(alias.name for alias in node.names)
+    return modules
+
+
+def _assert_isolated_import_succeeds(script: str) -> None:
+    env = os.environ.copy()
+    env.pop("CUDA_VISIBLE_DEVICES", None)
+    env.pop("HIP_VISIBLE_DEVICES", None)
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        cwd=REPO_ROOT,
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=120,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr
+
+
+def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
+    imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
+    assert "vllm_omni.config.config_factory" not in imported
+    assert "vllm_omni.config.pipeline_registry" not in imported
+
+
+def test_config_package_lazy_exports_still_resolve() -> None:
+    from vllm_omni.config import StageConfigFactory, register_pipeline
+    from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
+    from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
+
+    assert StageConfigFactory is DirectStageConfigFactory
+    assert register_pipeline is direct_register_pipeline
+
+
+def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
+    src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
+    init_start = src.index("def __init__(self) -> None:")
+    runtime_start = src.index("def init_diffusion_model_runner_runtime")
+    init_src = src[init_start:runtime_start]
+    runtime_src = src[runtime_start:]
+    assert "apply_minimax_h3_qwen3vl_patch" not in init_src
+    assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
+    assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
+
+
+def test_lora_config_import_does_not_load_pipeline_registry() -> None:
+    """Import LoRAConfig and diffusion.data in a fresh interpreter.
+
+    The NPU failure was ``data.py`` importing LoRAConfig while still loading,
+    then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
+    DiffusionOutput. This child process must not preload those exports.
+    """
+    _assert_isolated_import_succeeds(
+        """
+import sys
+
+from vllm_omni.config.lora import LoRAConfig
+
+if "vllm_omni.config.pipeline_registry" in sys.modules:
+    raise SystemExit("LoRAConfig import loaded pipeline_registry")
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+    raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
+
+from vllm_omni.diffusion.data import DiffusionOutput
+
+if DiffusionOutput is None or LoRAConfig is None:
+    raise SystemExit("expected LoRAConfig and DiffusionOutput")
+"""
+    )
+
+
+def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
+    """Resolving the pi0 topology must not import its runtime pipeline.
+
+    ``pipeline_pi0`` imports ``diffusion.data``. Importing it while
+    ``diffusion.data`` is only partially initialized closes the Ascend plugin's
+    circular-import loop, so the registry must depend only on a lightweight
+    pipeline-config module.
+    """
+    _assert_isolated_import_succeeds(
+        """
+import sys
+
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+    raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
+
+pipeline = resolve_pipeline_config("pi0")
+if pipeline is None or pipeline.model_type != "pi0":
+    raise SystemExit("expected the registered pi0 pipeline config")
+"""
+    )

tests/config/test_omni_config.py

@@ -5,11 +5,20 @@
 from __future__ import annotations
 
 from dataclasses import fields
+from inspect import Parameter, signature
 from pathlib import Path
 
+import msgspec
 import pytest
 from pydantic import ValidationError
+from pydantic.fields import FieldInfo
 from transformers import Qwen3OmniMoeConfig
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
 
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config import omni_config as omni_config_module
@@ -40,6 +49,7 @@ from vllm_omni.config.stage_config import (
     merge_pipeline_deploy,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
 from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -243,9 +253,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
     ],
     ids=["global", "ar-stage", "generation-stage"],
 )
-def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
-    with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
-        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
+    stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
+
+    assert stage.cache_config.cache_dtype == "fp8"
+    assert "cache_dtype" in stage.cache_config._omni_explicit_fields
+
+
+@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
+def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
+            )
+        ]
+    )
+
+    stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+
+    assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
+    assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
 
 
 def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
@@ -437,6 +467,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
         "connector_config",
         "runtime_config",
         "parallel_config",
+        "compilation_config",
+        "profiler_config",
         "quantization_config",
     }
     assert "diffusion_config" not in public_fields
@@ -455,10 +487,47 @@ def test_runtime_config_fields_match_structured_runtime_scope():
         "num_gpus",
         "log_level",
         "log_stats",
-        "profiler_config",
     }
 
 
+def test_upstream_config_compatible_fields_materialize_mapping_inputs():
+    compilation_config = {"backend": "eager"}
+    profiler_config = {"profiler": "cuda"}
+
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+        compilation_config=compilation_config,
+        profiler_config=profiler_config,
+    )
+
+    assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+    assert stage_config.compilation_config.backend == "eager"
+    assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+    assert stage_config.profiler_config.profiler == "cuda"
+
+
+def test_upstream_config_compatible_fields_validate_mapping_inputs():
+    with pytest.raises(ValidationError):
+        VllmOmniARStageConfig(
+            stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+            profiler_config={"delay_iterations": -1},
+        )
+
+
+def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
+    compilation_config = VllmCompilationConfig(backend="eager")
+    profiler_config = VllmProfilerConfig(profiler="cuda")
+
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+        compilation_config=compilation_config,
+        profiler_config=profiler_config,
+    )
+
+    assert stage_config.compilation_config is compilation_config
+    assert stage_config.profiler_config is profiler_config
+
+
 def test_sub_config_fields_match_structured_scopes():
     assert {f.name for f in fields(OmniStageModelConfig)} == {
         "model",
@@ -489,7 +558,6 @@ def test_sub_config_fields_match_structured_scopes():
         "enforce_eager",
         "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",
-        "compilation_config",
         "enable_multithread_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
@@ -499,29 +567,25 @@ def test_sub_config_fields_match_structured_scopes():
         "model_subdir",
         "tokenizer_subdir",
     }
-    assert {f.name for f in fields(OmniStageLoadConfig)} == {
+    vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
+    assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
+    assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
         "tokenizer",
-        "download_dir",
         "skip_tokenizer_init",
-        "load_format",
         "tokenizer_mode",
         "config_format",
         "skip_mm_profiling",
     }
-    assert {f.name for f in fields(OmniStageCacheConfig)} == {
-        "kv_cache_memory_bytes",
-        "gpu_memory_utilization",
-        "enable_prefix_caching",
+    assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
+    assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
+    assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
         "disable_hybrid_kv_cache_manager",
         "mm_processor_cache_gb",
         "mamba_ssm_cache_dtype",
     }
-    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
-        "max_num_seqs",
-        "max_num_batched_tokens",
+    assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
+    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
         "max_model_len",
-        "enable_chunked_prefill",
-        "async_scheduling",
     }
     assert {f.name for f in fields(OmniStageConnectorConfig)} == {
         "async_chunk",
@@ -530,17 +594,10 @@ def test_sub_config_fields_match_structured_scopes():
         "output_connectors",
         "input_connectors",
     }
-    assert {f.name for f in fields(OmniStageParallelConfig)} == {
-        "pipeline_parallel_size",
-        "data_parallel_size",
-        "tensor_parallel_size",
-        "enable_expert_parallel",
-        "world_size",
-    }
-    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
-        "pipeline_parallel_size",
-        "data_parallel_size",
-        "tensor_parallel_size",
+    vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
+    assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
+    assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
+    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
         "sequence_parallel_size",
         "ulysses_degree",
         "ring_degree",
@@ -554,20 +611,86 @@ def test_sub_config_fields_match_structured_scopes():
         "mask_sp_padding",
         "hsdp_shard_size",
         "hsdp_replicate_size",
-        "enable_expert_parallel",
-        "world_size",
     }
 
 
+def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
+    scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
+    assert scheduler_config.max_num_encoder_input_tokens == 4096
+    assert scheduler_config.encoder_cache_size == 4096
+
+    unresolved_scheduler_config = OmniStageSchedulerConfig()
+    assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
+    assert unresolved_scheduler_config.encoder_cache_size is None
+
+    parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+    assert parallel_config.data_parallel_index == 2
+
+    diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+    assert diffusion_parallel_config.data_parallel_index == 2
+
+
+@pytest.mark.parametrize(
+    "config_cls",
+    [
+        OmniStageLoadConfig,
+        OmniStageCacheConfig,
+        OmniStageSchedulerConfig,
+        OmniStageParallelConfig,
+        OmniStageDiffusionParallelConfig,
+    ],
+)
+def test_inherited_sub_configs_are_keyword_only(config_cls):
+    assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
+
+    with pytest.raises(TypeError):
+        config_cls(1)
+
+
+@pytest.mark.parametrize(
+    "config_cls",
+    [
+        OmniStageLoadConfig,
+        OmniStageCacheConfig,
+        OmniStageSchedulerConfig,
+        OmniStageParallelConfig,
+        OmniStageDiffusionParallelConfig,
+    ],
+)
+def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
+    config = config_cls()
+
+    uninitialized_fields = [
+        config_field.name
+        for config_field in fields(config)
+        if isinstance(getattr(config, config_field.name), FieldInfo)
+    ]
+    assert uninitialized_fields == []
+
+    msgspec.msgpack.encode(_serialize_stage_config(config))
+
+
+def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
+    omni_config = _from_pipeline_key("qwen3_tts")
+
+    for stage_config in omni_config.stage_configs:
+        payload = {
+            "stage_id": stage_config.stage_id,
+            "stage_config": _serialize_stage_config(stage_config),
+        }
+        msgspec.msgpack.encode(payload)
+
+
 def test_diffusion_parallel_config_fields_cover_legacy_surface():
     from vllm_omni.diffusion.data import DiffusionParallelConfig
 
     legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
     structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
     expected_upstream_fields = {"mask_sp_padding"}
+    vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
 
     assert legacy_fields | expected_upstream_fields <= structured_fields
-    assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
+    assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
 
 
 def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
@@ -610,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
 
 def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
     cfg = OmniStageDiffusionParallelConfig(
-        pipeline_parallel_size=2,
-        ulysses_degree=2,
+        ulysses_degree=4,
         use_hsdp=True,
         hsdp_shard_size=-1,
         hsdp_replicate_size=2,
@@ -622,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
 
 
 def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with TP"):
         OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with DP"):
         OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
 
@@ -1078,3 +1200,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
     cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
 
     assert cfg.quantization_config == quantization_config
+
+
+def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        quantization_config=quantization_config,
+    )
+
+    # Exercise the terminal construction performed by the future typed startup
+    # path without probing ports or loading remote model metadata.
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
+    cfg.enrich_config()
+
+    assert cfg.quantization_config is not None
+    assert cfg.quantization_config.get_name() == "int8"

tests/dfx/reliability/README.md

@@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal
 - Qwen3 tests: `test_reliability_qwen3_omni.py`
 - Wan2.2 tests: `test_reliability_wan22.py`
 - HunyuanImage tests: `test_reliability_hunyuan_image.py`
-- VoxCPM2 tests: `test_reliability_voxcpm2.py`
 
 ```bash
 pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -976,3 +976,68 @@ def test_voices_delete_invalid_requests(
             "err_message": err_message,
         }
     )
+
+
+# ─── POST /v1/audio/speech · MiniMax Music 3 ───
+
+# Text-to-music on the speech endpoint. Most of the speech contract does not
+# apply: there is no speaker, no reference audio and no temperature, and the
+# lyrics alone are not a request. Each case below pins a rejection so a caller
+# who assumes the usual TTS shape finds out immediately.
+_MINIMAX_MUSIC3_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="MiniMaxAI/MiniMax-Music3",
+            stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="minimax_music3",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
+_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
+# 25 frames per second, capped at six minutes of song.
+_MINIMAX_MUSIC3_MAX_FRAMES = 9000
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
+        pytest.param({"instructions": "   "}, ("instructions",), id="caption_blank"),
+        pytest.param({"input": "   "}, ("input",), id="lyrics_blank"),
+        pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
+        pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
+        pytest.param(
+            {"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
+            ("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
+            id="max_new_tokens_over_cap",
+        ),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
+def test_minimax_music3_speech_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body: dict[str, Any] = {
+        "model": omni_server.model,
+        "input": _MINIMAX_MUSIC3_LYRICS,
+        "instructions": _MINIMAX_MUSIC3_CAPTION,
+        "seed": 1,
+        "max_new_tokens": 250,
+        "response_format": "wav",
+    }
+    for key, value in overrides.items():
+        if value is None:
+            body.pop(key, None)
+        else:
+            body[key] = value
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -71,7 +71,7 @@ def _default_oom_device_spec() -> str:
 
 OOM_INJECTION_CONFIG = {
     "device": _default_oom_device_spec(),
-    "target_mem_ratio": 0.92,
+    "target_mem_ratio": 0.95,
     "hold_seconds": 0,
     "startup_timeout_sec": 20,
     "strict": False,

tests/dfx/reliability/test_reliability_voxcpm2.py

@@ -1,508 +0,0 @@
-"""
-VoxCPM2 reliability integration tests.
-
-Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
-"""
-
-from __future__ import annotations
-
-import concurrent.futures
-import os
-import time
-from pathlib import Path
-from typing import Any
-
-import pytest
-
-from tests.dfx.conftest import (
-    assert_fault_exception,
-    create_reliability_omni_server_params,
-    resolve_oom_device_spec,
-)
-from tests.dfx.reliability.helpers import (
-    FaultInjector,
-    assert_no_server_tree_process_residual_and_gpu_release,
-    extract_openai_error_contract_from_bytes,
-    get_health_raw,
-    inject_gpu_oom,
-    make_process_kill_fault_injector,
-    make_server_root_kill_fault_injector,
-    make_server_tree_kill_fault_injector,
-    post_json_raw,
-    run_fault_injection_with_rate_load,
-    stop_gpu_oom_hogs,
-    worker_residual_timeout_after_kill_signal,
-)
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
-from tests.helpers.runtime import OpenAIClientHandler
-
-RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
-    {
-        "test_name": "voxcpm2_reliability_default",
-        "server_params": {
-            "model": "openbmb/VoxCPM2",
-            "stage_config_name": "voxcpm2.yaml",
-            # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
-            # would otherwise override deploy YAML settings during server startup.
-            "server_args": ["--trust-remote-code", "--disable-log-stats"],
-        },
-    },
-]
-
-DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
-OOM_INJECTION_CONFIG = {
-    "target_mem_ratio": 0.96,
-    "hold_seconds": 0,
-    "startup_timeout_sec": 20,
-    "strict": False,
-}
-FAULT_ERROR_KEYWORDS = (
-    "the request failed",
-    "oom",
-    "out of memory",
-    "cuda",
-    "orchestrator",
-    "timeout",
-    "connection",
-    "500",
-    "503",
-)
-PROCESS_KILL_ERROR_KEYWORDS = (
-    "timeout",
-    "did not complete within",
-    "connection",
-    "engine",
-    "orchestrator",
-    "dead",
-    "internal",
-    "500",
-    "503",
-)
-RUNTIME_WORKER_PATTERN = "VLLM::"
-SERVE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-SERVE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-WORKER_SIGNAL_FAULT_PARAMS = [
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGTERM",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigterm",
-    ),
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGKILL",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigkill",
-    ),
-]
-
-VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
-INFLIGHT_INJECTION_REQUEST_RATE = 0.3
-INFLIGHT_INJECTION_REQUEST_COUNT = 10
-_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
-_SPEECH_FF_MAX_ELAPSED_SEC = 30
-
-
-def _speech_request_config(omni_server: Any) -> dict[str, Any]:
-    return {
-        "model": omni_server.model,
-        "input": "The weather is nice today, perfect for a walk in the park.",
-        "stream": False,
-        "timeout": 120.0,
-        "response_format": "wav",
-        "voice": "default",
-    }
-
-
-def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
-    openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
-
-
-def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
-    base = "The weather is nice today, perfect for a walk in the park. "
-    return {
-        "model": omni_server.model,
-        "input": base * 80,
-        "max_new_tokens": 4096,
-        "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
-        "stream": False,
-        "timeout": 300.0,
-        "response_format": "wav",
-        "voice": "default",
-    }
-
-
-def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
-    openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
-
-
-def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
-    payload = {
-        "model": model,
-        "input": "post-fault fast-fail check",
-        "voice": "default",
-    }
-    start = time.monotonic()
-    try:
-        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
-        assert status >= 500, (
-            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
-        )
-    except Exception:  # noqa: BLE001
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
-
-
-def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            if status == 503:
-                return
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-            return
-        time.sleep(0.5)
-    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_request_failure(
-    omni_server_function,
-    openai_client_function,
-) -> None:
-    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
-    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
-    handle = inject_gpu_oom(
-        device=device_spec,
-        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
-        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
-        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
-        strict=OOM_INJECTION_CONFIG["strict"],
-    )
-    try:
-        try:
-            _submit_oom_speech_request(openai_client_function, omni_server_function)
-        except Exception as exc:
-            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
-        else:
-            pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
-    finally:
-        stop_gpu_oom_hogs(handle)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_error_contract(
-    omni_server_function,
-) -> None:
-    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
-    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
-    handle = inject_gpu_oom(
-        device=device_spec,
-        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
-        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
-        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
-        strict=OOM_INJECTION_CONFIG["strict"],
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    oom_config = _oom_speech_request_config(omni_server_function)
-    payload = {
-        "model": oom_config["model"],
-        "input": oom_config["input"],
-        "voice": oom_config["voice"],
-        "max_new_tokens": oom_config["max_new_tokens"],
-        "ref_audio": oom_config["ref_audio"],
-    }
-    try:
-        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
-    finally:
-        stop_gpu_oom_hogs(handle)
-
-    assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
-    speech_error = extract_openai_error_contract_from_bytes(body)
-    assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
-    assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_request_failure(
-    omni_server_after_fault_function,
-    openai_client_function,
-) -> None:
-    try:
-        _submit_speech_request(openai_client_function, omni_server_after_fault_function)
-    except Exception as exc:
-        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
-    else:
-        pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_fast_fail_and_concurrent(
-    omni_server_after_fault_function,
-    openai_client_function,
-) -> None:
-    host = omni_server_after_fault_function.host
-    port = omni_server_after_fault_function.port
-    model = omni_server_after_fault_function.model
-    payload = {
-        "model": model,
-        "input": "Say hello in one short sentence.",
-        "voice": "default",
-    }
-
-    start = time.monotonic()
-    try:
-        status, body = post_json_raw(
-            host,
-            port,
-            "/v1/audio/speech",
-            payload,
-            timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
-        )
-        elapsed = time.monotonic() - start
-        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
-            f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
-        )
-        assert status >= 500, (
-            "[process_kill fast_fail] expected server-side failure after fault, "
-            f"got status={status}, body={body[:200]!r}"
-        )
-    except Exception:
-        elapsed = time.monotonic() - start
-        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
-            f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
-        )
-
-    start = time.monotonic()
-    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
-        futures = [
-            executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
-            for _ in range(3)
-        ]
-        done, pending = concurrent.futures.wait(
-            futures,
-            timeout=30,
-            return_when=concurrent.futures.ALL_COMPLETED,
-        )
-
-    elapsed = time.monotonic() - start
-    assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
-    assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
-
-    fault_observed = False
-    for future in done:
-        try:
-            future.result()
-        except Exception as exc:
-            fault_observed = True
-            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
-    assert fault_observed, (
-        "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_worker_with_load_request_failure(
-    omni_server_function,
-    openai_client_function,
-    fault_injector: FaultInjector,
-) -> None:
-    scenario = "kill_worker_with_load"
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: fault_injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=300.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )

tests/dfx/stability/scripts/test_stability_voxcpm2.py

@@ -1,69 +0,0 @@
-"""
-VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
-
-Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from tests.dfx.conftest import (
-    create_benchmark_indices,
-    create_test_parameter_mapping,
-    create_unique_server_params,
-    load_configs,
-)
-from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
-
-STABILITY_DIR = Path(__file__).resolve().parent.parent
-DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
-CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
-DEFAULT_NUM_PROMPTS_PER_BATCH = 20
-STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
-
-try:
-    BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-except FileNotFoundError:
-    BENCHMARK_CONFIGS = []
-
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
-server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
-def test_stability_voxcpm2(omni_server, stability_benchmark_params):
-    test_name = stability_benchmark_params["test_name"]
-    params = stability_benchmark_params["params"]
-    duration_sec = params.get("duration_sec", 300)
-    num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
-    request_rate = params.get("request_rate")
-    max_concurrency = params.get("max_concurrency")
-
-    bench_params = {
-        k: v
-        for k, v in params.items()
-        if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
-    }
-
-    result = run_stability_benchmark_loop(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=omni_server.model,
-        duration_sec=duration_sec,
-        params=bench_params,
-        request_rate=request_rate,
-        max_concurrency=max_concurrency,
-        result_dir=str(STABILITY_DIR),
-        num_prompts_per_batch=num_prompts_per_batch,
-        run_one_batch=_run_one_vllm_bench_batch,
-    )
-
-    assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
-    assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"

tests/diffusion/diffusion_kv/test_initialization.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+import vllm_omni.diffusion.diffusion_kv.initialization as initialization
+import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Executor:
+    def __init__(self) -> None:
+        self.specs = [{"layer0": object()}, {"layer0": object()}]
+        self.memory = [1024, 1024]
+        self.configs = None
+        self.profile_requests = None
+
+    def get_kv_cache_specs(self):
+        return self.specs
+
+    def determine_available_kv_memory(self, profile_requests):
+        self.profile_requests = profile_requests
+        return self.memory
+
+    def set_kv_cache_configs(self, configs) -> None:
+        self.configs = configs
+
+
+def _od_config(**overrides):
+    values = dict(
+        model="test-diffusion",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        dtype=torch.bfloat16,
+        max_model_len=64,
+        quantization_config=None,
+        tf_model_config=None,
+        enforce_eager=True,
+        is_moe=False,
+        additional_config={},
+        profiler_config=None,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.9,
+        kv_cache_memory_bytes=None,
+        max_num_batched_tokens=64,
+        num_gpus=1,
+    )
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(num_gpus=2)
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=4096),
+        max_in_flight_tokens=128,
+    )
+    worker_configs = [object(), object()]
+    scheduler_kv_cache_config = object()
+    profile_requests = [object()]
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
+    build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
+    monkeypatch.setattr(
+        initialization,
+        "build_native_kv_cache_configs",
+        build,
+    )
+    monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
+
+    result = initialization.initialize_diffusion_kv_control_plane(
+        executor,
+        od_config,
+        profile_requests=profile_requests,
+        device=torch.device("cpu"),
+    )
+
+    build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
+    assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
+    assert executor.configs == worker_configs
+    assert executor.profile_requests is profile_requests
+
+
+def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
+    monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
+
+    result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
+
+    assert result is None
+    assert executor.configs is None
+    get_device.assert_not_called()
+
+
+def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
+    executor = _Executor()
+    executor.memory = [1024]
+    od_config = _od_config(num_gpus=2)
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
+
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        initialization.initialize_diffusion_kv_control_plane(
+            executor,
+            od_config,
+            profile_requests=[object()],
+            device=torch.device("cpu"),
+        )
+
+
+def test_paged_control_plane_requires_prepared_profile_request() -> None:
+    with pytest.raises(ValueError, match="prepared memory profile requests"):
+        initialization.initialize_diffusion_kv_control_plane(
+            _Executor(),
+            _od_config(),
+            device=torch.device("cpu"),
+        )
+
+
+def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    num_blocks = 32
+
+    worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * num_blocks],
+    )
+
+    assert len(worker_configs) == 1
+    assert worker_configs[0].num_blocks == num_blocks
+    assert scheduler_kv_cache_config.num_blocks == num_blocks
+    assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
+
+
+def test_native_control_plane_config_initializes_scheduler() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 32],
+    )
+
+    scheduler = RequestScheduler()
+    scheduler.initialize(
+        od_config,
+        kv_cache_config=scheduler_kv_cache_config,
+        scheduler_block_size=spec.block_size,
+        hash_block_size=spec.block_size,
+        kv_vllm_config=vllm_config,
+    )
+
+    assert scheduler._diffusion_kv_manager is not None
+    scheduler.close()
+
+
+def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
+    od_config = _od_config(
+        diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=4,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.5,
+        kv_cache_memory_bytes=4096,
+        max_num_batched_tokens=64,
+    )
+    vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
+    original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
+    original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
+    original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
+    original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
+
+    diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
+
+    assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
+    assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
+    assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
+    assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
+
+
+def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(gpu_memory_utilization=0.42),
+    )
+
+    assert vllm_config.cache_config.gpu_memory_utilization == 0.42
+
+
+def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
+    quantization_config = build_quant_config(
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "ignore": [],
+        }
+    )
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=quantization_config,
+        )
+    )
+
+    assert model_config.dtype is torch.bfloat16
+    assert model_config.quantization == "modelopt"
+    assert model_config.quantization_config is quantization_config
+    assert model_config.is_quantized
+    assert model_config.is_moe is False
+    assert model_config.is_nvfp4_quantized() is False
+    assert model_config.original_max_model_len == model_config.max_model_len
+
+    unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=None,
+        )
+    )
+    assert not unquantized_model_config.is_quantized
+
+
+def test_minus_one_max_model_len_uses_model_limit() -> None:
+    od_config = SimpleNamespace(
+        max_model_len=-1,
+        tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+    )
+
+    assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
+
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        )
+    )
+    assert model_config.max_model_len == 4096
+    assert model_config.original_max_model_len == -1
+
+
+def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        ),
+    )
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+
+    initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 16],
+    )
+
+    assert vllm_config.model_config.original_max_model_len == -1
+    assert vllm_config.model_config.max_model_len == 256

tests/diffusion/diffusion_kv/test_manager.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+BLOCK_SIZE = 4
+
+
+def _config(num_blocks: int) -> KVCacheConfig:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+
+
+def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+    return DiffusionKVRequest(
+        f"{public_id}/diffusion-kv/{sequence_id}",
+        sequence_id=sequence_id,
+        prefix_len=4,
+        target_len=4,
+        seq_len=seq_len,
+        kv_contexts=kv_contexts,
+    )
+
+
+def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+    return DiffusionKVCacheManager(
+        _config(num_blocks),
+        max_model_len=max_model_len,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+    )
+
+
+def test_reserve_and_free_multi_cfg_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    requests = (_request("public", 0), _request("public", 1))
+
+    metadata = manager.reserve_request("public", requests)
+
+    assert metadata is not None
+    assert metadata.request_id == "public"
+    assert metadata.allocation_generation == 1
+    assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
+    assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+    assert all(sequence.context_ids == () for sequence in metadata.sequences)
+    assert metadata.contexts == ()
+    assert metadata == manager.get_metadata("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
+
+    manager.free_request("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
+    manager = _manager(3)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
+        manager.reserve_request(
+            "impossible",
+            tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
+        )
+
+    assert manager.has_request("impossible") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_temporary_capacity_pressure_returns_none() -> None:
+    manager = _manager(3)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+
+    allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
+
+    assert allocation is None
+    assert manager.has_request("waiting") is False
+
+
+def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    waiting = (_request("waiting", 0), _request("waiting", 1))
+
+    assert manager.reserve_request("waiting", waiting) is None
+    assert manager.has_request("waiting") is False
+
+    manager.free_request("running")
+    metadata = manager.reserve_request("waiting", waiting)
+
+    assert metadata is not None
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+
+
+def test_reallocated_request_gets_a_new_allocation_generation() -> None:
+    manager = _manager(3)
+    request = (_request("public", 0),)
+
+    first = manager.reserve_request("public", request)
+    assert first is not None
+    manager.free_request("public")
+    second = manager.reserve_request("public", request)
+
+    assert second is not None
+    assert second.allocation_generation == first.allocation_generation + 1
+
+
+def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    native_allocate = manager.native_manager.allocate_slots
+    num_calls = 0
+
+    def fail_second_allocation(*args, **kwargs):
+        nonlocal num_calls
+        num_calls += 1
+        if num_calls == 2:
+            raise RuntimeError("injected allocation failure")
+        return native_allocate(*args, **kwargs)
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
+
+    with pytest.raises(RuntimeError, match="injected allocation failure"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
+    manager = _manager(8)
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0, kv_contexts=(context,)),),
+        )
+
+
+def test_rejects_sequence_longer_than_admission_bound() -> None:
+    manager = _manager(8, max_model_len=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
+        manager.reserve_request("public", (_request("public", 0, seq_len=9),))
+
+
+@pytest.mark.parametrize(
+    "requests",
+    [
+        (_request("public", 1),),
+        (_request("public", 1), _request("public", 0)),
+        (_request("public", 0), _request("public", 2)),
+    ],
+)
+def test_rejects_ambiguous_sequence_order(requests) -> None:
+    manager = _manager(8)
+
+    with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
+        manager.reserve_request("public", requests)
+
+
+def test_close_releases_every_public_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    assert manager.reserve_request("first", (_request("first", 0),)) is not None
+    assert manager.reserve_request("second", (_request("second", 0),)) is not None
+
+    manager.close()
+
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+    assert manager.has_request("first") is False
+    assert manager.has_request("second") is False

tests/diffusion/diffusion_kv/test_request.py

@@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None:
 def _manager(*, enable_caching: bool) -> KVCacheManager:
     # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
     # already registered an out-of-tree spec, making the registry non-empty
-    # before FullAttentionSpec is registered. Initialize the native mappings
-    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
-    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
-    if register_all is not None:
-        register_all(None)
+    # before FullAttentionSpec is registered. Initialize the native mappings.
+    native_kv_managers.register_all_kvcache_specs(None)
 
     spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Backend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def _attention(*, enabled: bool) -> Attention:
+    attention = Attention.__new__(Attention)
+    nn.Module.__init__(attention)
+    attention.paged_kv_cache_role = "primary" if enabled else None
+    attention.paged_kv_cache_dtype = torch.bfloat16
+    attention.num_kv_heads = 2
+    attention.head_size = 8
+    attention.causal = False
+    attention.attn_backend = _Backend
+    return attention
+
+
+def _runner(attention: Attention) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=4),
+        model_config=SimpleNamespace(dtype=torch.float16),
+    )
+    pipeline = nn.Module()
+    pipeline.image_attention = attention
+    runner.pipeline = pipeline
+    return runner
+
+
+def test_runner_discovers_native_spec_from_loaded_attention() -> None:
+    specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
+
+    assert set(specs) == {"image_attention"}
+    spec = specs["image_attention"]
+    assert isinstance(spec, FullAttentionSpec)
+    assert spec.block_size == 4
+    assert spec.num_kv_heads == 2
+    assert spec.head_size == 8
+    assert spec.dtype is torch.bfloat16
+    assert spec.indexes_kv_by_block_stride is True
+    assert spec.non_causal is True
+
+
+def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
+    runner = _runner(_attention(enabled=False))
+
+    with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
+        runner.get_kv_cache_spec()
+
+
+def test_runner_retains_matching_rank_local_config() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
+    )
+
+    runner.set_kv_cache_config(config)
+
+    assert runner.kv_cache_config is config
+
+
+def test_runner_rejects_rank_local_config_for_different_layers() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
+    )
+
+    with pytest.raises(ValueError, match="layer mismatch"):
+        runner.set_kv_cache_config(config)
+
+
+def test_worker_selects_its_rank_local_config() -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 1
+    worker.od_config = SimpleNamespace(num_gpus=2)
+    worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
+    configs = [object(), object()]
+
+    worker.set_kv_cache_configs(configs)
+
+    assert worker.installed is configs[1]
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        worker.set_kv_cache_configs(configs[:1])
+
+
+def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
+    worker.init_snapshot = SimpleNamespace(free_memory=16384)
+    worker.requested_memory = 8192
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [4096]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=0,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=400,
+        after_profile=SimpleNamespace(free_memory=500),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+@pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
+def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=non_kv_cache_memory,
+        after_profile=SimpleNamespace(free_memory=100),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
+            f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
+        ),
+    ):
+        worker.determine_available_kv_memory(profile_request)
+
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=1200)
+    worker.requested_memory = 1000
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.8,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=650,
+        after_profile=SimpleNamespace(free_memory=550),
+    )
+    profiling_calls = []
+
+    def profile(snapshot, *, weights_memory):
+        profiling_calls.append((snapshot, weights_memory))
+        return nullcontext(profile_result)
+
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        profile,
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    assert profiling_calls == [(worker.init_snapshot, 250)]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.init_snapshot = None
+    worker.requested_memory = None
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
+    worker.model_runner = SimpleNamespace()
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+
+    with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
+        worker.determine_available_kv_memory(object())
+
+
+def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
+    gathered = []
+
+    def gather(local_result):
+        gathered.append(local_result)
+        return [local_result, (True, "peer-result")]
+
+    def fail_local_probe():
+        raise ValueError("local failure")
+
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
+
+    with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
+        diffusion_worker_module._run_and_gather_rank_values(
+            "test probe",
+            fail_local_probe,
+        )
+
+    assert gathered == [(False, "ValueError: local failure")]

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -85,19 +85,15 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
     assert new_req.diffusion_kv_metadata is metadata
 
 
-def test_paged_request_without_metadata_reaches_request_forward() -> None:
+def test_paged_request_without_metadata_fails_before_request_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = DiffusionOutput(output=None)
-    runner._execute_request_list = Mock(
-        return_value=SimpleNamespace(
-            runner_outputs=[SimpleNamespace(result=expected_output)],
-        )
-    )
+    runner._execute_request_list = Mock()
     req = SimpleNamespace(request_id="req-0")
 
-    assert runner.execute_model(req) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model(req)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_dense_request_rejects_metadata_before_request_forward() -> None:
@@ -141,18 +137,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
     assert calls == []
 
 
-def test_batch_path_allows_missing_metadata_before_forward() -> None:
+def test_batch_path_rejects_missing_metadata_before_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = object()
-    runner._execute_request_list = Mock(return_value=expected_output)
+    runner._execute_request_list = Mock()
     new_req = NewRequestData(
         request_id="req-0",
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -170,7 +166,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
     runner._execute_request_list.assert_not_called()
 
 
-def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
+def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.pipeline = object()
     runner._supports_step_mode = Mock(return_value=False)
@@ -179,10 +175,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    with pytest.raises(ValueError, match="does not support step execution"):
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
         runner.execute_stepwise(make_scheduler_output(new_req))
 
-    runner._supports_step_mode.assert_called_once()
+    runner._supports_step_mode.assert_not_called()
 
 
 def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -288,3 +284,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
     worker.execute_model(req, SimpleNamespace())
 
     assert calls == [(req, None)]
+
+
+def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
+    calls: list[tuple] = []
+
+    class ModelRunner:
+        def execute_model(self, req, **kwargs):
+            calls.append((req, kwargs))
+            return DiffusionOutput(output=None)
+
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = ModelRunner()
+    worker.lora_manager = None
+    worker._get_profiler = lambda: None
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
+        lambda: 1,
+    )
+    req_0 = SimpleNamespace(request_id="req-0")
+    req_1 = SimpleNamespace(request_id="req-1")
+    metadata_0 = make_metadata("req-0")
+    metadata_1 = make_metadata("req-1")
+    envelopes = [
+        NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
+        NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
+    ]
+
+    result = worker.execute_model(envelopes, SimpleNamespace())
+
+    assert calls == [
+        (
+            req_1,
+            {
+                "kv_prefetch_job": None,
+                "diffusion_kv_metadata": metadata_1,
+            },
+        )
+    ]
+    assert result["dp_rank"] == 1

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -74,14 +74,6 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
 
 
-@pytest.fixture(autouse=True)
-def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
-        lambda: None,
-    )
-
-
 @pytest.fixture(autouse=True)
 def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -229,9 +228,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -259,6 +257,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     ]
 
 
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
+    """A destroyed diffusion PP group must not block the next initialization."""
+
+    class _DestroyableGroup:
+        def __init__(self) -> None:
+            self.destroy_calls = 0
+
+        def destroy(self) -> None:
+            self.destroy_calls += 1
+
+    pipeline_group = _DestroyableGroup()
+    for name in ("_DP", "_CFG", "_SP"):
+        monkeypatch.setattr(parallel_state, name, None)
+    monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
+    monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
+    for name in ("_DP", "_PCP", "_TP", "_EP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+    monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
+
+    parallel_state.destroy_model_parallel()
+
+    assert pipeline_group.destroy_calls == 1
+    assert parallel_state.vllm_parallel_state._PP is None
+
+
 @pytest.mark.cpu
 @pytest.mark.core_model
 def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
@@ -295,9 +320,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)

tests/diffusion/distributed/test_hsdp.py

@@ -234,7 +234,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_tp(self):
         """Test that HSDP and Tensor Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with TP"):
             DiffusionParallelConfig(
                 tensor_parallel_size=2,
                 use_hsdp=True,
@@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_dp(self):
         """Test that HSDP and Data Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with DP"):
             DiffusionParallelConfig(
                 data_parallel_size=2,
                 use_hsdp=True,
@@ -264,158 +264,6 @@ class TestDiffusionParallelConfigHSDP:
         assert config.hsdp_shard_size == 2  # auto: 4 // 2
 
 
-class TestStandaloneHSDPDetection:
-    """Tests for standalone HSDP detection and dit_parallel_size calculation.
-
-    These tests verify the logic used in initialize_model_parallel() to detect
-    standalone HSDP mode and calculate effective parallel sizes.
-
-    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
-    """
-
-    @staticmethod
-    def compute_standalone_hsdp_params(
-        data_parallel_size: int = 1,
-        cfg_parallel_size: int = 1,
-        sequence_parallel_size: int = 1,
-        pipeline_parallel_size: int = 1,
-        tensor_parallel_size: int = 1,
-        fully_shard_degree: int = 1,
-        hsdp_replicate_size: int = 1,
-    ) -> dict:
-        """Compute standalone HSDP detection parameters.
-
-        This mirrors the logic in initialize_model_parallel().
-        """
-        dit_parallel_size = (
-            data_parallel_size
-            * cfg_parallel_size
-            * sequence_parallel_size
-            * pipeline_parallel_size
-            * tensor_parallel_size
-        )
-
-        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
-        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
-
-        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
-        if is_standalone_hsdp:
-            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
-        else:
-            effective_dit_parallel_size = dit_parallel_size
-
-        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
-
-        return {
-            "original_dit_parallel_size": dit_parallel_size,
-            "is_standalone_hsdp": is_standalone_hsdp,
-            "effective_dit_parallel_size": effective_dit_parallel_size,
-            "effective_dp_size": effective_dp_size,
-        }
-
-    def test_standalone_hsdp_basic(self):
-        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 4
-        assert result["effective_dp_size"] == 4
-
-    def test_standalone_hsdp_with_replicate(self):
-        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=2,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
-        assert result["effective_dp_size"] == 8
-
-    def test_combined_hsdp_sp_not_standalone(self):
-        """Test HSDP combined with SP is NOT detected as standalone.
-
-        This is a regression test for the bug where the condition
-        `dit_parallel_size == fully_shard_degree` incorrectly matched
-        combined modes like SP=4 + HSDP=4.
-        """
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=4,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4
-        assert result["is_standalone_hsdp"] is False
-        # Should NOT override dp_size for combined mode
-        assert result["effective_dp_size"] == 1  # original data_parallel_size
-
-    def test_combined_hsdp_cfg_not_standalone(self):
-        """Test HSDP combined with CFG is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_hsdp_pp_not_standalone(self):
-        """Test HSDP combined with PP is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            pipeline_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_no_hsdp_not_standalone(self):
-        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_multiple_parallelism_not_standalone(self):
-        """Test HSDP combined with multiple parallelism is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=2,
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4  # 2 * 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_standalone_hsdp_large_shard(self):
-        """Test standalone HSDP with large shard size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=8,
-            hsdp_replicate_size=1,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8
-        assert result["effective_dp_size"] == 8
-
-    def test_standalone_hsdp_large_replicate(self):
-        """Test standalone HSDP with large replicate size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=4,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
-        assert result["effective_dp_size"] == 16
-
-
 class TestHSDPShardConditions:
     """Tests for _hsdp_shard_conditions matching logic."""
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
@@ -91,6 +91,17 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
     return loader
 
 
+def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="",
+        dtype=torch.float32,
+        quantization_config="fp8",
+        parallel_config=DiffusionParallelConfig(data_parallel_size=2),
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+    )
+
+
 def test_strict_check_only_validates_source_prefix_parameters():
     model = _DummyPipelineModel(source_prefix="transformer.")
     loader = _make_loader_with_weights(["transformer.weight"])
@@ -394,6 +405,78 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    class UnsupportedOnlineMethod:
+        uses_meta_device = True
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnsupportedOnlineMethod()
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    with pytest.raises(ValueError, match="per-tensor FP8 linears"):
+        loader.load_model(load_device="cpu")
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/model_loader/test_fp8_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the Phase-I normalized online-FP8 cache contract."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    CACHE_COMPLETE_NAME,
+    CACHE_COMPONENT_SUBFOLDER,
+    CACHE_MANIFEST_NAME,
+    _source_fingerprint,
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _write_valid_cache(tmp_path):
+    source = tmp_path / "source.safetensors"
+    save_file({"weight": torch.ones(2, 2)}, str(source))
+
+    cache = tmp_path / "cache"
+    component = cache / CACHE_COMPONENT_SUBFOLDER
+    component.mkdir(parents=True)
+    shard = component / "part-00000.safetensors"
+    save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard))
+    (component / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": 4},
+                "weight_map": {"weight": shard.name},
+            }
+        )
+    )
+
+    manifest = {
+        "schema_version": 1,
+        "cache_kind": "normalized_fp8_runtime",
+        "component": "transformer",
+        "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+        "source_prefix": "transformer.",
+        "quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"},
+        "tensor_count": 1,
+        "source_files": [str(source)],
+        "source_fingerprint": _source_fingerprint({str(source)}),
+    }
+    cache.mkdir(exist_ok=True)
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+    (cache / CACHE_COMPLETE_NAME).write_text("phase1\n")
+    return cache, source, manifest
+
+
+def test_online_and_serialized_fp8_configs_are_classified():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    assert is_online_fp8_config(Fp8Config())
+    assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True))
+
+    cache_config = _transformer_quant_config()
+    assert not is_online_fp8_config(cache_config)
+    assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized
+    assert cache_config.resolve("text_encoder") is None
+
+
+def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path):
+    cache, source, expected = _write_valid_cache(tmp_path)
+
+    assert load_fp8_cache_manifest(cache) == expected
+
+    source.write_bytes(source.read_bytes() + b"changed")
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+
+    manifest["quantization"]["serialized"] = False
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+    index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / CACHE_MANIFEST_NAME).write_text("[]")
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path):
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    source_file = source_dir / "model.safetensors"
+    save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file))
+
+    model = nn.Module()
+    model._dit_modules = ("transformer",)
+    model.transformer = nn.Module()
+    model.transformer.block = nn.Module()
+    model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False)
+    model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+    model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+
+    source = SimpleNamespace(
+        model_or_path=str(source_dir),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    cache_dir = tmp_path / "cache"
+
+    manifest = build_normalized_fp8_cache(
+        model,
+        sources=(source,),
+        cache_dir=cache_dir,
+    )
+
+    assert manifest["cache_kind"] == "normalized_fp8_runtime"
+    tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors"))
+    scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors"))
+    assert tensors["block.weight"].dtype == torch.float8_e4m3fn
+    assert tuple(tensors["block.weight"].shape) == (3, 2)
+    assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5]))
+
+
+@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME])
+def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / missing).unlink()
+
+    assert load_fp8_cache_manifest(cache) is None

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -118,6 +118,14 @@ def _call_mgr(
     )
 
 
+def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
+    mgr = _make_cache_mgr()
+
+    assert isinstance(mgr, nn.Module)
+    assert dict(mgr.named_modules())["attn"] is mgr.attn
+    assert mgr.state_dict() == {}
+
+
 # ============================================================
 # Test 1: No AR KV — basic cache → reuse
 # ============================================================

tests/diffusion/models/pi0/__init__.py

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

tests/diffusion/models/pi0/test_pi0_parity.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle.
+
+Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit
+matching action chunks with LeRobot's ``PI0Policy`` when fed:
+  * the same weights (``lerobot/pi0_base``)
+  * the same pre-processed inputs (images, masks, tokens, state)
+  * the same initial noise tensor
+
+π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed
+``num_steps``), so the output is deterministic once the noise is fixed and
+``torch.allclose`` on the final action chunk is a valid oracle. This is the
+authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4).
+
+The OpenPI websocket online-serving e2e lives separately in
+``tests/e2e/online_serving/test_pi0_expansion.py``.
+
+Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni
+env), with the vllm-omni ``pi0`` package importable::
+
+    python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s
+
+Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env).
+The run is CPU/float32 with fixed defaults; the only override is
+``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the
+HF download; defaults to ``lerobot/pi0_base``).
+"""
+
+from __future__ import annotations
+
+import copy
+import importlib.util
+import os
+
+import pytest
+import torch
+
+# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it
+# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md).
+# Additionally gated on lerobot being importable.
+_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+
+# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ──
+DEVICE = "cpu"
+DTYPE_STR = "float32"
+ATOL = 1e-4
+NUM_STEPS = 10
+BATCH_SIZE = 2
+ACTION_DIM = 32
+STATE_DIM = 32
+ACTION_HORIZON = 50
+MAX_TOKEN_LEN = 48
+
+# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF
+# download; defaults to the HF repo id.
+MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base")
+
+
+def _resolve_checkpoint_dir() -> str:
+    """Return a local dir containing the pi0_base checkpoint (download if needed)."""
+    if os.path.isdir(MODEL_PATH):
+        return MODEL_PATH
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
+
+
+# ─── Dummy dataset stats (identity) ──────────────────────────────────
+def _dummy_dataset_stats() -> dict:
+    return {
+        "observation.state": {
+            "mean": torch.zeros(STATE_DIM),
+            "std": torch.ones(STATE_DIM),
+            "q01": torch.zeros(STATE_DIM),
+            "q99": torch.ones(STATE_DIM),
+        },
+        "action": {
+            "mean": torch.zeros(ACTION_DIM),
+            "std": torch.ones(ACTION_DIM),
+            "q01": torch.zeros(ACTION_DIM),
+            "q99": torch.ones(ACTION_DIM),
+        },
+        "images": {
+            cam: {
+                "mean": torch.zeros(3, 224, 224),
+                "std": torch.ones(3, 224, 224),
+                "q01": torch.zeros(3, 224, 224),
+                "q99": torch.ones(3, 224, 224),
+            }
+            for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+        },
+    }
+
+
+def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict:
+    """Reproducible dummy inputs — identical across both implementations."""
+    g = torch.Generator(device="cpu").manual_seed(0)
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
+        "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
+        "observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+            device
+        ),
+        "observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+            device
+        ),
+        "observation.images.right_wrist_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, generator=g, dtype=torch.float32
+        ).to(device),
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+# ─── LeRobot instantiation ────────────────────────────────────────────
+def _instantiate_lerobot():
+    from lerobot.policies.pi0 import PI0Policy
+    from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+
+    policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True)
+    policy.to(DEVICE)
+    policy.config.device = DEVICE
+    policy.eval()
+
+    pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
+    return policy, pre, post
+
+
+# ─── vllm-omni instantiation ──────────────────────────────────────────
+def _instantiate_vllm_omni():
+    """Build the vllm-omni π0 model in isolation (no pipeline, no engine)."""
+    from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction
+
+    cfg = Pi0Config(
+        max_action_dim=ACTION_DIM,
+        max_state_dim=STATE_DIM,
+        chunk_size=ACTION_HORIZON,
+        num_inference_steps=NUM_STEPS,
+        dtype=DTYPE_STR,
+    )
+    model = Pi0ForActionPrediction(cfg)
+    model.to(DEVICE).eval()
+    _load_lerobot_weights(model)
+    return model
+
+
+def _load_lerobot_weights(model):
+    """Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model.
+
+    ``Pi0ForActionPrediction.load_weights`` handles the leading ``model.``
+    prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass
+    the raw checkpoint dict.
+    """
+    import safetensors.torch
+
+    cache_dir = _resolve_checkpoint_dir()
+    path = os.path.join(cache_dir, "model.safetensors")
+    state = safetensors.torch.load_file(path)
+    model.load_weights(list(state.items()))
+
+
+# ─── Helpers to extract LeRobot's pre-processed inputs ────────────────
+def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
+    """Mimic what ``PI0Policy.predict_action_chunk`` feeds into
+    ``self.model.sample_actions``. We use these *exact* tensors for vllm-omni
+    so any divergence must come from the core model, not preprocessing.
+    """
+    images, img_masks = lerobot_policy._preprocess_images(processed_batch)
+    from lerobot.utils.constants import (
+        OBS_LANGUAGE_ATTENTION_MASK,
+        OBS_LANGUAGE_TOKENS,
+    )
+
+    lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS]
+    lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
+    state = lerobot_policy.prepare_state(processed_batch)
+    return images, img_masks, lang_tokens, lang_masks, state
+
+
+# ─── Shared fixed-noise sampler ───────────────────────────────────────
+def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
+    g = torch.Generator(device="cpu").manual_seed(42)
+    return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
+
+
+# ─── Main test ────────────────────────────────────────────────────────
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi0_vllm_omni_vs_lerobot():
+    print("\n[parity] Instantiating LeRobot…")
+    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
+
+    print("[parity] Instantiating vllm-omni…")
+    omni_model = _instantiate_vllm_omni()
+
+    print("[parity] Preparing shared inputs…")
+    raw_batch = _create_dummy_batch()
+    processed_batch = lerobot_pre(copy.deepcopy(raw_batch))
+    images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch)
+    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+    print(f"[parity] state.shape={state.shape}  lang_tokens.shape={lang_tokens.shape}")
+    print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})")
+    print(f"[parity] noise.shape={noise.shape}  dtype={noise.dtype}")
+
+    # ── LeRobot forward ──
+    print("[parity] Running LeRobot sample_actions…")
+    with torch.no_grad():
+        lerobot_actions = lerobot_policy.model.sample_actions(
+            images,
+            img_masks,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+    print(
+        f"[parity] LeRobot actions: shape={lerobot_actions.shape} "
+        f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}"
+    )
+
+    # ── vllm-omni forward ──
+    print("[parity] Running vllm-omni sample_actions…")
+    with torch.no_grad():
+        omni_actions = omni_model.sample_actions(
+            images=images,
+            image_masks=img_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            state=state,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+    print(
+        f"[parity] vllm-omni actions: shape={omni_actions.shape} "
+        f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}"
+    )
+
+    # ── Compare ──
+    diff = (lerobot_actions.float() - omni_actions.float()).abs()
+    max_diff = diff.max().item()
+    mean_diff = diff.mean().item()
+    print(f"[parity] |Δ| max={max_diff:.2e}  mean={mean_diff:.2e}  atol={ATOL:.1e}")
+    close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL)
+    print(f"[parity] torch.allclose(atol={ATOL}): {close}")
+
+    if not close:
+        print("\n[parity] ⚠️  Outputs diverge — running per-stage diagnostics…")
+        _diagnose_divergence(
+            lerobot_policy.model,
+            omni_model,
+            images,
+            img_masks,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise,
+        )
+
+    assert close, (
+        f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e}  mean_diff={mean_diff:.2e}"
+    )
+
+
+# ─── Per-stage divergence diagnostics ─────────────────────────────────
+@torch.no_grad()
+def _diagnose_divergence(
+    lerobot_flow_model,
+    omni_model,
+    images,
+    img_masks,
+    lang_tokens,
+    lang_masks,
+    state,
+    noise,
+):
+    """Localize a numerical mismatch to a specific pipeline stage:
+    1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
+    2. Prefix KV cache layer 0 — PaliGemma LM attention.
+    3. A single denoise_step velocity at t=1.0 — action expert forward.
+    """
+    from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+        make_att_2d_masks,
+        prepare_attention_masks_4d,
+    )
+
+    # ── Stage 1: prefix embeddings ──
+    lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix(
+        images, img_masks, lang_tokens, lang_masks
+    )
+    sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+    total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item()
+    print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e}   (shape={tuple(sg_prefix_embs.shape)})")
+    print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}")
+    print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}")
+
+    num_cams = len(images)
+    img_len = 256 * num_cams
+    lang_len = lr_prefix_embs.shape[1] - img_len
+    img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item()
+    lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item()
+    print(f"[diag]   image slice [:{img_len}] max|Δ| = {img_diff:.2e}   (num_cams={num_cams})")
+    print(f"[diag]   lang slice  [{img_len}:] max|Δ| = {lang_diff:.2e}   (len={lang_len})")
+
+    def _stats(name, t):
+        t = t.float()
+        print(
+            f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} "
+            f"std={t.std().item():.4f} min={t.min().item():+.2f} "
+            f"max={t.max().item():+.2f}"
+        )
+
+    _stats("LeRobot ", lr_prefix_embs)
+    _stats("vllm-omni", sg_prefix_embs)
+
+    # ── Stage 2: prefix KV cache ──
+    prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att)
+    prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1
+    prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d)
+
+    lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager"
+    _, lr_kv = lerobot_flow_model.paligemma_with_expert.forward(
+        attention_mask=prefix_att_4d,
+        position_ids=prefix_pos,
+        past_key_values=None,
+        inputs_embeds=[lr_prefix_embs, None],
+        use_cache=True,
+    )
+    _, sg_kv = omni_model.paligemma_with_expert.forward(
+        attention_mask=prefix_att_4d,
+        position_ids=prefix_pos,
+        past_key_values=None,
+        inputs_embeds=[sg_prefix_embs, None],
+        use_cache=True,
+    )
+
+    def _layer0_kv(cache):
+        if isinstance(cache, list):
+            return cache[0]
+        layer0 = cache.layers[0]
+        return layer0.keys, layer0.values
+
+    try:
+        lr_k, lr_v = _layer0_kv(lr_kv)
+        sg_k, sg_v = _layer0_kv(sg_kv)
+        dk = (lr_k.float() - sg_k.float()).abs().max().item()
+        dv = (lr_v.float() - sg_v.float()).abs().max().item()
+        print(f"[diag] prefix KV layer0  K max|Δ|={dk:.2e}  V max|Δ|={dv:.2e}")
+    except Exception as e:  # noqa: BLE001
+        print(f"[diag] could not extract prefix KV for comparison: {e}")
+
+    # ── Stage 3: a single denoise_step at t=1.0 ──
+    bsize = state.shape[0]
+    t = torch.ones(bsize, dtype=torch.float32, device=state.device)
+    lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t)
+    sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t)
+    print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}")
+
+
+if __name__ == "__main__":
+    test_pi0_vllm_omni_vs_lerobot()

tests/diffusion/models/pi0/test_pi0_units.py

@@ -0,0 +1,585 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU unit tests for the π0 (Pi-Zero) VLA model.
+
+Deliberately lightweight: pure utility code (config plumbing, mask building,
+sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
+that runs on CPU with no weights and no lerobot. The full LeRobot parity check
+lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
+
+    pytest tests/diffusion/models/pi0/test_pi0_units.py -v
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+    OPENPI_ATTENTION_MASK_VALUE,
+    Pi0ForActionPrediction,
+    _apply_norm,
+    _build_norm_buffers,
+    create_sinusoidal_pos_embedding,
+    get_gemma_config,
+    make_att_2d_masks,
+    prepare_attention_masks_4d,
+)
+from vllm_omni.diffusion.models.pi0.processor_pi0 import (
+    Pi0ImageProcessor,
+    build_model_inputs,
+    pil_image_to_tensor,
+    resize_with_pad,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ----------------------------------------------------------------------------
+# Pi0Config dataclass + checkpoint resolver
+# ----------------------------------------------------------------------------
+_LEROBOT_CFG = {
+    "type": "pi0",
+    "paligemma_variant": "gemma_2b",
+    "action_expert_variant": "gemma_300m",
+    "chunk_size": 50,
+    "max_action_dim": 32,
+    "max_state_dim": 32,
+    "num_inference_steps": 10,
+    "image_resolution": [224, 224],
+    "tokenizer_max_length": 48,
+    "dtype": "float32",
+    # Training-only keys that must be filtered out of the dataclass.
+    "optimizer_lr": 2.5e-5,
+    "scheduler_warmup_steps": 1000,
+    "input_features": {
+        "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.state": {"type": "STATE", "shape": [32]},
+    },
+}
+
+_EXPECTED_CAMERA_ORDER = [
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+]
+
+
+def test_config_camera_order_from_input_features():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
+
+
+def test_config_filters_training_only_keys():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert not hasattr(c, "optimizer_lr")
+    assert not hasattr(c, "scheduler_warmup_steps")
+
+
+def test_config_coerces_image_resolution_to_tuple():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert c.image_resolution == (224, 224)
+    assert isinstance(c.image_resolution, tuple)
+
+
+def test_config_non_square_resolution_raises():
+    with pytest.raises(ValueError):
+        Pi0Config(image_resolution=(224, 256))
+
+
+@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
+def test_config_malformed_resolution_raises(bad):
+    with pytest.raises(ValueError):
+        Pi0Config(image_resolution=bad)
+
+
+def test_config_defaults_and_empty():
+    c = Pi0Config.from_model_config(None)
+    assert c.chunk_size == 50
+    assert c.max_action_dim == 32
+    assert c.max_state_dim == 32
+    assert c.num_inference_steps == 10
+    assert c.image_resolution == (224, 224)
+    assert c.tokenizer_max_length == 48
+    assert c.max_cameras == 3
+    assert c.norm_stats is None
+    assert c.image_key_map == {}
+
+
+def test_config_deploy_yaml_override():
+    """Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
+    c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
+    assert c.dtype == "bfloat16"
+    assert c.num_inference_steps == 4
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: sinusoidal timestep embedding
+# ----------------------------------------------------------------------------
+def test_sinusoidal_embedding_shape():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
+    assert emb.shape == (3, 64)
+
+
+def test_sinusoidal_embedding_values_at_zero():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
+    assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
+    assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
+
+
+def test_sinusoidal_embedding_odd_dim_raises():
+    with pytest.raises(ValueError):
+        create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
+
+
+def test_sinusoidal_embedding_non_1d_raises():
+    with pytest.raises(ValueError):
+        create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
+
+
+def test_sinusoidal_embedding_distinct_timesteps_differ():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
+    for i in range(emb.shape[0]):
+        for j in range(i + 1, emb.shape[0]):
+            assert not torch.allclose(emb[i], emb[j])
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: attention mask construction
+# ----------------------------------------------------------------------------
+def test_make_att_2d_masks_bidirectional():
+    result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
+    assert result.all()
+
+
+def test_make_att_2d_masks_causal():
+    result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
+    expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
+    assert torch.equal(result[0], expected)
+
+
+def test_make_att_2d_masks_prefix_lm():
+    result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
+    assert result[0, 0, 0]
+    assert not result[0, 0, 3]
+    assert result[0, 3, 0]
+    assert not result[0, 3, 4]
+
+
+def test_make_att_2d_masks_respects_padding():
+    pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool)  # last token is padding
+    att = torch.zeros(1, 4, dtype=torch.long)  # full bidirectional
+    result = make_att_2d_masks(pad, att)
+    assert not result[0, :, 3].any(), "Nothing should attend to padded key"
+    assert not result[0, 3, :].any(), "Padded query should attend to nothing"
+
+
+def test_prepare_attention_masks_4d():
+    mask_2d = torch.tensor([[[True, False], [True, True]]])
+    mask_4d = prepare_attention_masks_4d(mask_2d)
+    assert mask_4d.shape == (1, 1, 2, 2)
+    assert mask_4d[0, 0, 0, 0].item() == 0.0
+    assert mask_4d[0, 0, 1, 1].item() == 0.0
+    expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
+    assert torch.equal(mask_4d[0, 0, 0, 1], expected)
+
+
+def test_attention_mask_value_finite_and_negative():
+    assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
+    assert OPENPI_ATTENTION_MASK_VALUE < -1e30
+
+
+# ----------------------------------------------------------------------------
+# Normalization round-trips
+# ----------------------------------------------------------------------------
+def test_build_norm_buffers_missing():
+    assert _build_norm_buffers(None, "state") is None
+    assert _build_norm_buffers({}, "state") is None
+    assert _build_norm_buffers({"state": None}, "state") is None
+
+
+def test_build_norm_buffers_mean_std():
+    stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
+    out = _build_norm_buffers(stats, "state")
+    assert out["mode"] == "mean_std"
+    assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
+    assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
+
+
+def test_build_norm_buffers_min_max():
+    stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
+    out = _build_norm_buffers(stats, "action")
+    assert out["mode"] == "min_max"
+    assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
+    assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
+
+
+def test_apply_norm_mean_std_forward_inverse():
+    stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
+    x = torch.tensor([[1.0, 0.3]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
+    x_back = _apply_norm(z, stats, inverse=True)
+    assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_min_max_forward_inverse():
+    stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
+    x = torch.tensor([[0.25, 0.5]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
+    x_back = _apply_norm(z, stats, inverse=True)
+    assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_preserves_padded_tail():
+    stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
+    x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z[:, :2], x[:, :2])
+    assert torch.allclose(z[:, 2:], x[:, 2:])
+
+
+def test_apply_norm_noop_when_stats_none():
+    x = torch.randn(2, 5)
+    assert torch.equal(_apply_norm(x, None, inverse=False), x)
+    assert torch.equal(_apply_norm(x, None, inverse=True), x)
+
+
+# ----------------------------------------------------------------------------
+# Gemma variant configs
+# ----------------------------------------------------------------------------
+def test_gemma_2b_dims():
+    c = get_gemma_config("gemma_2b")
+    assert c.width == 2048
+    assert c.depth == 18
+
+
+def test_gemma_300m_dims():
+    c = get_gemma_config("gemma_300m")
+    assert c.width == 1024
+    assert c.depth == 18
+
+
+def test_gemma_unknown_variant_raises():
+    with pytest.raises(ValueError):
+        get_gemma_config("gemma_7b")
+
+
+# ----------------------------------------------------------------------------
+# load_weights remap rules (tiny model; needs torch + transformers)
+# ----------------------------------------------------------------------------
+def _tiny_pi0_model():
+    """Construct a tiny π0 model for remap tests (full Gemma depth but small
+    action/state dims — the PaliGemma backbone is always full-size)."""
+    cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
+    return Pi0ForActionPrediction(cfg)
+
+
+@pytest.fixture(autouse=True)
+def _silence_pi0_loader(caplog):
+    # The remap tests feed a single synthetic weight; the loader's
+    # "missing params" warning is expected noise here.
+    logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
+    yield
+
+
+@pytest.mark.slow
+def test_remap_paligemma_submodules_to_nested_layout():
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
+    flat_key = "model." + target.replace(
+        "paligemma_with_expert.paligemma.model.vision_tower.",
+        "paligemma_with_expert.paligemma.vision_tower.",
+    )
+    payload = torch.full(params[target].shape, 0.12345)
+    before = params[target].detach().clone()
+    model.load_weights([(flat_key, payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "flat→nested remap broken"
+    assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_lm_head_remap_to_embed_tokens():
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+    assert target in params
+    payload = torch.full(params[target].shape, 0.5)
+    before = params[target].detach().clone()
+    model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
+    assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_vision_tower_version_robust_remap():
+    """transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
+    vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
+    load_weights must adapt so vision weights aren't silently dropped (which
+    would leave the vision encoder at random init)."""
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    # A representative vision-tower param as it exists on the model.
+    target = next(
+        k
+        for k in params
+        if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
+    )
+    vt = "paligemma_with_expert.paligemma.model.vision_tower."
+    rest = target[len(vt) :]
+    # Build the checkpoint key in whichever layout differs from the model's.
+    if rest.startswith("vision_model."):
+        ckpt_key = vt + rest[len("vision_model.") :]
+    else:
+        ckpt_key = vt + "vision_model." + rest
+    payload = torch.full(params[target].shape, 0.073)
+    before = params[target].detach().clone()
+    model.load_weights([("model." + ckpt_key, payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "vision_tower version-robust remap broken"
+    assert torch.allclose(after, payload)
+
+
+# ----------------------------------------------------------------------------
+# Image processor
+# ----------------------------------------------------------------------------
+def test_resize_with_pad_square():
+    out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+
+
+def test_resize_with_pad_wide():
+    """Wide source (H < W): padding applied to top/bottom (rows)."""
+    out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+    assert out[0, 0, 0, 112].item() == -1.0  # top row is padded
+
+
+def test_resize_with_pad_tall():
+    """Tall source (H > W): padding applied to left/right (cols)."""
+    out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+    assert out[0, 0, 112, 0].item() == -1.0  # left-most column padded
+
+
+def test_resize_with_pad_rejects_non_4d():
+    with pytest.raises(ValueError):
+        resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
+
+
+def test_pil_to_tensor():
+    import numpy as np
+    from PIL import Image
+
+    img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
+    t = pil_image_to_tensor(img)
+    assert t.shape == (1, 3, 10, 10)
+    expected = 128.0 / 255.0 * 2.0 - 1.0
+    assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
+
+
+def test_pil_to_tensor_rgba_to_rgb():
+    import numpy as np
+    from PIL import Image
+
+    rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
+    img = Image.fromarray(rgba, mode="RGBA")
+    t = pil_image_to_tensor(img)
+    assert t.shape == (1, 3, 8, 8)
+
+
+def test_make_empty_image():
+    proc = Pi0ImageProcessor(image_size=224)
+    empty = proc.make_empty_image()
+    assert empty.shape == (1, 3, 224, 224)
+    assert (empty == -1.0).all()
+
+
+def test_preprocess_single_resizes_ndarray():
+    import numpy as np
+
+    proc = Pi0ImageProcessor(image_size=224)
+    arr = np.full((120, 120, 3), 100, dtype=np.uint8)
+    t = proc.preprocess_single(arr)
+    assert t.shape == (1, 3, 224, 224)
+
+
+# ----------------------------------------------------------------------------
+# build_model_inputs: camera ordering, padding, masks, state
+# ----------------------------------------------------------------------------
+class _FakeTokenizer:
+    """Minimal stand-in for the PaliGemma tokenizer (no download)."""
+
+    def __call__(
+        self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
+    ):
+        ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
+        ids = ids[:max_length]
+        attn = [1] * len(ids)
+        while len(ids) < max_length:
+            ids.append(0)
+            attn.append(0)
+        return {"input_ids": ids, "attention_mask": attn}
+
+
+def test_build_model_inputs_camera_order_and_padding():
+    import numpy as np
+
+    config = Pi0Config.from_model_config(_LEROBOT_CFG)  # 3-camera config
+    device = torch.device("cpu")
+    # Provide only 2 of the 3 declared cameras.
+    robot_obs = {
+        "images": {
+            "observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
+            "observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
+        },
+        "state": [0.1] * 14,  # raw state shorter than max_state_dim=32
+        "prompt": "pick up the block",
+    }
+    images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
+        robot_obs, config, _FakeTokenizer(), device
+    )
+    assert len(images) == 3
+    assert all(img.shape == (1, 3, 224, 224) for img in images)
+    masks = [bool(m.item()) for m in image_masks]
+    assert masks == [True, True, False]
+    # The padded (3rd) camera must be all -1.
+    assert (images[2] == -1.0).all()
+    assert lang_tokens.shape == (1, 48)
+    assert lang_masks.shape == (1, 48)
+    assert state.shape == (1, 32)
+    # State padded with zeros beyond the 14 real entries.
+    assert torch.allclose(state[0, 14:], torch.zeros(18))
+
+
+def test_build_model_inputs_state_truncated():
+    config = Pi0Config.from_model_config(_LEROBOT_CFG)
+    robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
+    _, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
+    assert state.shape == (1, 32)
+
+
+def test_build_model_inputs_image_key_map():
+    import numpy as np
+
+    cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
+    robot_obs = {
+        "images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
+        "prompt": "go",
+    }
+    _, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
+    masks = [bool(m.item()) for m in image_masks]
+    assert masks == [True, False, False]  # only base_0_rgb (mapped) is present
+
+
+# ----------------------------------------------------------------------------
+# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
+#
+# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
+# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
+# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
+# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
+# output to golden values established at 5.3.0 and re-check them here: running on
+# the production transformers proves the port is version-stable and guards the
+# vision_tower / embed_scale version-robust paths against regressions.
+#
+#   PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
+#       pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
+# ----------------------------------------------------------------------------
+import os  # noqa: E402
+
+_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
+
+# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
+_GOLDEN_SUM = -67.808113
+_GOLDEN_STD = 0.252789
+_GOLDEN_MEAN = -0.021190
+_GOLDEN_EMBED_IMAGE_STD = 2.17603
+
+
+def _selfconsist_inputs(device):
+    # NB: draw order is load-bearing — the golden values were established with
+    # state drawn first, then the 3 images, then lang (shared RNG, seed 0).
+    g = torch.Generator(device="cpu").manual_seed(0)
+    b = 2
+    state = torch.randn(b, 32, generator=g)
+    images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
+    masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
+    lang = torch.randint(0, 1000, (b, 48), generator=g)
+    lmask = torch.ones(b, 48, dtype=torch.bool)
+    ng = torch.Generator(device="cpu").manual_seed(42)
+    noise = torch.randn(b, 50, 32, generator=ng)
+    return (
+        [im.to(device) for im in images],
+        [m.to(device) for m in masks],
+        lang.to(device),
+        lmask.to(device),
+        state.to(device),
+        noise.to(device),
+    )
+
+
+def _selfconsist_model(device):
+    import safetensors.torch
+
+    model = Pi0ForActionPrediction(
+        Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
+    )
+    state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
+    filled = model.load_weights(list(state.items()))
+    expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
+    missing = expected - (filled or set())
+    assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
+    return model.to(device).eval()
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_actions():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _selfconsist_model(device)
+    images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
+    actions = (
+        model.sample_actions(
+            images=images,
+            image_masks=masks,
+            lang_tokens=lang,
+            lang_masks=lmask,
+            state=state,
+            noise=noise,
+            num_steps=10,
+        )
+        .float()
+        .cpu()
+    )
+    s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
+    assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
+    assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
+    assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_embed_image():
+    """Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _selfconsist_model(device)
+    g = torch.Generator(device="cpu").manual_seed(7)
+    img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
+    std = model.paligemma_with_expert.embed_image(img).float().std().item()
+    assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
+        f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
+        "likely a vision_tower weight-load regression under this transformers version."
+    )

tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion
+config shim must not crash on missing ``head_dtype`` and the resulting
+LogitsProcessor must have ``head_dtype is None`` (= use model dtype)."""
+
+import os
+
+import pytest
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import (
+    cleanup_dist_env_and_memory,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+    SenseNovaU1ForCausalLM,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+from vllm_omni.transformers_utils.configs.sensenova_u1 import (
+    SenseNovaU1Config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _build_tiny_sensenova_u1(quant_config=None, prefix="model"):
+    """Build a minimal SenseNovaU1ForCausalLM for unit tests."""
+    config = SenseNovaU1Config(
+        llm_config={
+            "hidden_size": 64,
+            "num_attention_heads": 2,
+            "num_hidden_layers": 1,
+            "intermediate_size": 128,
+            "vocab_size": 32,
+            "max_position_embeddings": 128,
+            "max_position_embeddings_hw": 128,
+        },
+    )
+    return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix)
+
+
+def test_logits_processor_head_dtype_under_diffusion_shim():
+    # Initialize the distributed environment
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29543")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+
+    # Create a vLLM config with a wrapped diffusion config
+    fake_diff_config = _DiffusionVllmModelConfig(
+        model="sensenova-u1-test",
+        dtype=torch.bfloat16,
+        max_model_len=8192,
+        original_max_model_len=8192,
+    )
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+    )
+    vllm_config.modelconfig = fake_diff_config  # type: ignore[assignment]
+
+    # Build a tiny version of the Causal LM component
+    with set_current_vllm_config(vllm_config):
+        model = _build_tiny_sensenova_u1()
+
+    # Ensure that we have a set head_dtype attribute. Currently, the head_dtype
+    # is set to None
+    head_dtype = model.logits_processor.head_dtype
+    assert head_dtype is None
+    cleanup_dist_env_and_memory()

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -185,6 +185,79 @@ class TestDistributedLayerwiseOffloadHook:
         assert next_block.weight.stride() == expected_stride
         assert torch.equal(next_block.weight, expected)
 
+    def test_allgather_reconstructs_online_fp8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""
+
+        class OnlineFp8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical_weight = torch.tensor(
+                    [
+                        [1.0, 2.0, 3.0, 4.0],
+                        [5.0, 6.0, 7.0, 8.0],
+                        [9.0, 10.0, 11.0, 12.0],
+                    ],
+                    dtype=torch.float32,
+                ).to(torch.float8_e4m3fn)
+                self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
+                self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)
+
+        current_block = OnlineFp8Block()
+        rank0_block = OnlineFp8Block()
+        rank1_block = OnlineFp8Block()
+        expected_weight = rank0_block.weight.detach().clone()
+        expected_scale = rank0_block.weight_scale.detach().clone()
+        expected_stride = rank0_block.weight.stride()
+
+        rank0_hook = DistributedLayerwiseOffloadHook(
+            next_block=rank0_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=0,
+            copy_stream=DummyStream(),
+            comm_stream=DummyStream(),
+            pin_memory=False,
+        )
+        rank1_hook = DistributedLayerwiseOffloadHook(
+            next_block=rank1_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=1,
+            copy_stream=DummyStream(),
+            comm_stream=DummyStream(),
+            pin_memory=False,
+        )
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineFp8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert rank0_block.weight.stride() == expected_stride
+        torch.testing.assert_close(
+            rank0_block.weight.float(),
+            expected_weight.float(),
+            rtol=0,
+            atol=0,
+        )
+        torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
+
     def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         """DTensor wrapper should be preserved through prefetch/offload cycle."""
         current_block = TinyBlock(_make_values(1.0))

tests/diffusion/test_diffusion_config_propagation.py

@@ -92,6 +92,30 @@ class TestParallelConfigPropagation:
         od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
         assert od.num_gpus == 2
 
+    def test_dp_is_inferred_from_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
+        assert od.parallel_config.data_parallel_size == 2
+        assert od.parallel_config.world_size == 8
+
+    def test_explicit_dp_is_validated_against_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 2
+
+        with pytest.raises(ValueError, match="does not match WORLD-derived value"):
+            OmniDiffusionConfig.from_kwargs(
+                model="x",
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
+                num_gpus=8,
+            )
+
+    def test_hsdp_does_not_infer_ordinary_dp(self):
+        pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 1
+        assert od.parallel_config.world_size == 4
+
 
 class TestCreateDefaultDiffusion:
     """Verify engine_args structure from create_default_diffusion."""
@@ -130,6 +154,21 @@ class TestCreateDefaultDiffusion:
 
         assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 
+    def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
+        monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
+        od = _roundtrip_diffusion_config(
+            model="x",
+            kv_cache_memory_bytes=4096,
+            gpu_memory_utilization=0.75,
+            max_num_batched_tokens=2048,
+            max_model_len=4096,
+        )
+
+        assert od.kv_cache_memory_bytes == 4096
+        assert od.gpu_memory_utilization == 0.75
+        assert od.max_num_batched_tokens == 2048
+        assert od.max_model_len == 4096
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -161,9 +161,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
     assert exc_info.value is initialization_error
     custom_scheduler.initialize.assert_called_once_with(od_config)
     custom_scheduler.close.assert_called_once_with()
+
+
+def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        Mock(side_effect=RuntimeError("KV initialization failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="KV initialization failed"):
+        DiffusionEngine(SimpleNamespace())
+
     fake_executor.shutdown.assert_called_once_with()
 
 
+def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
+    events = []
+    expected_profile_requests = [object()]
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    kv_cache_config = object()
+    kv_vllm_config = object()
+    od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
+
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_init_executor",
+        lambda self, config: setattr(self, "executor", fake_executor),
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_prepare_diffusion_kv_profile_requests",
+        lambda self: events.append("prepare-profile") or expected_profile_requests,
+    )
+
+    def initialize(executor, config, *, profile_requests):
+        events.append("profile-workers")
+        assert executor is fake_executor
+        assert config is od_config
+        assert profile_requests is expected_profile_requests
+        return kv_cache_config, 16, 16, kv_vllm_config
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        initialize,
+    )
+
+    def init_scheduler(self, config, scheduler, *args, **kwargs):
+        events.append("initialize-scheduler")
+        self.scheduler = SimpleNamespace(close=Mock())
+
+    monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
+    monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
+
+    DiffusionEngine(od_config)
+
+    assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
+
+
 @pytest.mark.asyncio
 async def test_step_compatibility_wrapper_returns_final_batch() -> None:
     engine = _make_engine()

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -1,9 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 
 from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
     )
 
     assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
+
+
+def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
+    assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
+    assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
+
+
+def test_dense_mode_does_not_build_kv_profile_request() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
+
+    assert engine._prepare_diffusion_kv_profile_requests() is None
+    engine._make_dummy_request.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
+    [
+        (DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
+        (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
+        (DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
+    ],
+)
+def test_paged_kv_profile_requests_match_per_rank_batch(
+    execution_mode: DiffusionExecutionMode,
+    uses_dlo_dp: bool,
+    max_num_seqs: int,
+    expected_profile_requests: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.execution_mode = execution_mode
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        model_class_name="HunyuanImage3ForCausalMM",
+        max_num_seqs=max_num_seqs,
+        parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
+        enable_distributed_layerwise_offload=uses_dlo_dp,
+        dlo_use_allgather=True,
+    )
+    request = OmniDiffusionRequest(
+        prompt="profile",
+        request_id="profile-request",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+    engine._make_dummy_request = Mock(return_value=request)
+    prepared_layout = object()
+    scheduler_kv_state = (
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=3,
+            target_len=8,
+            seq_len=12,
+        ),
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/1",
+            sequence_id=1,
+            prefix_len=4,
+            target_len=8,
+            seq_len=13,
+        ),
+    )
+
+    def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        req.prepared_layout = prepared_layout
+        req.diffusion_kv_requests = scheduler_kv_state
+        return req
+
+    engine._prepare_request_for_admission = Mock(side_effect=preprocess)
+
+    result = engine._prepare_diffusion_kv_profile_requests()
+
+    assert result is not None
+    assert len(result) == expected_profile_requests
+    assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
+    assert all(profile_request.is_dummy_run() for profile_request in result)
+    assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
+    assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
+    assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
+    assert engine._diffusion_kv_profile_limits == (2, 13, 8)
+    assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
+    engine._make_dummy_request.assert_called_once_with(
+        height=1024,
+        width=1024,
+        guidance_scale=5.0,
+        num_image_inputs=3,
+    )
+    engine._prepare_request_for_admission.assert_called_once_with(request)
+
+
+@pytest.mark.parametrize(
+    ("num_sequences", "seq_len", "target_len"),
+    [
+        (3, 16, 8),
+        (1, 17, 8),
+        (1, 16, 9),
+    ],
+)
+def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
+    num_sequences: int,
+    seq_len: int,
+    target_len: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="too large",
+        request_id="too-large",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"too-large/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=target_len,
+                seq_len=seq_len,
+            )
+            for sequence_id in range(num_sequences)
+        ),
+    )
+
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(request)
+
+
+def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="fits",
+        request_id="fits",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"fits/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=8,
+                target_len=8,
+                seq_len=16,
+            )
+            for sequence_id in range(2)
+        ),
+    )
+
+    assert engine._prepare_request_for_admission(request) is request

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -473,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop():
         _stop_engine(engine)
 
 
+@pytest.mark.asyncio
+async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+
+    def fail_reservation(*args, **kwargs):
+        raise RuntimeError("native allocation bug")
+
+    engine.scheduler._diffusion_kv_manager = SimpleNamespace(
+        has_request=lambda request_id: False,
+        reserve_request=fail_reservation,
+        free_request=lambda request_id: None,
+    )
+    request = _make_request("kv-error")
+    request.diffusion_kv_requests = (object(),)
+    try:
+        response_stream = engine.async_add_req_and_stream_response(request)
+        output_queue = engine._out_streams[request.request_id]
+        for _ in range(300):
+            if not output_queue.empty():
+                break
+            await asyncio.sleep(0.01)
+        assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
+
+        output = await anext(response_stream)
+        assert output.error == "native allocation bug"
+        assert output.finished
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
 # ─────────────────────────── _RpcTask basics ───────────────────────────────
 
 

tests/diffusion/test_diffusion_model_runner.py

@@ -3,6 +3,7 @@
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -512,6 +513,84 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     assert runner.pipeline.last_req.num_reqs == 1
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    request = _make_request()
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    record_names = []
+    original_execute = runner._execute_request_list
+
+    def execute(*args, **kwargs):
+        record_names.append(kwargs["record_name"])
+        return original_execute(*args, **kwargs)
+
+    runner._execute_request_list = execute
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run([request])
+
+    assert runner.pipeline.forward_calls == 1
+    assert record_names == ["pipeline_memory_profile"]
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.step_execution = True
+    requests = [_make_request(), _make_request()]
+    requests[0].request_id = "profile-0"
+    requests[1].request_id = "profile-1"
+    for request in requests:
+        request.sampling_params.num_inference_steps = 1
+
+    observed_batch_rows = []
+    original_denoise_step = runner.pipeline.denoise_step
+
+    def denoise_step(input_batch, states):
+        observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
+        return original_denoise_step(input_batch, states)
+
+    runner.pipeline.denoise_step = denoise_step
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run(requests)
+
+    assert observed_batch_rows == [(2, 2)]
+    assert runner.state_cache == {}
+    assert runner.input_batch is None
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
 class _BatchPipeline:
     """Pipeline returning a configurable list of outputs from forward()."""
 

tests/diffusion/test_diffusion_request.py

@@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
 
     assert req.is_dummy_run()
     assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+    assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
 
 
 def test_non_dummy_request_is_not_identified_as_dummy_run():

tests/diffusion/test_diffusion_scheduler.py

@@ -8,12 +8,14 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]:
     return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
+def _initialize_paged_scheduler(
+    scheduler: BaseScheduler,
+    *,
+    num_blocks: int = 64,
+    max_num_seqs: int = 1,
+) -> None:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    scheduler.initialize(
+        SimpleNamespace(
+            diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+            max_model_len=64,
+            max_num_seqs=max_num_seqs,
+        ),
+        kv_cache_config=config,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+        ),
+    )
+
+
+def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+    request.diffusion_kv_requests = (
+        DiffusionKVRequest(
+            f"{request.request_id}/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=4,
+            target_len=4,
+            seq_len=seq_len,
+        ),
+    )
+
+
 class _StubScheduler:
     def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
@@ -400,18 +448,16 @@ class TestRequestScheduler:
         assert self.scheduler.has_requests() is False
 
     def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         prepared_layout = object()
         request.prepared_layout = prepared_layout
-        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
         kv_request = DiffusionKVRequest(
             "diffusion-kv/diffusion-kv/0",
             sequence_id=0,
             prefix_len=4,
             target_len=8,
             seq_len=16,
-            kv_contexts=(text_context,),
         )
         request.diffusion_kv_requests = (kv_request,)
 
@@ -421,12 +467,179 @@ class TestRequestScheduler:
 
         assert state is not None
         assert state.diffusion_kv_requests == (kv_request,)
-        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
         assert request.diffusion_kv_requests is None
         assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+        metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.request_id == request.request_id
+        assert len(metadata.sequences[0].block_ids[0]) == 4
+
+    def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
+        first_request = _make_request("first")
+        second_request = _make_request("second")
+        _attach_diffusion_kv(first_request)
+        _attach_diffusion_kv(second_request)
+        self.scheduler.add_request(first_request)
+        self.scheduler.add_request(second_request)
+
+        first_output = self.scheduler.schedule()
+
+        assert _new_ids(first_output) == ["first"]
+        assert first_output.num_waiting_reqs == 1
+
+        self.scheduler.update_from_output(first_output, _make_request_output("first"))
+        second_output = self.scheduler.schedule()
+
+        assert _new_ids(second_output) == ["second"]
+        assert second_output.num_waiting_reqs == 0
+        # The same output carries both the release and the newly admitted block
+        # table; the future Worker data plane must process finished ids first.
+        assert second_output.finished_req_ids == {"first"}
+        assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
+
+    def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
+        impossible = _make_request("impossible")
+        schedulable = _make_request("schedulable")
+        _attach_diffusion_kv(impossible, seq_len=8)
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "cannot fit even when the block pool is empty" in failed_state.error
+        assert sched_output.finished_req_ids == {"impossible"}
+        assert _new_ids(sched_output) == ["schedulable"]
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_request_output("schedulable"),
+        )
+
+        assert finished == {"impossible", "schedulable"}
+
+    def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
+        running = _make_request("running")
+        _attach_diffusion_kv(running)
+        self.scheduler.add_request(running)
+        assert _new_ids(self.scheduler.schedule()) == ["running"]
+
+        impossible = _make_request("impossible")
+        impossible.diffusion_kv_requests = tuple(
+            DiffusionKVRequest(
+                f"impossible/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            )
+            for sequence_id in range(4)
+        )
+        schedulable = _make_request("schedulable")
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "required_blocks=4, available_blocks=3" in failed_state.error
+        assert sched_output.finished_req_ids == {"impossible"}
+        assert _new_ids(sched_output) == ["schedulable"]
+
+    def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
+        _initialize_paged_scheduler(self.scheduler)
+        request = _make_request("native-error")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+
+        def raise_native_error(*args, **kwargs):
+            raise ValueError("native allocation bug")
+
+        monkeypatch.setattr(
+            manager,
+            "reserve_request",
+            raise_native_error,
+        )
+
+        sched_output = self.scheduler.schedule()
+
+        state = self.scheduler.get_request_state("native-error")
+        assert state is not None
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "native allocation bug"
+        assert sched_output.finished_req_ids == {"native-error"}
+        assert sched_output.scheduled_request_ids == []
+
+    def test_diffusion_kv_preemption_retains_allocation(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("preempted")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        metadata = manager.get_metadata("preempted")
+
+        assert self.scheduler.preempt_request("preempted") is True
+        resumed_output = self.scheduler.schedule()
+
+        assert resumed_output.scheduled_new_reqs == []
+        assert _cached_ids(resumed_output) == ["preempted"]
+        assert manager.get_metadata("preempted") == metadata
+
+    @pytest.mark.parametrize(
+        "status",
+        [
+            DiffusionRequestStatus.FINISHED_COMPLETED,
+            DiffusionRequestStatus.FINISHED_ABORTED,
+            DiffusionRequestStatus.FINISHED_ERROR,
+        ],
+    )
+    def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("terminal")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        assert manager.has_request("terminal") is True
+
+        self.scheduler.finish_requests("terminal", status)
+
+        assert manager.has_request("terminal") is False
 
     def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
 
         with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
@@ -443,7 +656,7 @@ class TestRequestScheduler:
         ],
     )
     def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         owner = request if owner_name == "request" else request.sampling_params
         setattr(owner, field_name, object())
@@ -921,6 +1134,45 @@ class TestDiffusionEngine:
         assert engine.scheduler.max_num_running_reqs == 1
         fake_executor_cls.assert_called_once_with(od_config)
 
+    def test_initializes_paged_scheduler_from_native_control_plane(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
+        fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
+        scheduler = mocker.Mock()
+        kv_cache_config = object()
+        kv_vllm_config = object()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+            lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
+        )
+
+        engine = DiffusionEngine(od_config, scheduler=scheduler)
+
+        assert engine.scheduler is scheduler
+        scheduler.initialize.assert_called_once_with(
+            od_config,
+            kv_cache_config=kv_cache_config,
+            scheduler_block_size=16,
+            hash_block_size=16,
+            kv_vllm_config=kv_vllm_config,
+        )
+
     def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
         scheduler = Scheduler()
         scheduler.initialize(SimpleNamespace())
@@ -987,6 +1239,21 @@ class TestDiffusionEngine:
         assert output.aborted is True
         assert output.abort_message == "Request req-finalize aborted."
 
+    def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = StepScheduler()
+        engine.scheduler.initialize(SimpleNamespace())
+
+        req_id = engine.scheduler.add_request(_make_request("req-error"))
+        engine.scheduler._finish_requests(
+            {req_id: DiffusionRequestStatus.FINISHED_ERROR},
+            {req_id: "KV request cannot fit"},
+        )
+
+        output = engine._finalize_finished_request(req_id)
+
+        assert output.error == "KV request cannot fit"
+
     @pytest.mark.asyncio
     async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)

tests/diffusion/test_diffusion_worker.py

@@ -13,11 +13,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from vllm_omni.diffusion.worker.diffusion_worker import (
-    DiffusionWorker,
-    _create_diffusion_worker_vllm_config,
-    _make_diffusion_vllm_model_config,
-)
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
@@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
     return worker
 
 
-def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
-    from types import SimpleNamespace
-
-    from vllm_omni.quantization import build_quant_config
-
-    od_config = SimpleNamespace(
-        model="dummy",
-        dtype=torch.bfloat16,
-        quantization_config=build_quant_config(
-            {
-                "quant_method": "modelopt",
-                "quant_algo": "FP8",
-                "ignore": [],
-            }
-        ),
-        tf_model_config=SimpleNamespace(),
-        enforce_eager=True,
-        is_moe=False,
-    )
-
-    model_config = _make_diffusion_vllm_model_config(od_config)
-
-    assert model_config.dtype is torch.bfloat16
-    assert model_config.quantization == "modelopt"
-    assert model_config.quantization_config is od_config.quantization_config
-    assert model_config.is_quantized()
-
-
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
@@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp:
         mock_buffer2.data.copy_.assert_not_called()
 
         assert result is True
-
-
-class TestWorkerVllmConfigAdditionalConfig:
-    """Test worker-side VllmConfig construction with additional_config."""
-
-    def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            return_value=mock_vllm_config,
-        )
-        od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
-        self, mocker: MockerFixture
-    ):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=[
-                TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
-                mock_vllm_config,
-            ],
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_count == 2
-        assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
-        mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=TypeError("additional_config values must be hashable"),
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        with pytest.raises(TypeError, match="hashable"):
-            _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -413,6 +413,26 @@ class TestRequestModeDispatch:
         assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
         executor.collective_rpc.assert_called_once()
 
+    def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
+        assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
+        assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
+
     @pytest.mark.parametrize(
         ("field", "value"),
         [

tests/e2e/online_serving/test_minimax_music3.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for MiniMax Music 3 text-to-music.
+
+This is a music model on the speech endpoint, so the request shape differs
+from every other TTS model here: ``input`` carries the lyrics, ``instructions``
+carries the caption that decides genre and arrangement, and both are required.
+
+Weekly rather than per-PR: the checkpoint is 27 GB across five components and a
+single request occupies two decode rows, because classifier-free guidance is
+mandatory for this model. Rejection cases live in
+``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "MiniMaxAI/MiniMax-Music3"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0
+
+# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The
+# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room
+# for the model ending a song before the cap.
+_FRAMES_10S = 250
+_FRAMES_30S = 750
+_MIN_BYTES_10S = 500_000
+_MIN_BYTES_30S = 1_500_000
+
+LYRICS = (
+    "[Verse]\nWalking down the empty street at midnight\n"
+    "Streetlights flicker like a broken dream\n"
+    "[Chorus]\nAnd I keep on walking\nTill the morning finds me"
+)
+
+CAPTION = (
+    "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, "
+    "soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. "
+    "Intimate bedroom production, gentle tape saturation, no bright cymbals."
+)
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="minimax_music3",
+    )
+]
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_001(omni_server, openai_client) -> None:
+    """Baseline: lyrics plus caption in, one ten-second song out.
+
+    Deploy Setting: minimax_music3.yaml
+    Input Modal: text (lyrics) + text (caption)
+    Output Modal: audio, 32 kHz stereo
+    Input Setting: stream=False, seed pinned, 250 frames
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": LYRICS,
+        "instructions": CAPTION,
+        "seed": 1,
+        "max_new_tokens": _FRAMES_10S,
+        "stream": False,
+        "response_format": "wav",
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "min_audio_bytes": _MIN_BYTES_10S,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_multi_window_002(omni_server, openai_client) -> None:
+    """A song long enough to span several acoustic windows.
+
+    The acoustic stage decodes overlapping 200-frame windows on a 100-frame
+    hop and crops the overlap back off, so a request that fits in one window
+    never exercises the stitching. 750 frames covers seven of them.
+
+    Deploy Setting: minimax_music3.yaml
+    Input Setting: stream=False, seed pinned, 750 frames
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": LYRICS,
+        "instructions": CAPTION,
+        "seed": 1,
+        "max_new_tokens": _FRAMES_30S,
+        "stream": False,
+        "response_format": "wav",
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "min_audio_bytes": _MIN_BYTES_30S,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_pi0_expansion.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving.
+
+Boots ``vllm serve --omni --deploy-config pi0.yaml`` and drives the real OpenPI
+websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses
+(handshake metadata → send observation → receive action chunk). Mirrors
+``tests/e2e/online_serving/test_dreamzero_expansion.py``. Needs a GPU + a
+pi0_base checkpoint; skipped in CI unless explicitly run.
+
+The in-process LeRobot parity oracle lives separately in
+``tests/diffusion/models/pi0/test_pi0_parity.py``.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import (
+    OmniServerParams,
+    get_open_port,
+    pi0_openpi_require_dependencies,
+    pi0_openpi_run_policy_session,
+    pi0_openpi_validate_session_result,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "lerobot/pi0_base"
+
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        port=8092,
+        server_args=[
+            "--deploy-config",
+            "vllm_omni/deploy/pi0.yaml",
+            "--served-model-name",
+            "pi0",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+        env_dict={
+            "ATTENTION_BACKEND": "torch",
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+            "VLLM_DISABLE_COMPILE_CACHE": "1",
+            "MASTER_PORT": str(get_open_port()),
+        },
+    )
+]
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_pi0_openpi_online(omni_server):
+    try:
+        pi0_openpi_require_dependencies()
+    except ModuleNotFoundError as exc:
+        pytest.skip(str(exc))
+
+    result = pi0_openpi_run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        prompt="pick up the red block and place it in the bin",
+        session_id="pi0-online-e2e",
+        num_steps=2,
+    )
+
+    # Asserts every returned chunk is [50, 32] + finite, and the handshake
+    # metadata matches pi0.yaml's policy_server_config.
+    pi0_openpi_validate_session_result(result)
+
+    metadata = result["metadata"]
+    assert tuple(metadata["image_resolution"]) == (224, 224)
+    assert metadata["needs_wrist_camera"] is True
+    assert metadata["needs_session_id"] is False
+    assert metadata["action_space"] == "joint_position"

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -452,11 +452,31 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, openai_client) -> None:
-    """
+    """Catastrophic-regression gate on one-word pronunciation.
+
+    Whether a one-word answer is pronounced intelligibly is sampled, not guaranteed: the
+    talker runs at temperature 0.9 (the upstream default, see qwen3_omni_moe.yaml stage 1),
+    so asserting it per request gates on a random variable, which is what the old 10-retry
+    loop was compensating for.
+
+    Sends a fixed 40 requests and requires at least 29 to transcribe as the expected
+    word. 85.6% is a predeclared minimum acceptable success probability, taken from a
+    measured HF reference run (77/90); vLLM-Omni itself measured 92.2%. Assuming 40
+    independent Bernoulli trials, >=29/40 has a nominal rejection probability of 0.9%
+    at p=.856, with 92.9% power at p=.60 and 55.9% at p=.70.
+
+    Those are conditional on p, not a measured false-failure rate for this CI: 77/90 is
+    one point estimate (Wilson 76.8%-91.4%), and at the low end of that interval the gate
+    would reject about 20% of runs. Pinning down a real rate would take repeated
+    report-only runs on the target hardware.
+
+    A drop to 70% is caught only about half the time, so this is a smoke gate for gross
+    pipeline breakage, not a quality tracker -- resolving moderate drift needs ~80+
+    samples and belongs in an accuracy benchmark. See #5395.
+
     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(),
@@ -474,17 +494,13 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
         "transcript_language": "en",
     }
 
-    # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py).
-    _similarity_assert_msg = "The audio content is not same as the text"
-    _max_retries = 10
-    for attempt in range(_max_retries):
-        try:
-            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
-            break
-        except AssertionError as e:
-            if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
-                raise
-            print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
+    # The judge runs per request so mismatches can be counted, which also transcribes
+    # concurrently where the previous request_num=5 path judged serially -- worth watching
+    # on tighter cards, since each judge loads whisper on the same device as the talker and
+    # code2wav stages.
+    openai_client.send_omni_request(
+        request_config, request_num=40, min_successes=29, max_concurrency=get_max_batch_size()
+    )
 
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)

tests/engine/test_async_omni_engine_input.py

@@ -371,3 +371,106 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er
         )
 
     assert stage_pool.get_bound_replica_id("req-error") is None
+
+
+def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture):
+    """A guided request is all-or-nothing.
+
+    Admitting the parent before its companion is built leaves an orphan: a
+    model whose guidance is mandatory never completes the pair, so the request
+    holds scheduler and KV capacity for the scheduler's whole hold budget and
+    then produces no audio. The build has to raise before anything is enqueued.
+    """
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin"))
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.request_queue = mocker.Mock()
+
+    with pytest.raises(ValueError, match="null twin"):
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture):
+    """The same holds when expansion succeeds but input processing does not."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt")
+    engine.request_queue = mocker.Mock()
+
+    with pytest.raises(RuntimeError, match="tokenizer rejected"):
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture):
+    """Construction and admission are separate steps."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
+    engine.request_queue = mocker.Mock()
+
+    companions = engine._build_cfg_companions(
+        parent_id="req",
+        original_prompt={"prompt": "positive"},
+        stage0_params=params,
+        sampling_params_list=[params],
+    )
+
+    assert len(companions) == 1
+    assert companions[0].companion_id == "req__cfg_text"
+    assert companions[0].parent_id == "req"
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture):
+    """An unguided request expands to no companions and admits normally."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: []
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.request_queue = mocker.Mock()
+
+    assert (
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+        == []
+    )
+    engine.request_queue.sync_q.put.assert_not_called()

tests/engine/test_stage_engine_args.py

@@ -8,8 +8,23 @@ from pathlib import Path
 
 import pytest
 from pydantic.fields import FieldInfo
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
+from vllm.engine.arg_utils import EngineArgs
 
-from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.omni_config import (
+    _LLM_STAGE_ENGINE_FIELDS,
+    OmniStageCacheConfig,
+    OmniStageLoadConfig,
+    OmniStageParallelConfig,
+    OmniStageSchedulerConfig,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+)
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
 from vllm_omni.config.stage_config import (
     DeployConfig,
@@ -53,6 +68,30 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
 
 _LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
 _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
+_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
+    {
+        "active_stream_window",
+        "codec_frame_rate_hz",
+        "custom_voice_dir",
+        "devices",
+        "disable_autocast",
+        "enable_multithread_weight_load",
+        "env",
+        "has_sampling_extra_args",
+        "log_level",
+        "log_stats",
+        "model_arch",
+        "model_subdir",
+        "num_gpus",
+        "num_replicas",
+        "num_weight_load_threads",
+        "omni_kv_config",
+        "parallel_config",
+        "subtalker_sampling_params",
+        "task_type",
+        "tokenizer_subdir",
+    }
+)
 
 
 @pytest.fixture(autouse=True)
@@ -197,6 +236,100 @@ def _legacy_and_typed_stages(
     return legacy_stages, omni_config
 
 
+def test_llm_stage_engine_field_schema_tracks_upstream_engine_args():
+    upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs))
+
+    assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS
+
+
+def test_typed_llm_projection_discovers_explicit_upstream_config_fields():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+        load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]),
+        cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"),
+        scheduler_config=OmniStageSchedulerConfig(
+            prefill_schedule_interval=2,
+            policy="priority",
+        ),
+        parallel_config=OmniStageParallelConfig(
+            data_parallel_master_ip="10.0.0.1",
+            enable_dbo=True,
+        ),
+    )
+
+    engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+    assert engine_args["ignore_patterns"] == ["*.bin"]
+    assert engine_args["block_size"] == 32
+    assert engine_args["kv_cache_dtype"] == "fp8"
+    assert engine_args["prefill_schedule_interval"] == 2
+    assert engine_args["scheduling_policy"] == "priority"
+    assert engine_args["data_parallel_address"] == "10.0.0.1"
+    assert engine_args["enable_dbo"] is True
+
+
+def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+    )
+
+    engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+    inherited_defaults = {
+        "ignore_patterns",
+        "block_size",
+        "kv_cache_dtype",
+        "prefill_schedule_interval",
+        "scheduling_policy",
+        "data_parallel_address",
+        "data_parallel_rank",
+        "data_parallel_size_local",
+        "data_parallel_rpc_port",
+        "enable_dbo",
+    }
+    assert inherited_defaults.isdisjoint(engine_args)
+
+
+def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+        scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"),
+    )
+
+    with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"):
+        stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+
+@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"])
+def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    compilation_config = VllmCompilationConfig(backend="eager")
+    profiler_config = VllmProfilerConfig(profiler="cuda")
+    stage_prefix = f"stage_{stage_id}_"
+    omni_config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=copy.deepcopy(deploy),
+        cli_overrides={
+            "model": model,
+            f"{stage_prefix}compilation_config": compilation_config,
+            f"{stage_prefix}profiler_config": profiler_config,
+        },
+    )
+
+    stage_config = omni_config.stage_by_id(stage_id)
+    typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model)
+
+    assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+    assert isinstance(typed_args["compilation_config"], VllmCompilationConfig)
+    assert typed_args["compilation_config"] is not stage_config.compilation_config
+    assert typed_args["compilation_config"].backend == "eager"
+
+    assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+    assert isinstance(typed_args["profiler_config"], VllmProfilerConfig)
+    assert typed_args["profiler_config"] is not stage_config.profiler_config
+    assert typed_args["profiler_config"].profiler == "cuda"
+
+
 def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
     pipeline, deploy, model = _engine_arg_inputs(tmp_path)
     legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
@@ -222,6 +355,13 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
         assert {name: typed_args[name] for name in legacy_args} == legacy_args
 
     thinker_args = typed_args_by_stage[0]
+    inherited_vllm_fields = {
+        field.name
+        for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig)
+        for field in fields(config_cls)
+    }
+    topology_projected_fields = {"scheduler_cls"}
+    assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args)
     assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
     assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
     assert thinker_args["stage_id"] == 0

tests/entrypoints/openai_api/test_serving_speech.py

@@ -33,6 +33,8 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     OpenAICreateAudioGenerateRequest,
     OpenAICreateSpeechRequest,
     SpeechBatchItem,
+    SpeechInputTokenDetails,
+    SpeechTokenUsage,
     StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
@@ -155,6 +157,7 @@ class TestAudioMixin:
 # Helper to create mock model output for endpoint tests
 def create_mock_audio_output_for_test(
     request_id: str = "speech-mock-123",
+    metrics: dict | None = None,
 ) -> OmniRequestOutput:
     class MockCompletionOutput:
         def __init__(self, index: int = 0):
@@ -181,11 +184,16 @@ def create_mock_audio_output_for_test(
     audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
     mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
 
-    return OmniRequestOutput.from_stage_output(
+    if metrics is None:
+        metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 5}}}
+
+    output = OmniRequestOutput.from_stage_output(
         mock_request_output,
         stage_id=0,
         final_output_type="audio",
     )
+    output.metrics = metrics
+    return output
 
 
 def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
@@ -233,6 +241,7 @@ def test_app(mocker: MockerFixture):
         models=mock_models,
         request_logger=mock_request_logger,
     )
+    speech_server._tts_tokenizer = lambda text, padding=False: {"input_ids": text.split()}
 
     # Skip TTS validation in tests (mock doesn't set up supported_speakers)
     speech_server._validate_tts_request = mocker.MagicMock(return_value=None)
@@ -370,6 +379,37 @@ class TestSpeechAPI:
         assert response.headers["content-type"] == "audio/wav"
         assert len(response.content) > 0
 
+    def test_create_speech_includes_token_usage_and_detail_headers(self, client):
+        payload = {
+            "input": "Hello world",
+            "model": "tts-model",
+            "voice": "alloy",
+            "response_format": "wav",
+        }
+        response = client.post("/v1/audio/speech", json=payload)
+        assert response.status_code == 200
+        assert response.headers["x-vllm-omni-input-tokens"] == "2"
+        assert response.headers["x-vllm-omni-output-tokens"] == "5"
+        assert response.headers["x-vllm-omni-total-tokens"] == "7"
+        assert response.headers["x-vllm-omni-input-text-tokens"] == "2"
+        assert response.headers["x-vllm-omni-input-audio-tokens"] == "0"
+
+    def test_build_speech_usage_headers_uses_usage_and_detail_field_names(self):
+        usage = SpeechTokenUsage(
+            input_tokens=7,
+            output_tokens=11,
+            total_tokens=18,
+            input_token_details=SpeechInputTokenDetails(text_tokens=3, audio_tokens=4),
+        )
+        headers = OmniOpenAIServingSpeech._build_speech_usage_headers(usage)
+        assert headers == {
+            "X-VLLM-OMNI-INPUT-TOKENS": "7",
+            "X-VLLM-OMNI-OUTPUT-TOKENS": "11",
+            "X-VLLM-OMNI-TOTAL-TOKENS": "18",
+            "X-VLLM-OMNI-INPUT-TEXT-TOKENS": "3",
+            "X-VLLM-OMNI-INPUT-AUDIO-TOKENS": "4",
+        }
+
     def test_create_speech_mp3_format(self, client):
         payload = {
             "input": "Hello world",
@@ -4019,9 +4059,7 @@ class TestTTSAsyncOffloading:
 
         assert adapter_model_type != legacy_tts_model_type
         assert any(
-            call.args
-            and call.args[0] == "TTS speech request %s: text=%r, model=%s"
-            and call.args[3] == adapter_model_type
+            call.args and call.args[0] == "TTS speech request %s: model=%s" and call.args[2] == adapter_model_type
             for call in log_info.call_args_list
         )
 

tests/entrypoints/openai_api/test_tts_detection.py

@@ -153,6 +153,7 @@ _PIPELINE_STAGES = [
     "indextts2_talker",
     "latent_generator",
     "llm",
+    "minimax_music3_ar",
     "ming_tts",
     "moss_tts",
     "moss_tts_codec",

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -2,16 +2,122 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for OpenAI-compatible video API encoding helpers."""
 
+import base64
+from io import BytesIO
+
+import httpx
 import numpy as np
 import pytest
 import torch
+from PIL import Image
+from vllm import envs
 
 from vllm_omni.diffusion.postprocess import rife_interpolator
 from vllm_omni.entrypoints.openai import video_api_utils
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _png_bytes() -> bytes:
+    buffer = BytesIO()
+    Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+def _install_http_transport(monkeypatch, handler, client_kwargs):
+    async_client = httpx.AsyncClient
+
+    def _client_factory(*args, **kwargs):
+        client_kwargs.append(kwargs.copy())
+        return async_client(*args, transport=httpx.MockTransport(handler), **kwargs)
+
+    monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
+    requested_paths = []
+    client_kwargs = []
+
+    def _handler(request):
+        requested_paths.append(request.url.path)
+        if request.url.path == "/redirect.png":
+            return httpx.Response(302, headers={"location": "/image.png"})
+        return httpx.Response(200, content=_png_bytes(), headers={"content-type": "image/png"})
+
+    monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+    _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+    image = await video_api_utils.decode_image_url("https://example.com/redirect.png")
+
+    assert image.size == (2, 1)
+    assert image.mode == "RGB"
+    assert requested_paths == ["/redirect.png", "/image.png"]
+    assert client_kwargs == [{"timeout": 60, "follow_redirects": True}]
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
+    requested_paths = []
+    client_kwargs = []
+
+    def _handler(request):
+        requested_paths.append(request.url.path)
+        return httpx.Response(302, headers={"location": "/image.png"})
+
+    monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", False)
+    _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+    with pytest.raises(InvalidInputReferenceError, match="redirect.*VLLM_MEDIA_URL_ALLOW_REDIRECTS"):
+        await video_api_utils.decode_image_url("https://example.com/redirect.png")
+
+    assert requested_paths == ["/redirect.png"]
+    assert client_kwargs == [{"timeout": 60, "follow_redirects": False}]
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_reports_http_status(monkeypatch):
+    client_kwargs = []
+
+    def _handler(request):
+        return httpx.Response(404)
+
+    monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+    _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+    with pytest.raises(InvalidInputReferenceError, match="server returned HTTP 404"):
+        await video_api_utils.decode_image_url("https://example.com/missing.png")
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_reports_connection_failure(monkeypatch):
+    client_kwargs = []
+
+    def _handler(request):
+        raise httpx.ConnectError("connection refused", request=request)
+
+    monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+    _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+    with pytest.raises(InvalidInputReferenceError, match="failed to download image"):
+        await video_api_utils.decode_image_url("https://example.com/unreachable.png")
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_keeps_data_urls_local(monkeypatch):
+    def _unexpected_http_client(*args, **kwargs):
+        pytest.fail("data URLs must not create an HTTP client")
+
+    monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _unexpected_http_client)
+    encoded_image = base64.b64encode(_png_bytes()).decode()
+
+    image = await video_api_utils.decode_image_url(f"data:image/png;base64,{encoded_image}")
+
+    assert image.size == (2, 1)
+    assert image.mode == "RGB"
+
+
 def _install_fake_video_mux(monkeypatch, mux_calls):
     def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
         mux_calls.append(

tests/entrypoints/openai_api/test_video_server.py

@@ -16,15 +16,17 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import av
+import httpx
 import numpy as np
 import pytest
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture
+from vllm import envs
 
 from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.openai import api_server
+from vllm_omni.entrypoints.openai import api_server, video_api_utils
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoGenerationRequest,
@@ -599,6 +601,45 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
     assert input_image.size == (40, 24)
 
 
+def test_i2v_video_generation_follows_allowed_image_redirect(test_client, mocker: MockerFixture, monkeypatch):
+    requested_paths = []
+    async_client = httpx.AsyncClient
+
+    def _handler(request):
+        requested_paths.append(request.url.path)
+        if request.url.path == "/redirect.png":
+            return httpx.Response(302, headers={"location": "/image.png"})
+        return httpx.Response(200, content=_make_test_image_bytes((40, 24)))
+
+    def _client_factory(*args, **kwargs):
+        return async_client(*args, transport=httpx.MockTransport(_handler), **kwargs)
+
+    monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+    monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A fox running through snow.",
+            "image_reference": json.dumps({"image_url": "https://example.com/redirect.png"}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    assert requested_paths == ["/redirect.png", "/image.png"]
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_image = engine.captured_prompt["multi_modal_data"]["image"]
+    assert isinstance(input_image, Image.Image)
+    assert input_image.size == (40, 24)
+
+
 def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
     video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -58,6 +58,23 @@ def test_default_stage_devices_from_sequence_parallel():
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
 
 
+def test_default_stage_devices_and_dp_from_num_gpus():
+    """Resolve omitted DP before deriving devices from an explicit GPU count."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "num_gpus": 8,
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.data_parallel_size == 2
+    assert parallel_config.world_size == 8
+    assert stage_cfg["engine_args"]["num_gpus"] == 8
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
+
+
 def test_default_stage_config_uses_parallel_size_kwargs():
     """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -76,8 +93,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
     assert parallel_config.enable_expert_parallel is True
 
 
-def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
-    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
+    """Keep omitted DP unresolved until the runtime WORLD size is known."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
         {
             "pipeline_parallel_size": None,
@@ -92,7 +109,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
     assert parallel_config.pipeline_parallel_size == 1
-    assert parallel_config.data_parallel_size == 1
+    assert parallel_config.data_parallel_size is None
     assert parallel_config.tensor_parallel_size == 1
     assert parallel_config.enable_expert_parallel is False
     assert stage_cfg["engine_args"]["enforce_eager"] is False
@@ -330,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-fp8-cache-dir",
+            "/tmp/dlo-fp8-cache",
         ]
     )
 
@@ -340,9 +359,11 @@ 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_fp8_cache_dir == "/tmp/dlo-fp8-cache"
     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_fp8_cache_dir"] == "/tmp/dlo-fp8-cache"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

tests/helpers/assertions.py

@@ -2,12 +2,16 @@
 
 import io
 import json
+import math
 import tempfile
 import threading
 import wave
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 if TYPE_CHECKING:
     from tests.helpers.runtime import DiffusionResponse
@@ -24,8 +28,109 @@ from tests.helpers.media import (
     preprocess_text,
 )
 
+_ResponseT = TypeVar("_ResponseT")
+
 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
+
+# Assertions that sample a random variable rather than testing whether serving works:
+# whether TTS audio is intelligible enough to transcribe back, and whether a preset voice's
+# timbre estimates as the expected gender. Both are decided by the model and the estimator,
+# so a caller measuring a success rate has to tell them apart from a serving failure. The
+# messages are constants because a caller matching the wording by hand would silently stop
+# matching if it were reworded, reclassifying every quality failure as a hard failure.
+AUDIO_MISMATCH_MESSAGE = "The audio content is not same as the text"
+GENDER_MISMATCH_MESSAGE = "estimated gender is"
+_QUALITY_FAILURE_MESSAGES = (AUDIO_MISMATCH_MESSAGE, GENDER_MISMATCH_MESSAGE)
+
+
+@dataclass(frozen=True)
+class SuccessRateGate:
+    """Gate a sampled quality check on how many requests pass, not on each one passing.
+
+    Built by the send helpers from their own arguments, so a test states a policy and never
+    touches this class. See ``collect_at_success_rate``.
+
+    Attributes:
+        min_successes: Fail below this many successes out of ``request_num``. Predeclare it
+            from a measured baseline; tuning it until CI passes defeats the gate.
+        concurrency: In-flight requests, in batches until the sample count is reached.
+            Defaults to sending all of them at once.
+    """
+
+    min_successes: int
+    concurrency: int | None = None
+
+
+def is_quality_failure(exc: BaseException) -> bool:
+    """True when ``exc`` is one of the sampled quality assertions above.
+
+    A caller absorbing a few of these into a success-rate budget must not absorb an HTTP
+    failure, missing audio or a text-keyword miss as well, or a broken server passes.
+    """
+    if not isinstance(exc, AssertionError):
+        return False
+    return any(message in str(exc) for message in _QUALITY_FAILURE_MESSAGES)
+
+
+def collect_at_success_rate(
+    sample: Callable[[], _ResponseT],
+    gate: SuccessRateGate,
+    *,
+    request_num: int,
+    report: Callable[[str], None] = print,
+) -> list[_ResponseT]:
+    """Send ``request_num`` samples and gate on how many passed, returning the successes.
+
+    ``sample`` sends one request and asserts it. Only the sampled quality assertions are
+    counted; anything else propagates immediately. Does not retry and does not exit early,
+    since both would bias the rate being measured.
+    """
+    if gate.min_successes > request_num:
+        raise ValueError(f"min_successes={gate.min_successes} exceeds request_num={request_num}, gate can never pass")
+    concurrency = min(gate.concurrency or request_num, request_num)
+    responses: list[_ResponseT] = []
+    failures: list[str] = []
+    lock = threading.Lock()
+
+    def _one_sample(_: int) -> None:
+        try:
+            response = sample()
+        except AssertionError as exc:
+            if not is_quality_failure(exc):
+                raise
+            with lock:
+                failures.append(str(exc))
+            return
+        with lock:
+            responses.append(response)
+
+    for start in range(0, request_num, concurrency):
+        size = min(concurrency, request_num - start)
+        with ThreadPoolExecutor(max_workers=size) as executor:
+            for _ in executor.map(_one_sample, range(size)):
+                pass
+
+    successes = request_num - len(failures)
+    report(f"success rate: {successes}/{request_num} succeeded")
+    assert successes >= gate.min_successes, (
+        f"success rate {successes}/{request_num} is below the "
+        f"{gate.min_successes}/{request_num} gate "
+        f"(95% CI {wilson_interval(successes, request_num)}). Failures:\n"
+        + "\n".join(f"  - {failure}" for failure in failures)
+    )
+    return responses
+
+
+def wilson_interval(successes: int, total: int, z: float = 1.96) -> str:
+    """Format a 95% Wilson score interval, so a rate failure shows how marginal it is."""
+    p = successes / total
+    denom = 1 + z * z / total
+    centre = (p + z * z / (2 * total)) / denom
+    half = z / denom * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total))
+    return f"{max(0.0, centre - half):.1%}-{min(1.0, centre + half):.1%}"
+
+
 # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
 # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
 # leading SFX token), which flakes the deterministic similarity gate. Short
@@ -474,7 +579,7 @@ def _assert_preset_voice_gender_from_audio(
     print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
     if estimated_gender != "unknown":
         assert estimated_gender == expected_gender, (
-            f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
+            f"{voice_name!r} is expected {expected_gender}, but {GENDER_MISMATCH_MESSAGE} {estimated_gender!r}"
         )
 
 
@@ -655,7 +760,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                     shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
                     longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
                     assert shorter_clean and (shorter_clean in longer_clean), (
-                        f"The audio content is not same as the text "
+                        f"{AUDIO_MISMATCH_MESSAGE} "
                         f"(short-text containment check failed: "
                         f"text={text_output!r}, transcript={transcript!r})"
                     )
@@ -666,7 +771,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                         text_output.lower(),
                     )
                     print(f"similarity is: {similarity}")
-                    assert similarity > similarity_threshold, "The audio content is not same as the text"
+                    assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE
             if audio_ref_text:
                 assert transcript is not None, "No audio transcript for reference-text validation"
                 audio_similarity = cosine_similarity_text(

tests/helpers/runtime.py

@@ -38,11 +38,13 @@ from vllm.distributed.parallel_state import (
 from vllm.logger import init_logger
 
 from tests.helpers.assertions import (
+    SuccessRateGate,
     assert_audio_speech_response,
     assert_diffusion_response,
     assert_http_error,
     assert_images_generations_response,
     assert_omni_response,
+    collect_at_success_rate,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.media import (
@@ -1891,8 +1893,20 @@ class OpenAIClientHandler:
             )
         ]
 
-    def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
-        """Chat completions via the OpenAI Python SDK (not raw HTTP)."""
+    def send_omni_request(
+        self,
+        request_config: dict[str, Any],
+        request_num: int = 1,
+        *,
+        min_successes: int | None = None,
+        max_concurrency: int | None = None,
+    ) -> list[OmniResponse]:
+        """Chat completions via the OpenAI Python SDK (not raw HTTP).
+
+        ``min_successes`` gates the batch on how many of the ``request_num`` requests pass
+        rather than requiring each one to, and returns only the successes; ``max_concurrency``
+        caps how many are in flight at once. See ``SuccessRateGate``.
+        """
         responses: list[OmniResponse] = []
         stream = request_config.get("stream", False)
         modalities = request_config.get("modalities", ["text", "audio"])
@@ -1923,22 +1937,6 @@ class OpenAIClientHandler:
         if extra_body:
             create_kwargs["extra_body"] = extra_body
 
-        if request_num == 1:
-            wall_start = time.perf_counter()
-            chat_completion = self.client.chat.completions.create(**create_kwargs)
-            resp = (
-                self._process_stream_omni_response(chat_completion, wall_start=wall_start)
-                if stream
-                else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
-            )
-            assert_omni_response(resp, request_config, run_level=self.run_level)
-            if resp.e2e_latency is not None:
-                self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
-            else:
-                self._print_client_stat("[omni] request#1 completed")
-            responses.append(resp)
-            return responses
-
         def _one():
             wall_start = time.perf_counter()
             chat_completion = self.client.chat.completions.create(**create_kwargs)
@@ -1948,6 +1946,33 @@ class OpenAIClientHandler:
                 else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
             )
 
+        if min_successes is not None:
+            gate = SuccessRateGate(min_successes=min_successes, concurrency=max_concurrency)
+
+            def _sample() -> OmniResponse:
+                resp = _one()
+                assert_omni_response(resp, request_config, run_level=self.run_level)
+                return resp
+
+            # Not _print_client_stat: the count is the gate's verdict, so it has to survive
+            # log_stats=False.
+            return collect_at_success_rate(
+                _sample,
+                gate,
+                request_num=request_num,
+                report=lambda line: print(f"[omni] {line}", flush=True),
+            )
+
+        if request_num == 1:
+            resp = _one()
+            assert_omni_response(resp, request_config, run_level=self.run_level)
+            if resp.e2e_latency is not None:
+                self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
+            else:
+                self._print_client_stat("[omni] request#1 completed")
+            responses.append(resp)
+            return responses
+
         with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
             futures = {executor.submit(_one): i + 1 for i in range(request_num)}
             for future in concurrent.futures.as_completed(futures):
@@ -3240,6 +3265,146 @@ def iter_omni_runner(
         print("OmniRunner stopped")
 
 
+# ─────────────────────────────────────────────────────────────────────
+# π0 (Pi-Zero) OpenPI websocket policy client
+#
+# Minimal client for the OpenPI realtime robot protocol
+# (``/v1/realtime/robot/openpi``): connect → read handshake metadata → send
+# observation → receive an action chunk. Used by the π0 online-serving e2e
+# (tests/e2e/online_serving/test_pi0_expansion.py). π0's observation is 3 RGB
+# cameras + proprioceptive state + a language prompt, built as synthetic numpy
+# frames here (no video assets).
+# ─────────────────────────────────────────────────────────────────────
+PI0_OPENPI_PATH = "/v1/realtime/robot/openpi"
+PI0_OPENPI_DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+# π0 / pi0_base camera identities (must match the server's image_feature_keys,
+# i.e. vllm_omni/deploy/pi0.yaml, whose image_key_map is empty).
+PI0_CAMERA_KEYS = (
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+)
+PI0_ACTION_HORIZON = 50
+PI0_ACTION_DIM = 32
+PI0_STATE_DIM = 32
+PI0_IMAGE_SIZE = 224
+
+
+def pi0_openpi_require_dependencies() -> None:
+    """Raise ModuleNotFoundError if the OpenPI websocket client deps are missing."""
+    missing = []
+    try:
+        import websockets.sync.client  # noqa: F401
+    except ImportError:
+        missing.append("websockets")
+    try:
+        from openpi_client import msgpack_numpy  # noqa: F401
+    except ImportError:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"π0 OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+def _pi0_decode_action_response(response: bytes | str) -> np.ndarray:
+    from openpi_client import msgpack_numpy
+
+    if isinstance(response, str):
+        raise RuntimeError(f"Inference failed: {response}")
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict) and decoded.get("type") == "error":
+        raise RuntimeError(f"Inference failed: {decoded.get('message', decoded)}")
+    return np.asarray(decoded, dtype=np.float32)
+
+
+def pi0_make_dummy_obs(*, prompt: str, session_id: str, image_size: int = PI0_IMAGE_SIZE) -> dict[str, Any]:
+    """A single π0 observation: 3 blank cameras (HWC uint8) + zero state + prompt."""
+    obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in PI0_CAMERA_KEYS}
+    obs["state"] = np.zeros(PI0_STATE_DIM, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def pi0_openpi_run_policy_session(
+    *,
+    host: str = "127.0.0.1",
+    port: int = 8000,
+    path: str = PI0_OPENPI_PATH,
+    prompt: str = PI0_OPENPI_DEFAULT_PROMPT,
+    session_id: str | None = None,
+    num_steps: int = 2,
+) -> dict[str, Any]:
+    """Connect, read handshake metadata, send ``num_steps`` observations."""
+    import uuid
+
+    import websockets.sync.client as websockets_client
+    from openpi_client import msgpack_numpy
+
+    session_id = session_id or str(uuid.uuid4())
+    uri = f"ws://{host}:{port}{path}"
+    packer = msgpack_numpy.Packer()
+    conn = websockets_client.connect(uri, compression=None, max_size=None, ping_interval=300, ping_timeout=3600)
+    try:
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        actions = []
+        for _ in range(num_steps):
+            payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id)
+            payload["endpoint"] = "infer"
+            conn.send(packer.pack(payload))
+            actions.append(_pi0_decode_action_response(conn.recv()))
+        return {"metadata": dict(metadata), "actions": actions, "session_id": session_id}
+    finally:
+        conn.close()
+
+
+def pi0_openpi_validate_session_result(
+    result: dict[str, Any],
+    *,
+    expected_action_horizon: int = PI0_ACTION_HORIZON,
+    expected_action_dim: int = PI0_ACTION_DIM,
+) -> None:
+    """Assert the handshake metadata + every returned action chunk for π0."""
+    metadata = result["metadata"]
+    required_keys = (
+        "image_resolution",
+        "needs_wrist_camera",
+        "needs_session_id",
+        "action_space",
+        "action_horizon",
+        "action_dim",
+    )
+    missing = [key for key in required_keys if key not in metadata]
+    if missing:
+        raise AssertionError(f"Missing π0 metadata keys: {missing}")
+
+    if tuple(metadata["image_resolution"]) != (PI0_IMAGE_SIZE, PI0_IMAGE_SIZE):
+        raise AssertionError(f"Unexpected image_resolution: {metadata['image_resolution']!r}")
+    if not metadata["needs_wrist_camera"]:
+        raise AssertionError("π0 test expects needs_wrist_camera=True")
+    if metadata["needs_session_id"]:
+        raise AssertionError("π0 is stateless; needs_session_id must be False")
+    if metadata["action_space"] != "joint_position":
+        raise AssertionError(f"Unexpected action_space: {metadata['action_space']!r}")
+    if int(metadata["action_horizon"]) != expected_action_horizon:
+        raise AssertionError(f"Unexpected action_horizon: {metadata['action_horizon']!r}")
+    if int(metadata["action_dim"]) != expected_action_dim:
+        raise AssertionError(f"Unexpected action_dim: {metadata['action_dim']!r}")
+
+    actions = result["actions"]
+    if not actions:
+        raise AssertionError("No actions returned from the server")
+    for index, action in enumerate(actions):
+        if action.shape != (expected_action_horizon, expected_action_dim):
+            raise AssertionError(
+                f"Action {index} shape mismatch: expected "
+                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+            )
+        if not np.isfinite(action).all():
+            raise AssertionError(f"Action {index} contains non-finite values")
+
+
 __all__ = [
     "DiffusionResponse",
     "HttpResponse",
@@ -3261,4 +3426,8 @@ __all__ = [
     "OpenAIClientHandler",
     "get_open_port",
     "dummy_messages_from_mix_data",
+    "pi0_openpi_require_dependencies",
+    "pi0_openpi_run_policy_session",
+    "pi0_openpi_validate_session_result",
+    "pi0_make_dummy_obs",
 ]

tests/helpers/tests/test_success_rate_helpers.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers import assertions
+from tests.helpers.assertions import (
+    AUDIO_MISMATCH_MESSAGE,
+    SuccessRateGate,
+    assert_omni_response,
+    collect_at_success_rate,
+    is_quality_failure,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_GATE = SuccessRateGate(min_successes=7, concurrency=2)
+
+
+def _fake_response(text: str) -> SimpleNamespace:
+    return SimpleNamespace(success=True, text_content=text, audio_content=None, audio_bytes=b"fake-wav")
+
+
+def _omni_response_failure(monkeypatch, *, text: str, transcript: str) -> AssertionError:
+    """Drive the real assert_omni_response to failure and hand back what it raised."""
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *a, **kw: transcript)
+    request_config = {"modalities": ["text", "audio"], "key_words": {"text": ["london"]}}
+    with pytest.raises(AssertionError) as excinfo:
+        assert_omni_response(_fake_response(text), request_config, run_level="full_model")
+    return excinfo.value
+
+
+def test_real_audio_mismatch_is_tolerated(monkeypatch):
+    # Ties the classifier to the actual producer: if assert_omni_response's wording drifts,
+    # this fails instead of the gate silently turning every quality failure into a hard one.
+    assert is_quality_failure(_omni_response_failure(monkeypatch, text="London", transcript="런던"))
+
+
+def test_real_keyword_failure_is_not_tolerated(monkeypatch):
+    assert not is_quality_failure(_omni_response_failure(monkeypatch, text="Paris", transcript="London"))
+
+
+@pytest.mark.parametrize("exc", [RuntimeError(AUDIO_MISMATCH_MESSAGE), ValueError("boom")])
+def test_non_assertion_failures_are_not_tolerated(exc):
+    # A transport error carrying the same text is still a serving failure, not a sample.
+    assert not is_quality_failure(exc)
+
+
+class _ScriptedSampler:
+    """One scripted outcome per call: raised if it is an exception, else returned."""
+
+    def __init__(self, outcomes: list) -> None:
+        self._outcomes = outcomes
+        self._lock = threading.Lock()
+        self.calls = 0
+
+    def __call__(self):
+        with self._lock:
+            outcome = self._outcomes[self.calls]
+            self.calls += 1
+        if isinstance(outcome, BaseException):
+            raise outcome
+        return outcome
+
+
+def _mismatch(detail: str) -> AssertionError:
+    return AssertionError(f"{AUDIO_MISMATCH_MESSAGE} (transcript={detail!r})")
+
+
+def test_tolerated_failures_are_counted_and_every_sample_is_sent():
+    sampler = _ScriptedSampler(["ok"] * 7 + [_mismatch("a"), _mismatch("b"), _mismatch("c")])
+
+    responses = collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
+
+    # No early exit once the gate is already satisfied: it would bias the measured rate.
+    assert sampler.calls == 10
+    assert responses == ["ok"] * 7
+
+
+def test_gate_failure_reports_the_count_and_every_failing_sample():
+    outcomes = ["ok"] * 6 + [_mismatch("first"), _mismatch("second"), _mismatch("third"), _mismatch("fourth")]
+
+    with pytest.raises(AssertionError) as excinfo:
+        collect_at_success_rate(_ScriptedSampler(outcomes), _GATE, request_num=10, report=lambda _: None)
+
+    message = str(excinfo.value)
+    assert "6/10" in message and "7/10 gate" in message
+    assert all(detail in message for detail in ("first", "second", "third", "fourth"))
+
+
+def test_serving_failure_is_not_absorbed_into_the_budget():
+    sampler = _ScriptedSampler(["ok", AssertionError("No audio output is generated"), *["ok"] * 8])
+
+    with pytest.raises(AssertionError, match="No audio output"):
+        collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
+
+    # Stopped on the serving failure rather than spending the remaining samples.
+    assert sampler.calls < 10
+
+
+def test_unreachable_gate_is_rejected_before_sending():
+    sampler = _ScriptedSampler(["ok"] * 3)
+
+    with pytest.raises(ValueError, match="can never pass"):
+        collect_at_success_rate(sampler, _GATE, request_num=3, report=lambda _: None)
+
+    assert sampler.calls == 0

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -24,6 +24,7 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
+    _apply_platform_overrides,
     load_deploy_config,
     merge_pipeline_deploy,
 )
@@ -184,6 +185,30 @@ class TestDeployTopology:
                 0.35,
             ]
 
+    @pytest.mark.parametrize(
+        "filename",
+        [
+            "minicpmo_4_5.yaml",
+            "minicpmo_4_5_2gpu.yaml",
+            "minicpmo_4_5_3gpu.yaml",
+        ],
+    )
+    def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None:
+        deploy = load_deploy_config(_DEPLOY_DIR / filename)
+        connector_extra = deploy.connectors["connector_of_shared_memory"]["extra"]
+        assert "code2wav_enable_npu_graph" not in connector_extra
+        assert "code2wav_max_npu_graphs" not in connector_extra
+
+        deploy = _apply_platform_overrides(deploy, platform="npu")
+        stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
+        assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+        assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+        assert stages[2].yaml_engine_args["enforce_eager"] is True
+        assert stages[2].yaml_engine_args["additional_config"] == {
+            "code2wav_enable_npu_graph": True,
+            "code2wav_max_npu_graphs": 32,
+        }
+
     def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None:
         pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
         talker = pipeline.get_stage(1)

tests/model_executor/models/test_minimax_music3_batch_payload.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 batched payload partitioning.
+
+The stage-0 payload is split per request by indexing list values: a list is
+indexed by batch row, anything else is handed to every request unchanged. So
+per-request metadata has to be carried as lists. Merging it into one dict made
+every request in a batch receive the last request's seed, window offset and
+terminal flag, which crosses seeds between songs, misplaces acoustic windows,
+and lets one request's terminal flag end another request's stream early.
+
+These drive ``make_omni_output`` directly with a hand-built batch, so no GPU
+and no engine are needed.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    MiniMaxMusic3TalkerForConditionalGeneration as Talker,
+)
+from vllm_omni.utils.mm_outputs import to_payload_element
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_WIDTH = 4
+
+
+def _info(*, seed: int, start: int, frames: int, last: bool) -> dict:
+    """One batch row's intermediate buffer, as ``postprocess`` leaves it."""
+    return {
+        "latent": torch.full((frames, _WIDTH), float(seed)),
+        "codes": {"audio": torch.full((frames, 8), seed)},
+        "meta": {
+            "audio_seed": seed,
+            "num_processed_tokens": start,
+            "codec_chunk_frames": frames,
+            "last_chunk": last,
+        },
+    }
+
+
+def _payload_for(output, row: int) -> dict:
+    """Partition the batch payload down to one request, as the runner does."""
+    flat = {}
+    for key, value in output.multimodal_outputs.items():
+        if isinstance(value, dict):
+            for qual, val in value.items():
+                flat[f"{key}.{qual}"] = val
+        else:
+            flat[key] = value
+    return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()}
+
+
+class TestBatchedMetadataPartitioning:
+    """Two requests in one batch must keep their own metadata."""
+
+    @staticmethod
+    def _two_request_batch():
+        infos = [
+            _info(seed=111, start=100, frames=3, last=False),
+            _info(seed=222, start=200, frames=5, last=True),
+        ]
+        talker = Talker.__new__(Talker)
+        return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+    def test_seeds_are_not_crossed(self) -> None:
+        """A song decoded with another request's seed is a different song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.audio_seed"] == 111
+        assert _payload_for(output, 1)["meta.audio_seed"] == 222
+
+    def test_window_offsets_stay_with_their_request(self) -> None:
+        """The start frame decides where a window lands in the song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100
+        assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200
+
+    def test_terminal_flag_does_not_leak_across_requests(self) -> None:
+        """One request finishing must not end another request's stream."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.last_chunk"] is False
+        assert _payload_for(output, 1)["meta.last_chunk"] is True
+
+    def test_span_length_matches_its_own_latent(self) -> None:
+        output = self._two_request_batch()
+        for row, frames in ((0, 3), (1, 5)):
+            entry = _payload_for(output, row)
+            assert entry["meta.codec_chunk_frames"] == frames
+            assert entry["latent"].shape[0] == frames
+
+    def test_latent_and_codes_follow_the_same_row(self) -> None:
+        """The tensors were already per-request; pin that they stay aligned."""
+        output = self._two_request_batch()
+        for row, seed in ((0, 111), (1, 222)):
+            entry = _payload_for(output, row)
+            assert entry["latent"][0, 0].item() == pytest.approx(float(seed))
+            assert int(entry["codes.audio"][0, 0]) == seed
+
+
+class TestPartialBatches:
+    """Only some rows carry a span on any given step."""
+
+    def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None:
+        infos = [{}, _info(seed=777, start=300, frames=2, last=True)]
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+        idle = _payload_for(output, 0)
+        assert idle["latent"].numel() == 0
+        assert idle["meta.audio_seed"] is None
+        assert idle["meta.last_chunk"] is None
+
+        active = _payload_for(output, 1)
+        assert active["meta.audio_seed"] == 777
+        assert active["meta.last_chunk"] is True
+
+    def test_batch_with_no_spans_emits_no_payload(self) -> None:
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}])
+        assert output.multimodal_outputs is None
+
+    def test_consumed_entries_are_cleared_from_the_buffer(self) -> None:
+        """The runner never clears this buffer, so a span left behind would be
+        re-emitted on every later step."""
+        infos = [_info(seed=1, start=0, frames=2, last=False)]
+        talker = Talker.__new__(Talker)
+        talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos)
+        assert "latent" not in infos[0]
+        assert "meta" not in infos[0]

tests/model_executor/models/test_minimax_music3_frame_state.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping.
+
+The talker emits overlapping windows while dropping frames it can no longer
+need, and flushes whatever is left as one span when the song ends. Getting the
+eviction boundary wrong by one frame silently corrupts the conditioning handed
+to the acoustic stage from the second window onward, which is audible but hard
+to attribute, so the lifecycle is simulated here frame by frame.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    NUM_CODEBOOKS,
+)
+from vllm_omni.model_executor.models.minimax_music3.frame_state import (
+    FrameBuffer,
+    RequestARState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# One scalar per frame instead of 32768 dims: the bookkeeping is what is under
+# test, and a frame's value doubles as its identity.
+_WIDTH = 4
+
+
+def _frame(index: int) -> torch.Tensor:
+    return torch.full((1, _WIDTH), float(index))
+
+
+def _identities(chunk: torch.Tensor) -> list[int]:
+    return [int(row[0].item()) for row in chunk]
+
+
+class TestFrameBuffer:
+    def test_slice_returns_the_requested_frames(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6]
+
+    def test_slice_is_relative_to_the_live_base(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(4)
+        assert buffer.base_frame == 4
+        assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7]
+
+    def test_slicing_a_discarded_range_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(5)
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(3, 6)
+
+    def test_slicing_past_the_end_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(4):
+            buffer.append(_frame(index))
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(2, 9)
+
+    def test_discard_is_clamped_at_both_ends(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(5):
+            buffer.append(_frame(index))
+        buffer.discard_before(-3)
+        assert buffer.base_frame == 0
+        buffer.discard_before(99)
+        assert buffer.base_frame == 5
+        assert buffer.end_frame == 5
+
+    def test_peak_tracks_the_high_water_mark(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(8)
+        assert buffer.peak_frames == 10
+
+
+class TestWindowLifecycle:
+    """Drive the state exactly as the talker does, one frame per step."""
+
+    @staticmethod
+    def _run(total_frames: int) -> tuple[list[list[int]], int]:
+        """Generate ``total_frames`` frames, returning the emitted spans.
+
+        Returns:
+            The frame identities of every emitted span, and the peak buffer
+            occupancy reached along the way.
+        """
+        state = RequestARState(seed=1, max_frames=total_frames)
+        spans: list[list[int]] = []
+        for index in range(total_frames):
+            state.frames.append(_frame(index))
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+            state.generated_frames += 1
+            window = state.ready_window()
+            if window is not None:
+                spans.append(_identities(state.take_window(window, is_final=False)))
+
+        # Terminal flush: everything not yet emitted, as one contiguous span.
+        start = state.emitted_windows * AR_CHUNK_HOP_FRAMES
+        end = state.frames.end_frame
+        if end > start:
+            spans.append(_identities(state.frames.slice(start, end)))
+        return spans, state.frames.peak_frames
+
+    def test_short_song_is_one_span(self) -> None:
+        spans, _ = self._run(50)
+        assert spans == [list(range(50))]
+
+    def test_every_frame_is_delivered_exactly_once_in_order(self) -> None:
+        """Spans overlap by design, but their union must be the whole song.
+
+        Reconstruct by taking each span's non-overlapping remainder, which is
+        what the acoustic stage does after cropping.
+        """
+        for total in (250, 437, 600, 901):
+            spans, _ = self._run(total)
+            rebuilt: list[int] = []
+            for span in spans:
+                rebuilt.extend(frame for frame in span if frame >= len(rebuilt))
+            assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames"
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        spans, _ = self._run(600)
+        for index, span in enumerate(spans[:-1]):
+            assert span[0] == index * AR_CHUNK_HOP_FRAMES
+            assert len(span) == AR_CHUNK_FRAMES
+
+    def test_mid_stream_windows_are_never_the_final_frame(self) -> None:
+        """A window is held back until a further frame proves it is not last."""
+        total = 600
+        spans, _ = self._run(total)
+        for span in spans[:-1]:
+            assert span[-1] < total - 1
+
+    def test_terminal_span_covers_the_tail(self) -> None:
+        total = 437
+        spans, _ = self._run(total)
+        assert spans[-1][-1] == total - 1
+
+    def test_terminal_span_is_bounded(self) -> None:
+        """The tail is at most one full window plus one hop."""
+        for total in (250, 301, 437, 600, 901, 1000):
+            spans, _ = self._run(total)
+            assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_buffer_stays_bounded_over_a_long_song(self) -> None:
+        """Peak occupancy must not grow with song length; that is the point
+        of evicting retired frames."""
+        _, peak_short = self._run(600)
+        _, peak_long = self._run(3000)
+        assert peak_long == peak_short
+        assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_emitted_window_count_matches_the_window_plan(self) -> None:
+        """What the stream emitted must reconcile with chunk_windows()."""
+        for total in (250, 437, 600):
+            spans, _ = self._run(total)
+            planned = chunk_windows(total)
+            # Every window is accounted for: the pre-terminal spans are one
+            # window each, and the terminal span carries whatever remains.
+            assert len(spans) <= len(planned)
+            assert spans[0][0] == planned[0].start
+
+
+class TestSamplingPositions:
+    def test_each_frame_reserves_one_slot_per_codebook(self) -> None:
+        """Codebook streams must never collide, or codes repeat across frames."""
+        state = RequestARState(seed=1, max_frames=10)
+        positions = [state.advance_position() for _ in range(5)]
+        assert positions == [0, 8, 16, 24, 32]
+
+    def test_draw_slots_are_globally_unique(self) -> None:
+        state = RequestARState(seed=1, max_frames=64)
+        slots: set[int] = set()
+        for _ in range(64):
+            base = state.advance_position()
+            for codebook in range(NUM_CODEBOOKS):
+                slot = base + codebook
+                assert slot not in slots
+                slots.add(slot)
+        assert len(slots) == 64 * NUM_CODEBOOKS
+
+
+class TestCodeRetention:
+    def test_codes_span_matches_the_frame_span(self) -> None:
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(10):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        codes = state.take_codes(2, 6)
+        assert codes is not None
+        assert codes.shape == (4, NUM_CODEBOOKS)
+        assert [int(row[0]) for row in codes] == [2, 3, 4, 5]
+
+    def test_out_of_range_span_is_dropped_not_raised(self) -> None:
+        """Codes are diagnostic; the audio path must not fail without them."""
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(4):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        assert state.take_codes(2, 99) is None
+        assert state.take_codes(0, 0) is None

tests/model_executor/models/test_minimax_music3_prompt.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.
+
+The prompt transformations are not cosmetic: they decide token positions, and
+token positions decide the audio. These tests pin the documented behaviours
+that a well-meaning cleanup would otherwise "fix".
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import (
+    chunk_windows,
+    crop_sample_bounds,
+    overlap_mel_length,
+)
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    BLEND_MEL_FRAMES,
+    DAV_HOP_SAMPLES,
+    HOP_MEL_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.prompt import (
+    SPECIAL_TOKEN_IDS,
+    build_cfg_null_token_ids,
+    build_prompt,
+    clean_caption,
+    normalize_lyrics,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestNormalizeLyrics:
+    def test_prepends_start_tag(self) -> None:
+        assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")
+
+    def test_lowercases_structure_tags(self) -> None:
+        assert "[verse]" in normalize_lyrics("[Verse]\nhello")
+        assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")
+
+    def test_text_after_a_leading_tag_is_dropped(self) -> None:
+        """Documented contract: a tag must sit on its own line.
+
+        Lyrics written next to a tag are silently lost, which is why the
+        serving docs tell callers to put tags on their own line.
+        """
+        out = normalize_lyrics("[Verse] Walking down the street")
+        assert "Walking" not in out
+        assert "[verse]" in out
+
+    def test_text_on_its_own_line_survives(self) -> None:
+        out = normalize_lyrics("[Verse]\nWalking down the street")
+        assert "Walking down the street" in out
+
+    def test_bracket_space_becomes_newline(self) -> None:
+        assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"
+
+
+class TestCleanCaption:
+    def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
+        """A caption must not be able to inject a real special token."""
+        assert clean_caption("<|bpm 92|>") == "bpm is 92"
+        assert "<|" not in clean_caption("<|bpm 92|>")
+
+    def test_lone_tag_keeps_its_inner_text(self) -> None:
+        assert clean_caption("<|instrumental|>") == "instrumental"
+
+    def test_markdown_is_stripped(self) -> None:
+        assert clean_caption("## Heading") == "Heading"
+        assert clean_caption("**bold**") == "bold"
+        assert clean_caption("- item") == "item"
+
+    def test_blank_lines_collapse_at_two_not_three(self) -> None:
+        """The source collapses runs of two or more, not three or more.
+
+        This shifts blank-line token positions and therefore the audio, so it
+        is pinned deliberately.
+        """
+        assert clean_caption("a\n\nb") == "a\nb"
+        assert clean_caption("a\n\n\n\nb") == "a\nb"
+
+
+class TestBuildPrompt:
+    def test_framing_tokens_are_present_and_ordered(self) -> None:
+        prompt = build_prompt("a bright pop song", "[Verse]\nhello")
+        assert prompt.startswith("<|im_start|><|caption_start|>")
+        assert prompt.endswith("<|im_end|><|audio_start|>")
+        assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")
+
+
+class TestBuildCfgNullTokenIds:
+    @staticmethod
+    def _prompt_ids(body_len: int) -> list[int]:
+        return (
+            [SPECIAL_TOKEN_IDS["<|im_start|>"]]
+            + list(range(1000, 1000 + body_len))
+            + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
+        )
+
+    def test_length_is_preserved(self) -> None:
+        """Equal length is what lets the two rows decode in lockstep."""
+        ids = self._prompt_ids(12)
+        assert len(build_cfg_null_token_ids(ids)) == len(ids)
+
+    def test_body_is_replaced_and_framing_survives(self) -> None:
+        ids = self._prompt_ids(12)
+        null = build_cfg_null_token_ids(ids)
+        assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
+        assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
+        assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
+        assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}
+
+    def test_input_is_not_mutated(self) -> None:
+        ids = self._prompt_ids(6)
+        before = list(ids)
+        build_cfg_null_token_ids(ids)
+        assert ids == before
+
+    def test_minimal_prompt_with_empty_body(self) -> None:
+        ids = [
+            SPECIAL_TOKEN_IDS["<|im_start|>"],
+            SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
+            SPECIAL_TOKEN_IDS["<|im_end|>"],
+            SPECIAL_TOKEN_IDS["<|audio_start|>"],
+        ]
+        assert build_cfg_null_token_ids(ids) == ids
+
+    def test_too_short_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="too short"):
+            build_cfg_null_token_ids([1, 2, 3])
+
+    def test_wrong_framing_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="framing"):
+            build_cfg_null_token_ids([9, 9, 9, 9, 9])
+
+
+class TestChunkWindows:
+    def test_no_frames_yields_nothing(self) -> None:
+        assert chunk_windows(0) == []
+
+    def test_short_song_is_one_final_window(self) -> None:
+        windows = chunk_windows(50)
+        assert len(windows) == 1
+        assert (windows[0].start, windows[0].end) == (0, 50)
+        assert windows[0].is_first and windows[0].is_last
+
+    def test_exactly_one_window(self) -> None:
+        windows = chunk_windows(AR_CHUNK_FRAMES)
+        assert len(windows) == 1
+        assert windows[0].is_last
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        windows = chunk_windows(450)
+        starts = [w.start for w in windows]
+        assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
+        assert windows[-1].is_last
+        assert not any(w.is_last for w in windows[:-1])
+
+    def test_windows_cover_every_frame(self) -> None:
+        frames = 437
+        windows = chunk_windows(frames)
+        covered: set[int] = set()
+        for window in windows:
+            covered.update(range(window.start, window.end))
+        assert covered == set(range(frames))
+
+    def test_final_window_is_clamped_to_the_song(self) -> None:
+        windows = chunk_windows(437)
+        assert windows[-1].end == 437
+
+    def test_only_the_first_window_is_first(self) -> None:
+        windows = chunk_windows(600)
+        assert windows[0].is_first
+        assert not any(w.is_first for w in windows[1:])
+
+
+class TestCropBounds:
+    def test_single_window_song_is_not_cropped(self) -> None:
+        window = chunk_windows(50)[0]
+        assert crop_sample_bounds(window) == (0, 0)
+
+    def test_first_window_keeps_its_head(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[0])
+        assert left == 0
+        assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
+
+    def test_last_window_keeps_its_tail(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[-1])
+        assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
+        assert right == 0
+
+    def test_middle_window_is_cropped_on_both_sides(self) -> None:
+        windows = chunk_windows(900)
+        middle = windows[len(windows) // 2]
+        left, right = crop_sample_bounds(middle)
+        assert left > 0 and right > 0
+
+    def test_overlap_is_half_a_window(self) -> None:
+        assert overlap_mel_length() == HOP_MEL_FRAMES // 2

tests/model_executor/models/test_minimax_music3_row_binding.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 row-to-request binding.
+
+vLLM compacts and reorders the persistent batch whenever a request finishes or
+is preempted, so a batch row number is only meaningful within the step that
+produced it. Every failure pinned here was a real defect: state keyed by row
+survived into a step where the row belonged to someone else, or state keyed by
+one spelling of the request id was released under another and leaked.
+
+These exercise the pure helpers and the binding arithmetic directly, so they
+need no GPU and no engine.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    DEFAULT_MAX_AUDIO_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    _CFG_UNCOND_SUFFIX,
+    _first_int,
+    _pair_id_from_request,
+    _request_key,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestRequestKey:
+    def test_engine_list_form_is_unwrapped(self) -> None:
+        """The engine injects global_request_id as a single-element list.
+
+        Keying state on the list's ``str()`` would produce "['req-1']", which
+        never matches the bare id that ``on_requests_finished`` delivers, so
+        every finished song would leak its frame buffer.
+        """
+        assert _request_key(["req-1"]) == "req-1"
+
+    def test_bare_string_is_unchanged(self) -> None:
+        assert _request_key("req-1") == "req-1"
+
+    def test_list_and_bare_forms_agree(self) -> None:
+        """Allocation and release must land on the same key."""
+        assert _request_key(["req-1"]) == _request_key("req-1")
+
+    def test_empty_and_missing_are_stable(self) -> None:
+        assert _request_key(None) == ""
+        assert _request_key([]) == ""
+
+    def test_tuple_form(self) -> None:
+        assert _request_key(("req-9",)) == "req-9"
+
+
+class TestPairIdDerivation:
+    def test_conditioned_request_is_its_own_pair(self) -> None:
+        assert _pair_id_from_request("req-1") == "req-1"
+
+    def test_companion_maps_to_its_partner(self) -> None:
+        assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1"
+
+    def test_pair_members_agree_on_the_pair_id(self) -> None:
+        cond = "req-42"
+        uncond = f"req-42{_CFG_UNCOND_SUFFIX}"
+        assert _pair_id_from_request(cond) == _pair_id_from_request(uncond)
+
+    def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None:
+        assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond"
+
+
+class TestFirstInt:
+    def test_serving_layer_list_wrapping(self) -> None:
+        """tts_params values arrive wrapped in a single-element list."""
+        assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_bare_value(self) -> None:
+        assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_missing_falls_back(self) -> None:
+        assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+    def test_nonsense_falls_back_rather_than_raising(self) -> None:
+        assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+
+def _feedback_plan(
+    request_ids: list[str],
+    computed: list[int],
+    scheduled: list[int],
+    prompt_tokens: dict[str, int],
+    has_codes: set[str],
+) -> list[int]:
+    """Mirror of the talker's feedback row selection.
+
+    Returns the flat token offsets whose embedding gets replaced by an audio
+    frame. Kept as a standalone function so the arithmetic can be checked
+    without constructing a model.
+    """
+    offsets: list[int] = []
+    running = 0
+    for count in scheduled:
+        offsets.append(running)
+        running += count
+
+    rows: list[int] = []
+    for row, request_id in enumerate(request_ids):
+        if request_id not in has_codes:
+            continue
+        if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0):
+            continue
+        rows.append(offsets[row])
+    return rows
+
+
+class TestFeedbackRowSelection:
+    """The guard that decides which token embeddings become audio frames."""
+
+    def test_decode_only_step_replaces_every_row(self) -> None:
+        ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"]
+        rows = _feedback_plan(
+            ids,
+            computed=[120, 120, 90, 90],
+            scheduled=[1, 1, 1, 1],
+            prompt_tokens={i: 50 for i in ids},
+            has_codes=set(ids),
+        )
+        assert rows == [0, 1, 2, 3]
+
+    def test_prefill_row_is_never_overwritten(self) -> None:
+        """A prompt token must keep its real embedding.
+
+        The earlier guard compared the flat token count against a row budget,
+        so a short prompt slipped through and had its first embeddings
+        replaced with a previous request's audio frame.
+        """
+        rows = _feedback_plan(
+            ["a", "b"],
+            computed=[0, 300],
+            scheduled=[8, 1],
+            prompt_tokens={"a": 8, "b": 40},
+            has_codes={"b"},
+        )
+        # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single
+        # decode token sits at flat offset 8.
+        assert rows == [8]
+
+    def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None:
+        """The exact shape that defeated the old token-count guard."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[0],
+            scheduled=[10],
+            prompt_tokens={"a": 10},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_chunked_prefill_midway_is_not_fed_back(self) -> None:
+        rows = _feedback_plan(
+            ["a"],
+            computed=[2048],
+            scheduled=[2048],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_final_prefill_chunk_is_not_fed_back(self) -> None:
+        """The step that completes a prompt still carries prompt tokens."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[4096],
+            scheduled=[904],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None:
+        """Tokens are a flat concatenation, so offsets are a running sum."""
+        ids = ["p", "d1", "d2"]
+        rows = _feedback_plan(
+            ids,
+            computed=[0, 500, 700],
+            scheduled=[64, 1, 1],
+            prompt_tokens={"p": 64, "d1": 40, "d2": 40},
+            has_codes={"d1", "d2"},
+        )
+        assert rows == [64, 65]
+
+    def test_request_without_codes_yet_is_skipped(self) -> None:
+        """The first decode step after prefill has no previous frame."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[50],
+            scheduled=[1],
+            prompt_tokens={"a": 50},
+            has_codes=set(),
+        )
+        assert rows == []
+
+    def test_reordered_batch_binds_by_request_not_position(self) -> None:
+        """The same two requests in swapped rows get swapped offsets.
+
+        This is the compaction case: holding codes in a row-indexed buffer
+        would have handed each request the other's frame.
+        """
+        prompt_tokens = {"a": 50, "b": 50}
+        first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"})
+        swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"})
+        assert first == [0]
+        assert swapped == [1]

tests/model_executor/test_cfg_pairing.py

@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the model-neutral CFG request-pairing scheduler patches.
+
+The patches keep a cond/uncond request pair step-locked inside one engine:
+admitted together, kept at equal ``num_computed_tokens``, finished together.
+They are installed onto a lightweight fake scheduler here; vLLM's real
+``Scheduler`` needs a full engine (and a GPU) to construct.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm_omni.model_executor.models.common import cfg_pairing
+from vllm_omni.model_executor.models.common.cfg_pairing import (
+    _equalize_cfg_pair_progress,
+    _patch_scheduler,
+    cfg_uncond_suffixes,
+    register_cfg_uncond_suffix,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ───────── fakes ─────────
+
+
+class _FakeRequest:
+    """Minimal stand-in for ``vllm.v1.request.Request`` (identity equality)."""
+
+    def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None:
+        self.request_id = request_id
+        self.sampling_params = SimpleNamespace(extra_args=extra_args)
+        self.num_computed_tokens = 0
+        self.step_tokens = step_tokens
+        self.skip_reading_prefix_cache = False
+
+
+class _FakeWaitingQueue(deque):
+    """FCFS waiting queue: vLLM's is also a ``deque`` subclass."""
+
+    def remove_requests(self, requests: list[_FakeRequest]) -> None:
+        keep = [req for req in self if not any(req is dropped for dropped in requests)]
+        self.clear()
+        self.extend(keep)
+
+    def prepend_request(self, request: _FakeRequest) -> None:
+        self.appendleft(request)
+
+
+@dataclass
+class _FakeSchedulerOutput:
+    num_scheduled_tokens: dict[str, int] = field(default_factory=dict)
+    total_num_scheduled_tokens: int = 0
+
+
+def _make_scheduler_cls() -> type:
+    """Fresh class per test so the in-place patches never stack."""
+
+    class FakeScheduler:
+        def __init__(self) -> None:
+            self.waiting = _FakeWaitingQueue()
+            self.running: list[_FakeRequest] = []
+            self.requests: dict[str, _FakeRequest] = {}
+            self.finished: list[set[str]] = []
+            self.observed_waiting: list[list[str]] = []
+            self.max_num_scheduled_tokens = 512
+            self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096)
+
+        def add_request(self, request: _FakeRequest) -> None:
+            self.requests[request.request_id] = request
+            self.waiting.append(request)
+
+        def finish_requests(self, request_ids: Any, finished_status: Any) -> None:
+            if request_ids is None:
+                ids = set(self.requests)
+            elif isinstance(request_ids, str):
+                ids = {request_ids}
+            else:
+                ids = set(request_ids)
+            self.finished.append(ids)
+            self.running = [req for req in self.running if req.request_id not in ids]
+            for req_id in ids:
+                self.requests.pop(req_id, None)
+
+        def schedule(self) -> _FakeSchedulerOutput:
+            self.observed_waiting.append([req.request_id for req in self.waiting])
+            while self.waiting:
+                self.running.append(self.waiting.popleft())
+            output = _FakeSchedulerOutput()
+            for req in self.running:
+                req.num_computed_tokens += req.step_tokens
+                output.num_scheduled_tokens[req.request_id] = req.step_tokens
+                output.total_num_scheduled_tokens += req.step_tokens
+            return output
+
+    return FakeScheduler
+
+
+@pytest.fixture
+def scheduler() -> Any:
+    cls = _make_scheduler_cls()
+    _patch_scheduler(cls)
+    return cls()
+
+
+@pytest.fixture(autouse=True)
+def _restore_suffix_registry():
+    original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES)
+    yield
+    cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear()
+    cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original)
+
+
+def _add_pair(
+    scheduler: Any,
+    pair_id: str | None = None,
+    base_id: str = "req-1",
+    suffix: str = "__cfg_uncond",
+) -> tuple[_FakeRequest, _FakeRequest]:
+    cond_extra: dict[str, Any] = {"cfg_role": "cond"}
+    uncond_extra: dict[str, Any] = {"cfg_role": "uncond"}
+    if pair_id is not None:
+        cond_extra["cfg_pair_id"] = pair_id
+        uncond_extra["cfg_pair_id"] = pair_id
+    cond = _FakeRequest(base_id, cond_extra)
+    uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra)
+    scheduler.add_request(cond)
+    scheduler.add_request(uncond)
+    return cond, uncond
+
+
+# ───────── pair-id derivation ─────────
+
+
+def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+
+    assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}
+    assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"}
+
+
+def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None:
+    register_cfg_uncond_suffix("::negative")
+    assert "::negative" in cfg_uncond_suffixes()
+
+    cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative")
+
+    assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}}
+
+
+def test_register_rejects_empty_suffix() -> None:
+    with pytest.raises(ValueError):
+        register_cfg_uncond_suffix("")
+
+
+def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None:
+    # Degenerate input: the companion never joins the cond pair, so the pair
+    # stays incomplete and the hold budget eventually releases it.
+    scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"}))
+
+    assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}}
+
+
+def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler, pair_id="pair-A")
+
+    assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}}
+    assert cond.request_id not in scheduler._cfg_pairs
+    assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A"
+
+
+def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+
+    assert cond.skip_reading_prefix_cache is True
+    assert uncond.skip_reading_prefix_cache is True
+
+
+# ───────── admission holding ─────────
+
+
+def test_lone_member_is_held_then_released(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting == [[]]  # the base scheduler saw nothing
+    assert [req.request_id for req in scheduler.waiting] == ["req-1"]
+    assert scheduler.running == []
+
+
+def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+    scheduler.schedule()
+    scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"]
+    assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"]
+    assert list(scheduler.waiting) == []
+
+
+def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2)
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+    scheduler.schedule()
+    scheduler.schedule()
+    assert scheduler.running == []
+    assert scheduler._cfg_pairs  # still hoping for the partner
+
+    scheduler.schedule()
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+    assert [req.request_id for req in scheduler.running] == ["req-1"]
+
+
+def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+    scheduler.add_request(_FakeRequest("other", None))
+    scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"]
+
+
+# ───────── progress equalization ─────────
+
+
+def _equalize_fixture(
+    cond_computed: int,
+    uncond_computed: int,
+    cond_sched: int,
+    uncond_sched: int,
+) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]:
+    cond = _FakeRequest("req-1", {"cfg_role": "cond"})
+    uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})
+    cond.num_computed_tokens = cond_computed
+    uncond.num_computed_tokens = uncond_computed
+    fake = SimpleNamespace(
+        _cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}},
+        requests={cond.request_id: cond, uncond.request_id: uncond},
+    )
+    output = _FakeSchedulerOutput(
+        num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched},
+        total_num_scheduled_tokens=cond_sched + uncond_sched,
+    )
+    return fake, output, cond, uncond
+
+
+def test_equalize_pulls_the_ahead_member_back() -> None:
+    fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert cond.num_computed_tokens == 6
+    assert uncond.num_computed_tokens == 6
+    assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4}
+    assert output.total_num_scheduled_tokens == 5
+
+
+def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None:
+    fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8)
+    assert output.total_num_scheduled_tokens == 8
+
+
+def test_equalize_skips_an_infeasible_rollback() -> None:
+    # Clawing back 4 tokens would leave the cond row with 0 scheduled tokens,
+    # which the scheduler output cannot express; leave the pair to _drop_split_pairs.
+    fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert cond.num_computed_tokens == 10
+    assert output.total_num_scheduled_tokens == 8
+
+
+# ───────── finish ─────────
+
+
+def test_finish_requests_finishes_the_partner(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+    scheduler.schedule()
+
+    scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED")
+
+    assert scheduler.finished == [{cond.request_id, uncond.request_id}]
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+    assert scheduler.running == []
+
+
+def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None:
+    _add_pair(scheduler)
+    scheduler.schedule()
+
+    scheduler.finish_requests(None, "FINISHED_ABORTED")
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+
+
+# ───────── non-CFG traffic ─────────
+
+
+def test_plain_requests_are_untouched(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("plain-1", None))
+    scheduler.add_request(_FakeRequest("plain-2", {"seed": 3}))
+
+    output = scheduler.schedule()
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler.observed_waiting == [["plain-1", "plain-2"]]
+    assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1}
+    assert output.total_num_scheduled_tokens == 2
+    # No pair ever existed, so the hold bookkeeping is never allocated.
+    assert not hasattr(scheduler, "_cfg_hold_counts")
+
+
+def test_default_suffix_matches_the_audex_stage_processor() -> None:
+    from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX
+
+    assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes()
+
+
+def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None:
+    from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied
+
+    assert apply_cfg_patches is cfg_pairing.apply_cfg_patches
+    assert cfg_patches_applied is cfg_pairing.cfg_patches_applied
+
+
+# ───────── install gate ─────────
+
+
+def test_gate_triggers_on_a_cfg_logits_processor() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    class AudexCFGLogitsProcessor:
+        pass
+
+    assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor]))
+    # vLLM keeps unresolved entries as dotted strings.
+    assert _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"])
+    )
+
+
+def test_gate_triggers_on_a_declared_cfg_role_default() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    assert _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale"))
+    )
+
+
+def test_gate_stays_off_for_non_cfg_stages() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=()))
+    assert not _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",))
+    )
+    assert not _stage_declares_cfg_pairs(SimpleNamespace())
+
+
+def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None:
+    from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys
+
+    assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == (
+        "cfg_role",
+        "cfg_scale",
+    )
+    assert _sampling_extra_args_keys({"temperature": 0.1}) == ()
+    assert _sampling_extra_args_keys(None) == ()

tests/model_tests/diffusion/test_alignment.py

@@ -77,6 +77,7 @@ EXCLUDED_MODELS = [
     "DreamZeroPipeline",
     "StableDiffusionXLPipeline",
     "Gr00tN1d7Pipeline",
+    "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
 ]

tests/platforms/npu/test_graph_tools.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.platforms.npu.graph_tools import (
+    CapturedDeviceGraph,
+    NPUExactGraphRunner,
+)
+from vllm_omni.platforms.npu.models import minicpmo_4_5_code2wav as code2wav_patch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_captured_graph_replay_clones_persistent_outputs():
+    static_input = torch.zeros(1)
+    static_output = torch.zeros(1)
+
+    class _Graph:
+        def replay(self):
+            static_output.copy_(static_input * 2)
+
+    graph = CapturedDeviceGraph(
+        graph=_Graph(),
+        static_inputs=(static_input,),
+        static_outputs=(static_output,),
+    )
+
+    first = graph.replay((torch.tensor([2.0]),))[0]
+    second = graph.replay((torch.tensor([3.0]),))[0]
+
+    torch.testing.assert_close(first, torch.tensor([4.0]))
+    torch.testing.assert_close(second, torch.tensor([6.0]))
+    assert first.data_ptr() != second.data_ptr()
+
+
+def test_capture_uses_vllm_global_graph_pool(monkeypatch):
+    runner = NPUExactGraphRunner()
+    global_pool = object()
+    seen_pools = []
+    synchronizations = 0
+
+    class _Graph:
+        def replay(self):
+            pass
+
+    @contextmanager
+    def graph_context(graph, *, pool):
+        del graph
+        seen_pools.append(pool)
+        yield
+
+    def synchronize():
+        nonlocal synchronizations
+        synchronizations += 1
+
+    fake_npu = SimpleNamespace(
+        NPUGraph=_Graph,
+        graph=graph_context,
+        synchronize=synchronize,
+    )
+    monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+    import vllm.platforms
+
+    monkeypatch.setattr(
+        vllm.platforms,
+        "current_platform",
+        SimpleNamespace(get_global_graph_pool=lambda: global_pool),
+    )
+
+    captured = runner.capture(
+        (torch.tensor([2.0]),),
+        lambda value: (value + 1,),
+    )
+
+    assert seen_pools == [global_pool]
+    assert synchronizations == 2
+    torch.testing.assert_close(captured.static_outputs[0], torch.tensor([3.0]))
+
+
+def test_exact_signature_dispatch_captures_then_replays(monkeypatch):
+    runner = NPUExactGraphRunner(max_graphs=2)
+    replay_count = 0
+
+    class _FunctionalGraph:
+        def __init__(self, compute):
+            self.compute = compute
+
+        def replay(self, inputs):
+            nonlocal replay_count
+            replay_count += 1
+            return tuple(output.detach().clone() for output in self.compute(*inputs))
+
+    monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+    monkeypatch.setattr(
+        runner,
+        "capture",
+        lambda inputs, compute: _FunctionalGraph(compute),
+    )
+
+    def compute(value):
+        return (value * 2,)
+
+    first = runner.run("unit", (torch.tensor([2.0]),), (False,), compute)
+    second = runner.run("unit", (torch.tensor([3.0]),), (False,), compute)
+
+    torch.testing.assert_close(first[0], torch.tensor([4.0]))
+    torch.testing.assert_close(second[0], torch.tensor([6.0]))
+    assert runner.stats == {"captures": 1, "failed": 0, "hits": 1}
+    assert replay_count == 1
+
+
+def test_capture_failure_is_fatal_for_stage_process(monkeypatch):
+    runner = NPUExactGraphRunner(
+        component_name="test component",
+        disable_config_hint="disable the test graph",
+    )
+    captures = 0
+    computes = 0
+
+    monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+
+    def fail_capture(inputs, compute):
+        nonlocal captures
+        captures += 1
+        raise RuntimeError("unsupported op")
+
+    monkeypatch.setattr(runner, "capture", fail_capture)
+
+    def compute(value):
+        nonlocal computes
+        computes += 1
+        return (value + 1,)
+
+    with pytest.raises(RuntimeError, match="restart the stage process"):
+        runner.run("unit", (torch.tensor([1.0]),), (False,), compute)
+    with pytest.raises(RuntimeError, match="cannot continue"):
+        runner.run(
+            "different-shape",
+            (torch.tensor([1.0, 2.0]),),
+            (True,),
+            compute,
+        )
+
+    assert captures == 1
+    assert computes == 1
+    assert runner.stats == {"captures": 0, "failed": 1, "hits": 0}
+
+
+def test_code2wav_runtime_disables_internal_format_and_jit(monkeypatch):
+    monkeypatch.delenv("ASCEND_LAUNCH_BLOCKING", raising=False)
+    compile_modes = []
+    fake_npu = SimpleNamespace(
+        config=SimpleNamespace(allow_internal_format=True),
+        set_compile_mode=lambda **kwargs: compile_modes.append(kwargs),
+    )
+    monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+    code2wav_patch.prepare_code2wav_graph_runtime()
+
+    assert fake_npu.config.allow_internal_format is False
+    assert compile_modes == [{"jit_compile": False}]
+
+
+def test_code2wav_runtime_rejects_launch_blocking(monkeypatch):
+    monkeypatch.setenv("ASCEND_LAUNCH_BLOCKING", "1")
+
+    with pytest.raises(RuntimeError, match="ASCEND_LAUNCH_BLOCKING=1"):
+        code2wav_patch.prepare_code2wav_graph_runtime()
+
+
+@pytest.mark.parametrize(
+    ("enabled", "max_graphs", "expected_prepared"),
+    [
+        (True, 7, 1),
+        (False, 7, 0),
+        (True, 0, 0),
+    ],
+)
+def test_code2wav_patch_reads_stage_additional_config(
+    monkeypatch,
+    enabled,
+    max_graphs,
+    expected_prepared,
+):
+    prepared = 0
+
+    class _Backend:
+        speech_window = torch.zeros(1)
+        flow = SimpleNamespace(training=False)
+
+    model = SimpleNamespace(
+        backend=None,
+        vllm_config=SimpleNamespace(
+            additional_config={
+                "code2wav_enable_npu_graph": enabled,
+                "code2wav_max_npu_graphs": max_graphs,
+            }
+        ),
+    )
+
+    def build_backend(instance):
+        instance.backend = _Backend()
+
+    def prepare_runtime():
+        nonlocal prepared
+        prepared += 1
+
+    monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+    monkeypatch.setattr(code2wav_patch, "prepare_code2wav_graph_runtime", prepare_runtime)
+    code2wav_patch._patched_build_backend(model)
+
+    assert prepared == expected_prepared
+    if expected_prepared:
+        graph_runner = code2wav_patch._backend_graph_runners[model.backend]
+        assert isinstance(graph_runner, NPUExactGraphRunner)
+        assert graph_runner.max_graphs == max_graphs
+    else:
+        assert model.backend not in code2wav_patch._backend_graph_runners
+
+
+@pytest.mark.parametrize("with_cache", [False, True])
+def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache):
+    calls = []
+
+    class _Backend:
+        _trt_stepper = None
+        _cfm_graph_wrapper = None
+
+    class _Runner:
+        def run(self, operation, inputs, constants, compute):
+            calls.append((operation, inputs, constants))
+            return compute(*inputs)
+
+    backend = _Backend()
+    runner = _Runner()
+    code2wav_patch._backend_graph_runners[backend] = runner
+
+    def graphable_estimator_step(
+        instance,
+        estimator,
+        *,
+        x,
+        mu,
+        time_embedding,
+        speakers,
+        cond,
+        cnn_cache,
+        att_cache,
+    ):
+        del instance, estimator, time_embedding, speakers, cond
+        cache_marker = x.new_tensor(0 if cnn_cache is None else 1)
+        return x + mu, cache_marker, cache_marker.clone()
+
+    monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None)
+    monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step)
+    value = torch.tensor([2.0])
+    cache = torch.tensor([1.0]) if with_cache else None
+    estimator = SimpleNamespace(t_embedder=lambda time: time + 1)
+    outputs = code2wav_patch._patched_estimator_step(
+        backend,
+        estimator,
+        x=value,
+        mu=value,
+        time=value,
+        speakers=value,
+        cond=value,
+        cnn_cache=cache,
+        att_cache=cache,
+    )
+
+    assert calls[0][0] == "cfm_estimator"
+    assert calls[0][2] == (with_cache,)
+    assert len(calls[0][1]) == (7 if with_cache else 5)
+    torch.testing.assert_close(calls[0][1][2], torch.tensor([[3.0]]))
+    torch.testing.assert_close(outputs[0], torch.tensor([4.0]))
+    torch.testing.assert_close(outputs[1], torch.tensor(1.0 if with_cache else 0.0))
+
+
+def test_code2wav_platform_wraps_flow_execution_context(monkeypatch):
+    entered = []
+
+    class _Backend:
+        pass
+
+    class _Context:
+        def __init__(self, device, require_math):
+            self.device = device
+            self.require_math = require_math
+
+        def __enter__(self):
+            entered.append((self.device, self.require_math))
+
+        def __exit__(self, exc_type, exc, traceback):
+            return False
+
+    backend = _Backend()
+    code2wav_patch._backend_graph_runners[backend] = NPUExactGraphRunner()
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_flow_execution_context",
+        lambda device, *, require_math: _Context(device, require_math),
+    )
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_original_setup_batch",
+        lambda instance, features, batch_size: ("setup", batch_size),
+    )
+    monkeypatch.setattr(
+        code2wav_patch,
+        "_original_decode_batch",
+        lambda instance, tokens, features, states, *, last_chunk, flush_encoder: (
+            "decode",
+            last_chunk,
+            flush_encoder,
+        ),
+    )
+    features = SimpleNamespace(speech_tokens=torch.zeros(1))
+
+    assert code2wav_patch._patched_setup_batch(backend, features, 3) == ("setup", 3)
+    assert code2wav_patch._patched_decode_batch(
+        backend,
+        torch.zeros(1),
+        features,
+        [],
+        last_chunk=True,
+    ) == ("decode", True, False)
+    assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)]

tests/platforms/npu/test_minicpmo_code2wav_npugraph.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ascend hardware validation for MiniCPM-o Code2Wav NPUGraph replay."""
+
+from math import prod
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
+    BatchedToken2Wav,
+)
+from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
+from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
+    _graphable_estimator_step,
+    prepare_code2wav_graph_runtime,
+)
+from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+    npu_token2wav_sdpa_context,
+)
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.omni,
+    *hardware_marks(res={"npu": "A3"}, num_cards=1),
+]
+
+
+class _Flow(nn.Module):
+    def __init__(self, estimator: nn.Module):
+        super().__init__()
+        self.decoder = SimpleNamespace(estimator=estimator)
+
+
+class _Token2Wav:
+    def __init__(self, estimator: nn.Module):
+        self.flow = _Flow(estimator).eval()
+        self.hift = nn.Identity().eval()
+        self.float16 = False
+        self.n_timesteps = 2
+        self.mel_cache_len = 1
+        self.source_cache_len = 2
+        self.speech_window = torch.ones(4, device="npu")
+
+
+def test_minicpmo_cfm_estimator_npugraph_matches_eager():
+    decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
+    prepare_code2wav_graph_runtime()
+
+    estimator = (
+        decoder_dit.DiT(
+            in_channels=6,
+            out_channels=2,
+            depth=2,
+            num_heads=2,
+            head_dim=8,
+            hidden_size=16,
+            mlp_ratio=2.0,
+        )
+        .npu()
+        .eval()
+    )
+    graph_runner = NPUExactGraphRunner(max_graphs=4)
+    adapter = BatchedToken2Wav(_Token2Wav(estimator))
+
+    def inputs(seed: int):
+        def make(shape: tuple[int, ...], offset: float):
+            values = torch.arange(
+                prod(shape),
+                device="npu",
+                dtype=torch.float32,
+            )
+            return torch.sin(values * 0.17 + seed + offset).reshape(shape)
+
+        return (
+            make((2, 2, 6), 0.0),
+            make((2, 2, 6), 0.5),
+            make((2, 1, 16), 1.0),
+            make((2, 1), 1.5),
+            make((2, 1, 6), 2.0),
+        )
+
+    def compute(*values, caches=None):
+        cnn_cache, att_cache = (None, None) if caches is None else caches
+        return _graphable_estimator_step(
+            adapter,
+            estimator,
+            x=values[0],
+            mu=values[1],
+            time_embedding=values[2],
+            speakers=values[3],
+            cond=values[4],
+            cnn_cache=cnn_cache,
+            att_cache=att_cache,
+        )
+
+    with torch.inference_mode(), npu_token2wav_sdpa_context(require_math=True):
+        warm_inputs = inputs(11)
+        warm_outputs = graph_runner.run(
+            "cfm_estimator",
+            warm_inputs,
+            (False,),
+            lambda *values: compute(*values),
+        )
+        replay_inputs = inputs(13)
+        expected = compute(*replay_inputs)
+        actual = graph_runner.run(
+            "cfm_estimator",
+            replay_inputs,
+            (False,),
+            lambda *values: compute(*values),
+        )
+        for eager, replayed in zip(expected, actual, strict=True):
+            torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+
+        cached_inputs = (*inputs(17), warm_outputs[1], warm_outputs[2])
+        graph_runner.run(
+            "cfm_estimator",
+            cached_inputs,
+            (True,),
+            lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+        )
+        replay_cached_inputs = (
+            *inputs(19),
+            warm_outputs[1] + 0.1,
+            warm_outputs[2] + 0.1,
+        )
+        expected_cached = compute(
+            *replay_cached_inputs[:5],
+            caches=(replay_cached_inputs[5], replay_cached_inputs[6]),
+        )
+        actual_cached = graph_runner.run(
+            "cfm_estimator",
+            replay_cached_inputs,
+            (True,),
+            lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+        )
+        for eager, replayed in zip(expected_cached, actual_cached, strict=True):
+            torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+        assert graph_runner.stats == {"captures": 2, "failed": 0, "hits": 2}

tests/platforms/npu/test_step_audio2_token2wav.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import importlib.util
 import sys
 import types
+from contextlib import contextmanager
 from pathlib import Path
 from types import SimpleNamespace
 
@@ -194,3 +195,71 @@ def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch)
 
     with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"):
         module.patch_step_audio2_hift_for_npu(SimpleNamespace())
+
+
+def test_sdpa_context_preserves_exceptions_from_token2wav_body(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module = _load_patch_module(monkeypatch)
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+    @contextmanager
+    def fake_math_context(*, require_available=False):
+        del require_available
+        yield
+
+    fake_attention_patch.npu_math_sdpa_context = fake_math_context
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with pytest.raises(RuntimeError, match="original graph failure"):
+        with module.npu_token2wav_sdpa_context():
+            raise RuntimeError("original graph failure")
+
+
+def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_patch_module(monkeypatch)
+    required = []
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+    @contextmanager
+    def fake_math_context(*, require_available=False):
+        required.append(require_available)
+        yield
+
+    fake_attention_patch.npu_math_sdpa_context = fake_math_context
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with module.npu_token2wav_sdpa_context(require_math=True):
+        pass
+
+    assert required == [True]
+
+
+def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_patch_module(monkeypatch)
+    fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+
+    def fail_patch():
+        raise RuntimeError("attention patch failed")
+
+    fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch
+    fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,)))
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+        fake_attention_patch,
+    )
+
+    with pytest.raises(RuntimeError, match="attention patch failed"):
+        with module.npu_token2wav_sdpa_context():
+            pass

tools/nightly/run_nightly_jobs.sh

@@ -43,7 +43,7 @@
 #
 #   stability (when included in TEST_TYPE):
 #     From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py
-#     model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six
+#     model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22 + hunyuan_image; all → all five
 #     LABEL_SUBSTR: if set, script path / job key / filename must contain it
 #
 # Requirements: bash, python3, PyYAML (pip install pyyaml)
@@ -588,7 +588,6 @@ def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -
 STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [
     ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)),
     ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)),
-    ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)),
     ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)),
     ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)),
     ("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)),

vllm_omni/benchmarks/data_modules/daily_omni_dataset.py

@@ -7,7 +7,8 @@ Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni
 
 Supports loading QA metadata from:
 - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
-- HuggingFace datasets (``dataset_path``): legacy online mode
+- HuggingFace Hub id (``dataset_path``): resolved with ``snapshot_download`` to the repo's
+  ``qa.json``, which reads the hub cache and therefore also works under ``HF_HUB_OFFLINE``
 
 Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir``
 is not set, the first request that needs on-disk media downloads that archive from the Hugging Face
@@ -24,9 +25,9 @@ Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``?
     This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``,
     ``random_seed``, ``self.data``) and implements **two explicit loaders**:
     ``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface``
-    (optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not**
+    (Hub ids, resolved to the repo's ``qa.json`` through the hub cache). The latter is **not**
     inheritance; it is the same Hub rows as before, factored into a helper so one class can
-    serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``.
+    serve both deployment modes without mandatory ``datasets``.
 
 Usage:
     from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset
@@ -38,7 +39,7 @@ Usage:
         random_seed=42,
     )
 
-    # HuggingFace mode (legacy, requires network)
+    # Hub id mode (works offline once the repo is in the HF hub cache)
     dataset = DailyOmniDataset(
         dataset_path="liarliar/Daily-Omni",
         dataset_split="train",
@@ -210,6 +211,68 @@ def _resolve_hf_cache_snapshot(path: Path) -> Path:
     return path
 
 
+#: QA metadata pulled by :func:`ensure_daily_omni_hub_root`. ``Videos.tar`` is deliberately left
+#: out: it is ~60GB and only :func:`ensure_daily_omni_hub_videos_dir` needs it, lazily, once a
+#: request actually asks for media.
+_DAILY_OMNI_QA_ALLOW_PATTERNS = ["qa.json", "QA.json"]
+
+
+def ensure_daily_omni_hub_root(repo_id: str) -> Path:
+    """Resolve a Daily-Omni Hub id to its snapshot directory, downloading QA metadata if needed.
+
+    Mirrors :func:`...seed_tts_dataset.resolve_seed_tts_root`: ``snapshot_download`` resolves the
+    repo through the **hub cache**, so an air-gapped run (``HF_HUB_OFFLINE=1``) keeps working as
+    long as the dataset was pre-staged there. ``datasets.load_dataset`` cannot do this — it only
+    consults the separate *datasets* cache, which a ``hf download`` never populates.
+
+    Raises:
+        ImportError: if ``huggingface_hub`` is not installed.
+        ValueError: if ``repo_id`` is empty.
+    """
+    rid = (repo_id or "").strip()
+    if not rid:
+        raise ValueError("repo_id is required to download Daily-Omni from Hugging Face")
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as e:
+        raise ImportError(
+            "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local "
+            "--daily-omni-qa-json / --daily-omni-video-dir mirror."
+        ) from e
+
+    cache = snapshot_download(
+        repo_id=rid,
+        repo_type="dataset",
+        allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS,
+    )
+    root = _resolve_hf_cache_snapshot(Path(cache).resolve())
+    logger.info("Daily-Omni Hub snapshot ready at %s (repo=%s)", root, rid)
+    return root
+
+
+def _daily_omni_qa_load_error(
+    repo_id: str | None,
+    snapshot_error: Exception | None,
+    datasets_error: Exception,
+) -> RuntimeError:
+    """Compose an actionable error after both QA-metadata sources failed.
+
+    Offline CI hits this when the dataset is missing from the hub cache *or* staged under a
+    different repo id, so the message names both attempts and the flags that bypass them.
+    """
+    parts = [f"Could not load Daily-Omni QA metadata for {repo_id!r}."]
+    if snapshot_error is not None:
+        parts.append(f"huggingface_hub snapshot_download failed: {type(snapshot_error).__name__}: {snapshot_error}")
+    parts.append(f"datasets.load_dataset failed: {type(datasets_error).__name__}: {datasets_error}")
+    parts.append(
+        "For offline runs, pre-stage the dataset in the HF hub cache "
+        "(`hf download --repo-type dataset <repo>`) and make sure the repo id matches, "
+        "or pass --daily-omni-qa-json / --daily-omni-video-dir pointing at a local mirror."
+    )
+    return RuntimeError(" ".join(parts))
+
+
 def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None:
     """Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local.
 
@@ -433,7 +496,8 @@ class DailyOmniDataset(BenchmarkDataset):
 
     QA metadata can be loaded from:
     - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
-    - HuggingFace datasets (``dataset_path``): legacy online mode
+    - HuggingFace Hub id (``dataset_path``): the repo's ``qa.json`` via ``snapshot_download``
+      (hub-cache backed, offline-safe); ``datasets.load_dataset`` only as a parquet fallback
 
     Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set,
     the first sample that needs on-disk media downloads that archive from the Hugging Face dataset
@@ -444,7 +508,8 @@ class DailyOmniDataset(BenchmarkDataset):
         qa_json_path: Path to local qa.json file (offline mode, preferred). When provided,
             ``dataset_path`` and ``dataset_split`` are ignored.
         dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if
-            ``qa_json_path`` is not provided (legacy online mode).
+            ``qa_json_path`` is not provided; resolved through the hub cache, so a pre-staged
+            repo also loads with ``HF_HUB_OFFLINE=1``.
         dataset_split: Dataset split to use (default: "train"). Used only in online mode.
         random_seed: Random seed for shuffling
         video_dir: Directory containing extracted video files (default: None; may be filled lazily
@@ -602,21 +667,42 @@ class DailyOmniDataset(BenchmarkDataset):
         self.data = _ListDatasetIterator(data)
 
     def _load_from_huggingface(self) -> None:
-        """Load QA rows via ``datasets.load_dataset`` (legacy / convenience path).
+        """Load QA rows for a Hub id, preferring the repo's ``qa.json`` over ``datasets``.
 
-        Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni``
-        and get the same parquet-backed rows as the Hub dataset card, with streaming (or
-        non-streaming if ``no_stream=True``) and shuffle.
+        The Hub repo ships ``qa.json`` (1,197 rows) next to ``Videos.tar``, so the primary path is
+        :func:`ensure_daily_omni_hub_root` → :meth:`_load_from_local_json`: ``snapshot_download``
+        reads the hub cache and therefore works under ``HF_HUB_OFFLINE``, the same contract as
+        Seed-TTS. ``datasets.load_dataset`` remains only as a fallback for mirrors
+        that publish parquet rows instead of ``qa.json``; it consults the *datasets* cache rather
+        than the hub cache and streams shards over HTTP, so it cannot serve air-gapped runs.
 
         This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base
         always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary
         API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing
         one ``sample()`` / request-building implementation for both.
         """
+        snapshot_error: Exception | None = None
+        try:
+            root = ensure_daily_omni_hub_root(self.dataset_path)
+        except Exception as e:  # noqa: BLE001 - any failure just falls through to `datasets`
+            snapshot_error = e
+        else:
+            qa_json = daily_omni_local_qa_json(root)
+            if qa_json is not None:
+                self.qa_json_path = qa_json
+                logger.info("Loading Daily-Omni QA metadata from Hub snapshot: %s", qa_json)
+                self._load_from_local_json()
+                return
+            logger.info(
+                "Hub snapshot %s ships no qa.json; falling back to `datasets.load_dataset`",
+                root,
+            )
+
         if load_dataset is None:
-            raise ImportError(
-                "datasets library is required for HuggingFace mode. "
-                "Install with: pip install datasets, or use local JSON mode instead."
+            raise _daily_omni_qa_load_error(
+                self.dataset_path,
+                snapshot_error,
+                ImportError("datasets library is not installed (pip install datasets)"),
             )
 
         load_kw: dict[str, Any] = {
@@ -626,7 +712,10 @@ class DailyOmniDataset(BenchmarkDataset):
         }
         if self.dataset_subset is not None:
             load_kw["name"] = self.dataset_subset
-        ds = load_dataset(self.dataset_path, **load_kw)
+        try:
+            ds = load_dataset(self.dataset_path, **load_kw)
+        except Exception as e:
+            raise _daily_omni_qa_load_error(self.dataset_path, snapshot_error, e) from e
         if not getattr(self, "disable_shuffle", False):
             ds = ds.shuffle(seed=self.random_seed)
         self.data = ds

vllm_omni/config/__init__.py

@@ -2,7 +2,6 @@
 Configuration module for vLLM-Omni.
 """
 
-from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.config.omni_config import (
@@ -22,7 +21,6 @@ from vllm_omni.config.omni_config import (
     VllmOmniGenerationStageConfig,
     VllmOmniOrchestratorConfig,
 )
-from vllm_omni.config.pipeline_registry import register_pipeline
 from vllm_omni.config.stage_config import (
     PIPELINE_WIDE_ENGINE_FIELDS,
     DeployConfig,
@@ -42,6 +40,15 @@ from vllm_omni.config.yaml_util import (
     to_dict,
 )
 
+# StageConfigFactory / register_pipeline pull pipeline_registry, which eagerly
+# imports PI0_PIPELINE → diffusion.data. Keep those lazy so
+# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
+# loading) cannot close a circular import through DiffusionOutput.
+_LAZY_ATTRS = {
+    "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
+    "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
+}
+
 __all__ = [
     # Legacy model-level configs.
     "LoRAConfig",
@@ -82,3 +89,19 @@ __all__ = [
     "merge_configs",
     "to_dict",
 ]
+
+
+def __getattr__(name: str):
+    spec = _LAZY_ATTRS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr_name = spec
+    import importlib
+
+    value = getattr(importlib.import_module(module_name), attr_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(__all__) | set(globals()))

vllm_omni/config/model.py

@@ -146,6 +146,11 @@ class OmniModelConfig(ModelConfig):
     task_type: str | None = None
     enable_sleep_mode: bool = False
     has_sampling_extra_args: bool = False
+    # Key names (not values) of the stage's default sampling ``extra_args``.
+    # Engine-core code runs before any request arrives, so this is the only
+    # place it can learn which request-shaping conventions a stage uses (e.g.
+    # ``cfg_role`` for classifier-free-guidance request pairs).
+    sampling_extra_args_keys: tuple[str, ...] = ()
 
     @property
     def registry(self):

vllm_omni/config/omni_config.py

@@ -12,12 +12,22 @@ from __future__ import annotations
 
 import copy
 from collections.abc import Mapping
-from dataclasses import dataclass, field, fields
+from dataclasses import InitVar, dataclass, field, fields
+from functools import wraps
+from inspect import Parameter, signature
 from pathlib import Path
 from typing import Any, Literal, TypeAlias, TypedDict, cast
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict, Field, model_validator
+from typing_extensions import Self
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
 from vllm.config.utils import config
+from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.config.stage_config import (
@@ -77,6 +87,49 @@ class _QuantizationEngineOverrides(TypedDict, total=False):
     quantization: str
 
 
+class _TrackExplicitConfigFields:
+    """Record constructor inputs without adding a serialized config field."""
+
+    @model_validator(mode="wrap")
+    @classmethod
+    def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any:
+        result = handler(value)
+        kwargs = getattr(value, "kwargs", None)
+        if kwargs is not None:
+            explicit_fields = frozenset(kwargs)
+        elif isinstance(value, cls):
+            explicit_fields = getattr(value, "_omni_explicit_fields", frozenset())
+        else:
+            explicit_fields = frozenset()
+        object.__setattr__(result, "_omni_explicit_fields", explicit_fields)
+        return result
+
+
+def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
+    """Make inherited Pydantic dataclass fields keyword-only as well."""
+    generated_init = cls.__init__
+    generated_signature = signature(cls)
+    keyword_only_signature = generated_signature.replace(
+        parameters=[
+            parameter.replace(kind=Parameter.KEYWORD_ONLY)
+            if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD}
+            else parameter
+            for parameter in generated_signature.parameters.values()
+        ]
+    )
+
+    @wraps(generated_init)
+    def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None:
+        if args:
+            raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)")
+        generated_init(self, **kwargs)
+
+    cls.__init__ = keyword_only_init
+    cls.__signature__ = keyword_only_signature
+    cls.__match_args__ = ()
+    return cls
+
+
 class _ModelEngineOverrides(TypedDict, total=False):
     model: str
     model_arch: str
@@ -106,7 +159,6 @@ class _ModelEngineOverrides(TypedDict, total=False):
     enforce_eager: bool
     max_cudagraph_capture_size: int
     enable_flashinfer_autotune: bool
-    compilation_config: dict[str, Any]
     enable_multithread_weight_load: bool
     num_weight_load_threads: int
     disable_autocast: bool
@@ -148,7 +200,6 @@ class _RuntimeEngineOverrides(TypedDict, total=False):
     num_gpus: int
     log_level: str
     log_stats: bool
-    profiler_config: dict[str, Any]
 
 
 class _ParallelConfigEngineOverrides(TypedDict, total=False):
@@ -192,6 +243,8 @@ class _StageEngineValues:
     runtime: _RuntimeEngineOverrides
     parallel: _ParallelEngineOverrides
     diffusion: _DiffusionEngineOverrides
+    compilation_config: Mapping[str, Any] | VllmCompilationConfig | None
+    profiler_config: Mapping[str, Any] | VllmProfilerConfig | None
 
 
 @dataclass(frozen=True)
@@ -341,7 +394,6 @@ class OmniStageModelConfig:
     enforce_eager: bool = False
     max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
     enable_flashinfer_autotune: bool | None = None
-    compilation_config: dict[str, Any] | None = None
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = Field(default=4, ge=1)
     disable_autocast: bool = False
@@ -352,21 +404,21 @@ class OmniStageModelConfig:
     tokenizer_subdir: str | None = None
 
 
-@config
-class OmniStageLoadConfig:
-    """Per-stage loading behavior and resolved tokenizer input."""
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
+    """vLLM loading behavior plus Omni stage-specific tokenizer inputs."""
 
     tokenizer: str | None = None
-    download_dir: str | None = None
     skip_tokenizer_init: bool = False
-    load_format: str = "auto"
     tokenizer_mode: str = "auto"
     config_format: str | None = None
     skip_mm_profiling: bool | None = None
 
 
-@config
-class OmniStageCacheConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig):
     """Per-stage engine cache and memory behavior.
 
     This is separate from ``_DiffusionConfigProjection.cache_config``, which configures
@@ -383,17 +435,26 @@ class OmniStageCacheConfig:
     mamba_ssm_cache_dtype: str | None = None
 
 
-@config
-class OmniStageSchedulerConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
     """Per-stage request scheduling behavior."""
 
+    # Upstream receives max_model_len only while materializing SchedulerConfig.
+    # Omni retains it as unresolved stage input until the owning engine process.
     max_num_seqs: int | None = Field(default=None, ge=1)
     max_num_batched_tokens: int | None = Field(default=None, ge=1)
     max_model_len: int | None = Field(default=None, ge=-1)
+    is_encoder_decoder: InitVar[bool] = False
     enable_chunked_prefill: bool | None = None
     async_scheduling: bool | None = None
 
-    def __post_init__(self) -> None:
+    def __post_init__(self, is_encoder_decoder: bool = False) -> None:
+        # Upstream initializes these derived fields in its terminal post-init.
+        # Keep them serializable here without running model-dependent checks.
+        self.max_num_encoder_input_tokens = self.max_num_batched_tokens
+        self.encoder_cache_size = self.max_num_batched_tokens
+
         if (
             self.max_num_batched_tokens is not None
             and self.max_num_seqs is not None
@@ -432,24 +493,73 @@ class OmniStageRuntimeConfig:
     num_gpus: int = Field(default=1, ge=1)
     log_level: str = "info"
     log_stats: bool = False
-    profiler_config: dict[str, Any] | None = None
 
 
-@config
-class OmniStageParallelConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
     """Common per-stage distributed parallelism behavior."""
 
-    pipeline_parallel_size: int = Field(default=1, ge=1)
-    data_parallel_size: int = Field(default=1, ge=1)
-    tensor_parallel_size: int = Field(default=1, ge=1)
-    enable_expert_parallel: bool = False
-    world_size: int = Field(default=1, ge=1, init=False)
+    # EngineArgs intentionally leaves these unresolved. The upstream terminal
+    # ParallelConfig defaults (rank 0, local size 1, port 29550, worker
+    # ``"auto"``) must not turn into explicit engine inputs in the head
+    # process merely because this transport class inherits ParallelConfig.
+    data_parallel_size_local: int | None = Field(default=None, ge=0)
+    data_parallel_rank: int | None = Field(default=None, ge=0)
+    data_parallel_rpc_port: int | None = None
+    worker_cls: str | None = None
+
+    @model_validator(mode="after")
+    def _validate_parallel_config(self) -> Self:
+        """Run upstream validation without resolving deferred engine inputs."""
+        deferred_values = (
+            self.data_parallel_size_local,
+            self.data_parallel_rank,
+            self.data_parallel_rpc_port,
+            self.worker_cls,
+            self.all2all_backend,
+            self.disable_custom_all_reduce,
+        )
+        if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size_local ({self.data_parallel_size_local}) "
+                f"must be <= data_parallel_size ({self.data_parallel_size})"
+            )
+        if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size:
+            raise ValueError(
+                f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})"
+            )
+        self.data_parallel_size_local = self.data_parallel_size
+        self.data_parallel_rank = 0
+        self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port
+        self.worker_cls = VllmParallelConfig.worker_cls
+        try:
+            VllmParallelConfig._validate_parallel_config(self)
+        finally:
+            (
+                self.data_parallel_size_local,
+                self.data_parallel_rank,
+                self.data_parallel_rpc_port,
+                self.worker_cls,
+                self.all2all_backend,
+                self.disable_custom_all_reduce,
+            ) = deferred_values
+        return self
 
     def __post_init__(self) -> None:
+        # Keep config construction transport-safe. Upstream runtime backend,
+        # rank, port, and platform resolution happens in the owning process.
+        self.data_parallel_index = self.data_parallel_rank
         self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size
 
+    @property
+    def world_size_across_dp(self) -> int:
+        # Omni's public world_size has historically included DP.
+        return self.world_size
 
-@config
+
+@_enforce_keyword_only_init
+@config(kw_only=True)
 class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     """Diffusion-stage distributed parallelism behavior."""
 
@@ -468,6 +578,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     hsdp_replicate_size: int = Field(default=1, ge=1)
 
     def __post_init__(self) -> None:
+        self.data_parallel_index = self.data_parallel_rank
         self.sequence_parallel_size = (
             self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
         )
@@ -489,12 +600,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             * self.cfg_parallel_size
         )
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
-                raise ValueError(
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
-                    f"data_parallel_size={self.data_parallel_size})"
-                )
+            incompatible = []
+            if self.tensor_parallel_size > 1:
+                incompatible.append("TP")
+            if self.data_parallel_size > 1:
+                incompatible.append("DP")
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 if other_parallel_world_size == 1:
                     raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")
@@ -577,6 +693,7 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_fp8_cache_dir: str | None = None
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -814,6 +931,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "omni_kv_config",
         "cfg_kv_collect_func",
         "max_num_seqs",
+        "kv_cache_memory_bytes",
+        "gpu_memory_utilization",
+        "max_num_batched_tokens",
+        "max_model_len",
         "enable_sleep_mode",
         "enforce_eager",
         "enable_multithread_weight_load",
@@ -839,6 +960,58 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP
     "stage_id",
 }
 
+
+def _upstream_engine_field_map(
+    config_cls: type[Any],
+    *,
+    aliases: Mapping[str, str] = {},
+    exclude: frozenset[str] = frozenset(),
+) -> dict[str, str]:
+    """Map reusable upstream config fields to their EngineArgs inputs."""
+    engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs))
+    return {
+        config_field.name: engine_name
+        for config_field in fields(config_cls)
+        if config_field.init
+        and config_field.name not in exclude
+        and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields
+    }
+
+
+_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig)
+_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmCacheConfig,
+    aliases={"cache_dtype": "kv_cache_dtype"},
+)
+_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmSchedulerConfig,
+    aliases={"policy": "scheduling_policy"},
+    # ``scheduler_cls`` is selected by the immutable stage topology, while
+    # ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the
+    # Omni schema.  Keep both out of the dynamic SchedulerConfig projection so
+    # one input cannot acquire two owners.
+    exclude=frozenset(
+        {
+            "scheduler_cls",
+            "disable_hybrid_kv_cache_manager",
+        }
+    ),
+)
+_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmParallelConfig,
+    aliases={"data_parallel_master_ip": "data_parallel_address"},
+    # Runtime owns worker selection. The private API-process fields are
+    # terminal vLLM internals rather than per-stage user inputs.
+    exclude=frozenset(
+        {
+            "distributed_executor_backend",
+            "worker_cls",
+            "_api_process_count",
+            "_api_process_rank",
+        }
+    ),
+)
+
 _QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__)
 _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
 _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
@@ -846,7 +1019,17 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
 _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
 _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
 _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
-_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
+_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values())
+_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(
+    f.name for f in fields(OmniStageDiffusionParallelConfig)
+) & frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}
+_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP)
+_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
 _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
 _COMMON_STAGE_ENGINE_FIELDS = (
     _QUANTIZATION_ENGINE_FIELDS
@@ -856,12 +1039,16 @@ _COMMON_STAGE_ENGINE_FIELDS = (
     | _SCHEDULER_ENGINE_FIELDS
     | _CONNECTOR_ENGINE_FIELDS
     | _RUNTIME_ENGINE_FIELDS
+    | _DIRECT_VLLM_CONFIG_ENGINE_FIELDS
+)
+_LLM_STAGE_ENGINE_FIELDS = (
+    _COMMON_STAGE_ENGINE_FIELDS
+    | _LLM_LOAD_ENGINE_FIELDS
+    | _LLM_CACHE_ENGINE_FIELDS
+    | _LLM_SCHEDULER_ENGINE_FIELDS
+    | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
+    | {"parallel_config"}
 )
-_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"}
-_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - {
-    "world_size"
-}
-_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
 _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
     _COMMON_STAGE_ENGINE_FIELDS
     | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
@@ -882,6 +1069,27 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
     StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
     StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
 }
+_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = {
+    StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS,
+    StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS,
+    StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP),
+}
+
+_LOAD_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _LOAD_ENGINE_FIELDS},
+    **_LOAD_CONFIG_ENGINE_FIELD_MAP,
+}
+_CACHE_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _CACHE_ENGINE_FIELDS},
+    **_CACHE_CONFIG_ENGINE_FIELD_MAP,
+}
+_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _SCHEDULER_ENGINE_FIELDS},
+    **_SCHEDULER_CONFIG_ENGINE_FIELD_MAP,
+}
+_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS}
+_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS}
+_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS}
 
 
 def _validate_stage_engine_override_ownership(
@@ -892,13 +1100,19 @@ def _validate_stage_engine_override_ownership(
     try:
         owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
         parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
+        parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type]
     except KeyError as exc:
         raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
 
     unowned_fields = set(overrides) - owner_fields
     parallel_config = overrides.get("parallel_config")
     if isinstance(parallel_config, Mapping):
-        unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields)
+        # Nested values traditionally use upstream config names, while flat
+        # EngineArgs uses aliases such as ``data_parallel_address``.  Accept
+        # either spelling at the boundary; the builder canonicalizes aliases
+        # before instantiating the inherited config.
+        accepted_parallel_names = parallel_config_fields | parallel_owner_fields
+        unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names)
     if unowned_fields:
         names = ", ".join(sorted(unowned_fields))
         raise ValueError(
@@ -969,17 +1183,28 @@ def _stage_engine_values(
         topology.execution_type,
         engine,
     )
+    if topology.execution_type in {
+        StageExecutionType.LLM_AR,
+        StageExecutionType.LLM_GENERATION,
+    }:
+        load_engine_fields = _LLM_LOAD_ENGINE_FIELDS
+        cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS
+        scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS
+    else:
+        load_engine_fields = _LOAD_ENGINE_FIELDS
+        cache_engine_fields = _CACHE_ENGINE_FIELDS
+        scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS
     return _StageEngineValues(
         quantization=cast(
             _QuantizationEngineOverrides,
             _select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS),
         ),
         model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
-        load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)),
-        cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)),
+        load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
+        cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
         scheduler=cast(
             _SchedulerEngineOverrides,
-            _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS),
+            _select_engine_overrides(engine, scheduler_engine_fields),
         ),
         connector=cast(
             _ConnectorEngineOverrides,
@@ -988,6 +1213,8 @@ def _stage_engine_values(
         runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)),
         parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
         diffusion=_DiffusionEngineOverrides.from_engine(engine),
+        compilation_config=_copy_value(engine.get("compilation_config")),
+        profiler_config=_copy_value(engine.get("profiler_config")),
     )
 
 
@@ -1042,6 +1269,8 @@ class BaseVllmOmniStageConfig:
     connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
     runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
     parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
+    compilation_config: VllmCompilationConfig | None = None
+    profiler_config: VllmProfilerConfig | None = None
     quantization_config: _QuantizationConfigType = None
 
     @property
@@ -1193,6 +1422,8 @@ def _build_common_stage_config_kwargs(
                 parallel_config,
             ),
             "parallel_config": parallel_config,
+            "compilation_config": _copy_value(engine.compilation_config),
+            "profiler_config": _copy_value(engine.profiler_config),
             "quantization_config": _copy_value(quantization_config),
         },
         input_proc,
@@ -1387,12 +1618,51 @@ def _build_load_config(
     return OmniStageLoadConfig(**kwargs)
 
 
+def _config_kwargs_from_engine_args(
+    engine: Mapping[str, Any],
+    field_map: Mapping[str, str],
+) -> dict[str, Any]:
+    """Convert EngineArgs names back to their upstream config field names."""
+    return {
+        config_name: _copy_value(engine[engine_name])
+        for config_name, engine_name in field_map.items()
+        if engine_name in engine and engine[engine_name] is not None
+    }
+
+
+def _normalize_config_mapping(
+    values: Mapping[str, Any],
+    field_map: Mapping[str, str],
+) -> dict[str, Any]:
+    """Canonicalize config and EngineArgs spellings to config field names."""
+    engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()}
+    normalized: dict[str, Any] = {}
+    source_names: dict[str, str] = {}
+    for name, value in values.items():
+        if value is None:
+            continue
+        config_name = name if name in field_map else engine_to_config.get(name, name)
+        if config_name in normalized:
+            raise ValueError(
+                f"Config mapping specifies both {source_names[config_name]!r} and {name!r} "
+                f"for upstream field {config_name!r}"
+            )
+        normalized[config_name] = _copy_value(value)
+        source_names[config_name] = name
+    return normalized
+
+
+def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]:
+    return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None}
+
+
 def _build_cache_config(
     deploy: DeployConfig,
     engine: _CacheEngineOverrides,
     execution_type: StageExecutionType,
 ) -> OmniStageCacheConfig:
-    kwargs = _config_kwargs(engine)
+    kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP)
+    kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS))
     if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None:
         kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching)
     if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION:
@@ -1406,7 +1676,8 @@ def _build_scheduler_config(
     engine: _SchedulerEngineOverrides,
     execution_type: StageExecutionType,
 ) -> OmniStageSchedulerConfig:
-    kwargs = _config_kwargs(engine)
+    kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP)
+    kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS))
     if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None:
         kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill)
     if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR:
@@ -1454,16 +1725,21 @@ def _build_parallel_config(
     config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig,
 ) -> OmniStageParallelConfig:
     parallel_config = _mapping_or_empty(engine.get("parallel_config"))
-    config_fields = {
-        config_field.name
-        for config_field in fields(config_cls)
-        if getattr(config_field.default, "init", None) is not False
-    }
-    kwargs = {
-        name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None
-    }
+    config_fields = frozenset(signature(config_cls).parameters)
+    is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig)
+    field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP
+    owned_parallel_fields = frozenset(field_map)
+    kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map)
+    kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields}
     kwargs = {name: value for name, value in kwargs.items() if name in config_fields}
-    kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields})
+    nested_kwargs = _normalize_config_mapping(parallel_config, field_map)
+    kwargs.update(
+        {
+            name: value
+            for name, value in nested_kwargs.items()
+            if name in owned_parallel_fields and name in config_fields
+        }
+    )
     if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None:
         kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size)
     if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None:

vllm_omni/config/pipeline_registry.py

@@ -33,6 +33,7 @@ from vllm.logger import init_logger
 from vllm_omni.config.stage_config import (
     PipelineConfig,
 )
+from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
 from vllm_omni.model_executor.models.audex.pipeline import (
     AUDEX_S2S_PIPELINE,
     AUDEX_THINKER_ONLY_PIPELINE,
@@ -84,6 +85,7 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import (
     MING_TTS_PIPELINE,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
+from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE
 from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_LOCAL_PIPELINE,
     MOSS_TTS_PIPELINE,
@@ -139,6 +141,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "lance": LANCE_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "Gr00tN1d7": GR00T_N1D7_PIPELINE,
+    "pi0": PI0_PIPELINE,
     "gepard": GEPARD_PIPELINE,
     "glm_image": GLM_IMAGE_PIPELINE,
     "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,
@@ -175,6 +178,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
     "minicpmo_4_5": MINICPMO_4_5_PIPELINE,
+    "minimax_music3": MINIMAX_MUSIC3_PIPELINE,
     "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
     "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
     "dynin_omni": DYNIN_OMNI_PIPELINE,

vllm_omni/config/stage_config.py

@@ -383,6 +383,7 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_fp8_cache_dir: str | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/data_entry_keys.py

@@ -91,6 +91,11 @@ class OmniPayloadMeta(TypedDict, total=False):
     ref_context_included: bool
     talker_prefill_offset: int
     omni_final_stage_id: int
+    # Per-request audio seed. Stages that draw their own noise (flow-matching
+    # / diffusion decoders) need the producing stage's seed to stay
+    # reproducible, and the producing stage's SamplingParams do not travel
+    # with the payload.
+    audio_seed: int
 
 
 class OmniPayload(TypedDict, total=False):
@@ -189,6 +194,7 @@ class MetaStruct(_StructBase):
     codec_left_context_frames: int | None = None
     code_flat_numel: int | None = None
     omni_final_stage_id: int | None = None
+    audio_seed: int | None = None
 
 
 class OmniPayloadStruct(_StructBase):