docs: publish PR 6279 benchmark figures

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

.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/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` | ✅︎ | ✅︎ | | ✅︎ | — |

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

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/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/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/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/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/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",

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

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_pi0 import PI0_PIPELINE
 from vllm_omni.model_executor.models.audex.pipeline import (
     AUDEX_S2S_PIPELINE,
     AUDEX_THINKER_ONLY_PIPELINE,
@@ -139,6 +140,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,

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -111,3 +111,9 @@ platforms:
         max_num_batched_tokens: 8192
         compilation_config:
           cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -93,3 +93,9 @@ platforms:
         max_num_batched_tokens: 8192
         compilation_config:
           cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -86,4 +86,20 @@ stages:
       max_tokens: 65536
       detokenize: true
 
-platforms: {}
+platforms:
+  npu:
+    stages:
+      - stage_id: 0
+        max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 1
+        max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 2
+        # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+        # estimator kernels through the NPU platform patch.
+        additional_config:
+          code2wav_enable_npu_graph: true
+          code2wav_max_npu_graphs: 32

vllm_omni/deploy/pi0.yaml

@@ -0,0 +1,52 @@
+# π0 (Pi-Zero) VLA deploy: single diffusion stage.
+#
+# Topology (PI0_PIPELINE) is declared in vllm_omni/diffusion/models/pi0/pipeline_pi0.py.
+# π0 takes multi-camera images + a language instruction + robot state and emits
+# a continuous action chunk [chunk_size, max_action_dim] via flow matching.
+# Actions are exposed through the OpenPI realtime serving API
+# (/v1/realtime/robot/openpi); they ride in multimodal_output["actions"].
+#
+# Default: one GPU, TP=1, eager (flow matching needs no CUDA graph), bsz 1
+# (OpenPI policy serving is one observation -> one action reply). float32 to
+# match the LeRobot pi0_base reference numerics.
+
+pipeline: pi0
+async_chunk: false
+distributed_executor_backend: mp
+dtype: float32
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    enforce_eager: true
+    model_class_name: Pi0Pipeline
+    model_config:
+      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
+      max_cameras: 3
+      # Camera identities (the order the model attends to). Serving maps raw obs
+      # keys onto these via image_key_map; leave image_key_map empty to match
+      # incoming keys verbatim.
+      image_feature_keys:
+        - observation.images.base_0_rgb
+        - observation.images.left_wrist_0_rgb
+        - observation.images.right_wrist_0_rgb
+      image_key_map: {}
+      # OpenPI handshake metadata advertised to the robot client.
+      policy_server_config:
+        image_resolution: [224, 224]
+        n_external_cameras: 1
+        needs_wrist_camera: true
+        needs_stereo_camera: false
+        needs_session_id: false
+        action_space: joint_position
+        action_horizon: 50
+        action_dim: 32
+        max_cameras: 3

vllm_omni/diffusion/data.py

@@ -1318,6 +1318,12 @@ class OmniDiffusionConfig:
                         self.update_multimodal_support()
                     else:
                         raise
+                elif cfg.get("type") == "pi0":
+                    # π0 (Pi-Zero) VLA — the LeRobot config.json uses ``type: "pi0"``.
+                    if self.model_class_name is None:
+                        self.model_class_name = "Pi0Pipeline"
+                    self.set_tf_model_config(TransformerConfig())
+                    self.update_multimodal_support()
                 elif architectures and len(architectures) == 1:
                     architecture = architectures[0]
                     from vllm_omni.diffusion.registry import DiffusionModelRegistry

vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py

@@ -203,6 +203,11 @@ def _apply_unquantized_gemm_with_weight(
     weight: torch.Tensor,
     bias: torch.Tensor | None,
 ) -> torch.Tensor:
+    # GEMM dispatch is selected from the host platform, not the tensor device.
+    # In CPU tests on a ROCm host that selects rocm_unquantized_gemm, which has
+    # no CPU kernel. Keep that combination on PyTorch's portable implementation.
+    if current_platform.is_rocm() and input_.device.type == "cpu":
+        return F.linear(input_, weight, bias)
     if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
         return linear_batch_invariant(input_, weight, bias)
     return dispatch_unquantized_gemm()(layer, input_, weight, bias)

vllm_omni/diffusion/models/pi0/__init__.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 (Pi-Zero) VLA model for vllm-omni.
+
+PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert + a
+flow-matching action head. Outputs a continuous action chunk
+``[horizon, action_dim]`` rather than tokens.
+"""
+
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+    GemmaVariantConfig,
+    PaliGemmaWithActionExpert,
+    Pi0ForActionPrediction,
+    create_sinusoidal_pos_embedding,
+    get_gemma_config,
+    make_att_2d_masks,
+    prepare_attention_masks_4d,
+)
+from vllm_omni.diffusion.models.pi0.pipeline_pi0 import (
+    Pi0Pipeline,
+    get_pi0_post_process_func,
+)
+from vllm_omni.diffusion.models.pi0.processor_pi0 import (
+    Pi0ImageProcessor,
+    build_model_inputs,
+    pil_image_to_tensor,
+    resize_with_pad,
+    tokenize_prompt,
+)
+
+__all__ = [
+    "Pi0Config",
+    "Pi0ForActionPrediction",
+    "PaliGemmaWithActionExpert",
+    "GemmaVariantConfig",
+    "get_gemma_config",
+    "create_sinusoidal_pos_embedding",
+    "make_att_2d_masks",
+    "prepare_attention_masks_4d",
+    "Pi0ImageProcessor",
+    "build_model_inputs",
+    "pil_image_to_tensor",
+    "resize_with_pad",
+    "tokenize_prompt",
+    "Pi0Pipeline",
+    "get_pi0_post_process_func",
+]

vllm_omni/diffusion/models/pi0/config.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config surface for the π0 (Pi-Zero) VLA model in vllm-omni.
+
+Only the parameters that actually shape runtime behaviour live here; the
+transformer dimensions (hidden size, head count, ...) are derived from
+``paligemma_variant`` / ``action_expert_variant`` inside the model via
+``get_gemma_config`` (see ``modeling_pi0``).
+
+The resolver reads the **raw LeRobot ``config.json``** (the field surface of
+``lerobot.policies.pi0.PI0Config``): only the recognized dataclass fields are
+used; any other keys in the file are ignored.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from dataclasses import fields as dataclass_fields
+from pathlib import Path
+from typing import Any
+
+# LeRobot / OpenPI observation key conventions.
+ACTION = "action"
+OBS_STR = "observation"
+OBS_STATE = OBS_STR + ".state"
+OBS_IMAGES = OBS_STR + ".images"
+
+
+@dataclass
+class Pi0Config:
+    """π0 VLA config (dataclass, not an HF ``PretrainedConfig``).
+
+    Mirrors the runtime-relevant subset of LeRobot ``PI0Config`` plus a couple
+    of serving-side knobs (``max_cameras``, ``image_feature_keys``,
+    ``image_key_map``, ``norm_stats``).
+    """
+
+    # Backbone variants — mapped to Gemma dimensions by ``get_gemma_config``.
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+
+    # Action chunk shape.
+    chunk_size: int = 50
+    max_action_dim: int = 32
+    max_state_dim: int = 32
+
+    # Flow-matching denoising schedule.
+    num_inference_steps: int = 10
+    # Sinusoidal timestep embedding periods (must match OpenPI/LeRobot).
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    # Image preprocessing. π0/SigLIP only support square inputs.
+    image_resolution: tuple[int, int] = (224, 224)
+    tokenizer_max_length: int = 48
+    # Number of camera slots the model attends to (real + padded).
+    max_cameras: int = 3
+
+    # Weight dtype the checkpoint was saved in; parity passes this through to
+    # match LeRobot's ``PI0Config.dtype``.
+    dtype: str = "float32"
+
+    # Per-dataset normalization stats, optional (schema matches LeRobot's
+    # ``NormalizerProcessorStep``; consumed by ``_build_norm_buffers``). ``None``
+    # means identity / pass-through (the ``lerobot/pi0_base`` case).
+    norm_stats: dict | None = None
+
+    # Ordered list of image feature keys (e.g. ``observation.images.base_0_rgb``)
+    # taken from the checkpoint's ``input_features``. This is the **camera order**
+    # that must be reproduced for LeRobot parity.
+    image_feature_keys: list[str] | None = None
+
+    # Optional map from raw OpenPI obs keys → ``image_feature_keys`` entries, so
+    # serving can translate wire obs (e.g. ``observation/exterior_image_0_left``)
+    # into the checkpoint's camera identities. Empty → identity match on key.
+    image_key_map: dict[str, str] = field(default_factory=dict)
+
+    # Stored for reference / camera-order derivation; not used at inference.
+    input_features: dict[str, Any] = field(default_factory=dict)
+    output_features: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # Coerce list → tuple (JSON has no tuples) and validate squareness.
+        res = self.image_resolution
+        if not isinstance(res, (tuple, list)) or len(res) != 2 or res[0] != res[1]:
+            raise ValueError(f"π0 expects a square image_resolution (H == W); got {res!r}.")
+        self.image_resolution = (int(res[0]), int(res[1]))
+
+        # Derive the camera order from input_features if not given explicitly.
+        if self.image_feature_keys is None and self.input_features:
+            self.image_feature_keys = [key for key in self.input_features if key.startswith(OBS_IMAGES + ".")]
+
+    # ------------------------------------------------------------------
+    # Constructors
+    # ------------------------------------------------------------------
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> Pi0Config:
+        """Build from a checkpoint directory's ``config.json``."""
+        checkpoint_dir = Path(checkpoint_dir)
+        config_path = checkpoint_dir / "config.json"
+        if not config_path.exists():
+            return cls()
+        with open(config_path, encoding="utf-8") as f:
+            raw = json.load(f)
+        return cls.from_model_config(raw)
+
+    @classmethod
+    def from_model_config(cls, model_config: dict[str, Any] | None) -> Pi0Config:
+        """Build from a config dict (LeRobot ``config.json`` or deploy yaml).
+
+        Keeps only the recognized dataclass fields (the LeRobot config also
+        carries many training-only keys) and coerces ``image_resolution`` to a
+        tuple.
+        """
+        if not model_config:
+            return cls()
+
+        raw = dict(model_config)
+        if "image_resolution" in raw:
+            raw["image_resolution"] = tuple(raw["image_resolution"])
+
+        allowed = {item.name for item in dataclass_fields(cls)}
+        filtered = {key: value for key, value in raw.items() if key in allowed}
+        return cls(**filtered)

vllm_omni/diffusion/models/pi0/modeling_pi0.py

@@ -0,0 +1,944 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only π0 (Pi-Zero) VLA math kernel for vllm-omni.
+
+A self-contained inference kernel: only the math that turns a robot observation
+into an action chunk, with no serving/request glue — validated bit-for-bit
+against the LeRobot ``PI0Policy`` reference (``max|Δ| < 1e-4``).
+
+π0 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert
+suffix + flow-matching head. Inference:
+
+1. Embed prefix (images + language) → bidirectional prefix tokens.
+2. Forward prefix through PaliGemma → a per-layer ``list[(k, v)]`` KV cache.
+3. For each Euler step ``t = 1.0, 1-dt, ..., 0``: embed the suffix
+   (state + noisy actions + timestep), run the action expert with cross
+   attention over the cached prefix K/V, predict velocity ``v_t``, and
+   integrate ``x_t = x_t + dt * v_t``.
+4. Return ``x_0`` as the action chunk ``(batch, action_horizon, action_dim)``.
+
+We walk Gemma decoder layers ourselves in ``_compute_layer_*`` so we never go
+through ``GemmaModel.forward`` (which has version-dependent mask/KV-cache
+behaviour). Targets modern transformers (PaliGemmaForConditionalGeneration with
+an inner PaliGemmaModel + DynamicCache).
+
+Reference implementations:
+   - OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py
+   - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py
+Weight source: https://huggingface.co/lerobot/pi0_base
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.auto import CONFIG_MAPPING
+from transformers.models.gemma.modeling_gemma import (
+    GemmaForCausalLM,
+    apply_rotary_pos_emb,
+)
+from transformers.models.paligemma.modeling_paligemma import (
+    PaliGemmaForConditionalGeneration,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+logger = logging.getLogger(__name__)
+
+# ──────────────────────────────────────────────────────────────────────
+# Constants
+# ──────────────────────────────────────────────────────────────────────
+DEFAULT_ACTION_DIM = 32
+DEFAULT_ACTION_HORIZON = 50
+DEFAULT_MAX_TOKEN_LEN = 48
+DEFAULT_NUM_INFERENCE_STEPS = 10
+DEFAULT_IMAGE_RESOLUTION = (224, 224)
+
+# Large negative value to fill masked-out positions in a float attention mask.
+# Matches OpenPI's constant exactly so that numerics line up during parity.
+# Ref: openpi/src/openpi/models/gemma.py
+OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Gemma variant configs (matches openpi/models/gemma.py get_config)
+# ──────────────────────────────────────────────────────────────────────
+
+
+class GemmaVariantConfig:
+    def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
+        self.width = width
+        self.depth = depth
+        self.mlp_dim = mlp_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+
+
+def get_gemma_config(variant: str) -> GemmaVariantConfig:
+    if variant == "gemma_2b":
+        return GemmaVariantConfig(2048, 18, 16384, 8, 1, 256)
+    elif variant == "gemma_300m":
+        return GemmaVariantConfig(1024, 18, 4096, 8, 1, 256)
+    else:
+        raise ValueError(f"Unknown variant: {variant}")
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Utility functions (match openpi/models_pytorch/pi0_pytorch.py)
+# ──────────────────────────────────────────────────────────────────────
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float = 4e-3,
+    max_period: float = 4.0,
+    device: torch.device = None,
+) -> torch.Tensor:
+    """Compute a sine/cosine positional embedding for scalar timesteps.
+
+    Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding
+    """
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    if time.ndim != 1:
+        raise ValueError("time tensor must be 1-D (batch_size,)")
+    if device is None:
+        device = time.device
+
+    # Use float64 for the log-linear sweep and for the inner products, to
+    # match the reference implementation's numerical behaviour exactly.
+    dtype = torch.float64
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None].to(dtype)
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
+    """Build a 2D attention mask from a padding mask and an autoregressive mask.
+
+    Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(f"att_masks must be 2-D, got {att_masks.ndim}-D")
+    if pad_masks.ndim != 2:
+        raise ValueError(f"pad_masks must be 2-D, got {pad_masks.ndim}-D")
+
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+def _build_norm_buffers(norm_stats: dict | None, key: str) -> dict[str, torch.Tensor] | None:
+    """Parse a ``norm_stats[key]`` entry into CPU tensors, or ``None``.
+
+    Matches LeRobot's ``NormalizationMode``:
+      - ``mean_std`` : ``forward = (x - mean) / std``, ``inverse = x * std + mean``
+      - ``min_max``  : forward maps into ``[-1, 1]``, inverse maps back.
+    """
+    if not norm_stats or not isinstance(norm_stats, dict):
+        return None
+    entry = norm_stats.get(key)
+    if not entry:
+        return None
+    mode = str(entry.get("mode", "mean_std")).lower()
+    if mode == "mean_std":
+        mean = entry.get("mean")
+        std = entry.get("std")
+        if mean is None or std is None:
+            return None
+        return {
+            "mode": mode,
+            "mean": torch.as_tensor(mean, dtype=torch.float32),
+            "std": torch.as_tensor(std, dtype=torch.float32),
+        }
+    if mode == "min_max":
+        lo = entry.get("min")
+        hi = entry.get("max")
+        if lo is None or hi is None:
+            return None
+        return {
+            "mode": mode,
+            "min": torch.as_tensor(lo, dtype=torch.float32),
+            "max": torch.as_tensor(hi, dtype=torch.float32),
+        }
+    return None
+
+
+def _apply_norm(
+    x: torch.Tensor,
+    stats: dict[str, torch.Tensor] | None,
+    inverse: bool,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    """Apply (un)normalization using the given stats. No-op if ``stats`` is None.
+
+    Broadcasts over leading dims; stats vectors align with the last dim. When
+    ``x`` has more last-dim entries than the stats (π0 pads state/action to
+    ``max_dim``), only the first ``len(stats)`` entries are transformed and
+    the padded tail is left untouched.
+    """
+    if stats is None:
+        return x
+    mode = stats["mode"]
+    if mode == "mean_std":
+        mean = stats["mean"].to(device=x.device, dtype=x.dtype)
+        std = stats["std"].to(device=x.device, dtype=x.dtype)
+        valid = mean.shape[0]
+        head = x[..., :valid]
+        if inverse:
+            head = head * std + mean
+        else:
+            head = (head - mean) / (std + eps)
+        if valid == x.shape[-1]:
+            return head
+        return torch.cat([head, x[..., valid:]], dim=-1)
+    if mode == "min_max":
+        lo = stats["min"].to(device=x.device, dtype=x.dtype)
+        hi = stats["max"].to(device=x.device, dtype=x.dtype)
+        valid = lo.shape[0]
+        head = x[..., :valid]
+        denom = (hi - lo).clamp_min(eps)
+        if inverse:
+            head = (head + 1.0) * 0.5 * denom + lo
+        else:
+            head = 2.0 * (head - lo) / denom - 1.0
+        if valid == x.shape[-1]:
+            return head
+        return torch.cat([head, x[..., valid:]], dim=-1)
+    return x
+
+
+def prepare_attention_masks_4d(att_2d_masks: torch.Tensor) -> torch.Tensor:
+    """Convert ``(B, S, S)`` bool masks to ``(B, 1, S, S)`` float masks.
+
+    ``True`` → 0.0 (attend), ``False`` → ``OPENPI_ATTENTION_MASK_VALUE``.
+    Ref: openpi PI0Pytorch._prepare_attention_masks_4d
+    """
+    att_2d_masks_4d = att_2d_masks[:, None, :, :]
+    return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Dual-backbone: PaliGemma + Action Expert
+# Ref: openpi/models_pytorch/gemma_pytorch.py PaliGemmaWithExpertModel
+# Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel
+# ──────────────────────────────────────────────────────────────────────
+
+
+def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Repeat KV heads for grouped-query attention.
+
+    ``(B, num_kv_heads, S, D) → (B, num_kv_heads * n_rep, S, D)``.
+    Matches ``transformers.models.gemma.modeling_gemma.repeat_kv``.
+    """
+    if n_rep == 1:
+        return hidden_states
+    b, nh, s, d = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(b, nh, n_rep, s, d)
+    return hidden_states.reshape(b, nh * n_rep, s, d)
+
+
+def _attend(query_states, key_states, value_states, attention_mask, num_kv_groups, scaling):
+    """Manual eager attention: ``softmax(Q Kᵀ · scale + mask) · V``.
+
+    Important: we slice ``attention_mask[..., :key_states.shape[-2]]`` so that
+    the same ``(B, 1, Q, prefix+suffix)`` mask produced by ``denoise_step``
+    can be re-used both in the prefix pass (K length = prefix) and in the
+    suffix pass (K length = prefix + suffix after we concat the cached prefix).
+    Stock transformers' ``eager_attention_forward`` adds the mask without
+    slicing and therefore breaks when the two dims disagree.
+    """
+    k = _repeat_kv(key_states, num_kv_groups)
+    v = _repeat_kv(value_states, num_kv_groups)
+    attn_weights = torch.matmul(query_states, k.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask[:, :, :, : k.shape[-2]]
+    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+    return torch.matmul(attn_weights, v)
+
+
+def _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma):
+    """Run one PaliGemma LM layer on the prefix only, returning the layer
+    output **and** the post-RoPE ``(k, v)`` so the caller can cache them for
+    the suffix pass."""
+    model = paligemma.model.language_model
+    layer = model.layers[layer_idx]
+    residual = hidden_states
+    x = layer.input_layernorm(hidden_states)
+
+    hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim)
+    q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2)
+    k = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2)
+    v = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2)
+
+    cos, sin = model.rotary_emb(v, position_ids)
+    q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
+
+    att = _attend(
+        q,
+        k,
+        v,
+        attention_mask,
+        num_kv_groups=layer.self_attn.num_key_value_groups,
+        scaling=1.0 / math.sqrt(layer.self_attn.head_dim),
+    )
+    att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim)
+
+    out = layer.self_attn.o_proj(att) + residual
+    after_resid = out
+    out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid
+    return out, (k, v)
+
+
+def _compute_layer_suffix_only(
+    layer_idx,
+    hidden_states,
+    prefix_kv,
+    attention_mask,
+    position_ids,
+    paligemma,
+    gemma_expert,
+):
+    """Run one action-expert layer on the suffix, attending to the cached
+    prefix K/V concatenated with freshly computed suffix K/V.
+
+    ``prefix_kv`` is a ``(k_prefix, v_prefix)`` tuple, shape
+    ``(B, num_kv_heads, prefix_len, head_dim)``, post-RoPE, produced by the
+    corresponding ``_compute_layer_prefix_only`` call. Bypasses
+    ``GemmaModel.forward`` entirely — no HF mask rebuild, no cache format
+    gymnastics.
+    """
+    layer = gemma_expert.model.layers[layer_idx]
+    residual = hidden_states
+    x = layer.input_layernorm(hidden_states)
+
+    hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim)
+    q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2)
+    k_suf = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2)
+    v_suf = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2)
+
+    # RoPE frequencies are shared between PaliGemma and the expert.
+    cos, sin = gemma_expert.model.rotary_emb(v_suf, position_ids)
+    q, k_suf = apply_rotary_pos_emb(q, k_suf, cos, sin, unsqueeze_dim=1)
+
+    # Concatenate cached prefix K/V (possibly different dtype) with suffix K/V.
+    k_prefix, v_prefix = prefix_kv
+    k = torch.cat([k_prefix.to(k_suf.dtype), k_suf], dim=2)
+    v = torch.cat([v_prefix.to(v_suf.dtype), v_suf], dim=2)
+
+    att = _attend(
+        q,
+        k,
+        v,
+        attention_mask,
+        num_kv_groups=layer.self_attn.num_key_value_groups,
+        scaling=1.0 / math.sqrt(layer.self_attn.head_dim),
+    )
+    att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim)
+
+    out = layer.self_attn.o_proj(att) + residual
+    after_resid = out
+    out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid
+    return out
+
+
+class PaliGemmaWithActionExpert(nn.Module):
+    """Dual-backbone transformer: PaliGemma (Gemma 2B) + Action Expert (Gemma 300M).
+
+    ``forward`` has two inference modes:
+      - **prefix_only**: ``inputs_embeds=[prefix, None]`` + ``use_cache=True``
+        → compute prefix hidden states + a layer-wise K/V cache.
+      - **suffix_only**: ``inputs_embeds=[None, suffix]`` + the cache from the
+        previous prefix pass → compute expert hidden states with cross-attention
+        over the concatenated (prefix, suffix) K/V.
+
+    Both modes walk Gemma decoder layers manually through ``_compute_layer_*``
+    instead of ``GemmaModel.forward`` — that way we own the attention mask
+    handling and the KV cache format (a plain ``list[(k, v)]`` one entry
+    per layer).
+
+    Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel
+    """
+
+    def __init__(self, vlm_config, action_expert_config):
+        super().__init__()
+
+        # Build HF PaliGemma config from the variant dims we expose.
+        # NOTE: we do NOT set use_adarms / adarms_cond_dim — those are
+        # OpenPI-only attributes used by π0.5, not π0.
+        vlm_config_hf = CONFIG_MAPPING["paligemma"]()
+        vlm_config_hf._vocab_size = 257152
+        vlm_config_hf.image_token_index = 257152
+        vlm_config_hf.text_config.hidden_size = vlm_config.width
+        vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
+        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
+        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
+        vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
+        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
+        vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
+        # transformers ≥ 5 uses ``dtype``; older versions still accept
+        # ``torch_dtype``. We set the canonical name; HF's PretrainedConfig
+        # normalizes one to the other internally.
+        vlm_config_hf.text_config.dtype = "float32"
+        vlm_config_hf.text_config.vocab_size = 257152
+        vlm_config_hf.vision_config.intermediate_size = 4304
+        vlm_config_hf.vision_config.projection_dim = 2048
+        vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
+        vlm_config_hf.vision_config.dtype = "float32"
+
+        action_expert_config_hf = CONFIG_MAPPING["gemma"](
+            head_dim=action_expert_config.head_dim,
+            hidden_size=action_expert_config.width,
+            intermediate_size=action_expert_config.mlp_dim,
+            num_attention_heads=action_expert_config.num_heads,
+            num_hidden_layers=action_expert_config.depth,
+            num_key_value_heads=action_expert_config.num_kv_heads,
+            vocab_size=257152,
+            hidden_activation="gelu_pytorch_tanh",
+            dtype="float32",
+        )
+
+        self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf)
+        self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf)
+        # The action expert doesn't embed tokens — it only consumes the
+        # suffix state/action embeddings we feed in.
+        self.gemma_expert.model.embed_tokens = None
+
+    def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Encode images with SigLIP vision tower + PaliGemma projector.
+
+        We run the two steps explicitly instead of calling
+        ``PaliGemmaModel.get_image_features`` because that convenience
+        method historically has returned either the projected tensor (old)
+        or the raw vision-tower output (new). Being explicit makes us
+        independent of which transformers release we're on.
+        """
+        # Shapes: pixel_values (B, 3, 224, 224) → SigLIP (B, 256, 1152)
+        #                                      → projector (B, 256, 2048)
+        vision_outputs = self.paligemma.model.vision_tower(pixel_values)
+        image_features = vision_outputs.last_hidden_state
+        return self.paligemma.model.multi_modal_projector(image_features)
+
+    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Embed language tokens with PaliGemma's embedding table, returning the
+        ``* sqrt(hidden)``-scaled embedding (PaliGemma/Gemma convention).
+
+        The scaling location moved across transformers releases:
+          * ≤ 5.3: ``embed_tokens`` is a plain ``nn.Embedding`` and the
+            ``* sqrt(hidden)`` normalizer lives inside ``GemmaModel.forward``
+            (which we bypass) — so we apply it explicitly here.
+          * ≥ 5.4: ``embed_tokens`` is a ``GemmaTextScaledWordEmbedding`` that
+            self-applies ``embed_scale = hidden_size ** 0.5`` — applying it again
+            would double-scale (≈45×). We detect this and skip the manual scale.
+        This makes ``embed_language_tokens`` return the canonical scaled embedding
+        on every transformers version.
+        """
+        embed_tokens = self.paligemma.model.language_model.embed_tokens
+        lang_emb = embed_tokens(tokens)
+        # If the embedding already self-scales (transformers ≥ 5.4), don't repeat it.
+        if getattr(embed_tokens, "embed_scale", None) is None:
+            lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
+        return lang_emb
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values=None,
+        inputs_embeds: list[torch.Tensor | None] | None = None,
+        use_cache: bool = False,
+    ):
+        """Dispatch to prefix_only / suffix_only and return
+        ``([prefix_out, suffix_out], past_key_values_or_None)``.
+        """
+        num_layers = self.paligemma.config.text_config.num_hidden_layers
+        pali_lm = self.paligemma.model.language_model
+        expert_lm = self.gemma_expert.model
+
+        if inputs_embeds[1] is None:
+            # Prefix-only: PaliGemma LM on (images + language) tokens; the
+            # per-layer post-RoPE K/V is collected into a list that the
+            # suffix pass will consume directly.
+            hidden_states = inputs_embeds[0]
+            kv_list: list[tuple[torch.Tensor, torch.Tensor]] = []
+            for layer_idx in range(num_layers):
+                hidden_states, kv = _compute_layer_prefix_only(
+                    layer_idx,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    paligemma=self.paligemma,
+                )
+                kv_list.append(kv)
+            hidden_states = pali_lm.norm(hidden_states)
+            return [hidden_states, None], (kv_list if use_cache else None)
+
+        if inputs_embeds[0] is not None:
+            raise ValueError(
+                "PaliGemmaWithActionExpert.forward only supports prefix-only "
+                "or suffix-only dispatch; got both inputs_embeds populated."
+            )
+        # Suffix-only: action expert on (state + noisy_actions + time),
+        # attending to the cached prefix K/V.
+        if not isinstance(past_key_values, list):
+            raise TypeError(
+                "suffix_only forward expects past_key_values to be the "
+                "list[(k, v)] produced by a previous prefix_only forward; "
+                f"got {type(past_key_values)}"
+            )
+        hidden_states = inputs_embeds[1]
+        for layer_idx in range(num_layers):
+            hidden_states = _compute_layer_suffix_only(
+                layer_idx,
+                hidden_states,
+                past_key_values[layer_idx],
+                attention_mask,
+                position_ids,
+                paligemma=self.paligemma,
+                gemma_expert=self.gemma_expert,
+            )
+        hidden_states = expert_lm.norm(hidden_states)
+        return [None, hidden_states], None
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Main π0 Model
+# ──────────────────────────────────────────────────────────────────────
+
+
+class Pi0ForActionPrediction(nn.Module):
+    """π0 VLA model for robot action prediction via flow matching.
+
+    Inference flow:
+      1. Embed prefix (images + language) → prefix tokens.
+      2. Forward prefix through PaliGemma → layer-wise KV cache.
+      3. For each denoising step ``t = 1.0, 1-dt, ..., 0``:
+         a. Embed suffix (state + x_t + timestep) → suffix tokens.
+         b. Forward suffix through the action expert with the prefix cache.
+         c. ``x_t = x_t + dt * v_t`` (Euler integration).
+      4. Return ``x_0`` as the predicted action chunk.
+    """
+
+    def __init__(
+        self,
+        config,
+        quant_config=None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        # ``quant_config`` is accepted for interface compatibility but unused —
+        # π0 runs in full precision for flow-matching parity.
+        del quant_config
+        self.config = config
+
+        self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM)
+        # ``max_state_dim`` is independent from ``max_action_dim`` — e.g. Aloha
+        # has state_dim=14, action_dim=32, both padded to their own max.  Fall
+        # back to action_dim for older configs that only expose a single size.
+        self.max_state_dim = getattr(config, "max_state_dim", self.action_dim)
+        self.action_horizon = getattr(config, "chunk_size", DEFAULT_ACTION_HORIZON)
+        self.num_inference_steps = getattr(config, "num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS)
+
+        paligemma_variant = getattr(config, "paligemma_variant", "gemma_2b")
+        action_expert_variant = getattr(config, "action_expert_variant", "gemma_300m")
+        vlm_config = get_gemma_config(paligemma_variant)
+        expert_config = get_gemma_config(action_expert_variant)
+        self.vlm_width = vlm_config.width
+        self.expert_width = expert_config.width
+
+        # Dual backbone
+        self.paligemma_with_expert = PaliGemmaWithActionExpert(vlm_config, expert_config)
+
+        # Action chunk projections (openpi pi0_pytorch.py).
+        self.action_in_proj = nn.Linear(self.action_dim, self.expert_width)
+        self.action_out_proj = nn.Linear(self.expert_width, self.action_dim)
+
+        # State is projected to the expert dim once (π0 only, not π0.5).
+        # Ref: lerobot modeling_pi0.py ``nn.Linear(config.max_state_dim, ...)``
+        self.state_proj = nn.Linear(self.max_state_dim, self.expert_width)
+
+        # Timestep + action fusion MLP (2W → W → W).
+        self.action_time_mlp_in = nn.Linear(2 * self.expert_width, self.expert_width)
+        self.action_time_mlp_out = nn.Linear(self.expert_width, self.expert_width)
+
+        # Optional mean/std normalization stats (LeRobot NormalizerProcessorStep
+        # semantics). Format on the config (all optional):
+        #     norm_stats = {
+        #         "state":  {"mode": "mean_std", "mean": [...], "std": [...]},
+        #         "action": {"mode": "mean_std", "mean": [...], "std": [...]},
+        #     }
+        # ``mode`` may also be ``"min_max"``.
+        #
+        # Missing norm_stats is logged at INFO (not WARNING) because the most
+        # common setup is "client pre-normalizes state and post-normalizes
+        # actions using dataset stats" — that's a supported mode, not an error.
+        self._state_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "state")
+        self._action_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "action")
+        if self._state_norm is None:
+            logger.info(
+                "π0: no state normalization stats on config.norm_stats — state values will pass through unchanged."
+            )
+        if self._action_norm is None:
+            logger.info(
+                "π0: no action normalization stats on config.norm_stats — "
+                "returned actions are in the model's normalized space."
+            )
+
+    # ── State / action normalization ─────────────────────────────────
+
+    def _normalize_state(self, state: torch.Tensor) -> torch.Tensor:
+        return _apply_norm(state, self._state_norm, inverse=False)
+
+    def _unnormalize_actions(self, actions: torch.Tensor) -> torch.Tensor:
+        return _apply_norm(actions, self._action_norm, inverse=True)
+
+    # ── Prefix embedding ─────────────────────────────────────────────
+
+    def embed_prefix(
+        self,
+        images: list[torch.Tensor],
+        image_masks: list[torch.Tensor],
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build the prefix embeddings, per-token padding mask, and AR mask.
+
+        Prefix tokens form a contiguous sequence
+        ``[img_cam_0..., img_cam_1..., ..., lang_tokens...]`` with
+        bidirectional attention; the returned ``att_masks`` are all zeros
+        because each token is free to attend to every other prefix token (the
+        suffix pass will put a causal boundary right before the state token).
+
+        Ref: openpi PI0Pytorch.embed_prefix
+        """
+        embs: list[torch.Tensor] = []
+        pad_masks: list[torch.Tensor] = []
+        att_masks: list[int] = []
+
+        for img, img_mask in zip(images, image_masks):
+            img_emb = self.paligemma_with_expert.embed_image(img)
+            bsize, num_img_embs = img_emb.shape[:2]
+            embs.append(img_emb)
+            pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs))
+            att_masks += [0] * num_img_embs
+
+        # ``embed_language_tokens`` already returns the ``* sqrt(hidden)``-scaled
+        # embedding (version-robust: self-scaled at transformers ≥ 5.4, manually
+        # scaled at ≤ 5.3). The image slice went through the projector and is
+        # NOT scaled. See embed_language_tokens for the version details.
+        lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens)
+
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks)
+        att_masks += [0] * lang_emb.shape[1]
+
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device)
+        att_masks = att_masks[None, :].expand(pad_masks.shape[0], -1)
+
+        return embs, pad_masks, att_masks
+
+    # ── Suffix embedding ─────────────────────────────────────────────
+
+    def embed_suffix(
+        self,
+        state: torch.Tensor,
+        noisy_actions: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build the suffix embeddings + masks: ``[state_token, action_tokens×H]``.
+
+        AR mask layout: ``[1, 1, 0, 0, ..., 0]`` — the state token is a causal
+        boundary (no later token attends backwards through it onto prefix
+        *by mistake*), the first action token starts a new causal block, and
+        the rest of the action tokens attend to each other bidirectionally.
+
+        Ref: openpi PI0Pytorch.embed_suffix
+        """
+        model_dtype = self.state_proj.weight.dtype
+        state = state.to(dtype=model_dtype)
+        noisy_actions = noisy_actions.to(dtype=model_dtype)
+        device = state.device
+
+        # State → (B, 1, W)
+        state_emb = self.state_proj(state)[:, None, :]
+        bsize = state_emb.shape[0]
+
+        # Sinusoidal timestep → (B, W) → expand across the horizon.
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=4e-3,
+            max_period=4.0,
+            device=device,
+        ).to(dtype=model_dtype)
+
+        # Fuse action and time via (2W → W → W) MLP with SiLU.
+        action_emb = self.action_in_proj(noisy_actions)  # (B, H, W)
+        action_time_emb = torch.cat([action_emb, time_emb[:, None, :].expand_as(action_emb)], dim=2)
+        action_time_emb = self.action_time_mlp_in(action_time_emb)
+        action_time_emb = F.silu(action_time_emb)
+        action_time_emb = self.action_time_mlp_out(action_time_emb)
+
+        embs = torch.cat([state_emb, action_time_emb], dim=1)
+        pad_masks = torch.ones(bsize, embs.shape[1], dtype=torch.bool, device=device)
+        att_masks = torch.tensor(
+            [1] + [1] + [0] * (self.action_horizon - 1),
+            dtype=embs.dtype,
+            device=device,
+        )[None, :].expand(bsize, -1)
+        return embs, pad_masks, att_masks
+
+    # ── Denoising step ───────────────────────────────────────────────
+
+    def denoise_step(
+        self,
+        state: torch.Tensor,
+        prefix_pad_masks: torch.Tensor,
+        past_key_values,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply one flow-matching denoising step: predict ``v_t`` from ``x_t``.
+
+        Uses the prefix KV cache from ``sample_actions`` and only runs the
+        action expert. Ref: openpi PI0Pytorch.denoise_step.
+        """
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep)
+
+        # Build the full (B, suffix_len, prefix_len + suffix_len) boolean mask:
+        #   * suffix queries can see every *valid* prefix key (padded camera
+        #     slots and padded language tokens are blocked).
+        #   * within the suffix, the state token is a causal boundary; action
+        #     tokens attend to each other bidirectionally.
+        batch_size = prefix_pad_masks.shape[0]
+        suffix_len = suffix_pad_masks.shape[1]
+        prefix_len = prefix_pad_masks.shape[1]
+
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+
+        # Position IDs continue from where the prefix's last valid token left off.
+        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
+        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
+
+        outputs_embeds, _ = self.paligemma_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, suffix_embs],
+            use_cache=False,
+        )
+
+        # Drop the state token, keep only the action tokens, project down.
+        suffix_out = outputs_embeds[1][:, -self.action_horizon :]
+        suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype)
+        return self.action_out_proj(suffix_out)
+
+    # ── Full action generation ───────────────────────────────────────
+
+    @torch.no_grad()
+    def sample_actions(
+        self,
+        images: list[torch.Tensor],
+        image_masks: list[torch.Tensor],
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+        state: torch.Tensor,
+        noise: torch.Tensor | None = None,
+        num_steps: int | None = None,
+    ) -> torch.Tensor:
+        """Generate an action chunk via iterative flow-matching denoising.
+
+        Convention: ``t=1`` is noise, ``t=0`` is the target — opposite of the
+        published π0 paper but matches both OpenPI and LeRobot.
+        Ref: openpi PI0Pytorch.sample_actions
+        """
+        if num_steps is None:
+            num_steps = self.num_inference_steps
+
+        bsize = state.shape[0]
+        device = state.device
+        if noise is None:
+            noise = torch.randn(
+                bsize,
+                self.action_horizon,
+                self.action_dim,
+                dtype=torch.float32,
+                device=device,
+            )
+
+        # 1. Prefix embeddings + mask building.
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, image_masks, lang_tokens, lang_masks
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
+        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
+
+        # 2. Forward prefix through PaliGemma LM, producing a list[(k, v)] cache.
+        _, past_key_values = self.paligemma_with_expert.forward(
+            attention_mask=prefix_att_2d_masks_4d,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None],
+            use_cache=True,
+        )
+
+        # 3. Euler-integrated denoising from t=1 down to t=0.
+        dt = -1.0 / num_steps
+        x_t = noise
+        for step in range(num_steps):
+            t = 1.0 + step * dt
+            time_tensor = torch.full((bsize,), t, dtype=torch.float32, device=device)
+            v_t = self.denoise_step(
+                state=state,
+                prefix_pad_masks=prefix_pad_masks,
+                past_key_values=past_key_values,
+                x_t=x_t,
+                timestep=time_tensor,
+            )
+            x_t = x_t + dt * v_t
+        return x_t
+
+    # ── Weight loading ───────────────────────────────────────────────
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        """Load weights from an ``lerobot/pi0_base`` safetensors checkpoint.
+
+        Checkpoint keys (after stripping the leading ``model.`` prefix) look
+        like ``paligemma_with_expert.paligemma.<sub>.*``. Modern transformers
+        nests those sub-modules one level deeper (``paligemma.model.<sub>``),
+        and PaliGemma ties ``lm_head.weight`` with ``embed_tokens.weight`` at
+        ``post_init`` time. Two remap rules make the load lossless:
+
+          1. ``paligemma.{vision_tower,multi_modal_projector,language_model}.*``
+             → ``paligemma.model.{...}.*``
+          2. ``paligemma.lm_head.weight``
+             → ``paligemma.model.language_model.embed_tokens.weight``
+             (the checkpoint does not store ``embed_tokens.weight`` on its
+             own — only the tied ``lm_head.weight`` copy — and
+             ``PaliGemmaForConditionalGeneration`` does not register
+             ``lm_head.weight`` as a Parameter when tied, so without this
+             rule the language embedding would silently remain at its
+             random init.)
+
+        On any remaining mismatch the loader logs a warning listing a sample
+        of the unmatched checkpoint keys and model params; downstream parity
+        tests surface these as a hard failure.
+        """
+        params_dict = dict(self.named_parameters())
+        buffers_dict = dict(self.named_buffers())
+
+        _PALIGEMMA_SUBMODULES = (
+            "vision_tower",
+            "multi_modal_projector",
+            "language_model",
+        )
+
+        def _remap(name: str) -> str:
+            # Strip the leading "model." that LeRobot's PI0Policy wrapper adds.
+            if name.startswith("model."):
+                name = name[len("model.") :]
+
+            # Historical MLP alias used by some early π0 checkpoints.
+            if name.startswith("time_mlp_in."):
+                name = "action_time_mlp_in." + name[len("time_mlp_in.") :]
+            elif name.startswith("time_mlp_out."):
+                name = "action_time_mlp_out." + name[len("time_mlp_out.") :]
+
+            # Nested PaliGemma layout.
+            for sub in _PALIGEMMA_SUBMODULES:
+                flat = f"paligemma_with_expert.paligemma.{sub}."
+                nested = f"paligemma_with_expert.paligemma.model.{sub}."
+                if name.startswith(flat) and not name.startswith(nested):
+                    return nested + name[len(flat) :]
+
+            # Tied lm_head → embed_tokens redirect (see docstring).
+            if name == "paligemma_with_expert.paligemma.lm_head.weight":
+                return "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+
+            return name
+
+        def _fix_vision_tower(name: str) -> str:
+            """Reconcile the SigLIP vision-tower nesting across transformers
+            versions. transformers ≤5.3 wraps the encoder in an extra
+            ``vision_tower.vision_model.*`` module; ≥5.4 flattens it to
+            ``vision_tower.*``. The ``lerobot/pi0_base`` checkpoint uses the
+            ``vision_model`` layout, so on newer transformers we must drop that
+            infix (else 437 vision params silently load nothing). Adaptive:
+            only rewrite when the rewritten key actually exists on the model.
+            """
+            vt = "paligemma_with_expert.paligemma.model.vision_tower."
+            if not name.startswith(vt):
+                return name
+            rest = name[len(vt) :]
+            if name in params_dict or name in buffers_dict:
+                return name  # model already expects this exact layout
+            if rest.startswith("vision_model."):
+                # Checkpoint has the infix; model (≥5.4) doesn't → strip it.
+                candidate = vt + rest[len("vision_model.") :]
+            else:
+                # Checkpoint lacks the infix; model (≤5.3) has it → insert it.
+                candidate = vt + "vision_model." + rest
+            return candidate if (candidate in params_dict or candidate in buffers_dict) else name
+
+        loaded = 0
+        skipped: list[str] = []
+        filled_params: set = set()
+        for name, loaded_weight in weights:
+            mapped = _fix_vision_tower(_remap(name))
+            if mapped in params_dict:
+                param = params_dict[mapped]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded += 1
+                filled_params.add(mapped)
+            elif mapped in buffers_dict:
+                buffers_dict[mapped].copy_(loaded_weight)
+                loaded += 1
+                filled_params.add(mapped)
+            else:
+                skipped.append(mapped)
+
+        # Reverse audit: any model param that got no checkpoint tensor at all
+        # would be running with random init. ``rotary_emb.inv_freq`` and
+        # friends are config-derived buffers, not stored in the checkpoint.
+        missing_params: list[str] = []
+        for pname in params_dict:
+            if pname in filled_params:
+                continue
+            if "rotary_emb" in pname or pname.endswith(".inv_freq"):
+                continue
+            missing_params.append(pname)
+
+        if missing_params or skipped:
+            parts: list[str] = []
+            if skipped:
+                parts.append(f"{len(skipped)} checkpoint key(s) had no home in the model (first 5: {skipped[:5]})")
+            if missing_params:
+                parts.append(f"{len(missing_params)} model param(s) received NO weight (first 5: {missing_params[:5]})")
+            logger.warning(
+                "π0 load_weights: %d tensors loaded — %s.",
+                loaded,
+                "; ".join(parts),
+            )
+        else:
+            logger.info(
+                "π0 load_weights: %d tensors loaded, 0 skipped, 0 missing.",
+                loaded,
+            )
+        return filled_params

vllm_omni/diffusion/models/pi0/pipeline_pi0.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 (Pi-Zero) VLA pipeline for vllm-omni.
+
+Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the
+DreamZero contract: the pipeline owns ALL preprocessing. It reads the raw robot
+observation from ``req.sampling_params.extra_args["robot_obs"]`` (delivered by
+the OpenPI realtime serving layer), builds model inputs, runs flow-matching
+denoising, and returns ``DiffusionOutput(output={"actions": ndarray})`` — which
+``diffusion_engine`` promotes to ``multimodal_output["actions"]`` for the client.
+
+π0 is stateless across calls (no KV reuse), so ``session_id`` / ``reset`` from
+the OpenPI protocol are accepted but ignored.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction
+from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+# Default tokenizer for the PaliGemma prefix (matches LeRobot Pi0).
+DEFAULT_PI0_TOKENIZER = "google/paligemma-3b-pt-224"
+
+
+def _pi0_post_process(x):
+    """Module-level identity post-process (picklable across the orchestrator's
+    multiprocess boundary — a local closure is not)."""
+    return x
+
+
+def get_pi0_post_process_func(od_config: OmniDiffusionConfig):
+    """π0 returns actions directly; post-processing is identity. The diffusion
+    engine resolves this via the registry (``_DIFFUSION_POST_PROCESS_FUNCS``) and
+    applies it engine-side, so the pipeline does NOT attach it to DiffusionOutput.
+    """
+    del od_config
+    return _pi0_post_process
+
+
+# Single-stage diffusion topology for online OpenPI serving.
+#
+# π0 is one diffusion stage (robot observation -> action chunk via flow
+# matching). It is registered in ``OMNI_PIPELINES`` (rather than relying on the
+# default single-stage fallback) because π0 is *online-served* via
+# ``vllm serve --deploy-config pi0.yaml``; that path
+# (``OmniStagePipelineFactory._create_from_registry``) requires a registered
+# pipeline so the deploy yaml's ``policy_server_config`` reaches the OpenPI
+# websocket layer. Offline-only models (e.g. ``internvla_a1``) need no topology.
+PI0_PIPELINE = PipelineConfig(
+    model_type="pi0",
+    model_arch="Pi0Pipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="action",
+            model_arch="Pi0Pipeline",
+        ),
+    ),
+)
+
+
+class Pi0Pipeline(nn.Module):
+    """π0 VLA pipeline: raw robot obs → continuous action chunk.
+
+    Registered as ``"Pi0Pipeline"`` in the diffusion registry. Weights are
+    self-loaded in ``__init__`` from the checkpoint's ``model.safetensors`` via
+    the kernel's ``load_weights`` (which handles the LeRobot key remaps).
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+        # Resolve od_config.model to a LOCAL directory. A bare HF repo id (e.g.
+        # the documented ``lerobot/pi0_base``) must be snapshot-downloaded first;
+        # otherwise config/tokenizer/weights would silently fall back to random
+        # init (this pipeline self-loads from a local dir and has no
+        # ``weights_sources`` for the framework loader to download from).
+        self.model_dir = self._resolve_model_dir(od_config.model)
+        self.config = self._build_config(od_config)
+
+        custom_args = od_config.custom_pipeline_args or {}
+        self.tokenizer_source = str(custom_args.get("tokenizer", self._resolve_tokenizer_source()))
+
+        # Torch dtype/device from od_config.
+        self._torch_dtype = self._resolve_dtype(od_config)
+        self._device = self._resolve_device(od_config)
+
+        self.tokenizer = self._load_tokenizer()
+        self.model = self._initialize_model()
+
+    # ------------------------------------------------------------------
+    # Construction helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _resolve_model_dir(model: str | None) -> str | None:
+        """Return a local directory for ``model``; download an HF repo id if needed.
+
+        Limits the snapshot to the files π0 actually reads (config, weights,
+        tokenizer) so a repo id doesn't pull unrelated artifacts.
+        """
+        if not model:
+            return None
+        if os.path.isdir(model):
+            return model
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(
+            repo_id=model,
+            allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"],
+        )
+
+    def _build_config(self, od_config: OmniDiffusionConfig) -> Pi0Config:
+        """Build Pi0Config from deploy-yaml model_config, falling back to the
+        checkpoint's config.json (raw LeRobot format)."""
+        if od_config.model_config:
+            config = Pi0Config.from_model_config(dict(od_config.model_config))
+            # The deploy yaml is authoritative, but must not silently drop
+            # checkpoint-derived fields it omits. Backfill the camera order and
+            # any normalization stats the yaml didn't specify.
+            if self.model_dir and (not config.image_feature_keys or config.norm_stats is None):
+                ckpt = Pi0Config.from_pretrained(self.model_dir)
+                if not config.image_feature_keys:
+                    config.image_feature_keys = ckpt.image_feature_keys
+                    if not config.input_features:
+                        config.input_features = ckpt.input_features
+                if config.norm_stats is None:
+                    config.norm_stats = ckpt.norm_stats
+            return config
+        if self.model_dir:
+            return Pi0Config.from_pretrained(self.model_dir)
+        return Pi0Config()
+
+    def _resolve_tokenizer_source(self) -> str:
+        """Prefer the checkpoint dir if it ships tokenizer files; else PaliGemma."""
+        if self.model_dir and os.path.isdir(self.model_dir):
+            if os.path.exists(os.path.join(self.model_dir, "tokenizer_config.json")):
+                return self.model_dir
+        return DEFAULT_PI0_TOKENIZER
+
+    @staticmethod
+    def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype:
+        dt = od_config.dtype
+        if isinstance(dt, torch.dtype):
+            return dt
+        return getattr(torch, str(dt).split(".")[-1], torch.float32)
+
+    @staticmethod
+    def _resolve_device(od_config: OmniDiffusionConfig) -> torch.device:
+        from vllm_omni.diffusion.distributed.utils import get_local_device
+
+        try:
+            return get_local_device()
+        except Exception:  # noqa: BLE001
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def _load_tokenizer(self):
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained(self.tokenizer_source)
+
+    def has_real_checkpoint(self) -> bool:
+        return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors"))
+
+    def _initialize_model(self) -> Pi0ForActionPrediction:
+        model = Pi0ForActionPrediction(self.config)
+        if self.has_real_checkpoint():
+            self._load_checkpoint(model)
+        else:
+            logger.info("Pi0Pipeline: no model.safetensors under %s; using random init.", self.model_dir)
+        model.to(device=self._device, dtype=self._torch_dtype)
+        model.eval()
+        return model
+
+    def _load_checkpoint(self, model: Pi0ForActionPrediction) -> None:
+        import safetensors.torch
+
+        path = os.path.join(self.model_dir, "model.safetensors")
+        logger.info("Pi0Pipeline: loading π0 weights from %s", path)
+        state = safetensors.torch.load_file(path)
+        model.load_weights(list(state.items()))
+
+    # ------------------------------------------------------------------
+    # Framework weight-loading hook
+    # ------------------------------------------------------------------
+    def load_weights(self, weights=()):  # noqa: D401
+        """No-op for the diffusion loader: π0 self-loads its checkpoint in
+        ``__init__`` (the kernel's ``load_weights`` handles the LeRobot remaps).
+        We expose no ``weights_sources``, so the loader passes an empty iterator
+        here; returning ``None`` skips its strict unloaded-weights check.
+        """
+        for _ in weights:
+            pass
+        return None
+
+    # ------------------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------------------
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
+        extra_args = getattr(req.sampling_params, "extra_args", None) or {}
+        robot_obs = extra_args.get("robot_obs")
+
+        if robot_obs is None:
+            # Dummy warmup path (no obs): return zeros so engine warmup/capture
+            # doesn't crash. Mirrors DreamZero's dummy-run handling.
+            first_prompt = req.prompts[0] if req.prompts else ""
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            num_steps = getattr(req.sampling_params, "num_inference_steps", None)
+            if prompt == "dummy run" or num_steps == 1:
+                logger.info("Pi0Pipeline: dummy warmup request without robot_obs — returning zeros.")
+                return DiffusionOutput(
+                    output={
+                        "actions": np.zeros(
+                            (self.config.chunk_size, self.config.max_action_dim),
+                            dtype=np.float32,
+                        )
+                    },
+                )
+            return DiffusionOutput(
+                error="Pi0Pipeline.forward requires sampling_params.extra_args['robot_obs'].",
+            )
+
+        images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
+            robot_obs, self.config, self.tokenizer, self._device
+        )
+
+        # State normalization (identity for pi0_base) inside the model.
+        state = self.model._normalize_state(state)
+
+        noise = extra_args.get("noise")
+        if noise is not None and not isinstance(noise, torch.Tensor):
+            noise = torch.as_tensor(noise, dtype=torch.float32, device=self._device)
+        elif isinstance(noise, torch.Tensor):
+            noise = noise.to(device=self._device, dtype=torch.float32)
+
+        num_steps = extra_args.get("num_inference_steps")
+
+        actions = self.model.sample_actions(
+            images=images,
+            image_masks=image_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            state=state,
+            noise=noise,
+            num_steps=num_steps,
+        )
+        actions = self.model._unnormalize_actions(actions)
+
+        # (B=1, horizon, action_dim) → (horizon, action_dim) numpy for the wire.
+        actions_np = actions.squeeze(0).float().cpu().numpy()
+
+        # Note: post-processing is applied engine-side via the registry
+        # (_DIFFUSION_POST_PROCESS_FUNCS), so we don't attach it here (a local
+        # closure wouldn't survive the orchestrator's multiprocess pickling).
+        return DiffusionOutput(output={"actions": actions_np})

vllm_omni/diffusion/models/pi0/processor_pi0.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Preprocessing for the π0 VLA model.
+
+Converts a raw robot observation (multi-camera images + language instruction +
+proprioceptive state) into the tensors that ``Pi0ForActionPrediction.sample_actions``
+consumes:
+
+  * images       : list of ``(1, 3, 224, 224)`` float tensors in ``[-1, 1]``,
+                   one per camera slot (real cameras first, then ``-1``-filled
+                   padding) in the checkpoint's ``image_feature_keys`` order
+  * image_masks  : list of ``(1,)`` bool tensors, ``True`` = real camera
+  * lang_tokens  : ``(1, L)`` long, PaliGemma tokenizer, ``\\n``-terminated
+  * lang_masks   : ``(1, L)`` bool
+  * state        : ``(1, max_state_dim)`` float32, zero-padded
+
+Image/prompt preprocessing matches OpenPI and LeRobot bit-for-bit — see the
+parity test (``test_pi0_parity.py``). A set of stateless helpers the diffusion
+pipeline calls directly (the DreamZero contract: the pipeline owns its
+preprocessing).
+
+Reference:
+  - OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad)
+  - OpenPI: openpi/src/openpi/models_pytorch/preprocessing_pytorch.py
+  - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py (Pi0NewLineProcessor)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+# Defaults straight from the π0 reference configs.
+PI0_IMAGE_SIZE = 224  # openpi/models/model.py IMAGE_RESOLUTION
+PI0_NUM_IMAGE_TOKENS = 256  # SigLIP So400m/14 on 224×224 → (224/14)**2
+PI0_IMAGE_TOKEN_INDEX = 257152  # openpi/models_pytorch/gemma_pytorch.py
+PI0_MAX_CAMERAS = 3  # openpi/models/model.py (3 camera slots)
+PI0_MAX_TOKEN_LEN = 48  # openpi/models/pi0_config.py
+
+
+def resize_with_pad(
+    images: torch.Tensor,
+    target_height: int,
+    target_width: int,
+    mode: str = "bilinear",
+) -> torch.Tensor:
+    """Resize ``(B, C, H, W)`` images to the target shape, preserving aspect
+    ratio with -1 padding on the short side.
+
+    Matches openpi ``image_tools.resize_with_pad_torch`` — the clamp to
+    [-1, 1] is what lets the padded region blend with SigLIP-normalized
+    pixels without adding signal at the boundary.
+    """
+    if images.ndim != 4:
+        raise ValueError(f"Expected 4-D (B,C,H,W), got {images.ndim}-D")
+    _, _, cur_h, cur_w = images.shape
+    ratio = max(cur_w / target_width, cur_h / target_height)
+    rh, rw = int(cur_h / ratio), int(cur_w / ratio)
+    align_corners = False if mode == "bilinear" else None
+    resized = F.interpolate(images, size=(rh, rw), mode=mode, align_corners=align_corners)
+    resized = resized.clamp(-1.0, 1.0)
+    ph, rem_h = divmod(target_height - rh, 2)
+    pw, rem_w = divmod(target_width - rw, 2)
+    return F.pad(resized, (pw, pw + rem_w, ph, ph + rem_h), value=-1.0)
+
+
+def pil_image_to_tensor(image: Image.Image) -> torch.Tensor:
+    """PIL → ``(1, C, H, W)`` float32 in ``[-1, 1]`` (SigLIP normalization)."""
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    arr = np.array(image, dtype=np.float32) / 255.0 * 2.0 - 1.0
+    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+
+
+class Pi0ImageProcessor:
+    """Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor."""
+
+    def __init__(self, image_size: int = PI0_IMAGE_SIZE):
+        self.image_size = image_size
+
+    def preprocess_single(self, image: Any) -> torch.Tensor:
+        """Accept a PIL image, an HWC uint8/float ndarray, or a CHW tensor and
+        return a ``(1, 3, image_size, image_size)`` float tensor in ``[-1, 1]``."""
+        t = self._to_tensor(image)
+        if t.shape[2] != self.image_size or t.shape[3] != self.image_size:
+            t = resize_with_pad(t, self.image_size, self.image_size)
+        return t
+
+    def _to_tensor(self, image: Any) -> torch.Tensor:
+        if isinstance(image, Image.Image):
+            return pil_image_to_tensor(image)
+        if isinstance(image, np.ndarray):
+            arr = image
+            # HWC → normalize uint8/[0,255] to [-1,1]; assume already [-1,1] if float.
+            if arr.ndim == 3 and arr.shape[-1] in (1, 3):
+                if np.issubdtype(arr.dtype, np.integer) or arr.max() > 1.0:
+                    arr = arr.astype(np.float32) / 255.0 * 2.0 - 1.0
+                else:
+                    arr = arr.astype(np.float32)
+                return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+            return torch.as_tensor(arr, dtype=torch.float32)
+        if isinstance(image, torch.Tensor):
+            t = image
+            if t.ndim == 3:  # (C, H, W) → (1, C, H, W)
+                t = t.unsqueeze(0)
+            return t.to(dtype=torch.float32)
+        raise TypeError(f"Unsupported image type for π0 preprocessing: {type(image)}")
+
+    def make_empty_image(self) -> torch.Tensor:
+        """Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot."""
+        return torch.full((1, 3, self.image_size, self.image_size), -1.0)
+
+
+def tokenize_prompt(tokenizer, text: str, max_token_len: int = PI0_MAX_TOKEN_LEN):
+    """Return ``(input_ids, attention_mask)`` lists, length ``max_token_len``.
+
+    Matches LeRobot's ``Pi0NewLineProcessor``: append ``\\n`` if the caller
+    didn't, then run the PaliGemma tokenizer with right-padding.
+    """
+    prompt = text or ""
+    if not prompt.endswith("\n"):
+        prompt = prompt + "\n"
+    enc = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_token_len,
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors=None,
+    )
+    return list(enc["input_ids"]), list(enc["attention_mask"])
+
+
+def _extract_images(robot_obs: dict, config) -> dict[str, Any]:
+    """Pull a ``{feature_key: image}`` map out of a raw robot obs.
+
+    Accepts two shapes:
+      * ``robot_obs["images"]`` is a dict keyed by camera name; or
+      * flat ``observation/<cam>`` style keys on ``robot_obs`` itself.
+    Keys are translated through ``config.image_key_map`` (raw obs key →
+    checkpoint feature key) so serving wire names map onto ``input_features``.
+    """
+    images = robot_obs.get("images")
+    if not isinstance(images, dict):
+        # Fall back to flat keys (anything that isn't a known scalar field).
+        images = {k: v for k, v in robot_obs.items() if k not in ("state", "prompt", "session_id", "reset", "images")}
+    key_map = getattr(config, "image_key_map", None) or {}
+    return {key_map.get(k, k): v for k, v in images.items()}
+
+
+def build_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device):
+    """Convert a raw robot observation into ``sample_actions`` inputs.
+
+    Reproduces the LeRobot camera ordering exactly: iterate
+    ``config.image_feature_keys`` (the ordered camera identities from the
+    checkpoint's ``input_features``); for each, use the supplied image (mask
+    ``True``) or a ``-1``-filled empty image (mask ``False``). Then tokenize
+    the prompt and zero-pad/truncate the state to ``max_state_dim``.
+
+    Returns ``(images, image_masks, lang_tokens, lang_masks, state)`` with a
+    leading batch dim of 1 on the tensors.
+    """
+    image_size = int(config.image_resolution[0])
+    img_proc = Pi0ImageProcessor(image_size=image_size)
+
+    feature_keys = config.image_feature_keys or []
+    if not feature_keys:
+        # No declared camera order — fall back to whatever the obs provides,
+        # capped at max_cameras (preserves insertion order).
+        provided = _extract_images(robot_obs, config)
+        feature_keys = list(provided.keys())[: config.max_cameras]
+
+    obs_images = _extract_images(robot_obs, config)
+
+    images: list[torch.Tensor] = []
+    image_masks: list[torch.Tensor] = []
+    for key in feature_keys[: config.max_cameras]:
+        img = obs_images.get(key)
+        if img is not None:
+            tensor = img_proc.preprocess_single(img).to(device=device)
+            mask = True
+        else:
+            tensor = img_proc.make_empty_image().to(device=device)
+            mask = False
+        images.append(tensor)
+        image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device))
+
+    # Pad up to max_cameras with empty slots if the checkpoint declares fewer
+    # cameras than the model attends to.
+    while len(images) < config.max_cameras:
+        images.append(img_proc.make_empty_image().to(device=device))
+        image_masks.append(torch.tensor([False], dtype=torch.bool, device=device))
+
+    prompt = robot_obs.get("prompt", "") or ""
+    ids, attn = tokenize_prompt(tokenizer, prompt, config.tokenizer_max_length)
+    lang_tokens = torch.tensor([ids], dtype=torch.long, device=device)
+    lang_masks = torch.tensor([attn], dtype=torch.bool, device=device)
+
+    state = _build_state(robot_obs.get("state"), config.max_state_dim, device)
+
+    return images, image_masks, lang_tokens, lang_masks, state
+
+
+def _build_state(raw_state, max_state_dim: int, device: torch.device) -> torch.Tensor:
+    """Zero-pad / truncate the raw state to ``(1, max_state_dim)`` float32."""
+    if raw_state is None:
+        return torch.zeros(1, max_state_dim, dtype=torch.float32, device=device)
+    state = torch.as_tensor(raw_state, dtype=torch.float32, device=device)
+    if state.dim() == 1:
+        state = state[None, :]
+    if state.shape[-1] < max_state_dim:
+        state = F.pad(state, (0, max_state_dim - state.shape[-1]))
+    elif state.shape[-1] > max_state_dim:
+        state = state[..., :max_state_dim]
+    return state

vllm_omni/diffusion/registry.py

@@ -181,6 +181,11 @@ _DIFFUSION_MODELS = {
         "pipeline_gr00t",
         "Gr00tN1d7Pipeline",
     ),
+    "Pi0Pipeline": (
+        "pi0",
+        "pipeline_pi0",
+        "Pi0Pipeline",
+    ),
     "LongCatImageEditPipeline": (
         "longcat_image",
         "pipeline_longcat_image_edit",
@@ -358,6 +363,11 @@ _NO_CACHE_ACCELERATION = {
     # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
     "NextStep11Pipeline",
     "AudioXPipeline",
+    # π0 is a flow-matching VLA with a self-contained sample_actions loop and no
+    # DiT-style ``.transformer`` block list, so cache_dit / tea_cache cannot apply
+    # to it; list it here so a stray cache_backend override disables gracefully
+    # instead of erroring.
+    "Pi0Pipeline",
     "LingBotWorldCausalDMDPipeline",
 }
 
@@ -560,6 +570,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LancePipeline": "get_lance_post_process_func",
     "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
+    "Pi0Pipeline": "get_pi0_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "FluxKontextPipeline": "get_flux_kontext_post_process_func",

vllm_omni/entrypoints/openai/serving_speech.py

@@ -80,6 +80,12 @@ from vllm_omni.utils.speaker_cache import (
 
 logger = init_logger(__name__)
 
+_SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS"
+_SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS"
+_SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
+_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
+_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
+
 # TTS Configuration
 #
 # The stage-key -> model-type mapping is NOT declared here: it is derived from
@@ -950,6 +956,24 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         details = self._compute_speech_input_details(request, tts_params)
         return build_speech_usage(details, output_tokens)
 
+    @staticmethod
+    def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]:
+        """Map speech usage into non-streaming response headers.
+
+        Returns an empty dict when usage is unavailable, allowing callers to
+        merge these headers with other optional response headers.
+        """
+        if usage is None:
+            return {}
+
+        return {
+            _SPEECH_USAGE_INPUT_TOKENS_HEADER: str(usage.input_tokens),
+            _SPEECH_USAGE_OUTPUT_TOKENS_HEADER: str(usage.output_tokens),
+            _SPEECH_USAGE_TOTAL_TOKENS_HEADER: str(usage.total_tokens),
+            _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER: str(usage.input_token_details.text_tokens),
+            _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER: str(usage.input_token_details.audio_tokens),
+        }
+
     def _estimate_fish_ref_code_len(self, ref_audio: object) -> int | None:
         """Estimate Fish Speech semantic token length from raw reference audio."""
         from vllm_omni.model_executor.models.fish_speech.dac_utils import (
@@ -3170,7 +3194,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # ``usage_out`` is an opt-in output channel: when a list is passed, the
         # computed SpeechTokenUsage is appended to it. The return stays a
         # 2-tuple so existing callers (and their test mocks) are unaffected;
-        # only the batch path, which surfaces per-item usage, opts in.
+        # batch and non-streaming response-header paths opt in when surfacing
+        # usage outside the raw audio body.
         request_id, generator, bytes_tts_params = await self._prepare_speech_generation(
             request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio
         )
@@ -3576,7 +3601,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 )
 
             collect: dict = {}
-            audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect)
+            usage_box: list[SpeechTokenUsage] = []
+            audio_bytes, media_type = await self._generate_audio_bytes(
+                request, request_id=request_id, usage_out=usage_box, collect=collect
+            )
             total_ms = (time.perf_counter() - request_start_s) * 1000.0
             logger.info(
                 "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
@@ -3584,7 +3612,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 total_ms,
                 len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
             )
-            headers = {}
+            headers = self._build_speech_usage_headers(usage_box[0] if usage_box else None)
             if collect.get("word_timestamps") is not None:
                 # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped).
                 ts_json = json.dumps(collect["word_timestamps"])

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -17,6 +17,7 @@ import httpx
 import numpy as np
 import torch
 from PIL import Image, UnidentifiedImageError
+from vllm import envs
 from vllm.multimodal.video import (
     VIDEO_LOADER_REGISTRY,
     VideoBackend,
@@ -172,11 +173,21 @@ async def decode_image_url(image_url: str) -> Image.Image:
         return _decode_base64_image(image_url, source="image_reference.image_url")
 
     if image_url.startswith(("http://", "https://")):
-        async with httpx.AsyncClient(timeout=60) as client:
+        allow_redirects = envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS
+        async with httpx.AsyncClient(timeout=60, follow_redirects=allow_redirects) as client:
             try:
                 response = await client.get(image_url)
                 response.raise_for_status()
-            except httpx.HTTPError as exc:
+            except httpx.HTTPStatusError as exc:
+                if exc.response.has_redirect_location and not allow_redirects:
+                    raise InvalidInputReferenceError(
+                        "Invalid image_reference.image_url: redirect response was rejected because "
+                        "VLLM_MEDIA_URL_ALLOW_REDIRECTS is disabled."
+                    ) from exc
+                raise InvalidInputReferenceError(
+                    f"Invalid image_reference.image_url: server returned HTTP {exc.response.status_code}."
+                ) from exc
+            except httpx.RequestError as exc:
                 raise InvalidInputReferenceError(
                     "Invalid image_reference.image_url: failed to download image."
                 ) from exc

vllm_omni/platforms/npu/__init__.py

@@ -1,6 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
 
 __all__ = ["NPUOmniPlatform"]
+
+
+def __getattr__(name: str):
+    if name != "NPUOmniPlatform":
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    globals()[name] = NPUOmniPlatform
+    return NPUOmniPlatform

vllm_omni/platforms/npu/graph_tools.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reusable exact-signature NPUGraph capture and replay helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class _ReplayGraph(Protocol):
+    def replay(self) -> None: ...
+
+
+@dataclass
+class CapturedDeviceGraph:
+    graph: _ReplayGraph
+    static_inputs: tuple[torch.Tensor, ...]
+    static_outputs: tuple[torch.Tensor, ...]
+
+    def replay(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+        with torch.inference_mode():
+            for static, current in zip(self.static_inputs, inputs, strict=True):
+                static.copy_(current)
+            self.graph.replay()
+            # Graph outputs are persistent and overwritten by the next replay.
+            # Clone before they become request-owned streaming cache entries.
+            return tuple(output.detach().clone() for output in self.static_outputs)
+
+
+def _tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]:
+    return tuple(value.shape), str(value.dtype), str(value.device)
+
+
+class NPUExactGraphRunner:
+    """Capture and replay tensor-only functions for exact NPU signatures."""
+
+    def __init__(
+        self,
+        *,
+        max_graphs: int = 32,
+        component_name: str = "device graph",
+        disable_config_hint: str = "disable graph capture",
+    ) -> None:
+        self.max_graphs = max(0, int(max_graphs))
+        self.component_name = component_name
+        self.disable_config_hint = disable_config_hint
+        self._enabled = self.max_graphs > 0
+        self._graphs: dict[tuple[object, ...], CapturedDeviceGraph] = {}
+        self._failed_keys: set[tuple[object, ...]] = set()
+        self._hits = 0
+        self._graph_pool: object | None = None
+
+    @staticmethod
+    def is_supported() -> bool:
+        npu = getattr(torch, "npu", None)
+        return npu is not None and all(
+            hasattr(npu, name)
+            for name in (
+                "NPUGraph",
+                "graph",
+                "is_current_stream_capturing",
+                "synchronize",
+            )
+        )
+
+    @staticmethod
+    def _stream_is_capturing() -> bool:
+        npu = getattr(torch, "npu", None)
+        is_capturing = getattr(npu, "is_current_stream_capturing", None)
+        if not callable(is_capturing):
+            return False
+        try:
+            return bool(is_capturing())
+        except (RuntimeError, TypeError):
+            return False
+
+    def _eligible(self, inputs: tuple[torch.Tensor, ...]) -> bool:
+        return (
+            self._enabled
+            and bool(inputs)
+            and all(value.device.type == "npu" for value in inputs)
+            and self.is_supported()
+            and not self._stream_is_capturing()
+        )
+
+    @property
+    def stats(self) -> dict[str, int]:
+        return {
+            "captures": len(self._graphs),
+            "failed": len(self._failed_keys),
+            "hits": self._hits,
+        }
+
+    def capture(
+        self,
+        inputs: tuple[torch.Tensor, ...],
+        compute: Callable[..., tuple[torch.Tensor, ...]],
+    ) -> CapturedDeviceGraph:
+        npu = torch.npu
+        static_inputs = tuple(value.detach().clone() for value in inputs)
+        npu.synchronize()
+        graph = npu.NPUGraph()
+        if self._graph_pool is None:
+            from vllm.platforms import current_platform
+
+            self._graph_pool = current_platform.get_global_graph_pool()
+        with torch.inference_mode(), npu.graph(graph, pool=self._graph_pool):
+            static_outputs = compute(*static_inputs)
+        npu.synchronize()
+        return CapturedDeviceGraph(
+            graph=graph,
+            static_inputs=static_inputs,
+            static_outputs=static_outputs,
+        )
+
+    def run(
+        self,
+        operation: str,
+        inputs: tuple[torch.Tensor, ...],
+        constants: tuple[object, ...],
+        compute: Callable[..., tuple[torch.Tensor, ...]],
+    ) -> tuple[torch.Tensor, ...]:
+        if self._failed_keys:
+            raise RuntimeError(
+                f"{self.component_name} cannot continue after a failed NPUGraph capture; "
+                f"restart the stage process and {self.disable_config_hint} before retrying."
+            )
+        if not self._eligible(inputs):
+            return compute(*inputs)
+
+        key = (
+            operation,
+            constants,
+            tuple(_tensor_signature(value) for value in inputs),
+        )
+        graph = self._graphs.get(key)
+        if graph is not None:
+            self._hits += 1
+            if self._hits == 1:
+                logger.info("%s started NPUGraph replay", self.component_name)
+            return graph.replay(inputs)
+
+        # Prime lazy kernels and allocator state before capture. The next call
+        # with the same exact tensor signature replays this graph.
+        eager_outputs = compute(*inputs)
+        if len(self._graphs) >= self.max_graphs:
+            logger.warning_once(
+                "%s reached the %d-entry NPUGraph limit; new tensor shapes will use eager execution.",
+                self.component_name,
+                self.max_graphs,
+            )
+            return eager_outputs
+        try:
+            self._graphs[key] = self.capture(inputs, compute)
+        except Exception as exc:
+            self._failed_keys.add(key)
+            self._enabled = False
+            logger.exception(
+                "%s failed to capture NPUGraph for %s; the torch-npu "
+                "allocator/RNG capture state may be invalid and this stage "
+                "process must be restarted.",
+                self.component_name,
+                operation,
+            )
+            raise RuntimeError(
+                f"{self.component_name} NPUGraph capture failed for {operation}; "
+                f"restart the stage process. To run eagerly, {self.disable_config_hint}."
+            ) from exc
+        logger.info(
+            "%s captured NPUGraph %d/%d for %s",
+            self.component_name,
+            len(self._graphs),
+            self.max_graphs,
+            operation,
+        )
+        return eager_outputs

vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py

@@ -18,7 +18,7 @@ This module:
 from __future__ import annotations
 
 from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import ExitStack, contextmanager
 
 import torch
 import torch.nn.functional as F
@@ -107,17 +107,20 @@ def _patched_attention_forward_chunk(
 
 
 @contextmanager
-def npu_math_sdpa_context() -> Iterator[None]:
+def npu_math_sdpa_context(*, require_available: bool = False) -> Iterator[None]:
     """Force SDPA MATH backend so Ascend does not call fused FA."""
-    try:
-        from torch.nn.attention import SDPBackend, sdpa_kernel
+    with ExitStack() as stack:
+        try:
+            from torch.nn.attention import SDPBackend, sdpa_kernel
 
-        with sdpa_kernel(SDPBackend.MATH):
-            yield
-    except Exception:
-        # Older torch / missing backend enum — just run as-is.
-        with nullcontext():
-            yield
+            stack.enter_context(sdpa_kernel(SDPBackend.MATH))
+        except Exception as exc:
+            if require_available:
+                raise RuntimeError(
+                    "MiniCPM-o Code2Wav NPU graph capture requires torch.nn.attention.sdpa_kernel(SDPBackend.MATH)"
+                ) from exc
+            logger.warning_once("MATH SDPA selection is unavailable; using the default Ascend SDPA backend.")
+        yield
 
 
 def _disable_upsample_encoder_compile() -> None:

vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inject MiniCPM-o Code2Wav NPUGraph acceleration on Ascend."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from contextlib import nullcontext
+from weakref import WeakKeyDictionary
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+_original_build_backend = None
+_original_estimator_step = None
+_original_setup_batch = None
+_original_decode_batch = None
+_backend_graph_runners: WeakKeyDictionary[object, NPUExactGraphRunner] = WeakKeyDictionary()
+_ENABLE_KEY = "code2wav_enable_npu_graph"
+_MAX_GRAPHS_KEY = "code2wav_max_npu_graphs"
+
+
+def _config_bool(value: object, default: bool) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
+def _graph_config(model: object) -> dict[str, object]:
+    config = getattr(getattr(model, "vllm_config", None), "additional_config", None)
+    return dict(config) if isinstance(config, Mapping) else {}
+
+
+def prepare_code2wav_graph_runtime() -> None:
+    """Select graph-capturable ACLNN kernels before Token2Wav is loaded."""
+    if os.environ.get("ASCEND_LAUNCH_BLOCKING") == "1":
+        raise RuntimeError(
+            "MiniCPM-o Code2Wav NPUGraph capture is incompatible with "
+            "ASCEND_LAUNCH_BLOCKING=1; unset it or set it to 0 before startup."
+        )
+    npu = torch.npu
+    npu.config.allow_internal_format = False
+    npu.set_compile_mode(jit_compile=False)
+    logger.info("Configured MiniCPM-o Code2Wav NPUGraph runtime (allow_internal_format=False, jit_compile=False)")
+
+
+def _flow_execution_context(device: torch.device, *, require_math: bool):
+    if device.type != "npu":
+        return nullcontext()
+    from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+        npu_token2wav_sdpa_context,
+    )
+
+    return npu_token2wav_sdpa_context(require_math=require_math)
+
+
+def _graphable_estimator_step(
+    backend,
+    estimator,
+    *,
+    x,
+    mu,
+    time_embedding,
+    speakers,
+    cond,
+    cnn_cache,
+    att_cache,
+):
+    """Run the CFM estimator body after host-backed timestep embedding."""
+    width = int(x.shape[-1])
+    speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
+    estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
+    cnn_out, att_out = backend._estimator_buffers(estimator, estimator_input, att_cache)
+    old_cnn = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
+    old_att = att_cache if att_cache is not None else [None] * len(estimator.blocks)
+    result = estimator.blocks_forward_chunk(
+        estimator_input,
+        time_embedding,
+        None,
+        old_cnn,
+        old_att,
+        cnn_out,
+        att_out,
+    )
+    return result, cnn_out, att_out
+
+
+def _patched_estimator_step(
+    self,
+    estimator,
+    *,
+    x,
+    mu,
+    time,
+    speakers,
+    cond,
+    cnn_cache,
+    att_cache,
+):
+    assert _original_estimator_step is not None
+    graph_runner = _backend_graph_runners.get(self)
+    if graph_runner is None or self._trt_stepper is not None or self._cfm_graph_wrapper is not None:
+        return _original_estimator_step(
+            self,
+            estimator,
+            x=x,
+            mu=mu,
+            time=time,
+            speakers=speakers,
+            cond=cond,
+            cnn_cache=cnn_cache,
+            att_cache=att_cache,
+        )
+    if (cnn_cache is None) != (att_cache is None):
+        raise ValueError("estimator CNN and attention caches must both be present or absent")
+
+    # The upstream embedder creates a frequency tensor on the host. Keep it
+    # outside capture while retaining the tensor-only estimator body in graph.
+    time_embedding = estimator.t_embedder(time).unsqueeze(1)
+    if cnn_cache is None:
+        return graph_runner.run(
+            "cfm_estimator",
+            (x, mu, time_embedding, speakers, cond),
+            (False,),
+            lambda step_x, step_mu, step_time, step_speakers, step_cond: _graphable_estimator_step(
+                self,
+                estimator,
+                x=step_x,
+                mu=step_mu,
+                time_embedding=step_time,
+                speakers=step_speakers,
+                cond=step_cond,
+                cnn_cache=None,
+                att_cache=None,
+            ),
+        )
+
+    return graph_runner.run(
+        "cfm_estimator",
+        (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache),
+        (True,),
+        lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att: _graphable_estimator_step(
+            self,
+            estimator,
+            x=step_x,
+            mu=step_mu,
+            time_embedding=step_time,
+            speakers=step_speakers,
+            cond=step_cond,
+            cnn_cache=step_cnn,
+            att_cache=step_att,
+        ),
+    )
+
+
+def _patched_setup_batch(self, features, batch_size):
+    assert _original_setup_batch is not None
+    with _flow_execution_context(
+        features.speech_tokens.device,
+        require_math=self in _backend_graph_runners,
+    ):
+        return _original_setup_batch(self, features, batch_size)
+
+
+def _patched_decode_batch(
+    self,
+    tokens,
+    features,
+    states,
+    *,
+    last_chunk,
+    flush_encoder=False,
+):
+    assert _original_decode_batch is not None
+    with _flow_execution_context(
+        tokens.device,
+        require_math=self in _backend_graph_runners,
+    ):
+        return _original_decode_batch(
+            self,
+            tokens,
+            features,
+            states,
+            last_chunk=last_chunk,
+            flush_encoder=flush_encoder,
+        )
+
+
+def _patched_build_backend(self) -> None:
+    if self.backend is not None:
+        return
+
+    config = _graph_config(self)
+    max_graphs = max(0, int(config.get(_MAX_GRAPHS_KEY, 32)))
+    graph_enabled = max_graphs > 0 and _config_bool(config.get(_ENABLE_KEY), False)
+    if graph_enabled:
+        # NPUOmniPlatform enables internal format for quantized LLM kernels.
+        # Code2Wav uses regular convolution kernels that must remain in the
+        # graph-capturable ACLNN path.
+        prepare_code2wav_graph_runtime()
+
+    assert _original_build_backend is not None
+    _original_build_backend(self)
+
+    graph_runner = None
+    if graph_enabled:
+        graph_runner = NPUExactGraphRunner(
+            max_graphs=max_graphs,
+            component_name="MiniCPM-o Code2Wav",
+            disable_config_hint=(
+                "set platforms.npu.stages[stage_id=2].additional_config.code2wav_enable_npu_graph=false"
+            ),
+        )
+        if self.backend.speech_window.device.type == "npu" and not graph_runner.is_supported():
+            raise RuntimeError(
+                "MiniCPM-o Code2Wav NPUGraph capture requires torch.npu "
+                "NPUGraph, graph, is_current_stream_capturing, and synchronize APIs."
+            )
+        if self.backend.flow.training:
+            raise ValueError("MiniCPM-o Code2Wav NPUGraph capture requires flow.eval()")
+        _backend_graph_runners[self.backend] = graph_runner
+
+    if graph_enabled:
+        logger.info(
+            "MiniCPM-o Code2Wav NPUGraph replay enabled (max_graphs=%d)",
+            max_graphs,
+        )
+
+
+def apply_minicpmo_4_5_code2wav_patch() -> None:
+    """Patch the generic Code2Wav backend builder with Ascend acceleration."""
+    global _PATCHED, _original_build_backend
+    global _original_decode_batch, _original_estimator_step, _original_setup_batch
+    if _PATCHED:
+        return
+
+    from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
+        BatchedToken2Wav,
+    )
+    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
+        MiniCPMO45Code2Wav,
+    )
+
+    _original_build_backend = MiniCPMO45Code2Wav._build_backend
+    _original_estimator_step = BatchedToken2Wav._estimator_step
+    _original_setup_batch = BatchedToken2Wav.setup_batch
+    _original_decode_batch = BatchedToken2Wav.decode_batch
+
+    MiniCPMO45Code2Wav._build_backend = _patched_build_backend  # type: ignore[method-assign]
+    BatchedToken2Wav._estimator_step = _patched_estimator_step  # type: ignore[method-assign]
+    BatchedToken2Wav.setup_batch = _patched_setup_batch  # type: ignore[method-assign]
+    BatchedToken2Wav.decode_batch = _patched_decode_batch  # type: ignore[method-assign]
+    _PATCHED = True
+    logger.debug("Applied NPU patch for MiniCPM-o 4.5 Code2Wav")

vllm_omni/platforms/npu/models/step_audio2_token2wav.py

@@ -13,7 +13,7 @@ Ascend-specific workarounds that must not live in the shared GPU model file:
 from __future__ import annotations
 
 from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from types import MethodType
 
 import numpy as np
@@ -113,20 +113,16 @@ def patch_step_audio2_hift_for_npu(hift: torch.nn.Module) -> None:
 
 
 @contextmanager
-def npu_token2wav_sdpa_context() -> Iterator[None]:
+def npu_token2wav_sdpa_context(*, require_math: bool = False) -> Iterator[None]:
     """Expand CosyVoice masks + force MATH SDPA to avoid FA 161001."""
-    try:
-        from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import (
-            apply_cosyvoice2_dit_attn_npu_patch,
-            npu_math_sdpa_context,
-        )
+    from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import (
+        apply_cosyvoice2_dit_attn_npu_patch,
+        npu_math_sdpa_context,
+    )
 
-        apply_cosyvoice2_dit_attn_npu_patch()
-        with npu_math_sdpa_context():
-            yield
-    except Exception:
-        with nullcontext():
-            yield
+    apply_cosyvoice2_dit_attn_npu_patch()
+    with npu_math_sdpa_context(require_available=require_math):
+        yield
 
 
 def _patched_ensure_models_loaded(self) -> None:

vllm_omni/platforms/npu/platform.py

@@ -36,6 +36,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         from vllm_ascend.utils import adapt_patch
 
         from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches
+        from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
+            apply_minicpmo_4_5_code2wav_patch,
+        )
         from vllm_omni.platforms.npu.models.minimax_h3 import apply_minimax_h3_qwen3vl_patch
         from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
             apply_qwen3_tts_code2wav_patch,
@@ -46,6 +49,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         adapt_patch(is_global_patch=True)
         apply_minimax_h3_qwen3vl_patch()
+        apply_minicpmo_4_5_code2wav_patch()
         apply_qwen3_tts_code2wav_patch()
         apply_qwen3_tts_tokenizer_v2_patch()
         apply_310p_patches()