Merge remote-tracking branch 'origin/main' into codex-api-server-user-guide

Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>

# Conflicts:
#	docs/user_guide/examples/online_serving/text_to_audio.md
#	examples/online_serving/stable_audio/README.md

.buildkite/amd/test-amd-merge.yml

@@ -72,7 +72,22 @@ steps:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
         - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-        - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+        - >-
+          pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+          --deselect='tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]'
+
+    # TODO: Diagnose unstable Stable Audio layerwise-offload memory savings on ROCm.
+    - label: "Diffusion Stable Audio Layerwise Offload Test (quarantined)"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: NonBlocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - >-
+          pytest -s -v
+          'tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]'
 
     - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2"
       agent_pool: mi300_4

.buildkite/cuda/test-merge.yml

@@ -85,6 +85,9 @@ steps:
           - vllm_omni/model_executor/models/common/alias_free_activation.py
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
+          - vllm_omni/entrypoints/openai/serving_speech.py
+          - vllm_omni/entrypoints/openai/tts_adapters/
+          - vllm_omni/engine/stage_pool.py
         commands:
           - |
             timeout 40m bash -c "
@@ -104,6 +107,9 @@ steps:
           - vllm_omni/model_executor/models/common/alias_free_activation.py
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
+          - vllm_omni/entrypoints/openai/serving_speech.py
+          - vllm_omni/entrypoints/openai/tts_adapters/
+          - vllm_omni/engine/stage_pool.py
           - .buildkite/common/scripts/run_cov_split.sh
           - pyproject.toml
         # 40 -> 60: the coverage split runs the model twice.

.buildkite/cuda/test-ready.yml

@@ -252,6 +252,9 @@ steps:
           - vllm_omni/model_executor/models/common/alias_free_activation.py
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
+          - vllm_omni/entrypoints/openai/serving_speech.py
+          - vllm_omni/entrypoints/openai/tts_adapters/
+          - vllm_omni/engine/stage_pool.py
         commands:
           - |
             timeout 20m bash -c "
@@ -260,25 +263,6 @@ steps:
             "
         mirror_hardwares: l4_1
 
-      - label: "TTS · Qwen3-TTS Base Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_qwen3_tts_base.py
-          - vllm_omni/model_executor/models/qwen3_tts/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/common/alias_free_activation.py
-          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
-          - vllm_omni/deploy/qwen3_tts.yaml
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
-              export VLLM_OMNI_USE_V2_RUNNER=1
-              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: l4_1
-
       - label: "Diffusion · Wan22 Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_wan22_t2v.py

docker/Dockerfile.npu

@@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \
       "stepaudio2-minicpmo>=0.1.1" \
       "step-audio2>=1.0.0"
 
+ENV VLLM_TARGET_DEVICE=npu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.a3

@@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \
       "stepaudio2-minicpmo>=0.1.1" \
       "step-audio2>=1.0.0"
 
+ENV VLLM_TARGET_DEVICE=npu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.ci

@@ -65,6 +65,7 @@ RUN python3 -m pip install -v --no-deps \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60
 
+ENV VLLM_TARGET_DEVICE=npu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.ci.a3

@@ -66,6 +66,7 @@ RUN python3 -m pip install -v --no-deps \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60
 
+ENV VLLM_TARGET_DEVICE=npu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docs/.nav.yml

@@ -5,14 +5,22 @@ nav:
     - getting_started/quickstart.md
     - getting_started/installation/*
   - Serving:
-    - OpenAI-Compatible API:
+    - API Server:
-      - Chat Completions API: serving/chat_completions_api.md
+      - Overview: serving/README.md
-      - Image Generation: serving/image_generation_api.md
+      - Core APIs:
-      - Image Edit: serving/image_edit_api.md
+        - Chat Completions: serving/chat_completions_api.md
-      - Text to Speech: serving/speech_api.md
+        - Text to Speech: serving/speech_api.md
-      - Videos: serving/videos_api.md
+        - Audio Generation: serving/audio_generate_api.md
-      - Audio Generation: serving/audio_generate_api.md
+        - Image Generation: serving/image_generation_api.md
-      - Streaming Video Input: serving/video_stream_api.md
+        - Image Edit: serving/image_edit_api.md
+        - Videos: serving/videos_api.md
+      - Streaming and Realtime:
+        - Realtime Audio: serving/realtime_api.md
+        - Full Duplex: serving/full_duplex_api.md
+        - Streaming Video Input: serving/video_stream_api.md
+        - Streaming Video Output: serving/streaming_video_output_api.md
+        - OpenPI Robot Policy: serving/openpi_api.md
+      - Standalone Experimental Servers: serving/standalone_servers.md
   - Examples:
     - examples/README.md
     - Offline Inference:

docs/examples/README.md

@@ -3,4 +3,6 @@
 vLLM-Omni's examples are split into two categories:
 
 - If you are using vLLM-Omni from within Python code, see the *Offline Inference* section.
-- If you are using vLLM-Omni from an HTTP application or client, see the *Online Serving* section.
+- If you are using vLLM-Omni from an HTTP application or client, start with
+  the [API Server guide](../serving/README.md) to choose an endpoint, then use
+  the *Online Serving* section for model-specific examples.

docs/getting_started/quickstart.md

@@ -91,7 +91,7 @@ if __name__ == "__main__":
 
 For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
 
-## Online Serving with OpenAI-Completions API
+## Online Serving with the API Server
 
 Text-to-image generation quickstart with vLLM-Omni:
 
@@ -100,20 +100,16 @@ vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091
 ```
 
 ```bash
-curl -s http://localhost:8091/v1/chat/completions \
+curl -s http://localhost:8091/v1/images/generations \
   -H "Content-Type: application/json" \
   -d '{
-    "messages": [
+    "prompt": "a cup of coffee on the table",
-      {"role": "user", "content": "a cup of coffee on the table"}
+    "size": "1024x1024",
-    ],
+    "response_format": "b64_json",
-    "extra_body": {
+    "seed": 42
-      "height": 1024,
+  }' | jq -r '.data[0].b64_json' | base64 -d > coffee.png
-      "width": 1024,
-      "num_inference_steps": 50,
-      "guidance_scale": 4.0,
-      "seed": 42
-    }
-  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > coffee.png
 ```
 
-For more details, please refer to [online serving](../user_guide/examples/online_serving/text_to_image.md).
+See the [API Server guide](../serving/README.md) to choose an endpoint. For
+model-specific details, refer to the [text-to-image online serving
+example](../user_guide/examples/online_serving/text_to_image.md).

docs/models/supported_models.md

@@ -21,7 +21,7 @@ th {
 |--------------|--------|-------------------|------------|---------|-----|-----------|--------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   | — |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
 | `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
@@ -62,7 +62,7 @@ th {
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
-| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | ✅︎ | — |
+| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
 | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |

docs/serving/README.md

@@ -0,0 +1,187 @@
+# API Server
+
+vLLM-Omni exposes OpenAI-compatible HTTP APIs plus vLLM-Omni extensions for
+image, audio, video, realtime, and robot-policy workloads. Use this page to
+choose an endpoint; follow the linked reference page for model-specific fields,
+wire protocols, and response details.
+
+Each server instance hosts one model. An endpoint is usable only when the
+loaded model supports that task.
+
+This guide focuses on vLLM-Omni's primary task endpoints. Compatible models
+may also expose standard routes inherited from vLLM, such as
+`POST /v1/completions`; those routes remain model- and task-dependent.
+
+## Start and Verify the Server
+
+Start a model with the shared serving command:
+
+```bash
+vllm serve <model> --omni --port 8091
+```
+
+Some models require additional flags or a deployment configuration. Use the
+model-specific guide when one is provided.
+
+After the server starts, check its health and served model name:
+
+```bash
+export VLLM_OMNI_BASE_URL=http://localhost:8091
+
+curl "$VLLM_OMNI_BASE_URL/health"
+curl "$VLLM_OMNI_BASE_URL/v1/models" | jq .
+```
+
+Use `http://localhost:8091/v1` as the `base_url` for an OpenAI SDK client. If
+the server was started with `--api-key`, include an `Authorization` header with
+`Bearer <api-key>` in requests.
+
+## Choose a Core HTTP Endpoint
+
+Prefer a task-specific endpoint when one matches your workload. Use Chat
+Completions for conversational or heterogeneous omni pipelines, rather than as
+the default wrapper for every generation task.
+
+| Task | Endpoint | Request | Response | Details |
+|------|----------|---------|----------|---------|
+| Conversation or multimodal understanding/generation | `POST /v1/chat/completions` | JSON | OpenAI-style JSON or SSE | [Chat Completions](chat_completions_api.md) |
+| Text-to-speech | `POST /v1/audio/speech` | JSON | Audio bytes or SSE | [Speech](speech_api.md) |
+| Sound, music, or ambient audio generation | `POST /v1/audio/generate` | JSON | Audio bytes | [Audio Generation](audio_generate_api.md) |
+| Text-to-image | `POST /v1/images/generations` | JSON | Base64 JSON or an image file | [Image Generation](image_generation_api.md) |
+| Image editing | `POST /v1/images/edits` | Multipart form | Base64 JSON or SSE | [Image Edit](image_edit_api.md) |
+| Video generation (recommended) | `POST /v1/videos` | Multipart form | Asynchronous job JSON | [Videos](videos_api.md) |
+| Video generation (blocking) | `POST /v1/videos/sync` | Multipart form | Video bytes | [Videos](videos_api.md#synchronous-response) |
+
+`POST /v1/videos/sync` is intended for tests and simple scripts. Use the
+asynchronous job API for production requests so clients can poll, download,
+and delete generated videos independently.
+
+## Minimal Requests
+
+The examples below assume a compatible model is already running. TTS voices,
+media inputs, generation controls, and output capabilities vary by model.
+
+=== "Chat"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/chat/completions" \
+      -H "Content-Type: application/json" \
+      -d '{
+        "messages": [{"role": "user", "content": "Describe vLLM-Omni briefly."}],
+        "modalities": ["text"]
+      }'
+    ```
+
+=== "Speech"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/audio/speech" \
+      -H "Content-Type: application/json" \
+      -d '{
+        "input": "Hello from vLLM-Omni.",
+        "voice": "vivian"
+      }' \
+      --output speech.wav
+    ```
+
+=== "Audio"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/audio/generate" \
+      -H "Content-Type: application/json" \
+      -d '{
+        "input": "Ocean waves on a quiet beach",
+        "audio_length": 5.0
+      }' \
+      --output audio.wav
+    ```
+
+=== "Image"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/images/generations" \
+      -H "Content-Type: application/json" \
+      -d '{
+        "prompt": "A small robot reading beside a window",
+        "size": "1024x1024",
+        "response_format": "file"
+      }' \
+      --output image.png
+    ```
+
+=== "Image edit"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/images/edits" \
+      -F "image=@input.png" \
+      -F "prompt=Turn the background into a snowy mountain" \
+      | jq -r '.data[0].b64_json' \
+      | base64 --decode > edited.png
+    ```
+
+=== "Video"
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/videos" \
+      -F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
+      -F "size=1280x720" \
+      -F "seconds=5"
+    ```
+
+    The response contains a video job `id`. Poll and download it with:
+
+    ```bash
+    curl "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>"
+    curl -L "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>/content" --output video.mp4
+    ```
+
+## Choose a Streaming or Realtime Endpoint
+
+These WebSocket APIs have different event schemas and cannot be used
+interchangeably.
+
+| Workload | Endpoint | Interaction model | Details |
+|----------|----------|-------------------|---------|
+| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
+| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
+| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
+| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
+| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
+| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
+
+All six routes are model- or configuration-dependent. In particular,
+`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the
+deployment explicitly enables duplex sessions. Clients should also verify the
+duplex capability payload because the query-parameter form falls back to the
+ordinary realtime handler when duplex is unavailable.
+
+## Related Endpoints
+
+| Purpose | Endpoints | Reference |
+|---------|-----------|-----------|
+| Discovery and readiness | `GET /health`, `GET /v1/models` | This page |
+| Batched conversations | `POST /v1/chat/completions/batch` | [Batch requests](chat_completions_api.md#batch-requests) |
+| Batched speech | `POST /v1/audio/speech/batch` | [Batch speech generation](speech_api.md#batch-speech-generation) |
+| TTS voice management | `GET/POST /v1/audio/voices`, `DELETE /v1/audio/voices/{name}` | [Voices](speech_api.md#voices-endpoint) |
+| Video job lifecycle | `GET /v1/videos`, `GET/DELETE /v1/videos/{video_id}`, `GET /v1/videos/{video_id}/content` | [Video endpoints](videos_api.md#endpoints) |
+| Release and restore stage memory | `POST /v1/omni/sleep`, `POST /v1/omni/wakeup` | [Sleep Mode](../features/sleep_mode.md) |
+
+## Standalone Experimental Servers
+
+PersonaPlex also provides a Moshi-compatible WebSocket server, while JoyVL
+provides a stateful interaction orchestrator in front of another model server.
+These are separate processes with their own routes; they are not additional
+paths on every `vllm serve --omni` instance. See [Standalone Experimental
+Servers](standalone_servers.md) before deploying either one.
+
+## Compatibility Notes
+
+- A registered path does not mean every model supports it. Unsupported tasks
+  return an OpenAI-style error or a WebSocket `error` event.
+- Use JSON for Chat, Speech, Audio Generation, and Image Generation. Use
+  `multipart/form-data` when uploading images, videos, or other files.
+- Generation can take minutes for large diffusion models. Configure client
+  timeouts accordingly; prefer asynchronous video jobs over long blocking
+  requests.
+- Model-specific parameters are vLLM-Omni extensions. Check the endpoint
+  reference before assuming an OpenAI SDK exposes them directly.

docs/serving/audio_generate_api.md

@@ -4,14 +4,14 @@ vLLM-Omni provides an API for text-to-audio generation using diffusion-based mod
 
 Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
 
-Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
 ## Quick Start
 
 ### Start the Server
 
 ```bash
-vllm-omni serve stabilityai/stable-audio-open-1.0 \
+vllm serve stabilityai/stable-audio-open-1.0 \
     --host 0.0.0.0 \
     --port 8091 \
     --gpu-memory-utilization 0.9 \
@@ -326,7 +326,7 @@ curl http://localhost:8091/health
 Enable debug logging:
 
 ```bash
-vllm-omni serve stabilityai/stable-audio-open-1.0 \
+vllm serve stabilityai/stable-audio-open-1.0 \
     --host 0.0.0.0 \
     --port 8091 \
     --gpu-memory-utilization 0.9 \

docs/serving/chat_completions_api.md

@@ -1,78 +1,110 @@
 # Chat Completions API
 
-vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
+Use `/v1/chat/completions` for conversational and multimodal pipelines that
-endpoint using diffusion models. This page explains how to pass generation
+follow the OpenAI Chat API schema. Depending on the loaded model, a request can
-parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
+contain text, image, audio, or video input and can return text, audio, images,
-models through this endpoint.
+or other model-specific output.
 
-!!! tip
+Some diffusion pipelines also support image generation and editing through
-    For dedicated endpoints that accept generation parameters as top-level
+Chat Completions. Prefer the dedicated [Image Generation
-    fields, see [Image Generation API](image_generation_api.md) and
+API](image_generation_api.md) or [Image Edit API](image_edit_api.md) for
-    [Image Edit API](image_edit_api.md).
+task-specific clients; their request fields and response contracts are more
+direct.
 
-## Passing Generation Parameters
+## Basic Request
-
-The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
-does not natively include diffusion-specific fields like `num_inference_steps`
-or `height`. How you pass these extra fields depends on your client.
-
-### curl / Python `requests`
-
-Wrap generation parameters inside an `"extra_body"` key in the JSON body:
 
 ```bash
-curl -s http://localhost:8091/v1/chat/completions \
+curl http://localhost:8091/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
     "messages": [
-      {"role": "user", "content": "A beautiful landscape painting"}
+      {"role": "user", "content": "Describe vLLM-Omni briefly."}
     ],
-    "extra_body": {
+    "modalities": ["text"]
-      "num_inference_steps": 50,
-      "seed": 42
-    }
   }'
 ```
 
-### OpenAI Python SDK
+Use `stream: true` for Server-Sent Events. Input media syntax, supported output
+`modalities`, and response choices depend on the model; see the
+[model-specific online serving examples](../user_guide/examples/online_serving/qwen3_omni.md).
 
-Use the `extra_body` **keyword argument**. The SDK automatically merges these
+!!! tip
-fields into the top-level request body:
+    Each server hosts one model. Query `GET /v1/models` for its served model
+    name when your client requires the `model` field.
 
-```python
+## vLLM Extension Parameters
-response = client.chat.completions.create(
+
-    model="Qwen/Qwen-Image",
+Send standard and vLLM-specific Chat parameters as top-level fields in a
-    messages=[{"role": "user", "content": "A beautiful landscape painting"}],
+direct HTTP request. When using the OpenAI Python SDK, pass vLLM-specific
-    extra_body={
+fields through the SDK's `extra_body` keyword; the SDK merges them into the
-        "num_inference_steps": 50,
+top-level JSON sent to the server.
-        "seed": 42,
+
-    },
+=== "curl"
-)
+
+    ```bash
+    curl http://localhost:8091/v1/chat/completions \
+      -H "Content-Type: application/json" \
+      -d '{
+        "messages": [{"role": "user", "content": "Write a short haiku."}],
+        "top_k": 40
+      }'
+    ```
+
+=== "OpenAI Python SDK"
+
+    ```python
+    from openai import OpenAI
+
+    client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+    response = client.chat.completions.create(
+        model="your-served-model",
+        messages=[{"role": "user", "content": "Write a short haiku."}],
+        extra_body={"top_k": 40},
+    )
+    ```
+
+??? note "Diffusion compatibility"
+
+    Diffusion models exposed through Chat Completions can accept supported
+    fields such as `num_inference_steps`, `seed`, `height`, and `width`. Send
+    them at the top level with direct HTTP, or in the SDK's `extra_body`
+    keyword argument.
+
+    A literal nested `"extra_body"` JSON object is accepted for compatibility,
+    but it is not recommended for new direct HTTP clients. Do not provide the
+    same parameter in multiple locations; the server returns a `400` error for
+    duplicate diffusion parameters. Prefer the dedicated [Image Generation
+    API](image_generation_api.md) or [Image Edit API](image_edit_api.md) when
+    either endpoint matches the task.
+
+## Batch Requests
+
+`POST /v1/chat/completions/batch` accepts the same shared generation fields,
+but `messages` is a list of conversations. The response contains one choice
+per conversation in input order.
+
+```bash
+curl http://localhost:8091/v1/chat/completions/batch \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      [{"role": "user", "content": "Summarize vLLM in one sentence."}],
+      [{"role": "user", "content": "Summarize vLLM-Omni in one sentence."}]
+    ],
+    "max_tokens": 64
+  }'
 ```
 
-!!! note "SDK `extra_body` vs. JSON `extra_body`"
+Batch Chat Completions does not support streaming, tools, beam search, or
-    These two `extra_body` usages look similar but work differently under the
+`n > 1`.
-    hood. The SDK flattens the dict into the top-level request JSON, while the
-    curl/requests approach sends it as a nested `"extra_body"` key. Both are
-    handled correctly by the server.
-
-!!! note "About the `ignored fields` warning"
-    You may see a log message like:
-
-    ```
-    WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
-    ```
-
-    This is **harmless**. It is emitted by vLLM's request validation layer
-    because these fields are not part of the standard OpenAI
-    `ChatCompletionRequest` schema. The fields are still stored internally
-    and correctly forwarded to the diffusion pipeline.
 
 ## Model-Specific Examples
 
-For complete examples with full request/response details, see the model-specific
+For complete examples with model-specific inputs and outputs, see:
-guides:
 
+- [Qwen3-Omni](../user_guide/examples/online_serving/qwen3_omni.md)
+- [Qwen2.5-Omni](../user_guide/examples/online_serving/qwen2_5_omni.md)
 - [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
 - [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
 - [GLM-Image](../user_guide/examples/online_serving/glm_image.md)

docs/serving/full_duplex_api.md

@@ -0,0 +1,116 @@
+# Full-Duplex WebSocket API
+
+vLLM-Omni provides an experimental full-duplex runtime for models that can
+continue receiving speech while producing speech. It adds persistent session
+state, model-specific turn policy, overlap handling, playback acknowledgement,
+and optional session resume.
+
+Full duplex is distinct from the turn-based [Realtime Audio API](realtime_api.md).
+
+## Choose an Endpoint
+
+| Endpoint | Protocol | Recommended use |
+|----------|----------|-----------------|
+| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
+| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
+
+Both endpoints use the same duplex engine and require the same model-side
+adapter. Prefer the Realtime projection unless the native lifecycle events are
+specifically required.
+
+## Enable Full Duplex
+
+The route becomes usable only when the deployment configuration explicitly
+sets:
+
+```yaml
+session_mode: duplex
+```
+
+The selected model pipeline must also provide a duplex serving adapter.
+Model-native deployments configure the corresponding engine runtime extension
+and control plane as part of their registered pipeline.
+
+!!! warning
+
+    `WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
+    enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
+    turn-based realtime handler when no duplex handler exists. Confirm that
+    `session.created.session.capabilities` is present before treating the
+    connection as full duplex.
+
+The current unified-runtime integrations are:
+
+- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5_duplex.yaml`;
+- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex.
+
+JoyVL is a separate HTTP interaction orchestrator and does not use these
+WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
+
+## MiniCPM-o Quick Start
+
+Start the duplex deployment:
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+  --deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
+  --trust-remote-code \
+  --port 8091
+```
+
+Stream a mono, PCM16, 16 kHz WAV file with the provided client:
+
+```bash
+python examples/online_serving/minicpmo/realtime_duplex_demo.py \
+  --url 'ws://localhost:8091/v1/realtime?duplex=1' \
+  --model openbmb/MiniCPM-o-4_5 \
+  --input-wav input_16k_mono.wav \
+  --ref-audio reference_voice.wav \
+  --output-dir /tmp/minicpmo-duplex
+```
+
+## Realtime Event Lifecycle
+
+A typical `/v1/realtime?duplex=1` session follows this lifecycle:
+
+1. Send `session.update` with the model, modalities, audio formats, and session
+   options.
+2. Wait for `session.created`; inspect `session.capabilities` instead of
+   assuming every model supports the same controls.
+3. Send `input_audio_buffer.append` events while microphone audio arrives.
+4. Send `input_audio_buffer.commit` at a user-turn boundary when required by
+   the model policy.
+5. Consume `response.created`, transcript deltas, `response.audio.delta`, and
+   `response.done` or `response.listen` events.
+6. Send `playback.ack` after audio has been played when the session advertises
+   playback acknowledgement support.
+7. Send `session.close` and wait for `session.closed`.
+
+Unlike the turn-based realtime endpoint, input may continue while a response
+is active. The server can emit `overlap.decision` to describe whether input was
+deferred, treated as a short acknowledgement, or used to interrupt output.
+
+## Capabilities and Model Differences
+
+The `session.created` payload includes capability fields such as
+`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
+`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
+runtime contract.
+
+For example, PersonaPlex supports native overlapping speech but currently
+advertises `supports_barge_in=false`; destructive output interruption and
+model-state rewind have not been validated for that integration. Capacity and
+session-resume behavior also depend on the selected deployment configuration.
+
+## Native Protocol
+
+`WS /v1/duplex` exposes lower-level lifecycle names including
+`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
+and `session.close`. It returns native session, input, response, overlap, and
+error events. This protocol is experimental and may evolve with the runtime;
+applications should use the provided Realtime client where possible.
+
+See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
+[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
+and [full-duplex runtime design](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/experimental/fullduplex/DESIGN.md)
+for model-specific validation and architecture details.

docs/serving/openpi_api.md

@@ -0,0 +1,109 @@
+# OpenPI Robot Policy WebSocket API
+
+Use `WS /v1/realtime/robot/openpi` for low-latency robot-policy inference with
+OpenPI-compatible clients. A persistent WebSocket carries observations to the
+loaded policy model and returns action arrays.
+
+Despite the `/realtime` prefix, this endpoint does not use the OpenAI Realtime
+event schema. It uses binary MessagePack frames with NumPy extensions.
+
+## Availability
+
+The endpoint is available only when the loaded diffusion policy configuration
+contains `policy_server_config`. The server sends that model-specific
+configuration to the client as the first binary frame.
+
+For example, DreamZero declares:
+
+```yaml
+policy_server_config:
+  image_resolution: [180, 320]
+  n_external_cameras: 2
+  needs_wrist_camera: true
+  needs_stereo_camera: false
+  needs_session_id: true
+  action_space: joint_position
+```
+
+Without this configuration, the WebSocket returns `Robot policy not available`
+and closes.
+
+## DreamZero Quick Start
+
+Start the bundled DreamZero deployment:
+
+```bash
+vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8091 \
+  --served-model-name dreamzero-droid \
+  --deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
+  --enforce-eager --disable-log-stats
+```
+
+Install the optional client dependencies, download the sample camera inputs,
+and run the client:
+
+```bash
+pip install openpi-client websockets opencv-python
+
+hf download YangshenDeng/vllm-omni-dreamzero-assets \
+  --repo-type dataset \
+  --local-dir outputs/dreamzero/assets
+
+python examples/online_serving/dreamzero/openpi_client.py \
+  --host 127.0.0.1 \
+  --port 8091 \
+  --video-dir outputs/dreamzero/assets
+```
+
+## Protocol
+
+The connection is request-response after an initial server handshake:
+
+```text
+connect
+  <- msgpack(policy_server_config)
+
+infer
+  -> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
+  <- msgpack(ndarray | dict[str, ndarray])
+
+reset
+  -> msgpack({"endpoint": "reset"})
+  <- msgpack({"status": "reset successful"})
+```
+
+If `endpoint` is omitted, the server treats the message as `infer`. Each
+inference message produces one action response after the engine request
+completes; action tokens or intermediate tensors are not streamed.
+
+NumPy arrays use the marker format implemented by `openpi-client`. The server
+also accepts the legacy vLLM NumPy marker representation on input. JSON text
+frames are not observation messages.
+
+## Sessions and Reset
+
+- The API layer's current-session and first-call counters are scoped to each
+  WebSocket connection.
+- `session_id` identifies model-side state across observations on that
+  connection. If omitted, it defaults to `default`.
+- The first inference for a session is sent to the model with `reset=true`.
+- Changing `session_id`, or sending the `reset` command, causes the next
+  inference to start with `reset=true`.
+- The policy pipeline owns observation transforms and persistent model state,
+  normally keyed by `session_id`; the API layer forwards the raw observation
+  dictionary.
+
+## Limits and Errors
+
+- Maximum inbound payload size is 64 MiB.
+- The server closes an idle connection after 30 seconds.
+- Invalid binary input returns a MessagePack
+  `{"type":"error","message":"Invalid request payload"}` response.
+- Inference failures return a generic `Internal inference error` without
+  exposing an internal traceback.
+- Unsupported NumPy object, structured, and complex dtypes are rejected.
+
+Observation keys, camera layout, state tensors, and action shapes are defined
+by the loaded policy rather than by this transport. Read the handshake before
+constructing observations. See the [DreamZero example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/dreamzero)
+for a complete OpenPI client and DROID simulation loop.

docs/serving/realtime_api.md

@@ -0,0 +1,82 @@
+# Realtime Audio WebSocket API
+
+Use `WS /v1/realtime` to stream PCM audio into a compatible omni model and
+receive incremental transcription and synthesized response audio. This is a
+turn-based streaming API: it starts one generation over an audio stream and
+finishes after the client closes that input stream.
+
+For a model that listens while it is already speaking, use the experimental
+[Full-Duplex API](full_duplex_api.md) instead.
+
+## Quick Start
+
+Start a compatible model. The currently documented example uses Qwen3-Omni:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+Run the provided client with a mono, 16-bit PCM, 16 kHz WAV file:
+
+```bash
+python examples/online_serving/qwen3_omni/openai_realtime_client.py \
+  --url ws://localhost:8091/v1/realtime \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --input-wav input_16k_mono.wav \
+  --output-wav response.wav \
+  --output-text response.txt
+```
+
+The client dependency is `websockets`.
+
+## Protocol
+
+Messages in both directions are JSON text frames. Audio carried inside an
+event is base64-encoded raw PCM16.
+
+| Direction | Event | Purpose |
+|-----------|-------|---------|
+| Server to client | `session.created` | Confirms that the WebSocket connection is ready |
+| Client to server | `session.update` | Selects and validates the served model |
+| Client to server | `input_audio_buffer.commit` with `final: false` | Starts generation over the incoming stream |
+| Client to server | `input_audio_buffer.append` | Appends base64 PCM16 audio |
+| Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input |
+| Server to client | `transcription.delta` | Carries incremental response text |
+| Server to client | `transcription.done` | Carries final text and token usage |
+| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio |
+| Server to client | `response.audio.done` | Marks the end of response audio |
+| Server to client | `error` | Reports an invalid event, model, or audio payload |
+
+A minimal client sends events in this order:
+
+```json
+{"type":"session.update","model":"Qwen/Qwen3-Omni-30B-A3B-Instruct"}
+{"type":"input_audio_buffer.commit","final":false}
+{"type":"input_audio_buffer.append","audio":"<base64-pcm16>"}
+{"type":"input_audio_buffer.commit","final":true}
+```
+
+The initial non-final commit intentionally starts generation before all audio
+has arrived. Continue sending `input_audio_buffer.append` events while the
+engine is consuming the stream.
+
+## Audio Handling
+
+- Input is mono PCM16 at 16 kHz for the Qwen3-Omni example.
+- `response.audio.delta.audio` contains base64-encoded PCM16 bytes.
+- Read `sample_rate_hz` from each audio event instead of assuming an output
+  rate. Qwen3-Omni output is typically 24 kHz.
+- Concatenate audio deltas in receive order to construct the output waveform.
+- Set `input_audio_buffer.commit.final` to `true` after the last input chunk so
+  the server can terminate the streaming request.
+
+## Availability and Limitations
+
+The path is registered on the unified API server, but it is usable only when
+the loaded pipeline implements realtime audio input and produces compatible
+audio output. Unsupported deployments return an `error` event. The endpoint
+does not provide duplex session resume, playback acknowledgement, overlap
+policy, or barge-in controls.
+
+See the [Qwen3-Omni online serving example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni)
+for concurrency options, chunk pacing, and per-delta audio debugging.

docs/serving/standalone_servers.md

@@ -0,0 +1,109 @@
+# Standalone Experimental Servers
+
+Most vLLM-Omni endpoints run on the unified server started with
+`vllm serve <model> --omni`. The experimental full-duplex package also contains
+two separate user-facing server processes. Their routes are not added to the
+unified API server.
+
+| Server | Transport | Role |
+|--------|-----------|------|
+| PersonaPlex compatibility server | WebSocket | Hosts the official PersonaPlex browser protocol and a raw-PCM alternative |
+| JoyVL interaction server | HTTP | Adds state, memory, proactive decisions, and delegation in front of a separate OpenAI-compatible model backend |
+
+## PersonaPlex Compatibility Server
+
+Use this server when a client requires the official Moshi/PersonaPlex wire
+protocol or the bundled browser UI. For the engine-owned unified duplex path,
+use [`WS /v1/realtime?duplex=1`](full_duplex_api.md) instead.
+
+Start it with:
+
+```bash
+HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
+  vllm_omni.experimental.fullduplex.personaplex.serving.server \
+  --port 8091
+```
+
+| Route | Purpose |
+|-------|---------|
+| `GET /` | Official PersonaPlex browser client, when its assets are available |
+| `GET /health` | Readiness check |
+| `WS /api/chat` | Moshi-compatible tagged binary protocol with Opus audio |
+| `WS /v1/audio/duplex` | Simpler float32 raw-PCM protocol |
+
+### `/api/chat` Protocol
+
+The first byte of every binary frame is a message tag:
+
+| Direction | Tag | Payload |
+|-----------|-----|---------|
+| Server to client | `0x00` | Ready handshake |
+| Client to server | `0x01` | Opus microphone audio at 24 kHz |
+| Server to client | `0x01` | Opus response audio |
+| Server to client | `0x02` | UTF-8 inner-monologue text |
+
+The optional `text_prompt` and `voice_prompt` query parameters select a persona
+and voice. The default server admits one conversation; use `--batch-size N` to
+configure multiple lockstep slots. Connections beyond capacity close with
+WebSocket code 1013.
+
+### `/v1/audio/duplex` Protocol
+
+Send `{"type":"open","persona":"...","voice":"..."}` as a JSON text frame,
+then exchange binary float32 PCM frames. The server returns generated float32
+audio as binary frames and text as `{"type":"text","text":"..."}`. Send
+`{"type":"close"}` and wait for `{"type":"done"}` to flush and finish.
+
+See the [PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex)
+for the browser and headless clients.
+
+## JoyVL Interaction Server
+
+JoyVL is an orchestration layer, not a model-serving engine. It calls a
+separate OpenAI-compatible backend and adds per-session frame history, memory,
+persona policy, and optional delegation.
+
+Start the model backend and orchestrator separately:
+
+```bash
+vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
+  --served-model-name JoyAI-VL-Interaction-Preview \
+  --port 8092 \
+  --max-model-len 131072 \
+  --enable-prefix-caching \
+  --limit-mm-per-prompt '{"image":256,"video":1}'
+
+python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
+  --port 8091 \
+  --main-backend-url http://127.0.0.1:8092/v1 \
+  --main-model JoyAI-VL-Interaction-Preview
+```
+
+| Method and route | Purpose |
+|------------------|---------|
+| `GET /health` | Readiness check |
+| `GET /v1/models` | Reports the configured interaction model |
+| `POST /v1/chat/completions` | Processes one frame or interaction tick |
+| `POST /reset` | Resets a session |
+| `POST /v1/streaming/reset` | Alias for `/reset` |
+| `POST /v1/streaming/persona` | Changes the session persona |
+
+Send one multimodal Chat Completions request per video frame, normally around
+one frame per second, and identify the session with `x-session-id`. Responses
+include an `interaction` block whose action is `silence`, `response`, or
+`delegate`.
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H 'content-type: application/json' \
+  -H 'x-session-id: demo' \
+  -d '{
+    "messages": [{"role": "user", "content": [
+      {"type": "text", "text": "Alert me if a fire appears"},
+      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
+    ]}]
+  }'
+```
+
+See the [JoyVL recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-VL-Interaction.md)
+for memory, personas, delegation backends, and the browser UI.

docs/serving/streaming_video_output_api.md

@@ -0,0 +1,92 @@
+# Streaming Video Output WebSocket API
+
+Use `WS /v1/realtime/video` to receive fragmented MP4 chunks while a compatible
+diffusion video model is generating. This endpoint streams generated video
+output; it is unrelated to [`WS /v1/video/chat/stream`](video_stream_api.md),
+which streams live video input into a conversational model.
+
+## Quick Start
+
+Start a model that supports step execution and enable diffusion streaming
+output:
+
+```bash
+vllm serve BestWishYsh/Helios-Distilled --omni \
+  --diffusion-streaming-output \
+  --port 8091
+```
+
+Install the client dependency and run the provided client:
+
+```bash
+pip install av websockets
+
+python examples/online_serving/streaming_video_generation/streaming_video_client.py \
+  --host 127.0.0.1 \
+  --port 8091 \
+  --model BestWishYsh/Helios-Distilled \
+  --prompt "A serene lakeside sunrise with mist over the water" \
+  --output output.mp4
+```
+
+## Protocol
+
+Control events are JSON text frames. Each `video.chunk_metadata` event is
+immediately followed by one binary WebSocket frame containing fragmented MP4
+bytes.
+
+| Direction | Event or frame | Purpose |
+|-----------|----------------|---------|
+| Client to server | `session.start` | Starts generation with a prompt and video parameters |
+| Server to client | `video.start` | Confirms the request ID, format, and accepted configuration |
+| Server to client | `video.chunk_metadata` | Describes the next binary media frame |
+| Server to client | Binary frame | Carries one fragmented MP4 chunk |
+| Client to server | `session.interaction` | Queues an optional prompt update during generation |
+| Server to client | `session.interaction.queued` | Confirms that an interaction was queued |
+| Client to server | `session.ping` | Refreshes the stall clock |
+| Server to client | `session.pong` | Acknowledges a ping |
+| Client to server | `session.stop` | Cancels the active generation |
+| Server to client | `session.done` | Reports completion or cancellation and the chunk count |
+| Server to client | `error` | Reports validation, generation, or timeout failures |
+
+Start a minimal session with:
+
+```json
+{
+  "type": "session.start",
+  "model": "BestWishYsh/Helios-Distilled",
+  "prompt": "A serene lakeside sunrise",
+  "format": "m4s",
+  "width": 640,
+  "height": 384,
+  "fps": 16,
+  "num_frames": 99
+}
+```
+
+Only `m4s` output is currently accepted. Concatenate the binary frames in
+receive order, then remux the fragmented stream if the target player requires
+a progressive MP4 file.
+
+## Mid-Generation Interaction
+
+Models that implement interaction updates can change the active prompt:
+
+```json
+{
+  "type": "session.interaction",
+  "interaction": {
+    "event_id": "sunlight-update",
+    "event": {"prompt": "Sunlight breaks through the mist"},
+    "transition_chunks": 2
+  }
+}
+```
+
+This feature is pipeline-dependent. An enabled endpoint does not imply that
+the loaded model supports midway prompt updates.
+
+The default start timeout is 10 seconds and the server reports a stall after
+about 60 seconds without generation progress or a `session.ping`. See the
+[streaming video generation example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/streaming_video_generation)
+for prompt-update scheduling and browser playback.

examples/online_serving/minicpmo/README.md

@@ -49,7 +49,7 @@ scenario below for live barge-in validation on the target GPU.
 Default:
 
 ```bash
-vllm-omni serve openbmb/MiniCPM-o-4_5 \
+vllm serve openbmb/MiniCPM-o-4_5 \
     --omni \
     --deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
     --trust-remote-code \

examples/online_serving/minicpmo/realtime_duplex_demo.py

@@ -155,6 +155,23 @@ def _response_in_progress(events: list[dict[str, object]]) -> bool:
     )
 
 
+async def _commit_input_after_playback_checkpoint(client: RealtimeDuplexClient) -> float:
+    """Commit input only after preserving already-played assistant history.
+
+    Native duplex can finish several responses while the user is still
+    streaming.  Under the ``ack_only`` playback policy, acknowledging those
+    responses after the final input commit is intentionally rejected: it could
+    append an old assistant message after the newly committed user message.
+    Checkpoint playback first so the server commits each assistant item at its
+    original history position.  A second acknowledgement after the commit can
+    then safely advance response-local playback that drained in the meantime.
+    """
+    await client.acknowledge_playback()
+    commit_sent_at_s = time.monotonic()
+    await client.commit()
+    return commit_sent_at_s
+
+
 def _event_count_after(
     events: list[dict[str, object]],
     event_type: str,
@@ -448,8 +465,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
             chunk_period_ms=_chunk_period_ms(client.events.events),
         )
         wait_for_post_commit_decision = False
-        commit_sent_at_s = time.monotonic()
+        commit_sent_at_s = await _commit_input_after_playback_checkpoint(client)
-        await client.commit()
         wait_error: str | None = None
         committed_index: int | None = None
         post_commit_decision: str | None = None

examples/online_serving/streaming_video_generation/README.md

@@ -37,7 +37,7 @@ During generation the client normally sends only `session.start` and then receiv
 ## Install Client Dependency
 
 ```bash
-pip install websockets
+pip install av websockets
 # For the Gradio demo:
 pip install 'vllm-omni[demo]' websockets
 ```

examples/online_serving/streaming_video_generation/streaming_video_client.py

@@ -8,7 +8,7 @@ the received bytes to disk when the session finishes or is interrupted.
 Image/reference input is intentionally not supported in this example yet.
 
 Requirements:
-    pip install websockets
+    pip install av websockets
 """
 
 import argparse

recipes/Boogu/Boogu-Image.md

@@ -1,11 +1,14 @@
 # Boogu-Image
 
-> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
+> Text-to-image and image-editing online serving
+(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
 
 ## Summary
 
 - Vendor: Boogu
-- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
+- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
+  `Boogu/Boogu-Image-0.1-Edit` (image editing), and
+  `Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
 - Task: Text-to-image generation and text-guided image editing (TI2I)
 - Mode: Online serving with the OpenAI-compatible API
 - Maintainer: Community
@@ -13,19 +16,23 @@
 ## When to use this recipe
 
 Use this recipe when you want a known-good starting point for serving
-`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
+`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
+`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
 `--diffusion-load-format diffusers`, and the upstream `boogu` package is not
 required).
 
 Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
-family. This recipe covers the Base text-to-image checkpoint, which pairs a
+family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
-Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
+with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
-Euler scheduler with time-shift. It handles photorealistic generation and
+time-shift. It handles photorealistic generation and Chinese/English text
-Chinese/English text rendering.
+rendering. The Edit and Edit-Turbo checkpoints use the same native
+`BooguImagePipeline` for image editing.
 
 ## References
 
 - Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
+- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
+- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
 - Project page: <https://boogu.org>
 - GitHub: <https://github.com/boogu-project/Boogu-Image>
 - Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
@@ -93,7 +100,7 @@ curl -s http://localhost:8091/v1/chat/completions \
 - **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
   to trim peak VRAM if needed.
 - **Key flags:**
-  - `--omni` — enables vLLM-Omni diffusion serving.
+    - `--omni` — enables vLLM-Omni diffusion serving.
 - **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
   `text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
   active whenever `guidance_scale > 1.0`.
@@ -111,13 +118,13 @@ the image-editing (TI2I) path activates automatically when a request carries a
 reference image. The Base text-to-image path is unaffected (no reference image
 is sent).
 
-#### Command
+### Command
 
 ```bash
 vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
 ```
 
-#### Verification
+### Verification
 
 Edit an image with `/v1/images/edits` (the model-card example — change a photo
 to a colored-pencil drawing). Diffusion parameters are plain multipart form
@@ -157,14 +164,14 @@ curl -s http://localhost:8091/v1/chat/completions \
   }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
 ```
 
-#### Notes
+### Notes
 
 - **Single reference image:** only one input image is supported for now (the
   upstream "Only support 1 reference image for now" limit).
 - **Guidance semantics:**
-  - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
+    - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
     `4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
-  - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
+    - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
     default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
     (3 model predictions per step), steering more strongly toward the reference
     image.
@@ -178,3 +185,69 @@ curl -s http://localhost:8091/v1/chat/completions \
   supports Boogu-Image-Edit directly
   (`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
   `--guidance-scale-2`).
+
+## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
+
+Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
+declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
+The upstream 1K hotfix is recommended for stable results and is pinned below
+to avoid silently loading the older checkpoint from the repository's default
+revision.
+
+### Command
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
+  --omni \
+  --revision hotfix-1k-20260708 \
+  --port 8091
+```
+
+### Verification
+
+Edit an image with the distilled four-step settings:
+
+```bash
+curl -s http://localhost:8091/v1/images/edits \
+  -F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \
+  -F image="@input.png" \
+  -F prompt="Change the style to a colored pencil drawing." \
+  -F num_inference_steps=4 \
+  -F guidance_scale=1.0 \
+  -F seed=42 \
+  | jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png
+```
+
+Or via chat completions:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": [
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
+        {"type": "text", "text": "Change the style to a colored pencil drawing."}
+      ]}
+    ],
+    "extra_body": {
+      "num_inference_steps": 4,
+      "guidance_scale": 1.0,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png
+```
+
+### Notes
+
+- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
+  hotfix over the 1.5K variant for more stable results.
+- **Recommended settings:** `num_inference_steps=4` and
+  `guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
+  behavior is already baked into the distilled checkpoint. A scale of `1.0`
+  avoids applying additional classifier-free guidance; unlike the regular Edit
+  checkpoint, it should not be increased to `5.0`.
+- **Reference images:** the same single-reference and `align_res` behavior as
+  the regular Edit checkpoint applies.
+- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
+  CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -69,11 +69,12 @@ reference-video preparation and MP4 output.
 
 For CUDA and ROCm deployments, non-streaming MP4 responses are encoded on the
 host CPU through PyAV/libx264 after generation. The response encoder selects the
-path automatically at runtime: inputs with a supported frame shape, common dtype,
+path automatically at runtime. The server-owned parallel converter accepts
-and RGB per-channel contiguous layout use direct planar frames; other inputs
+supported frame shapes and dtypes with either per-channel-contiguous or strided
-fall back to the legacy path before the PyAV container is opened. No CLI flag,
+RGB planes, including interleaved arrays materialized by output transport.
-model declaration, or user configuration is required. Streaming fMP4 output is
+Standalone callers without a parallel converter retain the legacy fallback for
-unchanged.
+strided planes. No CLI flag, model declaration, or user configuration is
+required. Streaming fMP4 output is unchanged.
 
 A community benchmark on 2x Xeon 8480C reported the following comparison
 between the legacy and direct planar paths
@@ -917,6 +918,72 @@ adapter and twice without it, then compare the four output digests. The adapter
 is bound and deterministic when each pair matches internally and the two pairs
 differ from each other.
 
+### FastH3 adapter
+
+[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step
+DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and
+schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused
+into the checkpoint at load time rather than switched per request, because it
+carries full-rank deltas that no LoRA layer can express.
+
+The bundle publishes four variants, so download one and point `--lora-path` at
+it; the repository root is refused rather than guessed at:
+
+```bash
+export FASTH3_DIR=/path/to/fasth3
+hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \
+  dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}"
+export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors"
+```
+
+Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server
+command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is
+correct even though FastH3 preview v1 distills T2VA only. Because the adapter is
+fused, `--lora-backend` does not apply and a request carrying a `lora=` field is
+rejected rather than served without the adapter it asked for.
+
+```bash
+-F 'num_inference_steps=4' \
+-F 'extra_params={"task":"t2va","duration":4.4}'
+```
+
+Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five
+sigma points bound four denoiser evaluations, and that count is what the step
+scheduler admits a request on. The server denoises on the release's own ladder,
+keeping H3's per-modality shifts at the checkpoint values, so a request that
+overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the
+student at noise levels it was never distilled at.
+
+Only a release that identifies itself as FastH3 is fused; any other
+`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact
+is then held to its own metadata: one that misdeclares its tensor counts or
+leaves a transformer block unedited is refused at startup instead of serving
+mostly base H3 weights on a four-step schedule. Offload is refused for the same
+reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
+`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
+
+> [!NOTE]
+> Only the **Dense / Data-Free** variant is supported today. The three VSA
+> variants carry `.set_weight` compression gates for modules that do not exist
+> in vLLM-Omni's dense H3 attention, and the release states dense attention is
+> not a supported substitute for them, so they are refused at startup rather
+> than silently run wrong. Video Sparse Attention for H3's packed
+> `[text | cond | audio | video]` sequence is follow-up work.
+
+Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
+1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
+
+| Adapter | Steps | End-to-end | Diffusion engine |
+| --- | ---: | ---: | ---: |
+| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s |
+| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s |
+
+The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE
+decoding and muxing are a fixed cost that dominates a clip this short; longer
+generations move the end-to-end figure toward the denoising one. Fusing the
+adapter does not measurably change startup: weight loading took 77.3 s with it
+against 85.8 s without.
+
 ## Key parameters
 
 | Parameter | Recommended value | Notes |

tests/benchmarks/benchmark_video_response_encoding.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Benchmark legacy and automatic MP4 response encoding on synthetic input."""
 
 from __future__ import annotations
@@ -11,13 +11,16 @@ import statistics
 import time
 from collections.abc import Callable
 from dataclasses import dataclass
+from functools import partial
 from pathlib import Path
+from typing import cast
 
 import numpy as np
 
 from vllm_omni.entrypoints.openai.video_api_utils import (
     _encode_video_bytes,
     _encode_video_bytes_legacy,
+    _PlanarFrameConverter,
 )
 
 
@@ -42,12 +45,19 @@ def _build_inputs(
     fps: int,
     audio_sample_rate: int,
     seed: int,
+    input_layout: str = "planar-backed",
 ) -> tuple[np.ndarray, np.ndarray]:
     rng = np.random.default_rng(seed)
     planar = rng.random((3, frames, height, width), dtype=np.float32)
     video = planar.transpose(1, 2, 3, 0)
-    if not all(video[..., channel].flags.c_contiguous for channel in range(3)):
+    if input_layout == "interleaved":
+        video = np.ascontiguousarray(video)
+        if not video.flags.c_contiguous:
+            raise RuntimeError("synthetic video is not C-contiguous")
+    elif input_layout == "planar-backed" and not all(video[..., channel].flags.c_contiguous for channel in range(3)):
         raise RuntimeError("synthetic video does not expose contiguous channel planes")
+    elif input_layout != "planar-backed":
+        raise ValueError(f"unsupported input layout: {input_layout}")
 
     sample_count = max(1, round(frames / fps * audio_sample_rate))
     timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
@@ -86,8 +96,8 @@ def _measure(
 
 def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]:
     selected = [record for record in records if record["label"] == label]
-    wall_values = [float(record["wall_ms"]) for record in selected]
+    wall_values = [cast(float, record["wall_ms"]) for record in selected]
-    cpu_values = [float(record["process_cpu_ms"]) for record in selected]
+    cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected]
     return {
         "runs": len(selected),
         "wall_ms": {
@@ -113,6 +123,16 @@ def _parse_args() -> argparse.Namespace:
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--rounds", type=_positive_int, default=3)
     parser.add_argument("--seed", type=int, default=20260817)
+    parser.add_argument(
+        "--input-layout",
+        choices=("planar-backed", "interleaved"),
+        default="planar-backed",
+    )
+    parser.add_argument(
+        "--frame-conversion-workers",
+        type=_positive_int,
+        default=8,
+    )
     parser.add_argument("--output-json", type=Path)
     args = parser.parse_args()
     if args.warmup < 0:
@@ -129,33 +149,41 @@ def main() -> None:
         fps=args.fps,
         audio_sample_rate=args.audio_sample_rate,
         seed=args.seed,
+        input_layout=args.input_layout,
     )
     baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
-    candidate = BenchmarkVariant("automatic", _encode_video_bytes)
+    converter = _PlanarFrameConverter(max_workers=args.frame_conversion_workers)
+    candidate = BenchmarkVariant(
+        "automatic",
+        partial(_encode_video_bytes, frame_converter=converter),
+    )
 
-    for _ in range(args.warmup):
+    try:
-        for variant in (baseline, candidate):
+        for _ in range(args.warmup):
-            _measure(
+            for variant in (baseline, candidate):
-                variant,
+                _measure(
-                video=video,
+                    variant,
-                audio=audio,
+                    video=video,
-                fps=args.fps,
+                    audio=audio,
-                audio_sample_rate=args.audio_sample_rate,
+                    fps=args.fps,
-            )
+                    audio_sample_rate=args.audio_sample_rate,
+                )
 
-    records: list[dict[str, object]] = []
+        records: list[dict[str, object]] = []
-    for round_index in range(args.rounds):
+        for round_index in range(args.rounds):
-        order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
+            order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
-        for variant in order:
+            for variant in order:
-            _, record = _measure(
+                _, record = _measure(
-                variant,
+                    variant,
-                video=video,
+                    video=video,
-                audio=audio,
+                    audio=audio,
-                fps=args.fps,
+                    fps=args.fps,
-                audio_sample_rate=args.audio_sample_rate,
+                    audio_sample_rate=args.audio_sample_rate,
-            )
+                )
-            record["round"] = round_index + 1
+                record["round"] = round_index + 1
-            records.append(record)
+                records.append(record)
+    finally:
+        converter.shutdown()
 
     output_hashes = {str(record["output_sha256"]) for record in records}
     if len(output_hashes) != 1:
@@ -165,8 +193,8 @@ def main() -> None:
     paired_savings_percent = []
     for round_index in range(args.rounds):
         pair = [record for record in records if record["round"] == round_index + 1]
-        baseline_ms = float(next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
+        baseline_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
-        candidate_ms = float(next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
+        candidate_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
         paired_savings_ms.append(baseline_ms - candidate_ms)
         paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
 
@@ -180,6 +208,8 @@ def main() -> None:
             "warmup": args.warmup,
             "rounds": args.rounds,
             "seed": args.seed,
+            "input_layout": args.input_layout,
+            "frame_conversion_workers": args.frame_conversion_workers,
             "video_shape": list(video.shape),
             "video_strides": list(video.strides),
             "codec_options": {"preset": "ultrafast", "threads": "0"},

tests/benchmarks/patch/test_patch.py

@@ -9,23 +9,98 @@ import asyncio
 import base64
 import json
 import time
+from argparse import Namespace
 from types import SimpleNamespace
 
 import pytest
 from pytest_mock import MockerFixture
 from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
 
+from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
+    SeedTTSSampleRequest,
+    SeedTTSTextSampleRequest,
+)
 from vllm_omni.benchmarks.patch.patch import (
     MixRequestFuncOutput,
     _apply_stage0_token_timings,
+    _attach_seed_tts_to_request_func_input,
     async_request_openai_chat_omni_completions,
     async_request_openai_realtime_duplex,
+    should_request_stage_metrics,
 )
 from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
 
 pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
 
 
+def _seed_tts_request_func_input() -> RequestFuncInput:
+    return RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="target text",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=2,
+        output_len=20,
+        extra_body={"modalities": ["text", "audio"]},
+    )
+
+
+def test_seed_tts_chat_request_carries_reference_audio_once() -> None:
+    reference_audio_url = "data:audio/wav;base64,AAAA"
+    sample = SeedTTSSampleRequest(
+        prompt="target text",
+        prompt_len=2,
+        expected_output_len=20,
+        multi_modal_data=None,
+        seed_tts_system_prompt="Clone the supplied voice.",
+        seed_tts_speech_extra={
+            "ref_audio": reference_audio_url,
+            "ref_text": "reference text",
+            "task_type": "Base",
+        },
+    )
+    request_func_input = _seed_tts_request_func_input()
+
+    _attach_seed_tts_to_request_func_input(sample, request_func_input)
+
+    assert request_func_input.omni_chat_messages == [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "Clone the supplied voice."}],
+        },
+        {"role": "user", "content": [{"type": "text", "text": "target text"}]},
+    ]
+    assert request_func_input.extra_body["ref_audio"] == reference_audio_url
+    serialized_request = json.dumps(
+        {"messages": request_func_input.omni_chat_messages, **request_func_input.extra_body},
+    )
+    assert serialized_request.count(reference_audio_url) == 1
+
+
+def test_seed_tts_text_chat_messages_do_not_add_reference_audio() -> None:
+    sample = SeedTTSTextSampleRequest(
+        prompt="target text",
+        prompt_len=2,
+        expected_output_len=20,
+        multi_modal_data=None,
+        seed_tts_system_prompt="Use the configured voice.",
+        seed_tts_speech_extra=None,
+    )
+    request_func_input = _seed_tts_request_func_input()
+
+    _attach_seed_tts_to_request_func_input(sample, request_func_input)
+
+    assert getattr(request_func_input, "seed_tts_row", False) is True
+    assert request_func_input.omni_chat_messages == [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "Use the configured voice."}],
+        },
+        {"role": "user", "content": [{"type": "text", "text": "target text"}]},
+    ]
+    assert "ref_audio" not in request_func_input.extra_body
+
+
 class MockResponse:
     """Mock aiohttp response for testing"""
 
@@ -274,6 +349,107 @@ def create_sse_chunk(data_dict):
     return f"data: {json.dumps(data_dict)}\n\n".encode()
 
 
+def test_chat_text_timing_metrics_request_stage_metrics():
+    args = Namespace(
+        backend="openai-chat-omni",
+        percentile_metrics="ttft,tpot,itl,e2el",
+        print_stage=False,
+        extra_body={},
+    )
+
+    assert should_request_stage_metrics(args) is True
+
+
+@pytest.mark.asyncio
+async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture):
+    """Engine timings recover TPOT when every text token arrives together."""
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "ABCD"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
+                "metrics": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 4,
+                            "vllm_itls_ms": [10.0, 11.0, 12.0],
+                            "vllm_tpot_ms": 11.0,
+                        }
+                    }
+                },
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+    mock_response = MockResponse(200, chunks)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    assert output.success is True
+    assert output.output_tokens == 4
+    assert output.itl == pytest.approx([0.010, 0.011, 0.012])
+    assert output.text_latency - output.ttft == pytest.approx(0.033)
+    assert output.tpot_measured is True
+
+
+@pytest.mark.asyncio
+async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "A"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "B"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+                "metrics": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 2,
+                            "vllm_itls_ms": [1000.0],
+                            "vllm_tpot_ms": 1000.0,
+                        }
+                    }
+                },
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+    mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    assert output.success is True
+    assert len(output.itl) == 1
+    assert 0.0 < output.itl[0] < 0.1
+
+
 # ============================================================================
 # output_tokens Tests
 # ============================================================================

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for SeedTTSTextDataset, SeedTTSTextSampleRequest, SeedTTSDesignDataset,
 and SeedTTSDesignSampleRequest.
 
@@ -234,34 +237,6 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
         assert req.seed_tts_utterance_id == "ok"
 
 
-def test_attach_sets_seed_tts_row_even_without_extra_body():
-    """seed_tts_row=True must be set for SeedTTSTextSampleRequest (no extra body)."""
-    from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSTextSampleRequest
-
-    req = SeedTTSTextSampleRequest(
-        prompt="hello world",
-        prompt_len=2,
-        expected_output_len=100,
-        multi_modal_data=None,
-        request_id="test-0",
-        seed_tts_speech_extra=None,
-        seed_tts_ref_wav_path="",
-    )
-    assert req.seed_tts_speech_extra is None
-    assert req.seed_tts_ref_wav_path == ""
-    # The fix ensures that even with speech_extra=None, the function
-    # sets seed_tts_row=True. We verify the source code has the fix.
-    import inspect
-
-    import vllm_omni.benchmarks.patch.patch as patch_mod
-
-    src = inspect.getsource(patch_mod._attach_seed_tts_to_request_func_input)
-    # seed_tts_row must be set BEFORE the 'if not ex: return' check
-    row_pos = src.index("seed_tts_row")
-    not_ex_pos = src.index("if not ex:")
-    assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
-
-
 def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
     from vllm_omni.benchmarks.data_modules import seed_tts_eval
 

tests/config/test_config_factory.py

@@ -41,6 +41,11 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+# Validation strings to match against in post init
+NO_STAGES_MATCH_STR = "no stages"
+NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage"
+
 
 @pytest.fixture(autouse=True)
 def _stable_test_platform(monkeypatch):
@@ -618,6 +623,7 @@ class TestPipelineDiscovery:
         p = PipelineConfig(
             model_type="custom_collide",
             hf_architectures=("SomeCollidingArch",),
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         assert p.hf_architectures == ("SomeCollidingArch",)
 
@@ -639,33 +645,63 @@ class TestStagePipelineConfig:
 
 
 class TestPipelineConfigNew:
-    def test_frozen(self):
+    def test_simple_init(self):
-        p = PipelineConfig(model_type="t", model_arch="A")
+        # Ensure pipeline config validates in post init.
-        with pytest.raises(AttributeError):
+        PipelineConfig(
-            p.model_type = "changed"
+            model_type="t",
+            model_arch="A",
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
 
-    def test_validate_valid(self):
+    def test_frozen(self):
         p = PipelineConfig(
+            model_type="t",
+            model_arch="A",
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        with pytest.raises(AttributeError):
+            p.model_type = "changed"
+
+    def test_validate_no_stages(self):
+        with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR):
+            PipelineConfig(model_type="t", model_arch="A")
+
+    def test_validate_no_terminal_stage(self):
+        """A pipeline with no ``final_output`` stage can never emit a result."""
+        with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR):
+            PipelineConfig(
+                model_type="t",
+                model_arch="A",
+                stages=(
+                    StagePipelineConfig(stage_id=0, model_stage="a"),
+                    StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
+                ),
+            )
+
+    def test_pipeline_can_have_final_output_in_any_stage(self):
+        """Ensure any stage may carry ``final_output``."""
+        PipelineConfig(
             model_type="t",
             model_arch="A",
             stages=(
-                StagePipelineConfig(stage_id=0, model_stage="a"),
+                StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),
                 StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
             ),
         )
-        assert p.validate() == []
-
-    def test_validate_no_stages(self):
-        p = PipelineConfig(model_type="t", model_arch="A")
-        assert any("no stages" in e.lower() for e in p.validate())
 
 
 class TestPipelineRegistration:
     def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
         deploy_key = "deploy_selected_pipeline"
         model_type_key = "hf_model_type_pipeline"
-        deploy_pipe = PipelineConfig(model_type=deploy_key)
+        deploy_pipe = PipelineConfig(
-        model_type_pipe = PipelineConfig(model_type=model_type_key)
+            model_type=deploy_key,
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        model_type_pipe = PipelineConfig(
+            model_type=model_type_key,
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
         register_pipeline(deploy_pipe)
         register_pipeline(model_type_pipe)
 
@@ -691,7 +727,11 @@ class TestPipelineRegistration:
         tmp_path,
     ):
         model_type_key = "registered_model_type_pipeline"
-        register_pipeline(PipelineConfig(model_type=model_type_key))
+        pipeline_cfg = PipelineConfig(
+            model_type=model_type_key,
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        register_pipeline(pipeline_cfg)
 
         deploy_path = tmp_path / "deploy.yaml"
         deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8")
@@ -726,6 +766,7 @@ class TestPipelineRegistration:
         pipe_cfg = PipelineConfig(
             model_type=pipeline_key,
             hf_architectures=("ArchitectureFallbackForTest",),
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         register_pipeline(pipe_cfg)
 
@@ -756,15 +797,18 @@ class TestPipelineRegistration:
             model_type="rejected_by_predicate",
             hf_architectures=(shared_arch,),
             hf_config_predicate=rejecting_predicate,
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         raises_cfg = PipelineConfig(
             model_type="raises_in_predicate",
             hf_architectures=(shared_arch,),
             hf_config_predicate=raising_predicate,
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         accept_cfg = PipelineConfig(
             model_type="accepted_by_predicate",
             hf_architectures=(shared_arch,),
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         register_pipeline(reject_cfg)
         register_pipeline(raises_cfg)
@@ -798,6 +842,7 @@ class TestPipelineRegistration:
             model_type="predicate_without_arch_match",
             hf_architectures=("DifferentArchitectureForTest",),
             hf_config_predicate=predicate,
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         register_pipeline(pipe_cfg)
 
@@ -821,6 +866,7 @@ class TestPipelineRegistration:
         resolved_cfg = PipelineConfig(
             model_type="callable_resolved_pipeline",
             hf_architectures=("CallableArchitectureForTest",),
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         seen_hf_configs = []
 
@@ -894,7 +940,10 @@ class TestPipelineRegistration:
     def test_pipeline_registration(self, clean_pipeline_registry):
         """Ensure that we can register and create a custom pipeline config."""
         new_model_type = "new_model_type"
-        pipe_cfg = PipelineConfig(model_type=new_model_type)
+        pipe_cfg = PipelineConfig(
+            model_type=new_model_type,
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
 
         # Register the new PipelineConfig
         assert new_model_type not in OMNI_PIPELINES
@@ -937,7 +986,10 @@ class TestPipelineRegistration:
         def custom_resolver(
             hf_config: FakeConfig,
         ) -> PipelineConfig:
-            return PipelineConfig(model_type=resolved_type)
+            return PipelineConfig(
+                model_type=resolved_type,
+                stages=SINGLE_STAGE_PIPE_CFG,
+            )
 
         # Register the new PipelineConfig
         assert new_model_type not in OMNI_PIPELINES
@@ -994,7 +1046,7 @@ class TestPipelineRegistration:
         pipe_cfg = PipelineConfig(
             model_type=pipeline_key,
             model_arch="DeployOnlyArch",
-            stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),),
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
         register_pipeline(pipe_cfg)
 
@@ -1027,7 +1079,11 @@ class TestPipelineRegistration:
 
     def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path):
         pipeline_key = "single_load_pipeline"
-        register_pipeline(PipelineConfig(model_type=pipeline_key))
+        pipeline_cfg = PipelineConfig(
+            model_type=pipeline_key,
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        register_pipeline(pipeline_cfg)
         deploy_path = tmp_path / "single_load.yaml"
         deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8")
 
@@ -1057,6 +1113,7 @@ class TestPipelineRegistration:
         pipeline = PipelineConfig(
             model_type="pipeline_with_default",
             default_deploy_config_name=default_name,
+            stages=SINGLE_STAGE_PIPE_CFG,
         )
 
         with patch(
@@ -1076,8 +1133,16 @@ class TestPipelineRegistration:
             OmniServingCapability.COMPLETIONS,
             "pipeline_a blocks completions",
         )
-        pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
+        pipe_a = PipelineConfig(
-        pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
+            model_type="detect_type",
+            endpoint_restrictions=(restriction,),
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        pipe_b = PipelineConfig(
+            model_type="override_type",
+            endpoint_restrictions=(),
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
         register_pipeline(pipe_a)
         register_pipeline(pipe_b)
 
@@ -1617,6 +1682,7 @@ stages:
                     model_stage="ar",
                     execution_type=StageExecutionType.LLM_AR,
                     requires_multimodal_data=True,
+                    final_output=True,
                 ),
             ),
         )
@@ -1637,6 +1703,7 @@ stages:
                     model_stage="ar",
                     execution_type=StageExecutionType.LLM_AR,
                     scheduler_cls=scheduler_cls,
+                    final_output=True,
                 ),
             ),
         )
@@ -1882,7 +1949,6 @@ class TestQwen3OmniPipeline:
         assert isinstance(p, PipelineConfig)
         assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
         assert len(p.stages) == 3
-        assert p.validate() == []
 
     def test_thinker(self):
         p = resolve_pipeline_config(
@@ -1936,7 +2002,6 @@ class TestQwen2_5OmniPipeline:
         assert isinstance(p, PipelineConfig)
         assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
         assert len(p.stages) == 3
-        assert p.validate() == []
 
     def test_thinker(self):
         p = resolve_pipeline_config("qwen2_5_omni")
@@ -1985,7 +2050,6 @@ class TestQwen3TTSPipeline:
         assert p is not None
         assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
         assert len(p.stages) == 2
-        assert p.validate() == []
 
     def test_talker_stage(self):
         p = resolve_pipeline_config("qwen3_tts")
@@ -2065,7 +2129,6 @@ class TestMingFlashOmniPipeline:
         assert isinstance(p, PipelineConfig)
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
         assert len(p.stages) == 2
-        assert p.validate() == []
 
     def test_thinker_stage(self):
         p = resolve_pipeline_config("ming_flash_omni")
@@ -2122,7 +2185,6 @@ class TestMingFlashOmniPipeline:
         assert isinstance(p, PipelineConfig)
         assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
         assert len(p.stages) == 1
-        assert p.validate() == []
 
     def test_tts_stage(self):
         p = resolve_pipeline_config("ming_flash_omni_tts")
@@ -2179,7 +2241,6 @@ class TestMingFlashOmniPipeline:
         assert isinstance(p, PipelineConfig)
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
         assert len(p.stages) == 1
-        assert p.validate() == []
 
     def test_thinker_only_stage(self):
         p = resolve_pipeline_config("ming_flash_omni_thinker_only")
@@ -2218,7 +2279,6 @@ class TestMingFlashOmniPipeline:
         assert p is not None
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
         assert len(p.stages) == 2
-        assert p.validate() == []
 
     def test_image_thinker_stage(self):
         s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0)

tests/config/test_omni_config.py

@@ -45,6 +45,7 @@ from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageDeployConfig,
     StageExecutionType,
+    StagePipelineConfig,
     load_deploy_config,
     merge_pipeline_deploy,
 )
@@ -1039,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
 
 
 def test_from_pipeline_config_accepts_pre_resolved_pipeline():
-    resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant")
+    resolved_pipeline = PipelineConfig(
+        model_type="callable_resolved_variant",
+        stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
+    )
 
     omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
 
@@ -1072,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p
     pipeline = PipelineConfig(
         model_type="pipeline_with_default",
         default_deploy_config_name=default_name,
+        stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
     )
     loaded_paths = []
 

tests/config/test_pipeline_registry.py

@@ -1,16 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for out of tree registration to OMNI_PIPELINES."""
 
 import pytest
 from transformers import PretrainedConfig
 
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
-from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
+from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
+    return PipelineConfig(
+        model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+    )
+
+
 @pytest.fixture
 def custom_resolver():
     """Build a reusable custom resolver for PipelineConfigs."""
@@ -22,7 +28,7 @@ def custom_resolver():
     def custom_resolver(
         hf_config: CustomConfigType,
     ) -> PipelineConfig:
-        return PipelineConfig(model_type="resolved_type")
+        return build_fake_pipeline_config("resolved_type")
 
     return custom_resolver
 
@@ -30,7 +36,7 @@ def custom_resolver():
 def test_register_pipeline_config(clean_pipeline_registry):
     """Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
     new_model_type = "new_model_type"
-    pipe_cfg = PipelineConfig(model_type=new_model_type)
+    pipe_cfg = build_fake_pipeline_config(new_model_type)
     assert new_model_type not in OMNI_PIPELINES
     register_pipeline(pipe_cfg)
     assert new_model_type in OMNI_PIPELINES
@@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
     """Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
     new_model_type = "new_model_type"
     unused_model_type = "foo"
-    pipe_cfg = PipelineConfig(model_type=unused_model_type)
+    pipe_cfg = build_fake_pipeline_config(unused_model_type)
     assert new_model_type not in OMNI_PIPELINES
     assert unused_model_type not in OMNI_PIPELINES
 

tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py

@@ -70,10 +70,14 @@ _PARAMS = [
         pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
         pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
         pytest.param(
-            {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
+            {"num_inference_steps": 0},
+            ("num_inference_steps", "greater than or equal to 1"),
+            id="num_inference_steps_zero",
         ),
         pytest.param(
-            {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
+            {"num_inference_steps": -1},
+            ("num_inference_steps", "greater than or equal to 1"),
+            id="num_inference_steps_negative",
         ),
         pytest.param(
             {"num_inference_steps": 6000},

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -157,7 +157,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
                 "ref_audio": None,
                 "ref_text": None,
             },
-            ("Invalid voice", "vivian"),
+            # NOTE: This is delegating down the TTS adapter, so the message is both specific
+            # to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
+            # (as opposed to having valida speakers, but getting an invalid speaker). it would
+            # be more ideal to validate that this is landing on `.validate()` on the adapter
+            # and ensuring that the returned message matches what we would get from calling the
+            # adapter validate directly.
+            ("no speakers configured",),
             id="customvoice_invalid_voice",
         ),
         pytest.param(

tests/diffusion/attention/test_attention_config.py

@@ -11,6 +11,7 @@ Tests cover:
 """
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -466,6 +467,120 @@ class TestCurrentDiffusionConfig:
 
 
 class TestAttentionInitUsesCurrentDiffusionConfig:
+    @pytest.mark.parametrize("mindiesd_available", [False, True])
+    def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
+        attention = Attention.__new__(Attention)
+        attention._scheduler_paged_kv = True
+        attention.paged_kv_cache_role = "primary"
+        attention.backend_pref = "FLASH_ATTN"
+        attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+        flash_output = torch.ones(1)
+        sdpa_output = torch.zeros(1)
+        dense_flash = Mock(return_value=flash_output)
+        sdpa = Mock(return_value=sdpa_output)
+        attention.attention = SimpleNamespace(forward=dense_flash)
+        attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+        monkeypatch.setattr(
+            layer_mod,
+            "get_forward_context",
+            lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
+        )
+        monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True)
+        monkeypatch.setattr(
+            layer_mod.current_omni_platform,
+            "supports_diffusion_dense_flash_attention",
+            lambda: mindiesd_available,
+        )
+
+        result = attention._run_local_attention(
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            AttentionMetadata(),
+        )
+
+        if mindiesd_available:
+            assert result is flash_output
+            dense_flash.assert_called_once()
+            sdpa.assert_not_called()
+        else:
+            assert result is sdpa_output
+            dense_flash.assert_not_called()
+            sdpa.assert_called_once()
+
+    def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
+        attention = Attention.__new__(Attention)
+        attention._scheduler_paged_kv = False
+        attention.paged_kv_cache_role = "primary"
+        attention.backend_pref = "FLASH_ATTN"
+        attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+        attention.attention = SimpleNamespace(
+            forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd"))
+        )
+        attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back")))
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        with pytest.raises(ModuleNotFoundError, match="mindiesd"):
+            attention._run_local_attention(
+                torch.zeros(1, dtype=torch.bfloat16),
+                torch.zeros(1, dtype=torch.bfloat16),
+                torch.zeros(1, dtype=torch.bfloat16),
+                AttentionMetadata(),
+            )
+
+    def test_paged_marker_does_not_change_dense_backend(self, monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _DenseBackend:
+            supports_paged_kv = False
+
+            @staticmethod
+            def get_name() -> str:
+                return "TORCH_SDPA"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        calls = []
+
+        def _fake_get_attn_backend_for_role(*args, **kwargs):
+            calls.append((args, kwargs))
+            return _DenseBackend, None
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY,
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attention = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                paged_kv_cache_role="primary",
+            )
+
+        assert len(calls) == 1
+        assert attention.attn_backend is _DenseBackend
+        assert attention.backend_pref is None
+
     def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
         class _FakeAttentionImpl:
             def __init__(self, **kwargs):

tests/diffusion/attention/test_flash_attn.py

@@ -353,8 +353,15 @@ def test_cross_attn_key_padding_vs_sdpa(k_len):
     print(f"Max absolute difference: {max_diff:.6f}")
     print(f"Mean absolute difference: {mean_diff:.6f}")
 
-    assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+    # Match the platform-specific comparison used by test_fa_vs_sdpa above.
-    assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+    if current_omni_platform.is_rocm():
+        # AITER FlashAttention and PyTorch SDPA use different BF16 reduction
+        # implementations on ROCm, so individual values can differ by one BF16
+        # step even though the outputs agree within the expected relative error.
+        torch.testing.assert_close(output_fa, output_sdpa, rtol=1e-2, atol=1.5e-2)
+    else:
+        assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+        assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
 
     print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
 

tests/diffusion/attention/test_piecewise_attn.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """End-to-end test for ``piecewise_attn`` (CPU).
 
 Verify that running attention in segments (causal outside full-attn spans,
@@ -98,9 +98,12 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
     query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
     key = query + 100
     value = query + 200
-    calls = []
+    calls: list[object] = []
 
-    def run_segment(segment_query, segment_key, segment_value, segment_batch):
+    output_buffer = torch.empty_like(query)
+
+    def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer):
+        assert segment_output_buffer is None
         calls.append(
             (
                 [segment.mode for segment in segment_batch.row_segments],
@@ -119,9 +122,13 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
         plan,
         plan.segments,
         run_segment,
+        output_buffer=output_buffer,
     )
 
+    assert output is output_buffer
     torch.testing.assert_close(output, query)
+    assert all(segment.query_range is None for segment in plan.segments)
+    assert not plan.segments_cover_query_contiguously
     assert calls == [
         (["causal", "causal"], [2, 4], [0, 1, 6, 7]),
         (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
@@ -129,6 +136,161 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
     ]
 
 
+def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch):
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)], [(2, 5)]],
+        query_offsets=[0, 0],
+        query_lens=[6, 6],
+        seq_lens=[6, 6],
+        device=DEVICE,
+    )
+    assert plan.homogeneous_batch_shape == (2, 6)
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    calls: list[torch.Tensor] = []
+
+    def fail_index_copy(*_args, **_kwargs):
+        raise AssertionError("homogeneous piecewise path must not use index_copy_")
+
+    monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy)
+
+    def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer):
+        assert segment_key is None
+        assert segment_value is None
+        assert segment_output_buffer is None
+        calls.append(segment_query.clone())
+        return segment_query + 1
+
+    output_buffer = torch.empty_like(query)
+    output = run_paged_piecewise_plan(
+        query,
+        None,
+        None,
+        plan,
+        plan.segments,
+        run_segment,
+        output_buffer=output_buffer,
+        use_homogeneous_batch=True,
+    )
+
+    assert output is output_buffer
+    torch.testing.assert_close(output, query + 1)
+    assert [call.shape[0] for call in calls] == [4, 6, 2]
+    torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]])
+    torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]])
+    torch.testing.assert_close(calls[2], query[[5, 11]])
+
+
+def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)], [(1, 4)]],
+        query_offsets=[0, 0],
+        query_lens=[6, 6],
+        seq_lens=[6, 6],
+        device=DEVICE,
+    )
+    assert plan.homogeneous_batch_shape is None
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    output = run_paged_piecewise_plan(
+        query,
+        None,
+        None,
+        plan,
+        plan.segments,
+        lambda segment_query, *_args: segment_query,
+        use_homogeneous_batch=True,
+    )
+    torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_slices_contiguous_segments():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+    key = query + 100
+    value = query + 200
+    expected_ranges = [(0, 2), (2, 5), (5, 6)]
+    calls: list[tuple[int, int]] = []
+
+    def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer):
+        assert segment_output_buffer is None
+        expected_range = expected_ranges[len(calls)]
+        start, end = expected_range
+        calls.append(expected_range)
+        assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr()
+        assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr()
+        assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr()
+        torch.testing.assert_close(segment_query, query[start:end])
+        return segment_query
+
+    output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment)
+
+    assert calls == expected_ranges
+    assert [segment.query_range for segment in plan.segments] == expected_ranges
+    assert plan.segments_cover_query_contiguously
+    torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+    output_buffer = torch.empty(6, 2, 3)
+    targets = []
+
+    def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer):
+        assert segment_output_buffer is not None
+        assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr()
+        targets.append(segment_output_buffer)
+        return segment_output_buffer.copy_(segment_query[..., :3])
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        run_segment,
+        output_buffer=output_buffer,
+    )
+
+    assert output is output_buffer
+    assert [target.shape[0] for target in targets] == [2, 3, 1]
+    torch.testing.assert_close(output, query[..., :3])
+
+
+def test_paged_piecewise_runner_returns_single_full_segment_directly():
+    plan = build_paged_piecewise_plan(
+        [[]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.randn(6, 2, 4)
+    runner_output = query + 1
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        lambda *_args: runner_output,
+    )
+
+    assert output is runner_output
+
+
 def test_paged_piecewise_runner_uses_runner_output_shape():
     plan = build_paged_piecewise_plan(
         [[(2, 5)]],

tests/diffusion/batching/test_diffusion_batching.py

@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
 This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
 diffusion model and validates that every caller receives its correct
 individual result.  When the underlying diffusion stage is configured with
-``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
+``max_num_seqs > 1``, the
 requests will be batched internally.
 
 Even without explicit batching config this test is useful for verifying
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
 ) -> None:
     """Run the same prompts sequentially then in parallel and print a comparison."""
 
-    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
         single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
         await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
         return
 
-    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
 
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
     fired via asyncio.gather. Each request_id must appear in the results."""
 
     async def _inner():
-        omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
+        omni = AsyncOmni(model=model_name, max_num_seqs=1)
         try:
             prompts = TEST_PROMPTS[:4]
             sp = _default_sampling_params()
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
     """
 
     async def _inner():
-        omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
+        omni = AsyncOmni(model=model_name, max_num_seqs=4)
         try:
             prompts = TEST_PROMPTS[:4]
             sp = _default_sampling_params()
@@ -528,7 +528,7 @@ if __name__ == "__main__":
         "--batch-size",
         type=int,
         default=1,
-        help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
+        help="Scheduler in-flight width (--max-num-seqs)",
     )
     parser.add_argument(
         "--mode",

tests/diffusion/conftest.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def restore_torch_default_dtype():
+    """Keep process-global dtype changes from leaking between diffusion tests."""
+    default_dtype = torch.get_default_dtype()
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(default_dtype)

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -485,8 +485,16 @@ def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest
     assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
     sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
     context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
-    assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
+    assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == (
-    assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
+        0,
+        9,
+        ((1, 2, 3), (4, 5)),
+    )
+    assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
+        1,
+        5,
+        ((6, 7), (8,)),
+    )
 
 
 def test_block_table_mutations_invalidate_prepared_attention_batches(

tests/diffusion/diffusion_kv/test_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
 
@@ -36,7 +36,7 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None:
         OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
 
 
-def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
+def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None:
     with pytest.raises(ValueError, match="does not support imported AR KV"):
         OmniDiffusionConfig.from_kwargs(
             diffusion_kv_mode="paged_scheduler",

tests/diffusion/diffusion_kv/test_initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -290,6 +290,7 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
         dtype=torch.bfloat16,
         non_causal=True,
     )
+    model_limit = vllm_config.model_config.max_model_len
 
     initialization.build_native_kv_cache_configs(
         vllm_config,
@@ -298,4 +299,8 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
     )
 
     assert vllm_config.model_config.original_max_model_len == -1
-    assert vllm_config.model_config.max_model_len == 256
+    # Sixteen physical pages provide ``16 * block_size`` tokens.  The native
+    # backend chooses the block geometry, so keep this assertion backend
+    # independent instead of assuming the generic 16-token block size.
+    expected_max_model_len = min(model_limit, spec.block_size * 16)
+    assert vllm_config.model_config.max_model_len == expected_max_model_len

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -20,12 +20,16 @@ from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRow,
     DiffusionPagedAttentionRowBinding,
+    DiffusionPagedAttentionRuntime,
 )
 from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
     override_paged_kv_adapter,
     set_forward_context,
+    set_forward_context_denoise_step_idx,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -122,8 +126,8 @@ class _FakeNativeImpl:
     def do_kv_cache_update(
         self,
         _layer,
-        key: torch.Tensor,
+        key: torch.Tensor | None,
-        value: torch.Tensor,
+        value: torch.Tensor | None,
         _kv_cache,
         slot_mapping: torch.Tensor,
     ) -> None:
@@ -155,6 +159,11 @@ def _make_adapter(
     capacity: int = 16,
 ) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
     events: list[tuple] = []
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "requires_diffusion_paged_kv_prewrite",
+        lambda: False,
+    )
     block_tables = _FakeBlockTables()
     block_tables.max_num_batched_tokens = capacity
     block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
@@ -301,6 +310,47 @@ def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch
     assert context.value_write.shape == (3, 8, 4)
 
 
+def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+    )
+    packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
+    query, key, value = packed_qkv.split(2 * 4, dim=-1)
+    query = query.reshape(1, 5, 2, 4)
+    key = key.reshape(1, 5, 2, 4)
+    value = value.reshape(1, 5, 2, 4)
+    assert not key.is_contiguous()
+    assert not value.is_contiguous()
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+    assert context.query.is_contiguous()
+    assert not context.key_write.is_contiguous()
+    assert not context.value_write.is_contiguous()
+    assert context.key_write.data_ptr() == key.data_ptr()
+    assert context.value_write.data_ptr() == value.data_ptr()
+
+
+def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+    )
+    query = torch.randn(1, 5, 2, 4)
+    key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+    value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+    assert key.shape == (1, 5, 2, 4)
+    assert key.stride(-1) != 1
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+    assert context.key_write.is_contiguous()
+    assert context.value_write.is_contiguous()
+
+
 def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, _ = _make_adapter(monkeypatch)
     layer.attn_backend.forward_includes_kv_cache_update = True
@@ -350,10 +400,48 @@ def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: p
     assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
     assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
     assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
+    for event in events[1:]:
+        build_kwargs = event[1]
+        assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
+        assert (
+            build_kwargs["slot_mappings"].untyped_storage().data_ptr()
+            == batch.slot_mappings.untyped_storage().data_ptr()
+        )
+
+
+def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, layer, _ = _make_adapter(monkeypatch)
+    layer.attn_backend.forward_includes_kv_cache_update = True
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "requires_diffusion_paged_kv_prewrite",
+        lambda: True,
+    )
+    adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
+        row_index=2,
+        max_seq_len=16,
+        block_ids=((3, 7, 9, 11),),
+    )
+    row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
+    batch = adapter.prepare_batch([row])
+    qkv = torch.randn(1, 4, 2, 4)
+    metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
+        output = _run_omni_paged_backend(context)
+
+    assert torch.equal(output, qkv)
+    assert layer.native_events == ["update", "forward", "forward", "forward"]
+    assert len(layer.updates) == 1
+    assert all(call[1] is None and call[2] is None for call in layer.calls)
 
 
 def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
+    # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
+    # output-buffer contract and is covered by the GPU adapter tests.
+    monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
     batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -450,6 +538,48 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     ]
 
 
+def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
+    build_id = 0
+
+    def build_metadata(**kwargs):
+        nonlocal build_id
+        shared_scheduler_metadata.fill_(build_id)
+        metadata = SimpleNamespace(
+            build_id=build_id,
+            causal=kwargs["causal"],
+            seq_lens=kwargs["seq_lens"].clone(),
+            query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
+            positions=kwargs["positions"].clone(),
+            slot_mappings=kwargs["slot_mappings"].clone(),
+            scheduler_metadata=shared_scheduler_metadata,
+            prefix_scheduler_metadata=None,
+        )
+        build_id += 1
+        return {"layer-0": metadata}
+
+    monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+    )
+    qkv = torch.randn(1, 4, 2, 4)
+    metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(
+            "layer-0",
+            qkv,
+            qkv,
+            qkv,
+            omni_attn_metadata=metadata,
+        )
+
+    segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
+    assert [item.item() for item in segment_schedules] == [1, 2, 3]
+    assert len({item.data_ptr() for item in segment_schedules}) == 3
+
+
 def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
     batch = adapter.prepare_batch(
@@ -851,6 +981,28 @@ def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: py
         adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
 
 
+def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
+            DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
+        ]
+    )
+    qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+    assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
+    torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
+    restored = context.restore_output(context.query)
+    assert restored.shape == qkv.shape
+    torch.testing.assert_close(restored[0, 0], qkv[0, 0])
+    assert torch.count_nonzero(restored[0, 1:]) == 0
+    torch.testing.assert_close(restored[1], qkv[1])
+
+
 def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch)
     batch = adapter.prepare_batch(
@@ -878,6 +1030,41 @@ def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None
         adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
 
 
+def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    metadata = DiffusionPagedAttentionMetadata(
+        prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
+        denoise_rows=(
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=2,
+                seq_len=5,
+                kv_start_pos=3,
+            ),
+        ),
+    )
+    runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
+
+    with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
+        set_forward_context_denoise_step_idx(0)
+        assert Attention._active_paged_kv_adapter() is adapter
+        assert adapter._active_batch.rows == metadata.prefill_rows
+
+        set_forward_context_denoise_step_idx(1)
+        assert Attention._active_paged_kv_adapter() is adapter
+        denoise_batch = adapter._active_batch
+        assert denoise_batch.rows == metadata.denoise_rows
+
+        set_forward_context_denoise_step_idx(2)
+        assert Attention._active_paged_kv_adapter() is adapter
+        assert adapter._active_batch is denoise_batch
+
+    assert adapter._active_batch is None
+
+
 def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
     row = DiffusionPagedAttentionRow(
@@ -992,7 +1179,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     selected_backends = []
-    impl_cls = Mock(return_value=SimpleNamespace(forward=Mock()))
+    specialized_backends = []
+    impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
     native_backend = SimpleNamespace(
         get_name=lambda: "ASCEND",
         indexes_kv_by_block_stride=lambda: True,
@@ -1003,16 +1191,28 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     original_backend_per_kind = {"full": object()}
     config = SimpleNamespace(
         attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
         model_config=SimpleNamespace(dtype=torch.float16),
         cache_config=SimpleNamespace(cache_dtype="auto"),
     )
 
     def select_backend(**_kwargs):
-        selected_backends.append((config.attention_config.backend, config.attention_config.backend_per_kind))
+        selected_backends.append(
+            (
+                config.attention_config.backend,
+                config.attention_config.backend_per_kind,
+                config.parallel_config.prefill_context_parallel_size,
+            )
+        )
         return native_backend
 
     monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
     monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+    monkeypatch.setattr(
+        adapter_module.current_omni_platform,
+        "get_diffusion_paged_kv_attn_backend",
+        lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+    )
     layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
     spec = FullAttentionSpec(
         block_size=16,
@@ -1031,9 +1231,11 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         ulysses_degree=2,
     )
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
+    assert specialized_backends == [(native_backend, 2)]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind
+    assert config.parallel_config.prefill_context_parallel_size == 2
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
@@ -1094,6 +1296,51 @@ def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
     assert torch.equal(output, torch.full_like(original, 6))
 
 
+def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
+    events: list[tuple[str, tuple[int, ...]]] = []
+
+    class Strategy:
+        name = "ulysses"
+
+        def pre_attention(self, query, key, value, metadata):
+            return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
+
+        def post_attention(self, output, _ctx):
+            events.append(("post", tuple(output.shape)))
+            return output
+
+    class Adapter:
+        def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
+            assert omni_attn_metadata is None
+            events.append(("kernel_input", tuple(query.shape)))
+            assert query[:, -1].tolist() == [[[3.0]]]
+            assert torch.equal(query, key)
+            assert torch.equal(query, value)
+            return SimpleNamespace(query=query)
+
+    class Backend:
+        def forward_paged(self, context):
+            return context.query
+
+    layer = Attention.__new__(Attention)
+    nn.Module.__init__(layer)
+    layer.prefix = "layer-0"
+    layer.paged_kv_cache_role = "primary"
+    layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
+    layer.attention = Backend()
+    layer.use_ring = False
+    layer._no_parallel_strategy = object()
+    layer._get_active_parallel_strategy = lambda: Strategy()
+    tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
+
+    with set_forward_context(), override_paged_kv_adapter(Adapter()):
+        get_forward_context().sp_padding_size = 1
+        output = layer._forward_impl(tensor, tensor, tensor)
+
+    assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
+    assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
+
+
 def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
     class Adapter:
         def prepare_layer_context(self, *args, **kwargs):
@@ -1145,3 +1392,15 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
 
     assert events == ["pre", "dense", "post"]
     assert output is qkv
+
+
+def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
+    layer = Attention.__new__(Attention)
+    nn.Module.__init__(layer)
+    layer._scheduler_paged_kv = True
+    layer.paged_kv_cache_role = "primary"
+    layer._get_active_parallel_strategy = lambda: object()
+    qkv = torch.zeros(1, 2, 2, 4)
+
+    with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
+        layer._forward_impl(qkv, qkv, qkv)

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -154,8 +154,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
 
     torch.manual_seed(1)
     query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
-    key = torch.randn_like(query)
+    packed_qkv = torch.randn(
-    value = torch.randn_like(query)
+        19,
+        3 * _NUM_HEADS * _HEAD_SIZE,
+        dtype=torch.float16,
+        device=device,
+    )
+    _, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
+    key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+    value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+    assert not key.is_contiguous()
+    assert not value.is_contiguous()
     omni_backend = FlashAttentionImpl(
         num_heads=_NUM_HEADS,
         head_size=_HEAD_SIZE,

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import nullcontext
 from types import SimpleNamespace
@@ -42,9 +42,10 @@ class _RunnerKVBackend:
         self.kv_cache_config = config
 
 
-def _attention(*, enabled: bool) -> Attention:
+def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention:
     attention = Attention.__new__(Attention)
     nn.Module.__init__(attention)
+    attention.prefix = prefix
     attention.paged_kv_cache_role = "primary" if enabled else None
     attention.paged_kv_cache_dtype = torch.bfloat16
     attention.num_kv_heads = 2
@@ -82,6 +83,23 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
     assert spec.non_causal is True
 
 
+def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None:
+    runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn"))
+
+    specs = runner.get_kv_cache_spec()
+
+    assert set(specs) == {"layers.0.self_attn.image_attn.attn"}
+
+
+def test_runner_rejects_duplicate_attention_prefixes() -> None:
+    runner = _runner(_attention(enabled=True, prefix="shared"))
+    second_attention = _attention(enabled=True, prefix="shared")
+    runner.pipeline.second_image_attention = second_attention
+
+    with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"):
+        runner.get_kv_cache_spec()
+
+
 def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
     runner = _runner(_attention(enabled=False))
 

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
 
 def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
     runner = object.__new__(DiffusionModelRunner)
-    runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+    runner.od_config = SimpleNamespace(
+        diffusion_kv_mode=mode,
+        parallel_config=SimpleNamespace(cfg_parallel_size=1),
+    )
     return runner
 
 
@@ -85,6 +88,51 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
     assert new_req.diffusion_kv_metadata is metadata
 
 
+def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+            DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    assert [
+        (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+        for row in attn_metadata.prefill_rows
+    ] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)]
+    assert [
+        (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+        for row in attn_metadata.denoise_rows
+    ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
+
+
+def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.od_config.parallel_config.cfg_parallel_size = 2
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank",
+        lambda: 1,
+    )
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+            DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1]
+    assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1]
+
+
 def test_paged_request_without_metadata_fails_before_request_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner._execute_request_list = Mock()

tests/diffusion/distributed/test_a2a_permute.py

@@ -43,21 +43,16 @@ def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatc
 
 @dataclass
 class _FakeAllocation:
+    """Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``."""
+
     size: int
     device: torch.device
 
+    def __post_init__(self) -> None:
+        self.buffer = torch.empty(self.size, dtype=torch.uint8)
 
-class _FakeHandle:
+    def __getitem__(self, item: slice) -> torch.Tensor:
-    rank = 0
+        return self.buffer[item]
-
-    def __init__(self, allocation: _FakeAllocation) -> None:
-        self.allocation = allocation
-        self.views: list[tuple[tuple[int, ...], torch.dtype]] = []
-
-    def get_buffer(self, rank, shape, dtype):
-        assert rank == self.rank
-        self.views.append((shape, dtype))
-        return self.allocation, shape, dtype
 
 
 def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
@@ -93,16 +88,20 @@ def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
         return allocation
 
     monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
-    monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda allocation, _group: _FakeHandle(allocation))
+    # The handle is only retained to keep the rendezvous registration alive.
+    monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object())
 
     device = torch.device("cuda:0")
     first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
     smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
     larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
 
-    assert first[0].size == 12
+    assert (first.shape, first.dtype) == ((2, 3), torch.float16)
-    assert smaller[0] is first[0]
+    assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16)
-    assert larger[0].size == 24
+    assert (larger.shape, larger.dtype) == ((4, 3), torch.float16)
+    # A sub-capacity request reuses the workspace; only growth allocates again.
+    assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr()
+    assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr()
     assert [allocation.size for allocation in allocations] == [12, 24]
     assert len(a2a_permute._SYMM_WORKSPACES) == 1
     assert len(all_reduce_calls) == 1

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """L1 unit tests for the native Boogu-Image pipeline.
 
@@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
     assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
 
 
+def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    revision = "hotfix-1k-20260708"
+    prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders")
+    scheduler_loader = mocker.patch(
+        f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=mock_dependencies["scheduler"],
+    )
+    mllm_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
+        return_value=mock_dependencies["mllm_wrapper"],
+    )
+    processor_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
+        return_value=mock_dependencies["processor"],
+    )
+    vae_loader = mocker.patch(
+        f"{_MODULE}.AutoencoderKL.from_pretrained",
+        return_value=mock_dependencies["vae"],
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        revision=revision,
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        num_gpus=1,
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    (source,) = pipeline.weights_sources
+    assert source.revision == revision
+    prefetch.assert_called_once_with(
+        "dummy-boogu",
+        ["scheduler", "vae", "mllm", "processor"],
+        local_files_only=True,
+        revision=revision,
+    )
+    for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader):
+        assert loader.call_args.kwargs["revision"] == revision
+
+
 def test_constructor_weights_sources(boogu_pipeline):
     (source,) = boogu_pipeline.weights_sources
     assert source.model_or_path == "dummy-boogu"

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
@@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend():
     pipeline._ensure_grouped_attention_backend_supported(2)
 
 
+def test_scheduler_paged_step_execution_is_rejected():
+    pipeline = _pipeline()
+    pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER
+    state = _state("paged-step", 0)
+
+    with pytest.raises(ValueError, match="request-level execution only"):
+        pipeline.prepare_encode(state)
+    with pytest.raises(ValueError, match="request-level execution only"):
+        pipeline.denoise_step(InputBatch.make_batch([state]))
+
+
 def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
     pipeline = _pipeline()
     pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for ImageKVCacheManager.
 
 Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
@@ -37,12 +37,21 @@ class MockAttention(nn.Module):
         super().__init__()
         self.paged_kv_active = False
         self.calls = []
+        self.paged_calls = []
+        self.paged_metadata = []
 
     def is_paged_kv_active(self):
         return self.paged_kv_active
 
     def forward(self, query, key, value, attn_metadata=None, **kwargs):
         self.calls.append((query, key, value))
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None:
+            self.paged_calls.append((query, key, value))
+            self.paged_metadata.append(attn_metadata)
+        if attn_metadata is not None and attn_metadata.joint_query is not None:
+            return torch.cat((attn_metadata.joint_query, query), dim=1)
         return query
 
 
@@ -105,6 +114,7 @@ def _call_mgr(
     shard_image_size=None,
     gen_timestep_scatter_index=None,
     position_ids=None,
+    full_attn_spans=None,
 ):
     query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
     attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -121,6 +131,7 @@ def _call_mgr(
         shard_image_size=shard_image_size,
         gen_timestep_scatter_index=gen_timestep_scatter_index,
         position_ids=position_ids,
+        full_attn_spans=full_attn_spans,
     )
 
 
@@ -132,6 +143,83 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    bs = 2
+    prefix_lens = torch.tensor([[2], [4]], dtype=torch.long)
+    first_q_len = 12
+    first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0)
+    full_attn_spans = [[(2, 12)], [(2, 12)]]
+
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        _call_mgr(
+            mgr,
+            bs,
+            q_len=first_q_len,
+            seq_len=first_q_len,
+            key_flat=first_k,
+            value_flat=first_v,
+            first_step=True,
+            gen_timestep_scatter_index=prefix_lens,
+            full_attn_spans=full_attn_spans,
+        )
+
+        assert mgr.image_kv_cache_map is None
+        assert mgr.image_kv_cache_lens is None
+        assert len(mgr.attn.paged_calls) == 1
+        assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+        assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans
+
+        current_q_len = IMAGE_TOKEN_LEN
+        current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0)
+        _call_mgr(
+            mgr,
+            bs,
+            q_len=current_q_len,
+            seq_len=prefix_lens[-1].item() + current_q_len,
+            key_flat=current_k,
+            value_flat=current_v,
+            position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1),
+            full_attn_spans=full_attn_spans,
+        )
+
+    assert len(mgr.attn.paged_calls) == 2
+    assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans
+
+
+def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    prefix_len = 4
+    prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0)
+    mgr._injected_ar_kv = [(prefix_key, prefix_value)]
+    current_key, current_value = _make_known_kv(3, base=1.0)
+
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        with pytest.raises(NotImplementedError, match="does not support imported AR KV"):
+            _call_mgr(
+                mgr,
+                bs=1,
+                q_len=3,
+                seq_len=7,
+                key_flat=current_key,
+                value_flat=current_value,
+                first_step=True,
+                gen_timestep_scatter_index=torch.tensor([[0]]),
+                full_attn_spans=[[(4, 7)]],
+            )
+
+
 @pytest.mark.parametrize(
     ("paged_kv_active", "expected_kv_heads"),
     [(False, NUM_HEADS), (True, NUM_KV_HEADS)],
@@ -155,6 +243,7 @@ def test_gqa_kv_stays_compressed_for_paged_attention(
         first_step=True,
         num_image_tokens=0,
         gen_timestep_scatter_index=_gen_timestep_index(1, 0),
+        full_attn_spans=[[(0, q_len)]] if paged_kv_active else None,
     )
 
     _, key_input, value_input = mgr.attn.calls[-1]

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -57,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
         return torch.zeros(1), torch.zeros(1)
 
     pipeline.diffuse = diffuse
-    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    decoded_audio = torch.zeros(1)
+    pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio))
     sampling = OmniDiffusionSamplingParams(
         quality="high",
         width=1344,
@@ -86,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
         num_inference_steps=50,
         extra_args={"task": "t2va", "aspect_ratio": "16:9"},
     )
-    assert output.output == pipeline.decode.return_value
+    assert output.output[0].shape == (1, 1, 1, 1, 3)
+    assert output.output[0].dtype == torch.uint8
+    assert output.output[1] is decoded_audio
 
 
 def test_pipeline_import_registry_and_component_discovery():
@@ -493,18 +496,41 @@ def test_joint_postprocess_is_multiprocessing_picklable():
 
     postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
     postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
-    video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+    # decode() hands over quantised frames already laid out as (B, T, H, W, C).
+    video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3)
     audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
 
     result = postprocess((video, audio), output_type="np")
 
     assert isinstance(result["video"], list)
+    assert result["video"][0].dtype == np.uint8
     assert result["video"][0].shape == (2, 4, 5, 3)
+    np.testing.assert_array_equal(result["video"][0], video[0].numpy())
     np.testing.assert_array_equal(result["audio"], audio.numpy())
     assert result["audio_sample_rate"] == 32000
     assert result["fps"] == 24
 
 
+def test_joint_video_output_is_quantized_before_transfer():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _minimax_h3_post_process,
+        _prepare_minimax_h3_video_output,
+    )
+
+    video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
+    expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
+    prepared = _prepare_minimax_h3_video_output(video)
+
+    assert prepared.dtype == torch.uint8
+    assert prepared.is_contiguous()
+    assert prepared.shape == (2, 2, 4, 5, 3)
+    torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
+
+    result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
+    assert result["video"][0].dtype == np.uint8
+    np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
+
+
 def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -572,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions():
 @pytest.mark.parametrize(
     "base_schedule",
     [
-        [0.9, 0.5, 0.0],
+        [1.5, 0.5, 0.0],
         [1.0, 0.5, 0.1],
         [1.0, 0.5, 0.5, 0.0],
         [1.0],
@@ -627,7 +653,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
         return torch.zeros(1), torch.zeros(1)
 
     pipeline.diffuse = diffuse
-    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1)))
     return pipeline
 
 

tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py

@@ -0,0 +1,498 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.minimax_h3.fasth3 import (
+    FASTH3_BASE_MODEL,
+    FASTH3_BASE_SCHEDULE,
+    FASTH3_DENOISE_STEPS,
+    FASTH3_FORMAT,
+    FastH3AdapterError,
+    FastH3WeightFusion,
+    resolve_fasth3_fusion,
+)
+from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
+from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
+from vllm_omni.errors import OmniClientError
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_RANK = 2
+_HIDDEN = 4
+_HEAD_DIM = 2
+_HEADS = 3
+_INNER = _HEAD_DIM * _HEADS  # attention inner size
+_FFN = 5
+
+
+def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """A rank-2 pair whose product is reproducible and non-symmetric."""
+    a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK)
+    b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK)
+    return a, b
+
+
+def _write_adapter(
+    path,
+    *,
+    tensors=None,
+    drop: str | None = None,
+    blocks: int = 1,
+    metadata: dict[str, str | None] | None = None,
+) -> None:
+    """Write an artifact in the published ``fastvideo-lora-v2`` shape."""
+    payload: dict[str, torch.Tensor] = {}
+    for block in range(blocks):
+        for suffix, (out_dim, in_dim) in {
+            "attn.to_q": (_INNER, _HIDDEN),
+            "attn.to_k": (_INNER, _HIDDEN),
+            "attn.to_v": (_INNER, _HIDDEN),
+            "attn.to_out.0": (_HIDDEN, _INNER),
+            "ff.net.0.proj": (2 * _FFN, _HIDDEN),
+            "ff.net.2": (_HIDDEN, _FFN),
+        }.items():
+            a, b = _factors(out_dim, in_dim)
+            payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a
+            payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b
+    payload.update(tensors or {})
+    if drop is not None:
+        payload.pop(drop, None)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # The published writer records what it emitted; mirror it so the fixtures
+    # are as self-describing as the real artifact.
+    declared = {
+        "format": FASTH3_FORMAT,
+        "finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1",
+        "base_model": FASTH3_BASE_MODEL,
+        "rank": str(_RANK),
+        "low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))),
+        "diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))),
+        "set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))),
+    }
+    # safetensors metadata is string-only, so a ``None`` override drops the key.
+    declared.update(metadata or {})
+    save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None})
+
+
+def _claim(path, **kwargs) -> FastH3WeightFusion | None:
+    """Whatever from_path decides, including declining the artifact."""
+    kwargs.setdefault("head_dim", _HEAD_DIM)
+    kwargs.setdefault("num_blocks", 1)
+    kwargs.setdefault("num_refiner_blocks", 0)
+    return FastH3WeightFusion.from_path(path, **kwargs)
+
+
+def _load(path, **kwargs) -> FastH3WeightFusion:
+    """The same, for the tests that only make sense on a claimed adapter."""
+    fusion = _claim(path, **kwargs)
+    assert fusion is not None, f"{path} was not claimed as a FastH3 adapter"
+    return fusion
+
+
+def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path):
+    plain = tmp_path / "peft" / "adapter_model.safetensors"
+    plain.parent.mkdir(parents=True)
+    save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain))
+    # No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to
+    # stay on the dynamic route rather than being fused.
+    assert _claim(plain.parent) is None
+    assert _claim(tmp_path / "missing") is None
+
+    claimed = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(claimed)
+    assert _claim(claimed.parent) is not None
+
+
+@pytest.mark.parametrize(
+    "identity",
+    [
+        # An ordinary H3 adapter out of FastVideo's own extraction tools.
+        {"finetuned_model": "someone/minimax-h3-style-lora"},
+        # Somebody else's adapter that merely names the student it imitates.
+        {"finetuned_model": "someone/fasth3-style-lora-for-h3"},
+        {"finetuned_model": ""},
+        {"finetuned_model": None},
+        # A FastH3 name over a different base model.
+        {"base_model": "Wan-AI/Wan2.2-TI2V-5B"},
+    ],
+)
+def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path):
+    """The fixture edits every block, so only the identity separates it from FastH3."""
+    path = tmp_path / "generic" / "adapter_model.safetensors"
+    _write_adapter(path, metadata=identity)
+
+    assert _claim(path.parent) is None
+
+
+def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path):
+    """Declining is for other people's adapters; a FastH3 file is held to its word."""
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, metadata={"diff_tensors": None})
+
+    with pytest.raises(FastH3AdapterError, match="omits diff_tensors"):
+        _claim(path.parent)
+
+
+def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path):
+    root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA"
+    for slug in ("dense-datafree", "vsa-datafree"):
+        _write_adapter(root / slug / "adapter_model.safetensors")
+    (root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"}))
+
+    with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"):
+        _load(root)
+    # One variant inside it is unambiguous.
+    assert _claim(root / "dense-datafree") is not None
+
+
+def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path):
+    # A plain PEFT LoRA saved as several shards must not hard-fail startup with a
+    # FastH3-specific message; it has to fall through to the dynamic LoRA route.
+    directory = tmp_path / "peft-sharded"
+    directory.mkdir()
+    for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"):
+        save_file(
+            {"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard)
+        )
+
+    assert _claim(directory) is None
+
+
+def test_low_rank_factors_reach_the_fused_projections(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path)
+    fusion = _load(path.parent)
+
+    qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32)
+    fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu()
+    # The checkpoint stores one head group at a time as [q, k, v], so the delta
+    # has to survive the loader's own unpacking as three separate projections.
+    q, k, v = torch.split(
+        _reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM),
+        [_INNER, _INNER, _INNER],
+    )
+    a, b = _factors(_INNER, _HIDDEN)
+    expected = b @ a
+    for got in (q, k, v):
+        assert torch.allclose(got, expected, atol=1e-5)
+
+
+def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path)
+    fusion = _load(path.parent)
+
+    fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu()
+    a, b = _factors(2 * _FFN, _HIDDEN)
+    value, gate = (b @ a).chunk(2, dim=0)
+    # The diffusers export is value-first; H3's fc1 is gate-first.
+    assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5)
+
+
+def test_diff_and_diff_b_edit_weights_and_biases(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(
+        path,
+        tensors={
+            "transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25),
+            "transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5),
+            "context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125),
+            "context_embedder.diff_b": torch.full((_HIDDEN,), 2.0),
+        },
+    )
+    fusion = _load(path.parent)
+
+    # A full-rank delta on an RMSNorm vector is exactly what a LoRA layer
+    # cannot express, and why this release is fused instead of switched.
+    assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25))
+    assert torch.allclose(
+        fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5)
+    )
+    assert torch.allclose(
+        fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(),
+        torch.full((_HIDDEN, _HIDDEN), 0.125),
+    )
+    assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0))
+
+
+def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)})
+    fusion = _load(path.parent)
+
+    untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3)
+    streamed = dict(
+        fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)])
+    )
+    assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0))
+    assert streamed["blocks.0.attn.q_norm.weight"] is untouched
+
+    # validate_fully_applied() releases the deltas, so a second stream would
+    # fuse nothing at all; it has to fail rather than serve base H3 weights.
+    with pytest.raises(FastH3AdapterError, match="already been fused"):
+        next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))]))
+
+
+def test_the_fused_result_keeps_the_parameter_dtype(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)})
+    fusion = _load(path.parent)
+
+    fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16))
+    assert fused.dtype == torch.bfloat16
+
+
+def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
+    dense = tmp_path / "dense" / "adapter_model.safetensors"
+    _write_adapter(dense)
+    assert _load(dense.parent).requires_vsa is False
+
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    assert _load(sparse.parent).requires_vsa is True
+
+
+def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
+    # Dropping it silently would load a model that is not the distilled student.
+    with pytest.raises(FastH3AdapterError, match="name no known"):
+        _load(path.parent)
+
+
+def test_an_unpaired_low_rank_factor_is_an_error(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight")
+
+    with pytest.raises(FastH3AdapterError, match="unpaired factor"):
+        _load(path.parent)
+
+
+def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path)
+    fusion = _load(path.parent)
+
+    fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN)))
+    # An unapplied delta is the failure that matters: the model would load and
+    # generate, just not as the distilled student.
+    with pytest.raises(FastH3AdapterError, match="never provided"):
+        fusion.validate_fully_applied()
+
+    for name, shape in (
+        ("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)),
+        ("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)),
+        ("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)),
+    ):
+        fusion.fuse(name, torch.zeros(shape))
+    fusion.validate_fully_applied()
+
+
+def _pipeline_stub(fusion=None):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._fasth3 = fusion
+    pipeline.od_config = SimpleNamespace()
+    pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0
+    return pipeline
+
+
+def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace:
+    fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None}
+    return SimpleNamespace(**{**fields, **kwargs})
+
+
+def _check_request(fusion, sampling) -> None:
+    fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0)
+
+
+def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion:
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, **kwargs)
+    return _load(path.parent)
+
+
+@pytest.mark.parametrize("task", ["fl2va", "ref2va"])
+def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path):
+    pipeline = _pipeline_stub(_fusion(tmp_path))
+
+    with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"):
+        pipeline._resolve_task(task, {})
+    assert pipeline._resolve_task("t2va", {}) == "t2va"
+    assert _pipeline_stub()._resolve_task(task, {}) == task
+
+
+@pytest.mark.parametrize("num_inference_steps", [None, 5, 50])
+def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path):
+    fusion = _fusion(tmp_path)
+
+    # The five sigma points bound four forwards, and a request states forwards:
+    # the interval contract the pinned schedules use, and the count the step
+    # scheduler admits a request on before any pipeline hook runs.
+    with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"):
+        _check_request(fusion, _sampling(num_inference_steps))
+    _check_request(fusion, _sampling())
+    # The pinned branch coerces the string form; this one cannot disagree.
+    _check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS)))
+
+
+def test_the_request_denoises_on_the_release_ladder(tmp_path):
+    pipeline = _pipeline_stub(_fusion(tmp_path))
+
+    positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling())
+
+    # dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the
+    # terminal 0.0: the pre-shift positions the student was distilled at, not
+    # the uniform ladder num_inference_steps would otherwise derive.
+    assert positions == (0.999, 0.749, 0.5, 0.25, 0.0)
+    # Five sigma points bound four transformer forwards, and forwards is the
+    # unit num_steps carries downstream (Cache-DiT, quality hints).
+    assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1
+    # H3's own per-modality shifts still apply on top, and are what turn those
+    # positions into the noise levels the four forwards run at.
+    assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx(
+        [0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6
+    )
+    assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx(
+        [0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6
+    )
+    # Without a fused adapter the same request keeps the undistilled ladder.
+    assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS)
+
+
+def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path):
+    pipeline = _pipeline_stub(_fusion(tmp_path))
+    pipeline.partition = "fl2va"
+
+    pipeline._fasth3.check_serving_contract(
+        partition=pipeline.partition,
+        od_config=pipeline.od_config,
+        video_shift=pipeline.default_video_shift,
+        audio_shift=pipeline.default_audio_shift,
+    )
+
+    # The ladder is stated pre-shift, so H3's 12/3 shifts are what place the
+    # student on the levels it was distilled at; the contract must not touch them.
+    assert pipeline.default_video_shift == 12.0
+    assert pipeline.default_audio_shift == 3.0
+
+
+@pytest.mark.parametrize(
+    "extra",
+    [{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}],
+)
+def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path):
+    fusion = _fusion(tmp_path)
+
+    with pytest.raises(OmniClientError, match="FastH3 requires"):
+        _check_request(fusion, _sampling(extra_args=extra))
+
+    _check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0}))
+
+
+def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path):
+    fusion = _fusion(tmp_path)
+
+    # The dynamic LoRA manager is skipped for a fused adapter, so an adapter
+    # asked for per request would be neither applied nor reported.
+    with pytest.raises(OmniClientError, match="per-request lora is unavailable"):
+        _check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1)))
+
+
+def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path):
+    fusion = _fusion(tmp_path)
+    for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"):
+        # A host-weight plan installs the transformer without load_weights(),
+        # so the fusion and its completeness check would both be skipped.
+        with pytest.raises(ValueError, match="cannot be combined with"):
+            fusion.check_serving_contract(
+                partition="fl2va",
+                od_config=SimpleNamespace(**{flag: True}),
+                video_shift=12.0,
+                audio_shift=3.0,
+            )
+
+
+def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0}
+    with pytest.raises(ValueError, match="Video Sparse Attention variant"):
+        _load(sparse.parent).check_serving_contract(partition="fl2va", **contract)
+
+    with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"):
+        _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
+
+
+def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
+    # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
+    # so a .diff aimed at one would otherwise be added transposed.
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))})
+    with pytest.raises(FastH3AdapterError, match="fused layout"):
+        _claim(path.parent)
+
+
+def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, blocks=2)
+
+    # Every tensor in this file is well formed and pairs up, so the fusion's
+    # own completeness check passes; only the model's depth reveals that a
+    # third of it would have been served as base H3 weights.
+    with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"):
+        _claim(path.parent, num_blocks=3)
+    with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"):
+        _claim(path.parent, num_blocks=1)
+    assert _claim(path.parent, num_blocks=2) is not None
+
+
+def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"})
+
+    # The writer records what it emitted, which is the file's one statement
+    # about its own completeness.
+    with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"):
+        _claim(path.parent)
+
+
+def test_only_a_configured_adapter_path_is_read(tmp_path):
+    path = tmp_path / "fasth3" / "adapter_model.safetensors"
+    _write_adapter(path)
+    transformer = SimpleNamespace(
+        arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0)
+    )
+    unbuilt = object()  # a transformer whose arch must not be touched
+
+    assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None
+    assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None
+    # Nothing configured, or more than one artifact, stays on the dynamic route
+    # without reading anything off the model.
+    assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None
+    assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None
+    assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None
+
+
+def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path):
+    assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True
+    assert _pipeline_stub().lora_is_fused is False
+
+
+def test_the_ladder_is_the_one_the_release_publishes():
+    # FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart.
+    assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0)
+    assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4

tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py

@@ -259,7 +259,11 @@ def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path):
         ("", "must not be empty"),
         (",,", "must not be empty"),
         ("1.0,oops,0.15,0.0", "is malformed"),
-        ("0.9,0.7,0.4,0.15,0.0", "is malformed"),
+        # A capped opening rung below 1.0 is a legal DMD2 schedule, so the
+        # out-of-range cases are a first position above 1.0 and a tail that
+        # never reaches clean latents.
+        ("1.5,0.7,0.4,0.15,0.0", "is malformed"),
+        ("1.0,0.7,0.4,0.15", "is malformed"),
     ],
 )
 def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message):

tests/diffusion/sched/test_dmd2_sigma_schedule.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import pytest
 
 from vllm_omni.diffusion.sched import DMD2SigmaSchedule
@@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries():
     [
         [],
         [1.0],
-        [0.9, 0.5, 0.0],
+        [1.5, 0.5, 0.0],
+        [0.0, 0.5, 0.0],
         [1.0, 0.5, 0.1],
         [1.0, 0.4, 0.6, 0.0],
         [1.0, 0.5, 0.5, 0.0],
@@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions):
         DMD2SigmaSchedule.from_positions(positions)
 
 
+def test_a_capped_first_position_is_a_schedule_not_a_malformed_one():
+    # A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise,
+    # so its opening rung sits just below 1.0 (FastH3: 999/1000).
+    schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0])
+
+    assert schedule.base_schedule[0] == 0.999
+    assert schedule.num_inference_steps == 4
+
+
 def test_one_base_schedule_drives_several_shift_scales():
     schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
 

tests/diffusion/test_diffusion_model_runner.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import contextmanager
 from types import SimpleNamespace
@@ -533,7 +533,14 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
         return original_execute(*args, **kwargs)
 
     runner._execute_request_list = execute
-    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    forward_context_calls = []
+
+    @contextmanager
+    def record_forward_context(*args, **kwargs):
+        forward_context_calls.append((args, kwargs))
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
     reset_peak_memory_stats = Mock()
     monkeypatch.setattr(
         model_runner_module.current_omni_platform,
@@ -549,6 +556,8 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
 
     assert runner.pipeline.forward_calls == 1
     assert record_names == ["pipeline_memory_profile"]
+    assert len(forward_context_calls) == 1
+    assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
     runner._validate_diffusion_kv_metadata.assert_not_called()
     reset_peak_memory_stats.assert_not_called()
     synchronize.assert_called_once_with()
@@ -575,7 +584,14 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
 
     runner.pipeline.denoise_step = denoise_step
     runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
-    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    forward_context_calls = []
+
+    @contextmanager
+    def record_forward_context(*args, **kwargs):
+        forward_context_calls.append((args, kwargs))
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
     reset_peak_memory_stats = Mock()
     monkeypatch.setattr(
         model_runner_module.current_omni_platform,
@@ -592,6 +608,8 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
     assert observed_batch_rows == [(2, 2)]
     assert runner.state_cache == {}
     assert runner.input_batch is None
+    assert len(forward_context_calls) == 1
+    assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
     runner._validate_diffusion_kv_metadata.assert_not_called()
     reset_peak_memory_stats.assert_not_called()
     synchronize.assert_called_once_with()

tests/diffusion/test_diffusion_streaming_output.py

@@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint:
                     model="test_model",
                     od_config=od_config,
                     metadata=metadata,
-                    batch_size=1,
                 )
 
     @staticmethod
@@ -526,6 +525,7 @@ class TestSupportedPipelines:
         runner = object.__new__(DiffusionModelRunner)
         runner.vllm_config = _make_vllm_config()
         runner.od_config = SimpleNamespace(
+            max_num_seqs=1,
             enable_cpu_offload=False,
             enable_layerwise_offload=False,
             enforce_eager=True,

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -47,7 +47,7 @@ def client(mock_engine):
     with patch.object(InlineStageDiffusionClient, "_enrich_config"):
         od_config = MagicMock(spec=OmniDiffusionConfig)
         od_config.streaming_output = False
-        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata)
         yield c
         c.shutdown()
 
@@ -234,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine):
                 model="test_model",
                 od_config=od_config,
                 metadata=metadata,
-                batch_size=1,
             )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1171,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation:
                 metadata,
                 "tcp://req",
                 "tcp://resp",
-                batch_size=1,
             )
 
     @pytest.mark.asyncio

tests/diffusion/test_prompt_update.py

@@ -387,7 +387,6 @@ class TestPromptUpdateIntegration:
                     model="test_model",
                     od_config=od_config,
                     metadata=metadata,
-                    batch_size=1,
                 )
 
     @staticmethod

tests/diffusion/test_stage_diffusion_proc.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import time
@@ -7,6 +7,8 @@ from dataclasses import asdict, dataclass
 
 import pytest
 
+import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc
+import vllm_omni.plugins as omni_plugins
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -26,6 +28,33 @@ BASE_INFER_STEPS = 10
 DELAY_BASE = 0.01
 
 
+def test_run_diffusion_proc_sets_lifecycle_before_loading_plugins(monkeypatch: pytest.MonkeyPatch) -> None:
+    events: list[str] = []
+
+    class StopProcessError(Exception):
+        pass
+
+    class TestStageDiffusionProc(StageDiffusionProc):
+        def __init__(self, model, od_config):
+            events.append("proc")
+            raise StopProcessError
+
+    monkeypatch.setattr(omni_plugins, "load_omni_general_plugins", lambda: events.append("plugins"))
+    monkeypatch.setattr(stage_diffusion_proc, "set_death_signal", lambda _: events.append("death_signal"))
+    monkeypatch.setattr(stage_diffusion_proc.signal, "signal", lambda *_: events.append("signal_handler"))
+
+    with pytest.raises(StopProcessError):
+        TestStageDiffusionProc.run_diffusion_proc(
+            model="test-model",
+            od_config=None,
+            handshake_address="test-address",
+            local_client=True,
+            headless=False,
+        )
+
+    assert events == ["death_signal", "signal_handler", "signal_handler", "plugins", "proc"]
+
+
 class MockDiffusionEngine:
     async def step_streaming(self, request):
         def simulate_step_delay(height, width, num_inference_steps) -> float:

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -2959,3 +2959,34 @@ def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter)
     )
 
     assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
+
+
+def test_save_async_boundary_holds_generation_without_request_counter(build_adapter):
+    """A boundary must not advance past a generation nothing can catch up to.
+
+    When the producing stage maintains ``_omni_segment_generation`` on the
+    request, a boundary may start the next segment's watermark at
+    ``generation + 1`` because the following frame will carry it. On paths
+    where the request has no such counter, that speculative advance makes every
+    subsequent frame look stale and starves the stream permanently
+    (vllm-project/vllm-omni#6816).
+    """
+    adapter, _ = build_adapter(stage_id=1)
+
+    request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen")
+    request.resumable = True
+    request._omni_segment_generation = None
+
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+    assert adapter._segment_generation["ext-nogen"] == 0
+
+    follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen")
+    follow_up.resumable = True
+    follow_up.num_computed_tokens = 3
+    follow_up._omni_segment_generation = None
+    queued_before = len(adapter._pending_save_reqs)
+
+    adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == queued_before + 1
+    assert adapter._pending_save_reqs[-1]["request"] is follow_up

tests/e2e/online_serving/test_boogu_image_edit.py

@@ -1,15 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
-Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions).
+Online serving tests for Boogu-Image-0.1-Edit and Edit-Turbo.
 
-The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base
+Both Edit checkpoints share the ``BooguImagePipeline`` class with the Base
 text-to-image checkpoint; the editing (TI2I) path activates when the request
-carries a reference image. Text guidance rides on ``guidance_scale`` (upstream
+carries a reference image. The regular Edit checkpoint uses text guidance via
-default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default
+``guidance_scale`` and optional image guidance via ``guidance_scale_2``.
-1.0 = off; a value > 1 enables the double-guidance path). Only a single
+Edit-Turbo is guidance-distilled and uses four steps with guidance scale 1.0.
-reference image is supported.
+Only a single reference image is supported.
 
 - ``test_single_image_to_image_001``: one reference image, text guidance only.
 - ``test_double_guidance_001``: one reference image, text + image guidance.
@@ -40,6 +40,10 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Change the style to a colored pencil drawing."
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 
+EDIT_MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
+
 
 def _get_diffusion_feature_cases(model: str):
     """Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
@@ -54,13 +58,44 @@ def _get_diffusion_feature_cases(model: str):
     ]
 
 
+def _get_single_edit_cases():
+    """Return the regular and distilled Edit server/request configurations."""
+    return [
+        pytest.param(
+            OmniServerParams(model=EDIT_MODEL),
+            {
+                "num_inference_steps": 2,
+                "guidance_scale": 5.0,
+            },
+            id="edit",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=EDIT_TURBO_MODEL,
+                server_args=["--revision", EDIT_TURBO_REVISION],
+            ),
+            {
+                "num_inference_steps": 4,
+                "guidance_scale": 1.0,
+            },
+            id="edit-turbo-hotfix-1k",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
 @pytest.mark.parametrize(
-    "omni_server",
+    ("omni_server", "sampling_params"),
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_single_edit_cases(),
-    indirect=True,
+    indirect=["omni_server"],
 )
-def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_single_image_to_image_001(
-    """Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``."""
+    omni_server: OmniServer,
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+):
+    """Single-reference text-guided edit smoke for regular and Turbo Edit."""
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -69,9 +104,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
         "model": omni_server.model,
         "messages": messages,
         "extra_body": {
-            "num_inference_steps": 2,
+            **sampling_params,
-            # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
-            "guidance_scale": 5.0,
             "seed": 42,
         },
     }
@@ -81,7 +114,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_diffusion_feature_cases(EDIT_MODEL),
     indirect=True,
 )
 def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -105,15 +138,20 @@ def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClien
     openai_client.send_diffusion_request(request_config)
 
 
-def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict:
+def _edits_request_config(
+    model: str,
+    *,
+    num_inference_steps: int = 2,
+    guidance_scale: float = 5.0,
+    guidance_scale_2: float | None = None,
+) -> dict:
     """Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params."""
     image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"])
     data = {
         "model": model,
         "prompt": EDIT_PROMPT,
-        "num_inference_steps": 2,
+        "num_inference_steps": num_inference_steps,
-        # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
+        "guidance_scale": guidance_scale,
-        "guidance_scale": 5.0,
         "seed": 42,
     }
     if guidance_scale_2 is not None:
@@ -133,19 +171,25 @@ def _assert_edits_response_has_image(responses) -> None:
 
 
 @pytest.mark.parametrize(
-    "omni_server",
+    ("omni_server", "sampling_params"),
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_single_edit_cases(),
-    indirect=True,
+    indirect=["omni_server"],
 )
-def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_images_edits_endpoint_001(
-    """Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint."""
+    omni_server: OmniServer,
-    responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model))
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+):
+    """Single-reference regular and Turbo Edit via ``/v1/images/edits``."""
+    responses = openai_client.send_images_edits_http_request(
+        _edits_request_config(omni_server.model, **sampling_params)
+    )
     _assert_edits_response_has_image(responses)
 
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_diffusion_feature_cases(EDIT_MODEL),
     indirect=True,
 )
 def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_boogu_image_edit_expansion.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
-L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I).
+L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I).
 
 The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
 text-guided editing and the text+image double-guidance path. This expansion
@@ -16,11 +16,12 @@ module adds the two edit paths smoke does not exercise:
   ``height`` / ``width``: asserts the output resolution follows the reference
   (upstream ``align_res``, on by default for a single reference).
 
-The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the
+Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
-TI2I path activates when the request carries a reference image. Only a single
+the TI2I path activates when the request carries a reference image. Only a
-reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``).
+single reference image is supported. The regular Edit cases stay smoke-depth
-CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this
+(``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
-model and are intentionally omitted.
+configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
+supported for these models and are intentionally omitted.
 
 From ``tests/``::
 
@@ -41,6 +42,8 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 
 MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
 EDIT_PROMPT = "Change the style to a colored pencil drawing."
 
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -57,6 +60,33 @@ def _get_diffusion_feature_cases(model: str):
     ]
 
 
+def _get_align_res_cases():
+    """Return regular and Turbo Edit configurations for align-res coverage."""
+    return [
+        pytest.param(
+            OmniServerParams(model=MODEL),
+            {
+                "num_inference_steps": 2,
+                "guidance_scale": 5.0,
+            },
+            id="edit",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=EDIT_TURBO_MODEL,
+                server_args=["--revision", EDIT_TURBO_REVISION],
+            ),
+            {
+                "num_inference_steps": 4,
+                "guidance_scale": 1.0,
+            },
+            id="edit-turbo-hotfix-1k",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
 @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
 def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
     """Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch."""
@@ -76,8 +106,16 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
     openai_client.send_diffusion_request(request_config)
 
 
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+@pytest.mark.parametrize(
-def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    ("omni_server", "sampling_params"),
+    _get_align_res_cases(),
+    indirect=["omni_server"],
+)
+def test_align_res_output_size(
+    omni_server: OmniServer,
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+) -> None:
     """Output resolution follows the reference image (upstream ``align_res``).
 
     No ``height`` / ``width`` is requested. The reference is 512x384 — both
@@ -92,8 +130,7 @@ def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAICli
         "model": omni_server.model,
         "messages": messages,
         "extra_body": {
-            "num_inference_steps": 2,
+            **sampling_params,
-            "guidance_scale": 5.0,
             "seed": 42,
         },
     }

tests/e2e/online_serving/test_minicpmo_4_5.py

@@ -1,12 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
 
 Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
 Talker, and Code2Wav stages.
 """
 
+import json
 import os
+import tempfile
+from copy import deepcopy
+from pathlib import Path
 
 import pytest
+from vllm.logger import DEFAULT_LOGGING_CONFIG
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -19,10 +27,36 @@ _MODEL = "openbmb/MiniCPM-o-4_5"
 _CI_DEPLOY = modify_stage_config(
     get_deploy_config_path("minicpmo_4_5.yaml"),
     updates={
-        "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}}
+        "stages": {
+            0: {"default_sampling_params.max_tokens": 64},
+            1: {
+                "default_sampling_params.max_tokens": 1024,
+                # Content-consistency assertions must not depend on a random
+                # codec trajectory: different valid TTS samples can append an
+                # audible tail even when Thinker text is identical.
+                "default_sampling_params.temperature": 0.0,
+            },
+        },
     },
 )
 
+_PROMPT_LOG_TEMP_DIR = tempfile.TemporaryDirectory(prefix="minicpmo45-prompt-selection-")
+_PROMPT_SELECTION_LOG = Path(_PROMPT_LOG_TEMP_DIR.name) / "prompt-selection.log"
+_PROMPT_LOGGING_CONFIG_PATH = Path(_PROMPT_LOG_TEMP_DIR.name) / "logging.json"
+_PROMPT_LOGGING_CONFIG = deepcopy(DEFAULT_LOGGING_CONFIG)
+_PROMPT_LOGGING_CONFIG["handlers"]["minicpmo_prompt_selection"] = {
+    "class": "logging.FileHandler",
+    "filename": str(_PROMPT_SELECTION_LOG),
+    "formatter": "vllm",
+    "level": "DEBUG",
+}
+_PROMPT_LOGGING_CONFIG["loggers"]["vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav"] = {
+    "handlers": ["minicpmo_prompt_selection"],
+    "level": "DEBUG",
+    "propagate": True,
+}
+_PROMPT_LOGGING_CONFIG_PATH.write_text(json.dumps(_PROMPT_LOGGING_CONFIG), encoding="utf-8")
+
 test_params = [
     pytest.param(
         OmniServerParams(
@@ -30,6 +64,10 @@ test_params = [
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
             server_args=["--trust-remote-code", "--async-chunk"],
+            env_dict={
+                "VLLM_CONFIGURE_LOGGING": "1",
+                "VLLM_LOGGING_CONFIG_PATH": str(_PROMPT_LOGGING_CONFIG_PATH),
+            },
         ),
         id="async_chunk",
     ),
@@ -74,6 +112,18 @@ _TTS_EXTRA_BODY = {
 }
 
 
+def _prompt_selection_log_offset() -> int:
+    return _PROMPT_SELECTION_LOG.stat().st_size if _PROMPT_SELECTION_LOG.exists() else 0
+
+
+def _read_prompt_selection_log(offset: int) -> str:
+    if not _PROMPT_SELECTION_LOG.exists():
+        return ""
+    with open(_PROMPT_SELECTION_LOG, "rb") as log_file:
+        log_file.seek(offset)
+        return log_file.read().decode("utf-8")
+
+
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.omni
@@ -126,6 +176,51 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config)
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_with_reference_audio(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()),
+        "stream": True,
+        "modalities": ["text", "audio"],
+        "extra_body": {
+            **_TTS_EXTRA_BODY,
+            "ref_audio": f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}",
+        },
+    }
+
+    log_offset = _prompt_selection_log_offset()
+    responses = openai_client.send_omni_request(request_config)
+    prompt_log = _read_prompt_selection_log(log_offset)
+
+    assert responses[0].success
+    assert responses[0].audio_bytes
+    assert "prompt_cache_id=runtime-ref-" in prompt_log
+    assert "minicpmo45_ref_" in prompt_log
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_with_default_reference(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()),
+        "stream": True,
+        "modalities": ["text", "audio"],
+        "extra_body": _TTS_EXTRA_BODY,
+    }
+
+    responses = openai_client.send_omni_request(request_config)
+
+    assert responses[0].success
+    assert responses[0].audio_bytes
+
+
 @pytest.mark.full_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)

tests/e2e/online_serving/test_nemotron_voicechat_duplex.py

@@ -18,6 +18,17 @@ from tests.helpers.stage_config import get_deploy_config_path
 MODEL = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_MODEL", "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B")
 DEPLOY_CONFIG = get_deploy_config_path("nemotron_labs_voicechat_duplex.yaml")
 TOKENIZER = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_LLM_PATH")
+
+# The client gives up after CLIENT_TIMEOUT_S. The server-side stage-input
+# deadline must fire first, otherwise a stalled stream surfaces as a bare
+# client TimeoutError with no server-side error at all. See #6816.
+CLIENT_TIMEOUT_S = 300
+SERVER_INPUT_WAIT_TIMEOUT_S = 240
+
+_SERVER_ENV: dict[str, str] = {"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S": str(SERVER_INPUT_WAIT_TIMEOUT_S)}
+if TOKENIZER:
+    _SERVER_ENV["NEMOTRON_VOICECHAT_LLM_PATH"] = TOKENIZER
+
 pytestmark = [pytest.mark.full_model, pytest.mark.omni]
 
 
@@ -29,7 +40,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.omni]
             OmniServerParams(
                 model=MODEL,
                 stage_config_path=DEPLOY_CONFIG,
-                env_dict={"NEMOTRON_VOICECHAT_LLM_PATH": TOKENIZER} if TOKENIZER else None,
+                env_dict=_SERVER_ENV,
             ),
             id="native-duplex",
         )
@@ -45,7 +56,7 @@ def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Pa
         (
             f"--url ws://{omni_server.host}:{omni_server.port}/v1/realtime --model {omni_server.model} "
             f"--input-wav {root / 'turn_taking.wav'} --input-channel 0 --max-frames 190 "
-            f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s 300 "
+            f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s {CLIENT_TIMEOUT_S} "
             f"--output-dir {tmp_path / 'native_duplex'}"
         ).split()
     )

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -55,7 +55,6 @@ tts_server_params = [
 
 
 @pytest.mark.advanced_model
-@pytest.mark.core_model
 @pytest.mark.tts
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)

tests/engine/test_arg_utils.py

@@ -286,6 +286,28 @@ def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
     }
 
 
+def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, monkeypatch, mocker):
+    """A successful remote tokenizer download must not enter the error path."""
+    tokenizer_dir = tmp_path / "CosyVoice-BlankEN"
+    tokenizer_dir.mkdir()
+    baseline_config = Mock()
+    warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning")
+
+    monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *args, **kwargs: str(tmp_path))
+    monkeypatch.setattr(OmniEngineArgs, "_patch_empty_hf_config", lambda *args, **kwargs: None)
+    monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config)
+    monkeypatch.setattr(
+        OmniModelConfig,
+        "from_vllm_model_config",
+        classmethod(lambda cls, model_config, **omni_kwargs: model_config),
+    )
+
+    args = OmniEngineArgs(model="remote/cosyvoice3", model_arch="CosyVoice3Model")
+    assert args.create_model_config() is baseline_config
+    assert args.tokenizer == str(tokenizer_dir)
+    warning.assert_not_called()
+
+
 def test_patch_missing_local_hf_config(tmp_path):
     """Native model bundles may only contain a checkpoints/ directory."""
     (tmp_path / "checkpoints").mkdir()

tests/engine/test_async_omni_engine_stage_init.py

@@ -172,7 +172,6 @@ def _make_stage_runtime() -> StageRuntime:
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -196,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
     engine.single_stage_mode = False
-    engine.diffusion_batch_size = 1
     engine.async_chunk = False
     engine.tokenizer = None
     engine._single_stage_id_filter = None
@@ -282,7 +280,6 @@ def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatc
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -371,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -399,16 +395,23 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
             os.environ[env_var] = old_env
 
 
-def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
+@pytest.mark.parametrize(
+    ("num_stages", "expected_inline"),
+    [(1, True), (2, False)],
+)
+def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(
+    monkeypatch,
+    num_stages,
+    expected_inline,
+):
     import vllm_omni.engine.stage_runtime as runtime_mod
     from vllm_omni.engine.stage_engine_startup import StageReplicaResources
 
     runtime = StageRuntime(
-        stage_configs=[types.SimpleNamespace()],
+        stage_configs=[types.SimpleNamespace()] * num_stages,
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=4,
         async_chunk=False,
     )
 
@@ -421,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
     def _capture_launch_diffusion_stage_replica(**kwargs):
         captured["stage_id"] = kwargs["metadata"].stage_id
         captured["stage_init_timeout"] = kwargs["stage_init_timeout"]
-        captured["batch_size"] = kwargs["batch_size"]
         captured["use_inline"] = kwargs["use_inline"]
         captured["omni_master_server"] = kwargs["omni_master_server"]
         return types.SimpleNamespace(), StageReplicaResources()
@@ -433,8 +435,7 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
     assert captured == {
         "stage_id": 0,
         "stage_init_timeout": 302,
-        "batch_size": 4,
+        "use_inline": expected_inline,
-        "use_inline": False,
         "omni_master_server": None,
     }
 
@@ -503,21 +504,20 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
     assert os.environ[runtime_env_var] == "parent-value"
 
 
-def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
+def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.engine.stage_init_utils as init_mod
 
-    od_config = types.SimpleNamespace(max_num_seqs=1)
+    od_config = types.SimpleNamespace(max_num_seqs=4)
     captured: dict[str, object] = {}
     monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
 
-    def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
+    def _capture_client(model, config, metadata, stage_init_timeout, use_inline):
         captured.update(
             model=model,
             config=config,
             metadata=metadata,
             stage_init_timeout=stage_init_timeout,
-            batch_size=batch_size,
             use_inline=use_inline,
         )
         return object()
@@ -531,27 +531,25 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_
         types.SimpleNamespace(),
         metadata,
         stage_init_timeout=12,
-        batch_size=4,
         use_inline=True,
     )
 
-    assert od_config.max_num_seqs == 1
+    assert od_config.max_num_seqs == 4
     assert captured == {
         "model": "dummy-model",
         "config": od_config,
         "metadata": metadata,
         "stage_init_timeout": 12,
-        "batch_size": 4,
         "use_inline": True,
     }
 
 
-def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
+def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
     import vllm_omni.engine.stage_engine_startup as startup_mod
 
-    od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1))
+    od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
     monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
     monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
     monkeypatch.setattr(
@@ -588,13 +586,12 @@ def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_s
         stage_config=types.SimpleNamespace(),
         metadata=types.SimpleNamespace(stage_id=0),
         stage_init_timeout=12,
-        batch_size=4,
         use_inline=False,
         omni_master_server=omni_master_server,
     )
 
     assert result is sentinel_client
-    assert od_config.max_num_seqs == 1
+    assert od_config.max_num_seqs == 4
     assert resources.manager is proc_manager
 
 
@@ -613,7 +610,6 @@ def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
         types.SimpleNamespace(),
         metadata,
         stage_init_timeout=12,
-        batch_size=4,
         use_inline=True,
     )
 
@@ -664,7 +660,6 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo
         stage_config=types.SimpleNamespace(),
         metadata=types.SimpleNamespace(stage_id=0),
         stage_init_timeout=12,
-        batch_size=4,
         use_inline=False,
         omni_master_server=omni_master_server,
     )
@@ -680,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -746,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
         log_stats=True,
     )
@@ -790,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
         log_stats=True,
     )
@@ -843,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -884,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=123,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -922,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys():
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=123,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
     plan = _make_llm_plan(
@@ -948,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -985,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
 
@@ -1024,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
         model="dummy-model",
         config_path="dummy-config",
         stage_init_timeout=1,
-        diffusion_batch_size=1,
         async_chunk=False,
     )
     stage_init_timeout = 302

tests/engine/test_single_stage_mode.py

@@ -575,7 +575,6 @@ class TestSingleStageInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
             single_stage_id_filter=stage_id_filter,
             omni_master_address="127.0.0.1",
@@ -807,7 +806,6 @@ class TestSingleStageInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
         )
         mocker.patch.object(runtime_mod, "prepare_engine_environment")
@@ -860,7 +858,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -928,7 +925,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -950,7 +946,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -1002,7 +997,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=2,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -1068,7 +1062,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=4,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -1120,7 +1113,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=4,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -1141,7 +1133,7 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+        od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
         mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mock_register = mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
@@ -1178,7 +1170,7 @@ class TestSingleStageReplicaInitialization:
                 os.environ[device_env_var] = prev_device_env
 
         assert result is sentinel_client
-        assert od_config.max_num_seqs is None
+        assert od_config.max_num_seqs == 4
         mock_register.assert_called_once_with(
             omni_master_address="127.0.0.1",
             omni_master_port=25000,
@@ -1203,7 +1195,6 @@ class TestSingleStageReplicaInitialization:
             request_address="tcp://127.0.0.1:26002",
             response_address="tcp://127.0.0.1:26003",
             proc_manager=mocker.ANY,
-            batch_size=4,
         )
         assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
 
@@ -1216,7 +1207,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=4,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
@@ -1258,7 +1248,6 @@ class TestSingleStageReplicaInitialization:
             model="fake-model",
             config_path="/fake/stages.yaml",
             stage_init_timeout=60,
-            diffusion_batch_size=4,
             async_chunk=False,
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",

tests/entrypoints/openai_api/test_duplex_handler.py

@@ -4425,6 +4425,167 @@ async def test_playback_ack_rejects_response_after_followup_user_commit():
     )
 
 
+@pytest.mark.asyncio
+async def test_precommit_zero_audio_ack_reserves_history_for_response_that_finishes_after_commit():
+    handler = OmniDuplexSessionHandler(
+        chat_service=FakeChatService(FakeEngineClient()),
+        config_timeout_s=0.1,
+        idle_timeout_s=1,
+    )
+    session = DuplexSession(
+        session_id="sid-zero-audio-checkpoint",
+        config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+    )
+    session.mark_user_input_activity()
+    response_id = session.begin_response()
+    item_id = f"item_{response_id}"
+    ws = TimedWebSocket()
+
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 0,
+            "committed_ms": 0,
+        },
+        ws.send_json,
+    )
+    session.commit_native_audio_input(transcript="later user input")
+    session.append_assistant_text("hello world")
+    session.mark_audio_sent(1000, text_chars=len("hello world"))
+    committed_message = session.end_response(commit_text=True, preserve_request=True)
+    session.register_history_item(item_id, committed_message)
+
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 1000,
+            "committed_ms": 1000,
+        },
+        ws.send_json,
+    )
+
+    assert not any(message.get("type") == "error" for message in ws.sent)
+    assert session.history == (
+        {"role": "assistant", "content": "hello world"},
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "audio_url",
+                    "audio_url": {"url": "native-duplex:input-audio"},
+                    "transcript": "later user input",
+                }
+            ],
+            "transcript": "later user input",
+        },
+    )
+
+
+@pytest.mark.asyncio
+async def test_precommit_partial_ack_extends_history_after_response_finishes():
+    handler = OmniDuplexSessionHandler(
+        chat_service=FakeChatService(FakeEngineClient()),
+        config_timeout_s=0.1,
+        idle_timeout_s=1,
+    )
+    session = DuplexSession(
+        session_id="sid-partial-checkpoint",
+        config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+    )
+    session.mark_user_input_activity()
+    response_id = session.begin_response()
+    item_id = f"item_{response_id}"
+    session.append_assistant_text("hello world")
+    session.mark_audio_sent(1000, text_chars=len("hello world"))
+    ws = TimedWebSocket()
+
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 500,
+            "committed_ms": 500,
+        },
+        ws.send_json,
+    )
+    session.commit_native_audio_input(transcript="later user input")
+    committed_message = session.end_response(commit_text=True, preserve_request=True)
+    session.register_history_item(item_id, committed_message)
+
+    assert session.history[0] == {"role": "assistant", "content": "hello"}
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 1000,
+            "committed_ms": 1000,
+        },
+        ws.send_json,
+    )
+
+    assert not any(message.get("type") == "error" for message in ws.sent)
+    assert session.history[0] == {"role": "assistant", "content": "hello world"}
+    assert session.history[1]["role"] == "user"
+
+
+@pytest.mark.asyncio
+async def test_explicit_truncate_seals_snapshot_against_later_playback_ack():
+    handler = OmniDuplexSessionHandler(
+        chat_service=FakeChatService(FakeEngineClient()),
+        config_timeout_s=0.1,
+        idle_timeout_s=1,
+    )
+    session = DuplexSession(
+        session_id="sid-hard-playback-truncate",
+        config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+    )
+    response_id = session.begin_response()
+    item_id = f"item_{response_id}"
+    session.append_assistant_text("hello world")
+    session.mark_audio_sent(1000, text_chars=len("hello world"))
+    session.end_response(commit_text=False, preserve_request=True)
+    ws = TimedWebSocket()
+
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 500,
+            "committed_ms": 500,
+            "truncate": True,
+        },
+        ws.send_json,
+    )
+    assert session.history == ({"role": "assistant", "content": "hello"},)
+
+    await handler._handle_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": response_id,
+            "item_id": item_id,
+            "played_ms": 1000,
+            "committed_ms": 1000,
+        },
+        ws.send_json,
+    )
+
+    assert not any(message.get("type") == "error" for message in ws.sent)
+    assert session.history == ({"role": "assistant", "content": "hello"},)
+
+
 @pytest.mark.asyncio
 async def test_cancel_does_not_append_old_assistant_after_followup_user():
     handler = OmniDuplexSessionHandler(

tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MiniCPM-o 4.5 chat reference audio transport."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import cast
+
+import numpy as np
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.inputs import TokensPrompt
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
+from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import (
+    llm2tts,
+    tts2code2wav_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _EngineClientStub:
+    stage_configs: list[dict[str, object]]
+
+
+@dataclass
+class _ModelConfigStub:
+    allowed_local_media_path: str = ""
+    allowed_media_domains: list[str] | None = None
+
+
+@dataclass
+class _ThinkerCompletionStub:
+    token_ids: list[int]
+    text: str
+    hidden_states: torch.Tensor
+    multimodal_output: dict[str, object] = field(default_factory=dict)
+
+
+@dataclass
+class _ThinkerRequestOutputStub:
+    request_id: str
+    prompt_token_ids: list[int]
+    outputs: list[_ThinkerCompletionStub]
+    multimodal_output: dict[str, object] | None = None
+
+
+@dataclass
+class _Code2WavRequestStub:
+    external_req_id: str
+    request_id: str
+    model_intermediate_buffer: dict[str, object]
+
+
+@dataclass
+class _TransferManagerStub:
+    request_payload: dict[str, object] = field(default_factory=dict)
+
+
+@pytest.fixture
+def serving_chat() -> OmniOpenAIServingChat:
+    return object.__new__(OmniOpenAIServingChat)
+
+
+def _stage_config(model_arch: str) -> dict[str, object]:
+    return {"engine_args": {"model_arch": model_arch}}
+
+
+def _chat_request(reference_audio_url: str) -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "target text"}],
+        ref_audio=reference_audio_url,
+    )
+
+
+@pytest.mark.asyncio
+async def test_chat_reference_audio_reaches_code2wav_transport(
+    serving_chat: OmniOpenAIServingChat,
+    mocker: MockerFixture,
+) -> None:
+    reference_audio_url = "data:audio/wav;base64,AAAA"
+    reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+    serving_chat.engine_client = _EngineClientStub(
+        [_stage_config("MiniCPMO45OmniForConditionalGeneration")],
+    )
+    serving_chat.model_config = _ModelConfigStub()
+    engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+    media_connector_class = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+    )
+    media_connector = media_connector_class.return_value
+    media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+    await serving_chat._attach_minicpmo45_reference_audio(
+        engine_prompt,
+        _chat_request(reference_audio_url),
+    )
+
+    media_connector_class.assert_called_once_with(
+        media_io_kwargs=None,
+        allowed_local_media_path="",
+        allowed_media_domains=None,
+    )
+    media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+    preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY]
+    assert isinstance(preserved_reference, tuple)
+    preserved_waveform, preserved_sample_rate = preserved_reference
+    assert isinstance(preserved_waveform, np.ndarray)
+    assert isinstance(preserved_sample_rate, int)
+    np.testing.assert_array_equal(preserved_waveform, reference_waveform)
+    assert preserved_sample_rate == 16000
+
+    thinker_output = _ThinkerRequestOutputStub(
+        request_id="request-0",
+        prompt_token_ids=[10],
+        outputs=[
+            _ThinkerCompletionStub(
+                token_ids=[20],
+                text="target text",
+                hidden_states=torch.zeros((2, 4)),
+            )
+        ],
+    )
+    talker_prompt = llm2tts([thinker_output], prompt=engine_prompt)[0]
+    talker_request = _Code2WavRequestStub(
+        external_req_id="request-0",
+        request_id="request-0",
+        model_intermediate_buffer=talker_prompt["model_intermediate_buffer"],
+    )
+    transfer_manager = _TransferManagerStub()
+
+    code2wav_payload = tts2code2wav_async_chunk(
+        transfer_manager,
+        {"codes": {"audio": torch.arange(7).reshape(-1, 1)}},
+        talker_request,
+        True,
+    )
+
+    assert code2wav_payload is not None
+    torch.testing.assert_close(code2wav_payload.codes.ref, torch.from_numpy(reference_waveform))
+    assert code2wav_payload.meta.ref_audio_sr == 16000
+
+
+@pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"])
+@pytest.mark.asyncio
+async def test_reference_audio_accepts_request_extra_mappings(
+    serving_chat: OmniOpenAIServingChat,
+    mocker: MockerFixture,
+    request_shape: str,
+) -> None:
+    reference_audio_url = "data:audio/wav;base64,AAAA"
+    reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+    serving_chat.engine_client = _EngineClientStub(
+        [_stage_config("MiniCPMO45OmniForConditionalGeneration")],
+    )
+    serving_chat.model_config = _ModelConfigStub()
+    engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+    if request_shape == "extra_body":
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "target text"}],
+            extra_body={"ref_audio": reference_audio_url},
+        )
+    else:
+        request = SimpleNamespace(model_extra={"ref_audio": reference_audio_url})
+    assert not hasattr(request, "ref_audio")
+    media_connector = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+    ).return_value
+    media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+    await serving_chat._attach_minicpmo45_reference_audio(engine_prompt, request)
+
+    media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+    preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY]
+    assert isinstance(preserved_reference, tuple)
+    np.testing.assert_array_equal(preserved_reference[0], reference_waveform)
+    assert preserved_reference[1] == 16000
+
+
+@pytest.mark.asyncio
+async def test_reference_audio_is_ignored_for_other_models(
+    serving_chat: OmniOpenAIServingChat,
+    mocker: MockerFixture,
+) -> None:
+    reference_audio_url = "data:audio/wav;base64,AAAA"
+    serving_chat.engine_client = _EngineClientStub([_stage_config("SomeOtherOmniForConditionalGeneration")])
+    engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+    media_connector_class = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+    )
+
+    await serving_chat._attach_minicpmo45_reference_audio(
+        engine_prompt,
+        _chat_request(reference_audio_url),
+    )
+
+    media_connector_class.assert_not_called()
+    assert MINICPMO45_REFERENCE_AUDIO_KEY not in cast(dict[str, object], engine_prompt)

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for OpenAI-compatible video API encoding helpers."""
 
 import base64
 import threading
 from io import BytesIO
+from typing import Any
 
 import av
 import httpx
@@ -41,7 +42,7 @@ def _install_http_transport(monkeypatch, handler, client_kwargs):
 @pytest.mark.asyncio
 async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
     requested_paths = []
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         requested_paths.append(request.url.path)
@@ -63,7 +64,7 @@ async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
 @pytest.mark.asyncio
 async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
     requested_paths = []
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         requested_paths.append(request.url.path)
@@ -81,7 +82,7 @@ async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_decode_image_url_reports_http_status(monkeypatch):
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         return httpx.Response(404)
@@ -95,7 +96,7 @@ async def test_decode_image_url_reports_http_status(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_decode_image_url_reports_connection_failure(monkeypatch):
-    client_kwargs = []
+    client_kwargs: list[dict[str, object]] = []
 
     def _handler(request):
         raise httpx.ConnectError("connection refused", request=request)
@@ -141,7 +142,7 @@ def _install_fake_video_mux(monkeypatch, mux_calls):
 
 
 def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
-    mux_calls = []
+    mux_calls: list[dict[str, Any]] = []
     _install_fake_video_mux(monkeypatch, mux_calls)
 
     frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)]
@@ -222,6 +223,19 @@ def test_channel_first_video_tensor_is_converted_to_channel_last():
     np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy())
 
 
+def test_contiguous_uint8_video_reaches_mux_without_conversion(monkeypatch):
+    video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
+
+    def fail_prepare(*args, **kwargs):
+        raise AssertionError("ready-to-encode uint8 video must not be normalized")
+
+    monkeypatch.setattr(video_api_utils, "_prepare_video_frames", fail_prepare)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames(video)
+
+    assert frames is video
+
+
 def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch):
     calls = []
 
@@ -506,6 +520,32 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
     assert len(coerce_calls) == 1
 
 
+def test_uint8_frames_reach_the_muxer_without_a_float_round_trip():
+    """uint8 is the muxer's own dtype, so normalisation must leave it alone.
+
+    MiniMax-H3 quantises on the accelerator to keep the response small across
+    its two process hops. Re-normalising to float here would undo that at the
+    last step and cost a full-size conversion in each direction.
+    """
+    video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
+
+    frames, frame_shape, common_dtype = video_api_utils._prepare_video_frames(video)
+
+    assert common_dtype == np.dtype(np.uint8)
+    assert frame_shape == (4, 5, 3)
+    np.testing.assert_array_equal(np.stack(frames), video)
+
+
+def test_non_uint8_integer_frames_are_still_normalised():
+    """Only uint8 is the muxer's dtype; other integer payloads still scale."""
+    video = np.full((2, 4, 5, 3), 255, dtype=np.int32)
+
+    frames, _, common_dtype = video_api_utils._prepare_video_frames(video)
+
+    assert np.issubdtype(common_dtype, np.floating)
+    np.testing.assert_allclose(frames[0], 1.0)
+
+
 @pytest.mark.parametrize(
     "video",
     [
@@ -538,9 +578,10 @@ def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch,
     monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
 
     prepared_frames, frame_shape, _ = video_api_utils._prepare_video_frames(video)
-    reference = np.rint(np.clip(np.stack([frame[..., :3] for frame in prepared_frames]), 0.0, 1.0) * 255.0).astype(
+    stacked = np.stack([frame[..., :3] for frame in prepared_frames])
-        np.uint8
+    # uint8 frames are the muxer's own dtype and reach it unchanged; everything
-    )
+    # else is still quantised out of the normalised [0, 1] range.
+    reference = stacked if stacked.dtype == np.uint8 else np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8)
 
     assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video"
     assert len(mux_inputs) == 1
@@ -549,6 +590,44 @@ def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch,
     np.testing.assert_array_equal(mux_inputs[0], reference)
 
 
+@pytest.mark.parametrize(
+    ("worker_count", "expected_path", "expected_reason"),
+    [
+        (None, "legacy_fallback", "non_contiguous_rgb_planes"),
+        (1, "legacy_fallback", "non_contiguous_rgb_planes"),
+        (8, "direct_planar", None),
+    ],
+)
+def test_interleaved_video_routing(monkeypatch, worker_count, expected_path, expected_reason):
+    routes = []
+
+    def fake_planar_mux(*args, **kwargs):
+        return b"planar-video"
+
+    def fake_compat_mux(frames, audio, **kwargs):
+        return b"legacy-video"
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux)
+    monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+    monkeypatch.setattr(video_api_utils, "_log_video_encoding_path", lambda **kwargs: routes.append(kwargs))
+
+    converter = None if worker_count is None else video_api_utils._PlanarFrameConverter(max_workers=worker_count)
+    try:
+        encoded = video_api_utils._encode_video_bytes(
+            np.zeros((2, 4, 6, 3), dtype=np.float32),
+            fps=12,
+            frame_converter=converter,
+        )
+    finally:
+        if converter is not None:
+            converter.shutdown()
+
+    assert encoded == (b"planar-video" if expected_path == "direct_planar" else b"legacy-video")
+    assert len(routes) == 1
+    assert routes[0]["selected_path"] == expected_path
+    assert routes[0].get("reason") == expected_reason
+
+
 @pytest.mark.parametrize(
     ("video", "reason"),
     [
@@ -690,7 +769,8 @@ def test_planar_bool_frames_match_bounded_compatible_output():
 )
 def test_planar_frame_rejects_undersized_plane(monkeypatch, plane_height, line_size):
     class FakePlane:
-        pass
+        height: int
+        line_size: int
 
     plane = FakePlane()
     plane.height = plane_height
@@ -758,6 +838,32 @@ def test_direct_and_compatible_paths_produce_equivalent_mp4(with_audio):
     np.testing.assert_array_equal(direct_frames, compatible_frames)
 
 
+@pytest.mark.parametrize("with_audio", [False, True])
+def test_parallel_interleaved_and_legacy_paths_produce_equivalent_mp4(with_audio):
+    rng = np.random.default_rng(7)
+    video = rng.random((3, 32, 32, 3), dtype=np.float32)
+    audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None
+    kwargs = {
+        "fps": 12,
+        "audio": audio,
+        "audio_sample_rate": 24000,
+        "video_codec_options": {"preset": "ultrafast", "threads": "1"},
+    }
+
+    converter = video_api_utils._PlanarFrameConverter(max_workers=8)
+    try:
+        parallel_bytes = video_api_utils._encode_video_bytes(
+            video,
+            **kwargs,
+            frame_converter=converter,
+        )
+    finally:
+        converter.shutdown()
+    legacy_bytes = video_api_utils._encode_video_bytes_legacy(video, **kwargs)
+
+    assert parallel_bytes == legacy_bytes
+
+
 def test_mux_closes_container_and_preserves_generator_error(monkeypatch):
     class GeneratorError(RuntimeError):
         pass

tests/entrypoints/test_omni_base_profiler.py

@@ -46,6 +46,16 @@ class TestOmniBaseProfiler:
             stage_ids=None,
         )
 
+    def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture):
+        """Legacy diffusion_batch_size callers receive an actionable error."""
+        mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine)
+        mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x)
+        mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize")
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"):
+            OmniBase(model="test-model", diffusion_batch_size=8)
+
     def test_start_profile_with_prefix(self, omni_base_instance, mock_engine):
         """Test that start_profile passes profile_prefix to collective_rpc."""
         omni_base_instance.start_profile(profile_prefix="test_trace")

tests/entrypoints/test_utils.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for vllm_omni.entrypoints.utils module."""
 
 import os
@@ -420,6 +423,7 @@ class TestLoadAndResolveStageConfigs:
         """Ensure that kwargs survive default stage creation."""
         engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
         kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
+        kwargs["revision"] = "pinned-revision"
         mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
         mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
 
@@ -433,6 +437,7 @@ class TestLoadAndResolveStageConfigs:
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
         assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
+        assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
 
     def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
         deploy_path = tmp_path / "qwen3_multi.yaml"

tests/examples/test_minicpmo_realtime_duplex_simple_demo.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import asyncio
 import base64
 import importlib.util
 import sys
@@ -172,6 +173,29 @@ def test_open_streaming_response_requires_post_commit_drain():
     )
 
 
+def test_final_input_commit_checkpoints_playback_first():
+    demo = _load_demo_module()
+    client = demo.RealtimeDuplexClient("ws://unused")
+    client.events.add(
+        {
+            "type": "response.created",
+            "response": {"id": "resp-zero-audio"},
+        }
+    )
+    calls = []
+
+    async def send(event):
+        calls.append(event)
+
+    client.send = send
+    commit_sent_at_s = asyncio.run(demo._commit_input_after_playback_checkpoint(client))
+
+    assert [event["type"] for event in calls] == ["playback.ack", "input_audio_buffer.commit"]
+    assert calls[0]["response_id"] == "resp-zero-audio"
+    assert calls[0]["played_ms"] == 0
+    assert isinstance(commit_sent_at_s, float)
+
+
 def _asymmetric_image():
     from PIL import Image
 

tests/model_executor/models/minicpmo_4_5/test_llm2tts.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the MiniCPM-o 4.5 stage 0 -> stage 1 bridge.
 
 Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2tts``:
@@ -25,6 +25,7 @@ import torch
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
     MiniCPMO45OmniForConditionalGeneration,
 )
+from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
 from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -313,6 +314,44 @@ class TestPromptAndMultiModal:
         )
         assert out[0]["multi_modal_data"] == mm
 
+    def test_reference_audio_is_added_to_stage_handoff(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+        waveform = torch.tensor([0.25, -0.5], dtype=torch.float32)
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={"multi_modal_data": {"audio": (waveform, 16000)}},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert intermediate_buffer["codes"]["ref"] == [0.25, -0.5]
+        assert intermediate_buffer["meta"]["ref_audio_sr"] == 16000
+
+    def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+        waveform = torch.tensor([0.125, -0.25], dtype=torch.float32)
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={MINICPMO45_REFERENCE_AUDIO_KEY: (waveform, 24000)},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert intermediate_buffer["codes"]["ref"] == [0.125, -0.25]
+        assert intermediate_buffer["meta"]["ref_audio_sr"] == 24000
+
+    def test_missing_reference_audio_does_not_add_stage_handoff_fields(self) -> None:
+        hidden_states = torch.zeros((2, _HIDDEN_DIM))
+
+        stage_outputs = llm2tts(
+            [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+            prompt={"multi_modal_data": {}},
+        )
+
+        intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+        assert "codes" not in intermediate_buffer
+        assert "ref_audio_sr" not in intermediate_buffer.get("meta", {})
+
     def test_internal_streaming_context_is_accepted(self) -> None:
         hidden = torch.zeros((2, _HIDDEN_DIM))
         out = llm2tts(

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -26,6 +26,7 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
+    StagePipelineConfig,
     _apply_platform_overrides,
     load_deploy_config,
     merge_pipeline_deploy,
@@ -65,7 +66,9 @@ class TestRegistryDeclaration:
         assert not hasattr(pipeline, "max_native_duplex_sessions")
 
     def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None:
-        pipeline = PipelineConfig(model_type="ordinary")
+        pipeline = PipelineConfig(
+            model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+        )
         assert pipeline.duplex_control_enabled is False
         assert pipeline.duplex_serving_adapter is None
         assert not hasattr(pipeline, "max_native_duplex_sessions")
@@ -80,10 +83,6 @@ class TestPipelineTopology:
         assert len(pipeline.stages) == 3
         assert [s.stage_id for s in pipeline.stages] == [0, 1, 2]
 
-    def test_topology_validates(self, pipeline: PipelineConfig) -> None:
-        # ``validate`` returns a list of structural errors; empty == valid.
-        assert pipeline.validate() == []
-
     def test_thinker_stage(self, pipeline: PipelineConfig) -> None:
         thinker = pipeline.get_stage(0)
         assert thinker is not None

tests/model_executor/models/test_nemotron_voicechat_registration.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
 
 import subprocess
@@ -59,7 +59,6 @@ def test_pipeline_registered_with_three_stages() -> None:
     assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
     assert code2wav.final_output and code2wav.final_output_type == "audio"
     assert code2wav.input_sources == (1,)
-    assert pipeline.validate() == []
 
 
 def test_architectures_registered() -> None:

tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections import defaultdict
 from types import SimpleNamespace
@@ -168,10 +168,14 @@ def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None:
 def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
     manager = _manager()
     request = _request("req")
-    request.additional_information = {
+    request.model_intermediate_buffer = {
         "codes": {"ref": [0.1, -0.1]},
         "meta": {"ref_audio_sr": 16000},
     }
+    request.additional_information = {
+        "codes": {"ref": [0.9]},
+        "meta": {"ref_audio_sr": 8000},
+    }
 
     payload = tts2code2wav_async_chunk(
         manager,
@@ -193,10 +197,33 @@ def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
     assert payload.meta.turn_end is True
 
 
+def test_first_chunk_falls_back_to_legacy_reference_fields() -> None:
+    manager = _manager()
+    request = _request("req")
+    request.model_intermediate_buffer = {
+        "meta": {"ref_audio_sr": 16000},
+    }
+    request.additional_information = {
+        "codes": {"ref": [0.9]},
+        "meta": {"ref_audio_sr": 8000},
+    }
+
+    payload = tts2code2wav_async_chunk(
+        manager,
+        _duplex_delta(*range(7), turn_end=True),
+        request,
+        True,
+    )
+
+    assert payload is not None
+    assert payload.codes.ref.tolist() == pytest.approx([0.9])
+    assert payload.meta.ref_audio_sr == 16000
+
+
 def test_full_payload_forwards_all_codes_and_request_metadata() -> None:
     manager = _manager()
     request = _request("req")
-    request.additional_information = {
+    request.model_intermediate_buffer = {
         "codes": {"ref": [0.1, -0.1]},
         "meta": {
             "ref_audio_sr": 16000,

tests/platforms/npu/test_diffusion_platform.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("available", [False, True])
+def test_dense_flash_capability_tracks_mindiesd(monkeypatch, available):
+    pytest.importorskip("vllm_ascend")
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    monkeypatch.setattr(
+        importlib.util,
+        "find_spec",
+        lambda name: object() if name == "mindiesd" and available else None,
+    )
+
+    assert NPUOmniPlatform.supports_diffusion_dense_flash_attention() is available
+
+
+def test_paged_config_uses_ascend_kernel_block_size() -> None:
+    pytest.importorskip("vllm_ascend")
+    from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
+
+    from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
+    NPUOmniPlatform.configure_diffusion_vllm_config(
+        vllm_config,
+        SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER),
+    )
+
+    assert vllm_config.cache_config.block_size == AscendAttentionBackend.get_supported_kernel_block_sizes()[0]
+
+
+def test_strict_ulysses_paged_backend_bypasses_ascend_pcp() -> None:
+    pytest.importorskip("vllm_ascend")
+    from vllm_ascend.attention.attention_v1 import (
+        AscendAttentionBackend,
+        AscendAttentionBackendImpl,
+        AscendAttentionMetadataBuilder,
+    )
+
+    from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+    assert (
+        NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+            AscendAttentionBackend,
+            ulysses_degree=1,
+        )
+        is AscendAttentionBackend
+    )
+    backend = NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+        AscendAttentionBackend,
+        ulysses_degree=2,
+    )
+    assert backend.get_impl_cls() is AscendAttentionBackendImpl
+    assert backend.get_builder_cls() is AscendAttentionMetadataBuilder

tests/platforms/test_record_device_event.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for OmniPlatform.record_device_event implementations."""
 
 import pytest
@@ -122,8 +122,8 @@ class TestNPUOmniPlatformRecordDeviceEvent:
         result = NPUOmniPlatform.record_device_event()
         assert result is None
 
-    def test_synchronizes_stream_then_records_event(self, mocker):
+    def test_synchronizes_stream_then_records_generic_event(self, mocker):
-        """NPU should sync stream first, then create and record an Event."""
+        """NPU should return an event consumable by the generic torch.Stream."""
         try:
             from vllm_omni.platforms.npu.platform import NPUOmniPlatform
         except ModuleNotFoundError:
@@ -133,14 +133,16 @@ class TestNPUOmniPlatformRecordDeviceEvent:
         mock_stream = mocker.MagicMock()
         mock_torch = mocker.MagicMock()
         mock_torch.npu.current_stream.return_value = mock_stream
-        mock_torch.npu.Event.return_value = mock_event
+        mock_torch.Event.return_value = mock_event
         mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch)
 
         result = NPUOmniPlatform.record_device_event()
 
         # Stream should be synced first (HCCL ordering)
         mock_stream.synchronize.assert_called_once()
-        # Then event should be created and recorded
+        # Then a backend-neutral event should be created and recorded.  The
-        mock_torch.npu.Event.assert_called_once()
+        # worker waits on it through the public torch.Stream wrapper.
+        mock_torch.Event.assert_called_once()
+        mock_torch.npu.Event.assert_not_called()
         mock_event.record.assert_called_once()
         assert result is mock_event

tests/utils/test_tracking_parser.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for TrackingArgumentParser and related utilities."""
 
 import argparse
@@ -35,7 +38,7 @@ _TEST_PIPELINE = PipelineConfig(
     stages=(
         StagePipelineConfig(stage_id=0, model_stage="stage0"),
         StagePipelineConfig(stage_id=1, model_stage="stage1"),
-        StagePipelineConfig(stage_id=2, model_stage="stage2"),
+        StagePipelineConfig(stage_id=2, model_stage="stage2", final_output=True),
     ),
 )
 

vllm_omni/benchmarks/patch/patch.py

@@ -119,6 +119,17 @@ def should_request_stage_metrics(args: Any) -> bool:
     if backend in _IMAGE_STAGE_METRICS_BACKENDS:
         return True
 
+    percentile_metrics = getattr(args, "percentile_metrics", ()) or ()
+    if isinstance(percentile_metrics, str):
+        percentile_metrics = percentile_metrics.split(",")
+    selected_metrics = {str(metric).strip().lower() for metric in percentile_metrics}
+    if backend == "openai-chat-omni" and selected_metrics.intersection({"tpot", "itl"}):
+        # Client receive timestamps normally provide these metrics. Under
+        # enough load, however, every text token can be coalesced into one
+        # event-loop read, leaving no positive interval. Request the engine's
+        # Stage 0 timings so that path still has an authoritative measurement.
+        return True
+
     extra_body = getattr(args, "extra_body", None) or {}
     modalities = extra_body.get("modalities") if isinstance(extra_body, dict) else None
     return backend == "openai-chat-omni" and "image" in (modalities or [])
@@ -854,6 +865,33 @@ def _update_output_stage_metrics_from_payload(
         output.stage_metrics.update(stage_snapshot)
 
 
+def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool:
+    """Apply native Stage 0 timings from a chat response snapshot."""
+    stage_metrics = output.stage_metrics
+    if not isinstance(stage_metrics, dict):
+        return False
+    stage0 = stage_metrics.get("0")
+    if not isinstance(stage0, dict):
+        return False
+
+    output_tokens = coerce_positive_int_scalar(stage0.get(defs.NUM_TOKENS_OUT))
+    expected_output_tokens = coerce_positive_int_scalar(output.output_tokens)
+    if output_tokens is None or expected_output_tokens is None or output_tokens != expected_output_tokens:
+        return False
+
+    return _apply_stage0_token_timings(
+        output,
+        [
+            {
+                "output_token_count": output_tokens,
+                "itls_ms": stage0.get(defs.VLLM_ITLS_MS),
+                "tpot_ms": stage0.get(defs.VLLM_TPOT_MS),
+            }
+        ],
+        expected_output_tokens=expected_output_tokens,
+    )
+
+
 def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]:
     if not isinstance(metrics, dict):
         return 0, 0.0, 0, 0.0
@@ -1121,6 +1159,14 @@ async def async_request_openai_chat_omni_completions(
 
                     output.latency = timestamp - st
                     output.generated_text = generated_text
+                    if output.output_tokens > 1 and not any(
+                        isinstance(value, int | float)
+                        and not isinstance(value, bool)
+                        and np.isfinite(value)
+                        and value > 0
+                        for value in output.itl
+                    ):
+                        _apply_chat_stage0_token_timings(output)
                     if output.itl:
                         # Align text_latency with ITL so TPOT formula and
                         # mean(ITL) are consistent.  Do NOT infer output_tokens

vllm_omni/config/pipeline_registry.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Pipeline registry and factory for vllm-omni.
 
 ``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig``
@@ -196,18 +196,13 @@ def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_typ
     since resolvers can return multiple different PipelineConfigs depending on the
     consumed config.
     """
-    errors: list[str] = []
     if isinstance(pipeline, PipelineConfig):
-        errors = pipeline.validate()
         model_type = model_type if model_type is not None else pipeline.model_type
-    else:
+    elif model_type is None:
-        if model_type is None:
+        raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
-            raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
 
     if model_type in OMNI_PIPELINES:
-        errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered")
+        logger.warning(f"Model type {model_type} is already registered; the old mapping will be clobbered")
-    if errors:
-        logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors)
     OMNI_PIPELINES[model_type] = pipeline
 
 

vllm_omni/config/stage_config.py

@@ -302,6 +302,17 @@ class PipelineConfig:
     # Global CLI spelling -> (stage id, stage-local spelling).
     stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict)
 
+    def __post_init__(self) -> None:
+        errors = self.get_validation_errors()
+
+        if errors:
+            # If we have multiple errors, put them on separate indented lines for readability
+            multi_sep = "\n\t"
+            error_str = multi_sep.join(errors)
+            if len(errors) > 1:
+                error_str = f"{multi_sep}{error_str}"
+            raise ValueError(f"PipelineConfig initialization failed with the following error(s): {error_str}")
+
     def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
         """Look up a stage by its ID."""
         for stage in self.stages:
@@ -309,7 +320,7 @@ class PipelineConfig:
                 return stage
         return None
 
-    def validate(self) -> list[str]:
+    def get_validation_errors(self) -> list[str]:
         """Return list of topology errors (empty if valid)."""
         errors: list[str] = []
         if not self.stages:
@@ -327,6 +338,11 @@ class PipelineConfig:
                     errors.append(f"Stage {stage.stage_id} references itself")
         if not any(not s.input_sources for s in self.stages):
             errors.append("No entry point (stage with empty input_sources)")
+        # Request completion and client-visible output are both driven by stages
+        # that declare ``final_output`` (see ``Orchestrator._route_output``). A
+        # pipeline without one can never emit a result, so every request hangs.
+        if not any(s.final_output for s in self.stages):
+            errors.append("No terminal stage (stage with final_output=True)")
         return errors
 
 

vllm_omni/deploy/wan2_2_ti2v.yaml

@@ -1,7 +1,6 @@
 # Wan2.2 TI2V-5B deploy.
 #
 # Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants.
-# Request-level batching uses scheduler max_num_seqs, not diffusion_batch_size.
 async_chunk: false
 trust_remote_code: true
 distributed_executor_backend: mp

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import os
 from functools import partial
@@ -7,11 +7,7 @@ from functools import partial
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.attention.backends.abstract import (
+from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
@@ -285,7 +281,9 @@ class FlashAttentionImpl(AttentionImpl):
         rank-local native cache context.  The adapter no longer performs the
         attention call itself; this method is the backend-owned execution
         boundary.  Its native layer wrapper keeps vLLM version-specific cache
-        and kernel details out of Omni's common ``Attention`` layer.
+        and kernel details out of Omni's common ``Attention`` layer. CUDA uses
+        the native vLLM FlashAttention writer/kernel contract. Ascend writes
+        the complete K/V span once before its piecewise FIA calls.
         """
 
         layer = getattr(paged_kv_context, "layer", None)
@@ -297,8 +295,18 @@ class FlashAttentionImpl(AttentionImpl):
         native_impl = getattr(layer, "impl", None)
         if native_impl is None:
             raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}")
-        if not layer.attn_backend.forward_includes_kv_cache_update:
+        prewrite_kv = current_omni_platform.requires_diffusion_paged_kv_prewrite()
-            native_impl.do_kv_cache_update(
+        read_kv_from_cache = prewrite_kv and layer.attn_backend.forward_includes_kv_cache_update
+        # The GPU/default path preserves native cache-update ownership. Ascend
+        # prewrites through vLLM-Ascend's normal-layout writer so piecewise FIA
+        # segments do not repeatedly scatter the same layer K/V.
+        if not layer.attn_backend.forward_includes_kv_cache_update or prewrite_kv:
+            cache_update = getattr(native_impl, "do_kv_cache_update", None)
+            if not callable(cache_update):
+                raise RuntimeError(
+                    f"Native attention implementation for {layer.layer_name!r} cannot update the KV cache"
+                )
+            cache_update(
                 layer,
                 paged_kv_context.key_write,
                 paged_kv_context.value_write,
@@ -308,33 +316,52 @@ class FlashAttentionImpl(AttentionImpl):
 
         def run_native_attention(
             query: torch.Tensor,
-            key: torch.Tensor,
+            key: torch.Tensor | None,
-            value: torch.Tensor,
+            value: torch.Tensor | None,
             native_metadata,
+            output: torch.Tensor | None,
         ) -> torch.Tensor:
-            output = torch.empty(
+            if output is None:
-                (query.shape[0], layer.num_heads, layer.head_size_v),
+                output = torch.empty(
-                dtype=query.dtype,
+                    (query.shape[0], layer.num_heads, layer.head_size_v),
-                device=query.device,
+                    dtype=query.dtype,
-            )
+                    device=query.device,
+                )
             return native_impl.forward(
                 layer,
                 query,
-                key,
+                None if read_kv_from_cache else key,
-                value,
+                None if read_kv_from_cache else value,
                 kv_cache,
                 native_metadata,
                 output,
             )
 
         if paged_kv_context.piecewise_plan is not None:
+            # Ascend FIA can execute identical CFG rows as one batched call
+            # per piece.  Keep that batch layout through the piecewise runner
+            # so it can concatenate row-major results instead of emitting an
+            # indexed ScatterUpdate for every segment.  The CUDA path keeps
+            # its existing output-buffer contract (including graph capture).
+            use_homogeneous_batch = (
+                current_omni_platform.is_npu() and paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None
+            )
+            output = None
+            if not use_homogeneous_batch:
+                output = torch.empty(
+                    (paged_kv_context.query.shape[0], layer.num_heads, layer.head_size_v),
+                    dtype=paged_kv_context.query.dtype,
+                    device=paged_kv_context.query.device,
+                )
             output = run_paged_piecewise_plan(
                 paged_kv_context.query,
-                paged_kv_context.key_write,
+                None if read_kv_from_cache else paged_kv_context.key_write,
-                paged_kv_context.value_write,
+                None if read_kv_from_cache else paged_kv_context.value_write,
                 paged_kv_context.piecewise_plan,
                 paged_kv_context.piecewise_native_metadata,
                 run_native_attention,
+                output_buffer=output,
+                use_homogeneous_batch=use_homogeneous_batch,
             )
         else:
             output = run_native_attention(
@@ -342,6 +369,7 @@ class FlashAttentionImpl(AttentionImpl):
                 paged_kv_context.key_write,
                 paged_kv_context.value_write,
                 paged_kv_context.native_metadata,
+                None,
             )
         return paged_kv_context.restore_output(output)
 
@@ -516,15 +544,8 @@ class FlashAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        try:
+        from mindiesd import attention_forward
-            from mindiesd import attention_forward
+
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
         # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
         # model marks extra["npu_attn_varlen"] and carries packed metadata, so
         # the padding document is excluded without reading or materializing
@@ -633,15 +654,8 @@ class FlashAttentionImpl(AttentionImpl):
             return None
         seq_q, seq_k = resolved
 
-        try:
+        from mindiesd import attention_forward_varlen
-            from mindiesd import attention_forward_varlen
+
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
         q = query.squeeze(0)  # [T, N, D] == TND
         k = key.squeeze(0)
         v = value.squeeze(0)
@@ -689,15 +703,8 @@ class FlashAttentionImpl(AttentionImpl):
         _, seq_k = resolved
         used_k = seq_k[0]  # real document length (first cumulative end)
 
-        try:
+        from mindiesd import attention_forward
-            from mindiesd import attention_forward
+
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
         # mindiesd always takes BSND input regardless of `layout`; the arg
         # selects the op-internal layout and laser only supports BNSD, so do
         # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.

vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Piecewise attention for mixed causal / full (bidirectional) masks.
 
@@ -92,6 +92,7 @@ def piecewise_attn(
 ):
     _check_homogeneous(full_attn_spans)
     spans = full_attn_spans[0]
+    ranges: tuple[tuple[int, int, int], ...]
     if query_ranges is None:
         query_len = query.shape[1]
         ranges = ((0, query_len, key.shape[1] - query_len),)
@@ -133,6 +134,11 @@ class PagedPiecewiseSegment:
 
     row_segments: tuple[Segment, ...]
     query_indices: torch.Tensor
+    query_range: tuple[int, int] | None
+    # Local [start, end) range shared by every row when the batch can retain
+    # its original [B, S] layout.  The packed query_indices remain available
+    # for metadata construction and the heterogeneous fallback.
+    local_query_range: tuple[int, int] | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -142,10 +148,16 @@ class PagedPiecewisePlan:
     spans: tuple[tuple[tuple[int, int], ...], ...]
     segments: tuple[PagedPiecewiseSegment, ...]
     num_query_tokens: int
+    segments_cover_query_contiguously: bool
+    # (batch_size, query_len) when all rows have the same local segment
+    # layout.  Native FIA still receives flattened row-major tokens, but the
+    # runner can restore outputs with one batch concat instead of one indexed
+    # scatter per segment.
+    homogeneous_batch_shape: tuple[int, int] | None = None
 
 
 PagedPiecewiseRunner = Callable[
-    [torch.Tensor, torch.Tensor, torch.Tensor, object],
+    [torch.Tensor, torch.Tensor | None, torch.Tensor | None, object, torch.Tensor | None],
     torch.Tensor,
 ]
 
@@ -164,7 +176,9 @@ def build_paged_piecewise_plan(
     if row_count == 0 or not (row_count == len(query_offsets) == len(query_lens) == len(seq_lens)):
         raise ValueError("Paged piecewise inputs must have one entry per row")
 
-    spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans)
+    spans: tuple[tuple[tuple[int, int], ...], ...] = tuple(
+        tuple((span_start, span_end) for span_start, span_end in row_spans) for row_spans in full_attn_spans
+    )
 
     packed_offsets = [0]
     for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True):
@@ -188,46 +202,206 @@ def build_paged_piecewise_plan(
     for row_segments in zip(*segments_by_row, strict=True):
         if len({segment.mode for segment in row_segments}) != 1:
             raise ValueError("Paged piecewise rows must use the same attention mode per segment")
-        query_indices = []
+        query_indices: list[int] = []
+        local_ranges: list[tuple[int, int]] = []
         for row_index, (segment, query_offset) in enumerate(zip(row_segments, query_offsets, strict=True)):
             local_start = segment.q_start - query_offset
             local_end = segment.q_end - query_offset
+            local_ranges.append((local_start, local_end))
             query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end))
+        query_range = None
+        if query_indices and query_indices == list(range(query_indices[0], query_indices[-1] + 1)):
+            query_range = (query_indices[0], query_indices[-1] + 1)
+        local_query_range = local_ranges[0] if len(set(local_ranges)) == 1 else None
         packed_segments.append(
             PagedPiecewiseSegment(
                 row_segments=tuple(row_segments),
                 query_indices=torch.tensor(query_indices, dtype=torch.long, device=device),
+                query_range=query_range,
+                local_query_range=local_query_range,
             )
         )
 
-    return PagedPiecewisePlan(spans, tuple(packed_segments), packed_offsets[-1])
+    covered = 0
+    for segment in packed_segments:
+        if segment.query_range is None or segment.query_range[0] != covered:
+            break
+        covered = segment.query_range[1]
+    segments_cover_query_contiguously = covered == packed_offsets[-1]
+    homogeneous_batch_shape: tuple[int, int] | None = None
+    if (
+        row_count > 1
+        and len(set(query_lens)) == 1
+        and all(segment.local_query_range is not None for segment in packed_segments)
+    ):
+        # A common local range is sufficient even when rows have different
+        # global offsets or KV lengths; those values remain row-specific in
+        # native metadata.
+        homogeneous_batch_shape = (row_count, int(query_lens[0]))
+    return PagedPiecewisePlan(
+        spans=spans,
+        segments=tuple(packed_segments),
+        num_query_tokens=packed_offsets[-1],
+        segments_cover_query_contiguously=segments_cover_query_contiguously,
+        homogeneous_batch_shape=homogeneous_batch_shape,
+    )
+
+
+def _run_homogeneous_paged_piecewise_plan(
+    query: torch.Tensor,
+    key: torch.Tensor | None,
+    value: torch.Tensor | None,
+    plan: PagedPiecewisePlan,
+    segment_metadata: Sequence[object],
+    segment_runner: PagedPiecewiseRunner,
+    output_buffer: torch.Tensor | None,
+) -> torch.Tensor:
+    """Run a homogeneous batch while retaining the old [B, S] layout.
+
+    The native FIA contract is flattened TND, so each segment is flattened
+    only for the kernel call.  Results are reshaped back to [B, L] and
+    concatenated along the sequence dimension, which avoids the per-segment
+    ``index_copy_`` used by heterogeneous packed rows.
+    """
+
+    batch_shape = plan.homogeneous_batch_shape
+    assert batch_shape is not None
+    batch_size, query_len = batch_shape
+    if query.shape[0] != batch_size * query_len:
+        raise ValueError(
+            f"Homogeneous paged piecewise plan expects {batch_size * query_len} query tokens, got {query.shape[0]}"
+        )
+    query_rows = query.reshape(batch_size, query_len, *query.shape[1:])
+    key_rows = None if key is None else key.reshape(batch_size, query_len, *key.shape[1:])
+    value_rows = None if value is None else value.reshape(batch_size, query_len, *value.shape[1:])
+
+    segment_outputs: list[torch.Tensor] = []
+    for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
+        local_range = segment.local_query_range
+        if local_range is None:
+            raise ValueError("Homogeneous paged piecewise execution requires a shared local segment range")
+        start, end = local_range
+        segment_query = query_rows[:, start:end].contiguous().reshape(-1, *query.shape[1:])
+        segment_key = None if key_rows is None else key_rows[:, start:end].contiguous().reshape(-1, *key.shape[1:])
+        segment_value = (
+            None if value_rows is None else value_rows[:, start:end].contiguous().reshape(-1, *value.shape[1:])
+        )
+        segment_output = segment_runner(segment_query, segment_key, segment_value, metadata, None)
+        expected_tokens = batch_size * (end - start)
+        if segment_output.shape[0] != expected_tokens:
+            raise ValueError(
+                f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+                f"for a {expected_tokens}-token homogeneous segment"
+            )
+        segment_outputs.append(segment_output.reshape(batch_size, end - start, *segment_output.shape[1:]))
+
+    if not segment_outputs:
+        result = torch.empty_like(query)
+    elif len(segment_outputs) == 1:
+        result = segment_outputs[0].reshape(plan.num_query_tokens, *segment_outputs[0].shape[2:])
+    else:
+        result_rows = torch.cat(segment_outputs, dim=1)
+        result = result_rows.reshape(plan.num_query_tokens, *result_rows.shape[2:])
+    if output_buffer is not None:
+        if output_buffer.shape != result.shape:
+            raise ValueError(
+                f"Paged piecewise output buffer shape {tuple(output_buffer.shape)} does not match "
+                f"homogeneous result shape {tuple(result.shape)}"
+            )
+        output_buffer.copy_(result)
+        return output_buffer
+    return result
 
 
 def run_paged_piecewise_plan(
     query: torch.Tensor,
-    key: torch.Tensor,
+    key: torch.Tensor | None,
-    value: torch.Tensor,
+    value: torch.Tensor | None,
     plan: PagedPiecewisePlan,
     segment_metadata: Sequence[object],
     segment_runner: PagedPiecewiseRunner,
+    output_buffer: torch.Tensor | None = None,
+    *,
+    use_homogeneous_batch: bool = False,
 ) -> torch.Tensor:
-    """Run and scatter piecewise attention over packed Q/K/V write tensors."""
+    """Run piecewise attention with view fast paths for contiguous segments.
+
+    ``use_homogeneous_batch`` is an opt-in for native backends that can keep
+    identical rows in the legacy batch layout.  The default remains the
+    indexed packed path so existing GPU/heterogeneous contracts are unchanged.
+    """
 
     if query.shape[0] != plan.num_query_tokens:
         raise ValueError(f"Paged piecewise plan has {plan.num_query_tokens} query tokens, got {query.shape[0]}")
+    if output_buffer is not None and output_buffer.shape[0] != plan.num_query_tokens:
+        raise ValueError(
+            f"Paged piecewise output buffer has {output_buffer.shape[0]} tokens; expected {plan.num_query_tokens}"
+        )
 
-    output = None
+    if use_homogeneous_batch and plan.homogeneous_batch_shape is not None:
+        return _run_homogeneous_paged_piecewise_plan(
+            query,
+            key,
+            value,
+            plan,
+            segment_metadata,
+            segment_runner,
+            output_buffer,
+        )
+
+    output = output_buffer
+    contiguous_outputs = []
     for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
         indices = segment.query_indices
+        segment_output_buffer = None
+        if segment.query_range is None:
+            segment_query = query.index_select(0, indices)
+            segment_key = None if key is None else key.index_select(0, indices)
+            segment_value = None if value is None else value.index_select(0, indices)
+        else:
+            start, end = segment.query_range
+            segment_query = query[start:end]
+            segment_key = None if key is None else key[start:end]
+            segment_value = None if value is None else value[start:end]
+            if output is not None:
+                segment_output_buffer = output[start:end]
         segment_output = segment_runner(
-            query.index_select(0, indices),
+            segment_query,
-            key.index_select(0, indices),
+            segment_key,
-            value.index_select(0, indices),
+            segment_value,
             metadata,
+            segment_output_buffer,
         )
+        if segment_output.shape[0] != indices.shape[0]:
+            raise ValueError(
+                f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+                f"for a {indices.shape[0]}-token segment"
+            )
+        if segment_output_buffer is not None and segment_output is not segment_output_buffer:
+            if segment_output.shape != segment_output_buffer.shape:
+                raise ValueError(
+                    f"Paged piecewise runner returned shape {tuple(segment_output.shape)} "
+                    f"for output buffer shape {tuple(segment_output_buffer.shape)}"
+                )
+            segment_output_buffer.copy_(segment_output)
+            segment_output = segment_output_buffer
+        if plan.segments_cover_query_contiguously:
+            if output_buffer is None:
+                contiguous_outputs.append(segment_output)
+            continue
         if output is None:
             output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:]))
-        output.index_copy_(0, indices, segment_output)
+        if segment.query_range is None:
+            output.index_copy_(0, indices, segment_output)
+        else:
+            start, end = segment.query_range
+            output[start:end].copy_(segment_output)
 
+    if output_buffer is not None and plan.segments_cover_query_contiguously:
+        return output_buffer
+    if contiguous_outputs:
+        if len(contiguous_outputs) == 1:
+            return contiguous_outputs[0]
+        return torch.cat(contiguous_outputs, dim=0)
     assert output is not None
     return output

vllm_omni/diffusion/attention/layer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright (c) Microsoft Corporation and Jiarui Fang
 # SPDX-License-Identifier: Apache-2.0
 # DeepSpeed Team & Jiarui Fang
@@ -23,6 +23,7 @@ from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
 from vllm_omni.diffusion.forward_context import (
     get_forward_context,
@@ -44,6 +45,8 @@ def _try_extract_layer_index(prefix: str) -> int | None:
 
 
 class Attention(nn.Module):
+    _scheduler_paged_kv = False
+
     def __init__(
         self,
         num_heads: int,
@@ -111,6 +114,44 @@ class Attention(nn.Module):
             role_category=role_category,
             allow_trtllm_default=allow_trtllm_default,
         )
+        scheduler_paged_kv = (
+            config is not None
+            and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+        )
+        self._scheduler_paged_kv = scheduler_paged_kv
+        if (
+            scheduler_paged_kv
+            and paged_kv_cache_role is not None
+            and spec is None
+            and not attn_backend_cls.supports_paged_kv
+        ):
+            # FLASH_ATTN is an Omni selector, not a device-specific kernel.
+            # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
+            # attention on NPU. A platform may legitimately default dense
+            # attention to SDPA (for example when an optional diffusion FA
+            # package is absent), but a paged layer must advertise the
+            # capability so the Worker can register its native cache view.
+            # Resolve the selector only for this marked layer; unmarked dense
+            # attention keeps the platform default unchanged.
+            from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+            dense_backend_name = attn_backend_cls.get_name()
+            paged_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+            attn_backend_cls, spec = get_attn_backend_for_role(
+                role=role,
+                head_size=head_size,
+                attention_config=paged_attention_config,
+                role_category=role_category,
+                allow_trtllm_default=allow_trtllm_default,
+            )
+            logger.info(
+                "Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
+                "does not support Scheduler-owned KV",
+                role,
+                attn_backend_cls.get_name(),
+                dense_backend_name,
+            )
         parallel_config = getattr(config, "parallel_config", None)
         allgather_degree = getattr(parallel_config, "allgather_degree", 1)
         # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
@@ -321,6 +362,17 @@ class Attention(nn.Module):
         # Get the appropriate parallel strategy based on SP active state
         strategy = self._get_active_parallel_strategy()
         paged_adapter = self._active_paged_kv_adapter()
+        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+        if (
+            self._scheduler_paged_kv
+            and self.paged_kv_cache_role is not None
+            and paged_adapter is None
+            and not in_kv_memory_profile
+        ):
+            raise RuntimeError(
+                "Scheduler-paged diffusion attention reached model forward without an active Worker adapter. "
+                "Only the startup KV memory profile may execute before paged KV initialization."
+            )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
             raise NotImplementedError(
@@ -343,9 +395,39 @@ class Attention(nn.Module):
         # For Ring: Concat joint_q
         query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)
 
-        # 2. Kernel execution stays inside the selected Omni backend.  The
+        # Scheduler rows describe the logical sequence, while strict Ulysses
-        # Worker adapter only prepares the native page-table context after
+        # may append synthetic tokens solely to make the image shard divisible.
-        # SP has produced rank-local Q/K/V tensors.
+        # Remove those tokens after the all-to-all and put zero placeholders
+        # back before the reverse all-to-all.
+        paged_sp_padding = 0
+        paged_sp_padding_offset = query.shape[1]
+        if use_paged_attention and strategy is not self._no_parallel_strategy and strategy.name == "ulysses":
+            forward_ctx = get_forward_context() if is_forward_context_available() else None
+            paged_sp_padding = int(getattr(forward_ctx, "sp_padding_size", 0))
+            if paged_sp_padding:
+                joint_len = int(getattr(ctx, "joint_len", 0))
+                joint_strategy = str(getattr(ctx, "joint_strategy", "front"))
+                paged_sp_padding_offset = query.shape[1] - (joint_len if joint_strategy == "rear" else 0)
+                padding_start = paged_sp_padding_offset - paged_sp_padding
+                if padding_start < 0:
+                    raise ValueError(
+                        "Paged Ulysses padding exceeds the post-all-to-all sequence: "
+                        f"padding={paged_sp_padding}, sequence={query.shape[1]}"
+                    )
+
+                def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor:
+                    return torch.cat(
+                        (tensor[:, :padding_start], tensor[:, paged_sp_padding_offset:]),
+                        dim=1,
+                    ).contiguous()
+
+                query = _remove_paged_sp_padding(query)
+                key = _remove_paged_sp_padding(key)
+                value = _remove_paged_sp_padding(value)
+
+        # 2. This is the shared GPU/NPU boundary. The Worker adapter prepares
+        # the native page-table context after SP has produced rank-local Q/K/V;
+        # backend resolution below selects CUDA or Ascend execution.
         if use_paged_attention:
             assert paged_adapter is not None
             paged_kv_context = paged_adapter.prepare_layer_context(
@@ -363,6 +445,11 @@ class Attention(nn.Module):
             else:
                 out = self._run_local_attention(query, key, value, attn_metadata)
 
+        if paged_sp_padding:
+            padding_start = paged_sp_padding_offset - paged_sp_padding
+            output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:]))
+            out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1)
+
         # 3. Post-processing (Reverse Communication)
         # For Ulysses: AllToAll Output, and AllGather Joint Output
         out = strategy.post_attention(out, ctx)
@@ -371,7 +458,7 @@ class Attention(nn.Module):
 
     @staticmethod
     def _active_paged_kv_adapter():
-        """Return the opaque Worker adapter installed for this forward."""
+        """Return the Worker adapter selected by Runner-owned metadata."""
 
         if not is_forward_context_available():
             return None
@@ -392,7 +479,25 @@ class Attention(nn.Module):
             )
             return self.sdpa_fallback.forward(query, key, value, attn_metadata)
 
-        # Fallback to standard attention
+        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+        # NPU only: the startup KV-capacity profile needs tensor shapes, not a
+        # paged attention result. If dense MindIE-SD is absent, SDPA provides
+        # that profile forward. Formal GPU/NPU paged requests never use this
+        # branch because their Worker adapter is active.
+        if (
+            self._scheduler_paged_kv
+            and self.paged_kv_cache_role is not None
+            and in_kv_memory_profile
+            and current_omni_platform.is_npu()
+            and self.attn_backend.get_name() == "FLASH_ATTN"
+            and not current_omni_platform.supports_diffusion_dense_flash_attention()
+        ):
+            logger.warning_once(
+                "The startup KV memory profile is using SDPA because MindIE-SD is unavailable. "
+                "Formal paged requests still use the platform-native paged attention backend."
+            )
+            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
         return self.attention.forward(query, key, value, attn_metadata)
 
     def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None:

vllm_omni/diffusion/data.py

@@ -1086,10 +1086,9 @@ class OmniDiffusionConfig:
             "need_recv_cache", False
         ):
             raise ValueError(
-                "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
+                "paged_scheduler Diffusion KV does not support imported AR KV; "
-                "disable need_recv_cache until connector-aware admission is implemented"
+                "disable need_recv_cache until connector-aware import is implemented"
             )
-
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
         if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import copy
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any
 
@@ -22,9 +22,9 @@ from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
-    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionMetadata,
     DiffusionPagedAttentionRowBinding,
-    PreparedDiffusionPagedAttentionBatch,
+    DiffusionPagedAttentionRuntime,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -32,13 +32,6 @@ DiffusionKVIdentity = tuple[str, int | None, str | None]
 DiffusionKVSnapshot = tuple[object, ...]
 
 
-@dataclass(frozen=True)
-class _DiffusionKVRequestState:
-    generation: int
-    snapshot: DiffusionKVSnapshot
-    row_token_lens: tuple[tuple[DiffusionKVIdentity, int], ...]
-
-
 @dataclass(frozen=True)
 class _DiffusionKVRowInstall:
     identity: DiffusionKVIdentity
@@ -46,6 +39,13 @@ class _DiffusionKVRowInstall:
     block_ids: tuple[tuple[int, ...], ...]
 
 
+@dataclass(frozen=True)
+class _DiffusionKVRequestState:
+    generation: int
+    snapshot: DiffusionKVSnapshot
+    row_installs: tuple[_DiffusionKVRowInstall, ...]
+
+
 class DiffusionKVModelRunnerBackend:
     """Native paged-KV state and BlockTable operations for a model runner."""
 
@@ -503,7 +503,7 @@ class DiffusionKVModelRunnerBackend:
             self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
                 generation=metadata.allocation_generation,
                 snapshot=snapshot,
-                row_token_lens=current_state.row_token_lens,
+                row_installs=current_state.row_installs,
             )
             return False
 
@@ -530,7 +530,7 @@ class DiffusionKVModelRunnerBackend:
         self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
             generation=metadata.allocation_generation,
             snapshot=snapshot,
-            row_token_lens=tuple((install.identity, install.token_len) for install in installs),
+            row_installs=tuple(installs),
         )
         return True
 
@@ -559,24 +559,11 @@ class DiffusionKVModelRunnerBackend:
             raise RuntimeError("paged_scheduler native attention adapter is not initialized")
         return self.paged_attention_adapter
 
-    def prepare_paged_attention_batch(
+    def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata):
-        self,
+        """Activate Runner-owned request metadata for an internal denoise loop."""
-        rows: Sequence[DiffusionPagedAttentionRow],
-    ) -> PreparedDiffusionPagedAttentionBatch:
-        """Build native page-table metadata for one Diffusion forward.
 
-        The scheduler allocation payload describes capacity, while the model
+        runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata)
-        integration supplies the current ``query_len``/``kv_start_pos`` span
+        return runtime, runtime.activate()
-        for each row.  Keeping this boundary explicit prevents the Worker from
-        guessing model-specific text/image layout.
-        """
-
-        return self.get_paged_attention_adapter().prepare_batch(rows)
-
-    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
-        """Return the context manager that exposes a prepared batch to Omni Attention."""
-
-        return self.get_paged_attention_adapter().activate(batch)
 
     def _resolve_paged_attention_row(
         self,
@@ -588,11 +575,12 @@ class DiffusionKVModelRunnerBackend:
         identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
         request_state = self._diffusion_kv_request_states.get(request_id)
         if request_state is not None:
-            for installed_identity, token_len in request_state.row_token_lens:
+            for install in request_state.row_installs:
-                if installed_identity == identity:
+                if install.identity == identity:
                     return DiffusionPagedAttentionRowBinding(
                         row_index=row_index,
-                        max_seq_len=token_len,
+                        max_seq_len=install.token_len,
+                        block_ids=install.block_ids,
                     )
         raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
 

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -34,6 +34,7 @@ class DiffusionPagedAttentionRowBinding:
 
     row_index: int
     max_seq_len: int
+    block_ids: tuple[tuple[int, ...], ...] = ()
 
 
 DiffusionKVRowResolver = Callable[
@@ -100,6 +101,10 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         finally:
             attention_config.backend = previous_backend
             attention_config.backend_per_kind = previous_backend_per_kind
+        attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
+            attn_backend,
+            ulysses_degree=ulysses_degree,
+        )
         canonical_spec = replace(
             spec,
             num_kv_heads=num_kv_heads,
@@ -152,9 +157,11 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
             raise RuntimeError(
                 f"Native paged-attention backend {self.attn_backend.get_name()!r} has no forward implementation"
             )
-        if not self.attn_backend.forward_includes_kv_cache_update and not callable(
+        needs_explicit_cache_update = (
-            getattr(impl, "do_kv_cache_update", None)
+            not self.attn_backend.forward_includes_kv_cache_update
-        ):
+            or current_omni_platform.requires_diffusion_paged_kv_prewrite()
+        )
+        if needs_explicit_cache_update and not callable(getattr(impl, "do_kv_cache_update", None)):
             raise RuntimeError(
                 f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache"
             )
@@ -204,6 +211,30 @@ class DiffusionPagedAttentionRow:
         return (self.request_id, self.sequence_id, self.context_id)
 
 
+@dataclass(frozen=True)
+class DiffusionPagedAttentionMetadata:
+    """Runner-owned row layouts for one request-level denoise loop."""
+
+    prefill_rows: tuple[DiffusionPagedAttentionRow, ...]
+    denoise_rows: tuple[DiffusionPagedAttentionRow, ...]
+
+    def __post_init__(self) -> None:
+        if not self.prefill_rows or not self.denoise_rows:
+            raise ValueError("Paged attention metadata requires both prefill and denoise rows")
+        prefill_identities = tuple(row.identity for row in self.prefill_rows)
+        denoise_identities = tuple(row.identity for row in self.denoise_rows)
+        if prefill_identities != denoise_identities:
+            raise ValueError(
+                "Paged attention prefill and denoise rows must use the same ordered identities: "
+                f"prefill={prefill_identities!r}, denoise={denoise_identities!r}"
+            )
+
+    def rows_for_step(self, step_idx: int | None) -> tuple[DiffusionPagedAttentionRow, ...]:
+        if type(step_idx) is not int or step_idx < 0:
+            raise ValueError(f"Paged attention requires a non-negative denoise step index, got {step_idx!r}")
+        return self.prefill_rows if step_idx == 0 else self.denoise_rows
+
+
 @dataclass(frozen=True)
 class PreparedDiffusionPagedAttentionBatch:
     """Native metadata shared by all paged attention layers in one forward."""
@@ -236,13 +267,81 @@ class DiffusionPagedAttentionContext:
     piecewise_native_metadata: tuple[Any, ...]
     query_token_shape: tuple[int, ...]
     query_has_head_dims: bool
+    output_scatter_indices: torch.Tensor | None = None
 
     def restore_output(self, output: torch.Tensor) -> torch.Tensor:
+        if self.output_scatter_indices is not None:
+            flat_output = output.reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+            restored = output.new_zeros(
+                (*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
+            ).reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+            restored.index_copy_(0, self.output_scatter_indices, flat_output)
+            if self.query_has_head_dims:
+                return restored.reshape(
+                    *self.query_token_shape,
+                    self.layer.num_heads,
+                    self.layer.head_size_v,
+                )
+            return restored.reshape(
+                *self.query_token_shape,
+                self.layer.num_heads * self.layer.head_size_v,
+            )
         if self.query_has_head_dims:
             return output.reshape(*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
         return output.reshape(*self.query_token_shape, self.layer.num_heads * self.layer.head_size_v)
 
 
+class DiffusionPagedAttentionRuntime:
+    """Activate Runner-prepared rows lazily as the denoise loop advances."""
+
+    def __init__(
+        self,
+        adapter: DiffusionPagedAttentionAdapter,
+        metadata: DiffusionPagedAttentionMetadata,
+    ) -> None:
+        self.adapter = adapter
+        self.metadata = metadata
+        self._active_phase: bool | None = None
+        self._activation: Any | None = None
+        self._closed = False
+
+    def ensure_active(self, step_idx: int | None) -> DiffusionPagedAttentionAdapter:
+        if self._closed:
+            raise RuntimeError("Paged attention runtime is already closed")
+        is_prefill = step_idx == 0
+        rows = self.metadata.rows_for_step(step_idx)
+        if self._active_phase == is_prefill:
+            return self.adapter
+
+        self._deactivate()
+        batch = self.adapter.prepare_batch(rows)
+        activation = self.adapter.activate(batch)
+        activation.__enter__()
+        self._activation = activation
+        self._active_phase = is_prefill
+        return self.adapter
+
+    def _deactivate(self) -> None:
+        activation = self._activation
+        self._activation = None
+        self._active_phase = None
+        if activation is not None:
+            activation.__exit__(None, None, None)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._deactivate()
+        self._closed = True
+
+    @contextmanager
+    def activate(self) -> Iterator[DiffusionPagedAttentionRuntime]:
+        try:
+            yield self
+        finally:
+            self.close()
+
+
 class DiffusionPagedAttentionAdapter:
     """Translate diffusion rows into vLLM metadata for an Omni paged backend."""
 
@@ -545,15 +644,16 @@ class DiffusionPagedAttentionAdapter:
         )
 
     @staticmethod
-    def _validate_token_layout(
+    def _token_packing_indices(
         token_shape: tuple[int, ...],
         batch: PreparedDiffusionPagedAttentionBatch,
         *,
         name: str,
-    ) -> None:
+        device: torch.device,
+    ) -> torch.Tensor | None:
         if len(token_shape) == 1:
             if token_shape[0] == batch.num_tokens:
-                return
+                return None
             raise ValueError(
                 f"Paged attention {name} token count must match the prepared write batch: "
                 f"tokens={token_shape[0]}, prepared={batch.num_tokens}"
@@ -561,11 +661,23 @@ class DiffusionPagedAttentionAdapter:
         if len(token_shape) == 2:
             batch_size, tokens_per_row = token_shape
             query_lens = tuple(row.query_len for row in batch.rows)
-            if batch_size == len(batch.rows) and all(query_len == tokens_per_row for query_len in query_lens):
+            if batch_size != len(batch.rows) or any(query_len > tokens_per_row for query_len in query_lens):
-                return
+                raise ValueError(
-            raise ValueError(
+                    f"Paged attention batched {name} layout must match prepared rows for the current write: "
-                f"Paged attention batched {name} layout must match prepared rows for the current write: "
+                    f"shape={token_shape}, row_query_lens={query_lens}"
-                f"shape={token_shape}, row_query_lens={query_lens}"
+                )
+            if all(query_len == tokens_per_row for query_len in query_lens):
+                return None
+            return torch.cat(
+                [
+                    torch.arange(
+                        row_index * tokens_per_row,
+                        row_index * tokens_per_row + query_len,
+                        dtype=torch.long,
+                        device=device,
+                    )
+                    for row_index, query_len in enumerate(query_lens)
+                ]
             )
         raise ValueError(
             f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; "
@@ -626,19 +738,39 @@ class DiffusionPagedAttentionAdapter:
                 query_offsets.append(query_offsets[-1] + query_len)
             query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
             seq_lens_cpu = torch.tensor(seq_lens_values, dtype=torch.int32)
-            metadata_by_segment.append(
+            if packed_segment.query_range is None:
-                self._build_native_metadata(
+                positions = batch.positions.index_select(0, packed_segment.query_indices)
-                    query_lens=query_lens,
+                slot_mappings = batch.slot_mappings.index_select(-1, packed_segment.query_indices)
-                    query_start_loc=query_start_loc_cpu.to(self.device),
+            else:
-                    query_start_loc_cpu=query_start_loc_cpu,
+                start, end = packed_segment.query_range
-                    seq_lens=seq_lens_cpu.to(self.device),
+                positions = batch.positions[start:end]
-                    seq_lens_cpu=seq_lens_cpu,
+                slot_mappings = batch.slot_mappings[..., start:end]
-                    positions=batch.positions.index_select(0, packed_segment.query_indices),
+            segment_metadata = self._build_native_metadata(
-                    block_tables=batch.block_tables,
+                query_lens=query_lens,
-                    slot_mappings=batch.slot_mappings.index_select(-1, packed_segment.query_indices),
+                query_start_loc=query_start_loc_cpu.to(self.device),
-                    causal=(row_segments[0].mode == "causal"),
+                query_start_loc_cpu=query_start_loc_cpu,
-                )
+                seq_lens=seq_lens_cpu.to(self.device),
+                seq_lens_cpu=seq_lens_cpu,
+                positions=positions,
+                block_tables=batch.block_tables,
+                slot_mappings=slot_mappings,
+                causal=(row_segments[0].mode == "causal"),
             )
+            # FA3's full-CUDA-graph metadata builder reuses one persistent
+            # scheduler buffer across builds. Piecewise attention prepares all
+            # segments before executing any of them, so each segment needs its
+            # own snapshot rather than a view overwritten by the next build.
+            snapshotted = set()
+            for native_metadata in segment_metadata.values():
+                metadata_id = id(native_metadata)
+                if metadata_id in snapshotted:
+                    continue
+                snapshotted.add(metadata_id)
+                for field_name in ("scheduler_metadata", "prefix_scheduler_metadata"):
+                    scheduler_metadata = getattr(native_metadata, field_name, None)
+                    if isinstance(scheduler_metadata, torch.Tensor):
+                        setattr(native_metadata, field_name, scheduler_metadata.clone())
+            metadata_by_segment.append(segment_metadata)
         self._active_piecewise_native_metadata = tuple(metadata_by_segment)
         return self._active_piecewise_native_metadata
 
@@ -678,9 +810,36 @@ class DiffusionPagedAttentionAdapter:
             head_size=layer.head_size_v,
             name="value",
         )
-        self._validate_token_layout(query_token_shape, batch, name="query")
+        if key_token_shape != query_token_shape:
-        self._validate_token_layout(key_token_shape, batch, name="key")
+            if len(key_token_shape) == 1:
-        self._validate_token_layout(value_token_shape, batch, name="value")
+                raise ValueError(
+                    "Paged attention key token count must match the query layout: "
+                    f"query={query_token_shape}, key={key_token_shape}"
+                )
+            raise ValueError(
+                "Paged attention batched key layout must match prepared rows and the query layout: "
+                f"query={query_token_shape}, key={key_token_shape}"
+            )
+        if value_token_shape != query_token_shape:
+            if len(value_token_shape) == 1:
+                raise ValueError(
+                    "Paged attention value token count must match the query layout: "
+                    f"query={query_token_shape}, value={value_token_shape}"
+                )
+            raise ValueError(
+                "Paged attention batched value layout must match prepared rows and the query layout: "
+                f"query={query_token_shape}, value={value_token_shape}"
+            )
+        output_scatter_indices = self._token_packing_indices(
+            query_token_shape,
+            batch,
+            name="Q/K/V",
+            device=query.device,
+        )
+        if output_scatter_indices is not None:
+            query_flat = query_flat.index_select(0, output_scatter_indices)
+            key_flat = key_flat.index_select(0, output_scatter_indices)
+            value_flat = value_flat.index_select(0, output_scatter_indices)
         if query.device != self.device or key.device != self.device or value.device != self.device:
             raise ValueError(
                 f"Paged attention Q/K/V must be on {self.device}; "
@@ -711,10 +870,18 @@ class DiffusionPagedAttentionAdapter:
                 f"expected {batch.num_tokens}"
             )
 
-        # vLLM's writer consumes slot_mapping directly, so a span that starts
+        # vLLM's CUDA writer consumes the source token strides directly, so
+        # preserve strided QKV-projection views instead of materializing an
+        # extra K/V copy here. Platform writers with stricter contracts (for
+        # example, Ascend PA scatter) make their inputs contiguous locally.
+        # The native CUDA writer still assumes heads and head dimensions are
+        # packed within each token; materialize unusual inner layouts only.
+        if key_flat.stride(-1) != 1 or (key_flat.shape[-2] > 1 and key_flat.stride(-2) != key_flat.shape[-1]):
+            key_flat = key_flat.contiguous()
+        if value_flat.stride(-1) != 1 or (value_flat.shape[-2] > 1 and value_flat.stride(-2) != value_flat.shape[-1]):
+            value_flat = value_flat.contiguous()
+        # The writer also consumes slot_mapping directly, so a span that starts
         # inside a physical block naturally writes only the requested suffix.
-        key_flat = key_flat.contiguous()
-        value_flat = value_flat.contiguous()
         try:
             native_metadata = batch.attn_metadata[layer_name]
         except KeyError as exc:
@@ -743,6 +910,7 @@ class DiffusionPagedAttentionAdapter:
             piecewise_native_metadata=piecewise_native_metadata,
             query_token_shape=query_token_shape,
             query_has_head_dims=query_has_head_dims,
+            output_scatter_indices=output_scatter_indices,
         )
 
     @staticmethod

vllm_omni/diffusion/distributed/a2a_permute.py

@@ -142,7 +142,7 @@ def _get_symm_buffer(
                 stream_id=stream_id,
             )
             _SYMM_WORKSPACES[key] = workspace
-        return workspace.handle.get_buffer(workspace.handle.rank, symm_shape, dtype)
+        return workspace.allocation[:required_bytes].view(dtype).view(symm_shape)
 
 
 def clear_a2a_permute_workspaces() -> None:
@@ -172,7 +172,9 @@ def ulysses_qkv_fwd(x: torch.Tensor, group_name: str, world_size: int) -> torch.
     rows = B * s_local
     symm_in = _get_symm_buffer((rows, p, lc), x.dtype, x.device, group_name)
     # (B, S_local, H, D) -> (rows, p, lc); H is row-major so column block r = heads [r*Hl:(r+1)*Hl]
-    symm_in.copy_(x.reshape(rows, p, lc))
+    # Fused QKV projections produce row-strided views. Stage complete rows with
+    # the copy engine instead of materializing the view with TensorIterator.
+    torch.ops.a2ap.copy_rows(x.view(rows, p, lc), symm_in)
     out = torch.empty(p, rows, lc, device=x.device, dtype=x.dtype)
     torch.ops.a2ap.all_to_all_permute(symm_in, out, 1, 0, group_name)
     # (p, rows, lc) -> (B, S_global, H_local, D), sequence ordered rank-major

vllm_omni/diffusion/distributed/csrc/a2a_permute.cu

@@ -7,6 +7,7 @@
 // cached per group — the PR's keyed-devcomm manager API is not in 2.11).
 
 #include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
 #include <c10/macros/Macros.h>
 #include <ATen/native/cuda/MemoryAccess.cuh>
 #include <ATen/cuda/CUDAContext.h>
@@ -29,6 +30,37 @@ using namespace c10d::symmetric_memory;
 #error "Build requires NCCL >= 2.28 with symmetric-memory device support"
 #endif
 
+void copy_rows(const at::Tensor& input, at::Tensor& out) {
+  TORCH_CHECK(input.is_cuda() && out.is_cuda(), "a2a_permute: copy_rows requires CUDA tensors");
+  TORCH_CHECK(input.device() == out.device(), "a2a_permute: copy_rows device mismatch");
+  TORCH_CHECK(input.scalar_type() == out.scalar_type(), "a2a_permute: copy_rows dtype mismatch");
+  TORCH_CHECK(input.dim() == 3, "a2a_permute: copy_rows input must be 3-D");
+  TORCH_CHECK(out.sizes() == input.sizes(), "a2a_permute: copy_rows shape mismatch");
+  TORCH_CHECK(out.is_contiguous(), "a2a_permute: copy_rows output must be contiguous");
+  TORCH_CHECK(
+      input.stride(2) == 1 && input.stride(1) == input.size(2),
+      "a2a_permute: copy_rows requires contiguous rows");
+
+  const int64_t row_elements = input.size(1) * input.size(2);
+  TORCH_CHECK(input.stride(0) >= row_elements, "a2a_permute: copy_rows input rows overlap");
+  const size_t rows = static_cast<size_t>(input.size(0));
+  const size_t row_bytes = static_cast<size_t>(row_elements) * input.element_size();
+  const size_t input_pitch =
+      static_cast<size_t>(input.stride(0) * input.element_size());
+
+  c10::cuda::CUDAGuard guard(input.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  C10_CUDA_CHECK(cudaMemcpy2DAsync(
+      out.data_ptr(),
+      row_bytes,
+      input.data_ptr(),
+      input_pitch,
+      row_bytes,
+      rows,
+      cudaMemcpyDeviceToDevice,
+      stream));
+}
+
 // ---- kernel (verbatim from PR) ----------------------------------------------
 __device__ inline void copy_bytes_vec16_aligned(
     const char* src, char* dst, size_t nbytes, size_t tid, size_t stride) {
@@ -227,8 +259,10 @@ void all_to_all_permute(
 } // namespace
 
 TORCH_LIBRARY_FRAGMENT(a2ap, m) {
+  m.def("copy_rows(Tensor input, Tensor(a!) out) -> ()");
   m.def("all_to_all_permute(Tensor input, Tensor(a!) out, int scatter_dim, int gather_dim, str group_name) -> ()");
 }
 TORCH_LIBRARY_IMPL(a2ap, CUDA, m) {
+  m.impl("copy_rows", TORCH_FN(copy_rows));
   m.impl("all_to_all_permute", TORCH_FN(all_to_all_permute));
 }

vllm_omni/diffusion/forward_context.py

@@ -29,11 +29,19 @@ class ForwardContext:
     vllm_config: VllmConfig | None = None
     omni_diffusion_config: OmniDiffusionConfig | None = None
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
+    # Runner-owned paged execution metadata/runtime. Attention resolves the
+    # active Worker adapter from it; model code must not construct BlockTable
+    # rows or activate the runtime directly.
+    paged_kv_runtime: object | None = None
     # Active Worker-side paged KV adapter.  The adapter is installed only for
     # the duration of a paged forward; dense forwards leave this as ``None``.
     # Keep the field opaque here to avoid coupling the common context module to
     # the diffusion_kv implementation.
     paged_kv_adapter: Any | None = None
+    # Startup-only memory profiling runs before Scheduler-owned pages exist.
+    # Attention layers use this explicit scope to distinguish that probe from
+    # a malformed paged request whose Worker adapter was not activated.
+    in_diffusion_kv_memory_profile: bool = False
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
     denoise_timestep: float | None = None
@@ -154,6 +162,8 @@ def create_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -161,6 +171,8 @@ def create_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -186,6 +198,8 @@ def set_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    paged_kv_runtime: object | None = None,
+    in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
 ):
@@ -197,6 +211,8 @@ def set_forward_context(
         vllm_config=vllm_config,
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
+        paged_kv_runtime=paged_kv_runtime,
+        in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
     )
@@ -246,6 +262,11 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:
         _forward_context.denoise_step_idx = step_idx
+        if step_idx is not None:
+            paged_kv_runtime = getattr(_forward_context, "paged_kv_runtime", None)
+            ensure_active = getattr(paged_kv_runtime, "ensure_active", None)
+            if callable(ensure_active):
+                ensure_active(step_idx)
 
 
 def set_forward_context_denoise_timestep(timestep: float | None) -> None:

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -48,7 +48,6 @@ class InlineStageDiffusionClient(StageClientBase):
         model: str,
         od_config: OmniDiffusionConfig,
         metadata: StageMetadata,
-        batch_size: int = 1,
     ) -> None:
         self.model = model
         self.od_config = od_config
@@ -63,7 +62,6 @@ class InlineStageDiffusionClient(StageClientBase):
         self.prompt_transform_func = None
         self.prompt_expand_func = None
         self.engine_input_source = metadata.engine_input_source
-        self.batch_size = batch_size
 
         self._enrich_config()
         self._engine = DiffusionEngine.make_engine(self.od_config)
@@ -79,10 +77,10 @@ class InlineStageDiffusionClient(StageClientBase):
         self._engine.executor.register_failure_callback(self._mark_engine_dead)
 
         logger.info(
-            "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (max_num_seqs=%d)",
             self.stage_id,
             self.replica_id,
-            self.batch_size,
+            self.od_config.max_num_seqs,
         )
 
     def _enrich_config(self) -> None:

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Native vLLM-Omni pipeline for Boogu-Image-0.1.
 
@@ -220,7 +220,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
             DiffusersPipelineLoader.ComponentSource(
                 model_or_path=od_config.model,
                 subfolder="transformer",
-                revision=None,
+                revision=od_config.revision,
                 prefix="transformer.",
                 fall_back_to_pt=True,
             )
@@ -233,10 +233,18 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
         # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
         # race; prefetch the whole component set before any from_pretrained.
         boogu_subfolders = ["scheduler", "vae", "mllm", "processor"]
-        prefetch_subfolders(model, boogu_subfolders, local_files_only=local_files_only)
+        prefetch_subfolders(
+            model,
+            boogu_subfolders,
+            local_files_only=local_files_only,
+            revision=od_config.revision,
+        )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            model, subfolder="scheduler", local_files_only=local_files_only
+            model,
+            subfolder="scheduler",
+            local_files_only=local_files_only,
+            revision=od_config.revision,
         )
 
         mllm = from_pretrained_with_prefetch(
@@ -246,6 +254,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
             prefetch_list=boogu_subfolders,
             local_files_only=local_files_only,
             torch_dtype=od_config.dtype,
+            revision=od_config.revision,
         )
         # Upstream reuses the full VLM as an optional instruction rewriter and
         # encodes with its inner model (no ``lm_head``); the rewriter is not
@@ -255,7 +264,10 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
         self.mllm = mllm.to(self._execution_device)
 
         self.processor = Qwen3VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+            model,
+            subfolder="processor",
+            local_files_only=local_files_only,
+            revision=od_config.revision,
         )
 
         self.vae = from_pretrained_with_prefetch(
@@ -264,6 +276,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
             subfolder="vae",
             prefetch_list=boogu_subfolders,
             local_files_only=local_files_only,
+            revision=od_config.revision,
         ).to(self._execution_device)
 
         self.transformer = BooguImageTransformer2DModel(od_config=od_config)

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import inspect
 import logging
@@ -1069,18 +1069,98 @@ class ImageKVCacheManager(nn.Module):
         )
         return key, value
 
-    def forward(
+    def clear_legacy_prompt_kv_cache(self) -> None:
+        """Drop model-owned dense KV state before Scheduler-paged execution."""
+
+        self.image_kv_cache_map = None
+        self.image_kv_cache_lens = None
+
+    def _forward_paged(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        *,
+        first_step: bool,
+        query_lens: list[int],
+        seq_lens: list[int],
+        shard_image_size: int | None,
+        full_attn_spans: list[list[tuple[int, int]]],
+        uncond_cfg_prefill: bool,
+    ) -> torch.Tensor:
+        """Run Hunyuan attention against Worker-managed Scheduler pages.
+
+        This model boundary is shared by GPU and NPU: it describes Hunyuan's
+        Q/K/V layout and mixed-attention spans, while the common Attention
+        layer resolves the actual CUDA or Ascend paged kernel.
+        """
+
+        if uncond_cfg_prefill:
+            raise RuntimeError("Hunyuan negative-CFG prefill must run before paged row activation")
+        if self._injected_ar_kv is not None:
+            raise NotImplementedError("Hunyuan Scheduler-paged KV does not support imported AR KV")
+        if not query_lens or len(seq_lens) != len(query_lens):
+            raise ValueError("Hunyuan Scheduler-paged KV requires matching query and sequence lengths")
+        if len(full_attn_spans) != len(query_lens):
+            raise ValueError("Hunyuan Scheduler-paged KV requires one full-attention span row per sequence")
+
+        self.clear_legacy_prompt_kv_cache()
+        bs = len(query_lens)
+        q_len = query_lens[0]
+        seq_len = seq_lens[0]
+        assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}"
+
+        head_num_per_rank = query.shape[1]
+        kv_head_num_per_rank = key.shape[1]
+        head_dim = query.shape[2]
+
+        query = query.reshape(bs, q_len, head_num_per_rank, head_dim)
+        key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
+        value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
+
+        joint_text_query = joint_text_key = joint_text_value = None
+        if self.sp_size > 1 and first_step:
+            if shard_image_size is None or shard_image_size <= 0:
+                raise ValueError("Hunyuan paged Ulysses requires a positive local image shard size")
+            local_prompt_len = seq_len - shard_image_size
+            joint_query_len = query.shape[1] - shard_image_size
+            if local_prompt_len != joint_query_len:
+                raise ValueError(
+                    "Hunyuan paged Ulysses prompt layout mismatch: "
+                    f"key_prompt={local_prompt_len}, query_prompt={joint_query_len}"
+                )
+            joint_text_query = query[:, :joint_query_len]
+            joint_text_key = key[:, :local_prompt_len]
+            joint_text_value = value[:, :local_prompt_len]
+            query = query[:, joint_query_len:]
+            key = key[:, local_prompt_len:]
+            value = value[:, local_prompt_len:]
+
+        if joint_text_query is None:
+            attn_metadata = AttentionMetadata(full_attn_spans=full_attn_spans)
+        else:
+            attn_metadata = AttentionMetadata(
+                joint_query=joint_text_query,
+                joint_key=joint_text_key,
+                joint_value=joint_text_value,
+                joint_strategy="front",
+                full_attn_spans=full_attn_spans,
+            )
+        attn_output = self.attn(query, key, value, attn_metadata)
+        return attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
+
+    def _forward_dense_legacy(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
-        self.image_token_len = kwargs.get("num_image_tokens")
+        """Run the device-neutral, model-owned dense prompt-KV compatibility path."""
+
         first_step = kwargs.get("first_step")
         uncond_cfg_prefill = kwargs.get("uncond_cfg_prefill", False)
-
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
         shard_image_size = kwargs.get("shard_image_size") if self.sp_size > 1 else None
@@ -1109,8 +1189,7 @@ class ImageKVCacheManager(nn.Module):
                 key = key[:, :0, :, :]
                 value = value[:, :0, :, :]
         elif first_step:
-            self.image_kv_cache_map = None  # reset first
+            self.clear_legacy_prompt_kv_cache()
-            self.image_kv_cache_lens = None
             key, value = self._cache_prompt_kv(
                 key,
                 value,
@@ -1134,7 +1213,7 @@ class ImageKVCacheManager(nn.Module):
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
 
-        if not keep_kv_compressed and not self.attn.is_paged_kv_active():
+        if not keep_kv_compressed:
             key = repeat_kv(key, repeat_num)
             value = repeat_kv(value, repeat_num)
             if self.sp_size > 1:
@@ -1163,6 +1242,34 @@ class ImageKVCacheManager(nn.Module):
         attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
         return attn_output
 
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        self.image_token_len = kwargs.get("num_image_tokens")
+        if self.attn.is_paged_kv_active():
+            full_attn_spans = kwargs.get("full_attn_spans")
+            if full_attn_spans is None:
+                raise ValueError("Hunyuan Scheduler-paged KV requires full_attn_spans metadata")
+            return self._forward_paged(
+                query,
+                key,
+                value,
+                first_step=bool(kwargs.get("first_step")),
+                query_lens=kwargs.get("query_lens"),
+                seq_lens=kwargs.get("seq_lens"),
+                shard_image_size=kwargs.get("shard_image_size") if self.sp_size > 1 else None,
+                full_attn_spans=full_attn_spans,
+                uncond_cfg_prefill=kwargs.get("uncond_cfg_prefill", False),
+            )
+        if attention_mask is None:
+            raise ValueError("Hunyuan dense attention requires an attention mask")
+        return self._forward_dense_legacy(query, key, value, attention_mask, **kwargs)
+
 
 @dataclass
 class CausalMMOutputWithPast(CausalLMOutputWithPast):
@@ -2706,7 +2813,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         # List[List[...]] per-sample metadata indexed along the CFG batch dim
         if isinstance(model_kwargs.get("full_attn_spans"), list):
             model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"][s.start : s.stop]
-
         # custom_pos_emb: tuple of (cos, sin)
         if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None:
             cos, sin = model_kwargs["custom_pos_emb"]

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import copy
 import logging
@@ -72,7 +72,6 @@ _STEP_OUTPUT_SIZE = "hunyuan_output_size"
 _STEP_COT_TEXT_LIST = "hunyuan_cot_text_list"
 _STEP_AR_KV = "hunyuan_ar_kv"
 _STEP_PROMPT_KV = "hunyuan_prompt_kv"
-
 _HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil"
 
 
@@ -466,6 +465,11 @@ class HunyuanImage3Pipeline(
         return self._pipeline
 
     def _validate_step_request(self, state: "StepRequestState") -> None:
+        if self._uses_scheduler_paged_kv():
+            raise ValueError(
+                "HunyuanImage3 paged_scheduler currently supports request-level execution only; "
+                "disable step execution or use dense_legacy."
+            )
         prompt = state.prompt
         sampling = state.sampling
         if prompt is None:
@@ -663,6 +667,12 @@ class HunyuanImage3Pipeline(
             prefix_lens.append(int(prompt_kv[0]["lens"][branch].item()))
         return prefix_lens
 
+    def _uses_scheduler_paged_kv(self) -> bool:
+        return (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+        )
+
     def _merge_step_model_inputs(
         self,
         states: list["StepRequestState"],
@@ -693,6 +703,7 @@ class HunyuanImage3Pipeline(
                         values,
                         row_branches,
                         prefix_lens or [None] * len(row_branches),
+                        strict=True,
                     )
                 ]
                 continue
@@ -1871,8 +1882,7 @@ class HunyuanImage3Pipeline(
         row_branches = [branch for branch in range(cfg_factor) for _ in states]
         return row_state_indexes, row_branches
 
-    @staticmethod
+    def _validate_step_group_states(self, states: list["StepRequestState"]) -> tuple[bool, int]:
-    def _validate_step_group_states(states: list["StepRequestState"]) -> tuple[bool, int]:
         if not states:
             raise ValueError("HunyuanImage3 denoise_step received an empty group.")
 
@@ -2035,6 +2045,11 @@ class HunyuanImage3Pipeline(
         states = list(input_batch.states)
         if not states:
             raise ValueError("HunyuanImage3 denoise_step received an empty batch.")
+        if self._uses_scheduler_paged_kv():
+            raise ValueError(
+                "HunyuanImage3 paged_scheduler currently supports request-level execution only; "
+                "disable step execution or use dense_legacy."
+            )
         self._ensure_grouped_attention_backend_supported(len(states))
         outputs: dict[str, torch.Tensor] = {}
         for group in self._split_step_groups(states):

vllm_omni/diffusion/models/minimax_h3/fasth3.py

@@ -0,0 +1,679 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""FastVideo FastH3: a four-step DMD2 student of MiniMax-H3.
+
+FastH3 replaces H3's 49 denoiser evaluations with four. It ships as an adapter
+over the base checkpoint rather than as a full release, so it reuses H3's text
+encoder, video VAE, audio VAE, tokenizers and schedulers unchanged.
+
+The artifact is *not* a PEFT LoRA, and it is not request-switchable. Its own
+metadata states the reconstruction as::
+
+    W = W_base + lora_B @ lora_A; then .diff/.diff_b added and .set_weight assigned
+
+so besides rank-64 factors it carries full-rank ``.diff``/``.diff_b`` deltas for
+RMSNorm weights, biases, patch projections and the final layer - none of which a
+LoRA layer can express - and the VSA variants add ``.set_weight`` tensors for
+compression gates that do not exist in the base transformer at all. The adapter
+is therefore fused into the checkpoint stream at load time, before the weights
+are sharded, which is also what the release's model card requires.
+
+The low-rank factors carry no alpha: the reconstruction adds ``lora_B @ lora_A``
+directly, i.e. a scale of exactly 1.
+
+Two checkpoint spellings meet here. The adapter is written in the diffusers
+namespace (``transformer_blocks.0.attn.to_q``) while vLLM-Omni loads H3's native
+one (``blocks.0.attn.qkv_proj``), whose attention and MLP projections are fused.
+Every mapping and layout convention below was verified tensor by tensor against
+the released full checkpoint (``FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree``):
+``W_base + delta`` reproduces it to bf16 rounding.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+from safetensors import safe_open
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+    from .minimax_h3_transformer import MiniMaxH3DiTModel
+
+logger = init_logger(__name__)
+
+# FastVideo's generic adapter container, not a FastH3 marker: their tools emit
+# it for ordinary H3 adapters too.
+FASTH3_FORMAT = "fastvideo-lora-v2"
+FASTH3_MANIFEST = "adapter_manifest.json"
+# The release identity: the distilled student the adapter came from, over the
+# base it edits. The name is held to FastVideo's own namespace as well as the
+# student's name, so an unrelated adapter that merely mentions FastH3 is not
+# claimed.
+FASTH3_BASE_MODEL = "MiniMaxAI/MiniMax-H3"
+_FASTH3_IDENTITY_KEY = "finetuned_model"
+_FASTH3_IDENTITY_NAMESPACE = "fastvideo/"
+_FASTH3_IDENTITY_MARKER = "fasth3"
+
+# The rectified-flow positions the student was distilled at, and the ladder the
+# server samples on. The release states them as `dmd_denoising_steps`
+# [999, 749, 500, 250] (`sampling.base_timesteps` in the bundle manifest):
+# timestep indices out of 1000, i.e. pre-shift positions, closed here with the
+# terminal 0.0 every rectified-flow schedule ends on. The opening rung is 0.999
+# rather than 1.0 because training capped the noise level there
+# (`max_timestep_ratio`).
+#
+# These are positions, not final sigmas: H3's schedulers apply their own
+# per-modality shift on top (12 for video, 3 for audio), which is what
+# reproduces the levels the student saw. Nothing here overrides those shifts.
+FASTH3_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.749, 0.5, 0.25, 0.0))
+# Its five points bound four transformer forwards, one per sigma interval. That
+# count is the one a request states, the one Cache-DiT is refreshed with and the
+# one step execution admits a request on - the same interval contract H3's
+# pinned checkpoint schedules and native LoRAs already use.
+FASTH3_DENOISE_STEPS = FASTH3_BASE_SCHEDULE.num_inference_steps
+# Preview v1 distills the text-to-video-and-audio path only.
+FASTH3_SUPPORTED_TASKS = frozenset({"t2va"})
+
+_LORA_A = ".lora_A.weight"
+_LORA_B = ".lora_B.weight"
+_DIFF = ".diff"
+_DIFF_B = ".diff_b"
+_SET_WEIGHT = ".set_weight"
+
+# Adapter module prefix -> the native parameter it edits, minus the
+# ``.weight``/``.bias`` suffix.
+_MODEL_LEVEL_TARGETS = {
+    "proj_in": "video_patch_proj",
+    "proj_out": "final_layer.video_out",
+    "audio_proj_in": "audio_patch_proj",
+    "audio_proj_out": "final_layer.audio_out",
+    "context_embedder": "condition_proj",
+    "time_embedder.linear_1": "time_embedder.proj_in",
+    "time_embedder.linear_2": "time_embedder.proj_out",
+    "norm_out.linear": "final_layer.adaln_proj.linear",
+    "norm_out.norm": "final_layer.norm",
+}
+
+# Per-block adapter suffix -> (native suffix, how a delta enters the native
+# parameter). H3 stores attention as one grouped QKV matrix and the MLP as one
+# fused gate/up matrix, so those deltas need placing rather than adding.
+_PLAIN, _QKV_Q, _QKV_K, _QKV_V, _SWAP_HALVES = "plain", "q", "k", "v", "swap_halves"
+_QKV_SLOTS = (_QKV_Q, _QKV_K, _QKV_V)
+_BLOCK_TARGETS = {
+    "attn.to_q": ("attn.qkv_proj", _QKV_Q),
+    "attn.to_k": ("attn.qkv_proj", _QKV_K),
+    "attn.to_v": ("attn.qkv_proj", _QKV_V),
+    "attn.to_out.0": ("attn.out_proj", _PLAIN),
+    "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
+    "ff.net.2": ("mlp.fc2", _PLAIN),
+    "adaln_proj.linear": ("adaln_proj.linear", _PLAIN),
+    "norm1": ("norm1", _PLAIN),
+    "norm2": ("norm2", _PLAIN),
+}
+
+# Adapter block prefix -> native block prefix.
+_BLOCK_PREFIXES = (
+    ("token_refiner.refiner_blocks.", "token_refiner.blocks."),
+    ("transformer_blocks.", "blocks."),
+)
+
+
+class FastH3AdapterError(ValueError):
+    """The artifact is a FastH3 adapter, but it cannot be applied as one."""
+
+
+@dataclass
+class _ParamPatch:
+    """Everything the adapter contributes to one native parameter."""
+
+    # layout -> (lora_A, lora_B). A grouped QKV parameter collects three.
+    low_rank: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = field(default_factory=dict)
+    diff: torch.Tensor | None = None
+    layout: str = _PLAIN
+
+
+def _swap_halves(tensor: torch.Tensor) -> torch.Tensor:
+    """Exchange the two halves of a fused gate/up matrix.
+
+    The diffusers export stores the feed-forward projection value-first while
+    H3's native ``mlp.fc1`` is gate-first, so a delta computed in the diffusers
+    layout has to be swapped before it can be added to the native parameter.
+    """
+    if tensor.shape[0] % 2:
+        raise FastH3AdapterError(f"fused gate/up delta must split evenly, got {tuple(tensor.shape)}")
+    first, second = tensor.chunk(2, dim=0)
+    return torch.cat((second, first), dim=0)
+
+
+def _place_in_grouped_qkv(deltas: Mapping[str, torch.Tensor], *, head_dim: int) -> torch.Tensor:
+    """Interleave per-projection deltas into H3's grouped QKV layout.
+
+    The checkpoint stores one head group at a time as ``[q, k, v]``, which is
+    what :func:`_reorder_grouped_qkv_to_qkv` unpacks on the way in. A delta
+    built from the separate diffusers projections has to be folded back into
+    that order.
+    """
+    missing = sorted(set(_QKV_SLOTS) - set(deltas))
+    if missing:
+        raise FastH3AdapterError(f"grouped QKV delta is missing its {missing} projections")
+    parts = []
+    for slot in _QKV_SLOTS:
+        delta = deltas[slot]
+        if delta.shape[0] % head_dim:
+            raise FastH3AdapterError(
+                f"QKV {slot} delta rows {delta.shape[0]} are not a multiple of head_dim {head_dim}"
+            )
+        parts.append(delta.reshape(delta.shape[0] // head_dim, head_dim, *delta.shape[1:]))
+    groups = parts[0].shape[0]
+    if any(part.shape[0] != groups for part in parts):
+        raise FastH3AdapterError("QKV projections disagree on the number of head groups")
+    return torch.cat(parts, dim=1).reshape(groups * 3 * head_dim, *parts[0].shape[2:])
+
+
+def _resolve_native_target(module: str) -> tuple[str, str, tuple[str, int] | None] | None:
+    """Map an adapter module path to ``(native path, layout, block)``.
+
+    ``block`` is the ``(native block prefix, index)`` the module sits in, or
+    ``None`` for a model-level one. Coverage is checked per block, so the index
+    has to survive the mapping rather than being folded into the path.
+    """
+    native = _MODEL_LEVEL_TARGETS.get(module)
+    if native is not None:
+        return native, _PLAIN, None
+    for adapter_prefix, native_prefix in _BLOCK_PREFIXES:
+        if not module.startswith(adapter_prefix):
+            continue
+        remainder = module[len(adapter_prefix) :]
+        index, _, suffix = remainder.partition(".")
+        if not index.isdigit():
+            return None
+        target = _BLOCK_TARGETS.get(suffix)
+        if target is None:
+            return None
+        native_suffix, layout = target
+        return f"{native_prefix}{index}.{native_suffix}", layout, (native_prefix, int(index))
+    return None
+
+
+def _split_adapter_key(name: str) -> tuple[str, str] | None:
+    """Split an adapter tensor name into ``(module path, role)``."""
+    for marker, role in (
+        (_LORA_A, "lora_a"),
+        (_LORA_B, "lora_b"),
+        (_DIFF_B, "diff_b"),
+        (_DIFF, "diff"),
+        (_SET_WEIGHT, "set_weight"),
+    ):
+        if name.endswith(marker):
+            return name[: -len(marker)], role
+    return None
+
+
+def _is_fasth3_release(metadata: Mapping[str, str]) -> bool:
+    """Whether this is a FastH3 release rather than merely FastVideo's format.
+
+    Claiming on the container alone would fuse an ordinary H3 adapter into the
+    checkpoint and put every request on a four-step schedule it was never
+    distilled for, and so would a bare ``fasth3`` substring: an adapter someone
+    else named after the student is not the student.
+    """
+    if metadata.get("format") != FASTH3_FORMAT:
+        return False
+    identity = metadata.get(_FASTH3_IDENTITY_KEY, "").lower()
+    if not identity.startswith(_FASTH3_IDENTITY_NAMESPACE) or _FASTH3_IDENTITY_MARKER not in identity:
+        return False
+    # Only held against a file that states it.
+    base_model = metadata.get("base_model")
+    return base_model is None or base_model.casefold() == FASTH3_BASE_MODEL.casefold()
+
+
+def _check_declared_counts(
+    metadata: Mapping[str, str],
+    counted: Mapping[str, int],
+    *,
+    weights_path: Path,
+) -> None:
+    """Hold the artifact to the tensor counts its own metadata declares.
+
+    The writer of the ``fastvideo-lora-v2`` format records how many tensors of
+    each kind it emitted, which is the one statement in the file about its own
+    completeness. A truncated or partially re-exported artifact is otherwise
+    indistinguishable from a small one, so a claimed file has to carry all of
+    them rather than opt out by omission.
+    """
+    for key, seen in counted.items():
+        declared = metadata.get(key)
+        if declared is None:
+            raise FastH3AdapterError(
+                f"{weights_path} declares the FastH3 identity but omits {key}; "
+                "a claimed adapter has to state what it emitted"
+            )
+        try:
+            expected = int(declared)
+        except (TypeError, ValueError) as exc:
+            raise FastH3AdapterError(f"{weights_path} declares a non-numeric {key}={declared!r}") from exc
+        if expected != seen:
+            raise FastH3AdapterError(
+                f"{weights_path} declares {key}={expected} but carries {seen}; the adapter is incomplete "
+                "and would leave most of the transformer on base H3 weights"
+            )
+
+
+def _check_block_coverage(
+    seen: Mapping[str, set[int]],
+    *,
+    expected: Mapping[str, int],
+    weights_path: Path,
+) -> None:
+    """Every block of the model this adapter is loaded against must be edited.
+
+    The release drops tensors training left unchanged, so per-parameter
+    coverage is legitimately sparse - but a distilled student touches every
+    block, so a block with no edits at all means the artifact does not match
+    this model.
+    """
+    for prefix, count in expected.items():
+        indices = seen.get(prefix, set())
+        wanted = set(range(count))
+        if indices == wanted:
+            continue
+        missing = sorted(wanted - indices)
+        extra = sorted(indices - wanted)
+        raise FastH3AdapterError(
+            f"{weights_path} edits {len(indices)} of the model's {count} {prefix}* blocks "
+            f"(missing={missing[:5]}, unknown={extra[:5]}); it is not an adapter for this checkpoint"
+        )
+
+
+class FastH3WeightFusion:
+    """Fuse a FastH3 adapter into the H3 checkpoint stream as it is loaded."""
+
+    def __init__(
+        self,
+        *,
+        source: Path,
+        patches: Mapping[str, _ParamPatch],
+        head_dim: int,
+        requires_vsa: bool,
+    ) -> None:
+        self._source = source
+        self._patches = dict(patches)
+        self._head_dim = head_dim
+        self.requires_vsa = requires_vsa
+        self._applied: set[str] = set()
+        self._device: torch.device | None = None
+
+    @property
+    def source(self) -> Path:
+        return self._source
+
+    @property
+    def base_schedule(self) -> tuple[float, ...]:
+        """The rectified-flow positions this student samples on.
+
+        The fused checkpoint is a four-step student, so the ladder comes from
+        the release rather than from the many-step teacher's metadata or from
+        the uniform one ``num_inference_steps`` would otherwise derive.
+        """
+        return FASTH3_BASE_SCHEDULE.base_schedule
+
+    @classmethod
+    def from_path(
+        cls,
+        path: str | Path,
+        *,
+        head_dim: int,
+        num_blocks: int,
+        num_refiner_blocks: int,
+    ) -> FastH3WeightFusion | None:
+        """Build a fusion from an adapter file or directory, else ``None``.
+
+        Returning ``None`` keeps every other ``--lora-path`` artifact on the
+        dynamic LoRA route; only a file carrying the FastH3 release identity is
+        claimed here.
+
+        The block counts are the model's, and the artifact has to cover them:
+        claiming a partial adapter would switch the server onto the four-step
+        contract while most of the transformer still held base H3 weights.
+        """
+        weights_path = _resolve_adapter_file(path)
+        if weights_path is None:
+            return None
+
+        patches: dict[str, _ParamPatch] = {}
+        gate_tensors: list[str] = []
+        unmapped: list[str] = []
+        blocks_seen: dict[str, set[int]] = {}
+        counted = {"low_rank_tensors": 0, "diff_tensors": 0}
+        with safe_open(weights_path, framework="pt", device="cpu") as checkpoint:
+            metadata = checkpoint.metadata() or {}
+            if not _is_fasth3_release(metadata):
+                return None
+            for name in checkpoint.keys():
+                split = _split_adapter_key(name)
+                if split is None:
+                    unmapped.append(name)
+                    continue
+                module, role = split
+                if role == "set_weight":
+                    # A VSA compression gate: a module the base transformer does
+                    # not have, so there is nothing to fuse it into.
+                    gate_tensors.append(name)
+                    continue
+                target = _resolve_native_target(module)
+                if target is None:
+                    unmapped.append(name)
+                    continue
+                native_module, layout, block = target
+                if block is not None:
+                    blocks_seen.setdefault(block[0], set()).add(block[1])
+                counted["diff_tensors" if role in ("diff", "diff_b") else "low_rank_tensors"] += 1
+                native_param = f"{native_module}.{'bias' if role == 'diff_b' else 'weight'}"
+                patch = patches.setdefault(native_param, _ParamPatch(layout=layout))
+                tensor = checkpoint.get_tensor(name)
+                if role in ("diff", "diff_b"):
+                    if patch.diff is not None:
+                        raise FastH3AdapterError(f"duplicate {role} for {native_param}")
+                    patch.diff = tensor
+                else:
+                    a, b = patch.low_rank.get(layout, (None, None))
+                    patch.low_rank[layout] = (tensor, b) if role == "lora_a" else (a, tensor)
+
+        if unmapped:
+            raise FastH3AdapterError(
+                f"FastH3 adapter at {weights_path} has {len(unmapped)} tensors that name no known "
+                f"H3 parameter: {sorted(unmapped)[:5]}"
+            )
+        _check_declared_counts(
+            metadata,
+            {**counted, "set_weight_tensors": len(gate_tensors)},
+            weights_path=weights_path,
+        )
+        _check_block_coverage(
+            blocks_seen,
+            expected={"blocks.": num_blocks, "token_refiner.blocks.": num_refiner_blocks},
+            weights_path=weights_path,
+        )
+        for native_param, patch in patches.items():
+            for slot, (a, b) in patch.low_rank.items():
+                if a is None or b is None:
+                    raise FastH3AdapterError(f"FastH3 adapter has an unpaired factor for {native_param} slot {slot!r}")
+            # Only the low-rank factors are placed into H3's fused QKV and
+            # gate/up layouts; a full-rank delta is added as it comes, so one
+            # aimed at a fused parameter would silently land transposed.
+            if patch.diff is not None and patch.layout != _PLAIN:
+                raise FastH3AdapterError(
+                    f"FastH3 adapter carries a full-rank delta for {native_param}, which H3 stores in the "
+                    f"{patch.layout!r} fused layout; this loader can only place low-rank factors there"
+                )
+
+        fusion = cls(
+            source=weights_path,
+            patches=patches,
+            head_dim=head_dim,
+            requires_vsa=bool(gate_tensors),
+        )
+        logger.info(
+            "FastH3 adapter %s: rank=%s, parameters patched=%d, low-rank=%s, diff=%s, set_weight=%d",
+            weights_path,
+            metadata.get("rank", "?"),
+            len(patches),
+            metadata.get("low_rank_tensors", "?"),
+            metadata.get("diff_tensors", "?"),
+            len(gate_tensors),
+        )
+        return fusion
+
+    def _compute_device(self, weight: torch.Tensor) -> torch.device:
+        """Where to reconstruct a delta.
+
+        H3's per-block modulation projection is 96768x2688, so rebuilding all
+        343 patched parameters is a few TFLOP of rank-64 products. On CPU that
+        adds minutes to a load that already has a startup deadline, so the
+        accelerator does the arithmetic whenever there is one.
+        """
+        if weight.device.type != "cpu":
+            return weight.device
+        if self._device is None:
+            # Ask the platform rather than PyTorch's global accelerator
+            # registry, so an out-of-tree backend controls its own placement.
+            self._device = current_omni_platform.get_torch_device()
+        return self._device
+
+    @staticmethod
+    def _widen(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
+        """Move to ``device``, then widen to float32.
+
+        Asking ``Tensor.to`` for a device and a dtype at once converts on the
+        host and ships twice the bytes; splitting it moves bfloat16 and widens
+        on the accelerator.
+        """
+        return tensor.to(device, non_blocking=True).to(torch.float32)
+
+    def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor:
+        """Return ``weight`` with this adapter's contribution added."""
+        patch = self._patches.get(name)
+        if patch is None:
+            return weight
+        self._applied.add(name)
+
+        device = self._compute_device(weight)
+
+        delta: torch.Tensor | None = None
+        if patch.low_rank:
+            if patch.layout in _QKV_SLOTS:
+                per_slot = {
+                    slot: self._widen(b, device) @ self._widen(a, device) for slot, (a, b) in patch.low_rank.items()
+                }
+                delta = _place_in_grouped_qkv(per_slot, head_dim=self._head_dim)
+            else:
+                a, b = patch.low_rank[patch.layout]
+                if patch.layout == _SWAP_HALVES:
+                    # Permuting the rows of B permutes the rows of the product,
+                    # so swap the rank-64 factor instead of the full delta.
+                    b = _swap_halves(b)
+                delta = self._widen(b, device) @ self._widen(a, device)
+        if patch.diff is not None:
+            diff = self._widen(patch.diff, device)
+            delta = diff if delta is None else delta + diff
+        if delta is None:
+            return weight
+        if delta.shape != weight.shape:
+            raise FastH3AdapterError(
+                f"FastH3 delta for {name} has shape {tuple(delta.shape)}, parameter is {tuple(weight.shape)}"
+            )
+        # Leave the result on the compute device. These weights are bound for
+        # the accelerator anyway, so returning them to host memory would pay a
+        # device-to-host copy of the whole checkpoint only for the loader to
+        # send it straight back: measured at 152s against 15s for 60 GiB of
+        # patched projections, against 17s for the unavoidable upload alone.
+        # Fold the base weight into the freshly built delta in place. Promoting
+        # the weight to float32 on its own would allocate two more buffers the
+        # size of the parameter, and H3's largest patched projection is 0.5 GiB.
+        return delta.add_(weight.to(device, non_blocking=True)).to(weight.dtype)
+
+    def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]:
+        """Fuse every streamed checkpoint tensor on its way into the model."""
+        if self._applied:
+            # ``validate_fully_applied`` released the deltas, so a second stream
+            # would fuse nothing and then pass its own completeness check: the
+            # server would serve base H3 weights on the student's ladder.
+            raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint")
+        for name, weight in weights:
+            yield name, self.fuse(name, weight)
+
+    def validate_fully_applied(self) -> None:
+        """Close the fusion: every edit must have met its parameter.
+
+        A silently unapplied delta is the failure mode that matters here: the
+        model would load and generate, just not as the distilled student. The
+        weights are loaded once, so the mapped payloads are dropped afterwards
+        rather than held for the life of the process.
+        """
+        missing = sorted(set(self._patches) - self._applied)
+        if missing:
+            raise FastH3AdapterError(
+                f"FastH3 adapter edits {len(missing)} parameters the checkpoint never provided: {missing[:5]}"
+            )
+        for patch in self._patches.values():
+            patch.low_rank.clear()
+            patch.diff = None
+
+    def check_serving_contract(
+        self,
+        *,
+        partition: str,
+        od_config: Any,
+        video_shift: float,
+        audio_shift: float,
+    ) -> None:
+        """Hold a starting server to the ladder this student was trained on."""
+        if partition == "ref2va":
+            raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition")
+        offloads = [
+            flag
+            for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload")
+            if getattr(od_config, flag, False)
+        ]
+        if offloads:
+            # A host-weight plan installs the transformer without going through
+            # load_weights(), which is where the fusion and its completeness
+            # check live. Serving base H3 weights under a four-step schedule
+            # would otherwise degrade output with nothing to signal it.
+            raise ValueError(
+                f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with "
+                f"{sorted(offloads)}. Serve it without offload."
+            )
+        if self.requires_vsa:
+            raise ValueError(
+                f"{self.source} is a Video Sparse Attention variant of FastH3, and its "
+                "compression gates have no counterpart in vLLM-Omni's dense H3 attention. Use the "
+                "release's Dense variant until VSA lands for H3's packed sequence."
+            )
+        logger.info(
+            "FastH3 adapter active: sigma points %s for %d transformer forwards, "
+            "flow_shift=%g, audio_flow_shift=%g, tasks=%s",
+            list(self.base_schedule),
+            FASTH3_DENOISE_STEPS,
+            video_shift,
+            audio_shift,
+            sorted(FASTH3_SUPPORTED_TASKS),
+        )
+
+    def check_task(self, task: str) -> None:
+        """Refuse a task this preview never distilled."""
+        if task not in FASTH3_SUPPORTED_TASKS:
+            raise OmniClientError(
+                f"FastH3 preview v1 distills {sorted(FASTH3_SUPPORTED_TASKS)} only, got task={task!r}"
+            )
+
+    def check_request(self, sampling: Any, *, video_shift: float, audio_shift: float) -> None:
+        """Refuse a request that would sample the student off its rungs."""
+        if sampling.lora_request is not None:
+            # The adapter is already in the weights and the dynamic LoRA manager
+            # is skipped, so nothing would apply the requested one. Serving the
+            # request anyway would quietly ignore it.
+            raise OmniClientError(
+                f"this server fused {self.source} into the checkpoint at startup, so per-request "
+                "lora is unavailable; drop the lora field"
+            )
+        if int(sampling.num_inference_steps or 0) != FASTH3_DENOISE_STEPS:
+            raise OmniClientError(
+                f"FastH3 is a four-step student and requires num_inference_steps={FASTH3_DENOISE_STEPS} "
+                "(one transformer forward per sigma interval)"
+            )
+        # The checkpoint's per-modality shifts turn the release's positions into
+        # the noise levels the student was distilled at, so a request that moves
+        # them samples where it was never trained.
+        extra = sampling.extra_args or {}
+        for key, expected in (("flow_shift", video_shift), ("audio_flow_shift", audio_shift)):
+            try:
+                requested = float(extra.get(key, expected))
+            except (TypeError, ValueError) as exc:
+                raise OmniClientError(f"FastH3 requires {key}={expected:g}") from exc
+            if not math.isclose(requested, expected):
+                raise OmniClientError(f"FastH3 requires {key}={expected:g}, got {requested:g}")
+
+
+def resolve_fasth3_fusion(od_config: Any, transformer: MiniMaxH3DiTModel) -> FastH3WeightFusion | None:
+    """Claim ``--lora-path`` when it points at a FastH3 adapter.
+
+    FastH3 rewrites RMSNorm weights and biases, so it cannot be expressed as a
+    request-switchable LoRA and is fused into the checkpoint instead. Any other
+    artifact returns None here and stays on the dynamic LoRA route.
+    """
+    lora_path = getattr(od_config, "lora_path", None)
+    if isinstance(lora_path, (list, tuple)):
+        if len(lora_path) != 1:
+            return None
+        lora_path = lora_path[0]
+    if not lora_path:
+        return None
+    # Placing a delta into the fused QKV parameter needs the head size, and
+    # completeness is judged against the model's depth, so the architecture is
+    # only read once an adapter is actually configured.
+    arch = transformer.arch
+    return FastH3WeightFusion.from_path(
+        lora_path,
+        head_dim=arch.attention_head_dim,
+        num_blocks=arch.num_layers,
+        num_refiner_blocks=arch.token_refiner_num_layers,
+    )
+
+
+def _safetensors_metadata(path: Path) -> Mapping[str, str]:
+    """The header metadata of a safetensors file, or ``{}`` if unreadable."""
+    try:
+        with safe_open(path, framework="pt", device="cpu") as checkpoint:
+            return checkpoint.metadata() or {}
+    except Exception:  # noqa: BLE001 - a path that will not open is simply not ours
+        return {}
+
+
+def _resolve_adapter_file(path: str | Path) -> Path | None:
+    """Find the single adapter file at ``path``, or ``None``."""
+    candidate = Path(path)
+    if candidate.is_file():
+        return candidate if candidate.suffix == ".safetensors" else None
+    if not candidate.is_dir():
+        return None
+    named = candidate / "adapter_model.safetensors"
+    if named.is_file():
+        return named
+    files = sorted(candidate.glob("*.safetensors"))
+    if len(files) == 1:
+        return files[0]
+    # The published repository bundles four variants under one root, marked by an
+    # adapter_manifest.json. Such a bundle is ambiguous rather than loadable - but
+    # only hold that against a directory that actually carries FastH3 artifacts;
+    # an unrelated multi-shard LoRA stays on the dynamic route via ``None``.
+    if (candidate / FASTH3_MANIFEST).is_file() or any(
+        _is_fasth3_release(_safetensors_metadata(file)) for file in files
+    ):
+        raise FastH3AdapterError(
+            f"{candidate} holds several FastH3 adapters; point --lora-path at one variant "
+            "(for example dense-datafree/adapter_model.safetensors)"
+        )
+    return None
+
+
+__all__ = [
+    "FASTH3_BASE_MODEL",
+    "FASTH3_BASE_SCHEDULE",
+    "FASTH3_DENOISE_STEPS",
+    "FASTH3_FORMAT",
+    "FASTH3_SUPPORTED_TASKS",
+    "FastH3AdapterError",
+    "FastH3WeightFusion",
+    "resolve_fasth3_fusion",
+]