[Doc] Add Qwen-Image-2.1 recipe to recipes index (#8070)

Signed-off-by: SamitHuang <285365963@qq.com>

.buildkite/amd/test-amd-nightly.yml

@@ -64,3 +64,21 @@ steps:
         commands:
           - export VLLM_ROCM_USE_AITER=1
           - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py::test_text_to_text_audio_001 -m "full_model and rocm and MI325 and cards_2" --run-level "full_model"
+
+  - group: ":speaker: CosyVoice3 MI300 Nightly"
+    steps:
+      # Exercise all sync/async and Chinese/English cases with the platform's
+      # default attention backend. On ROCm this retains end-to-end AITER
+      # coverage without putting its cold JIT build on the READY critical path.
+      - label: "CosyVoice3-TTS E2E Test"
+        agent_pool: mi300_1
+        depends_on: amd-build
+        mirror_hardwares: [amdproduction]
+        grade: NonBlocking
+        timeout_in_minutes: 90
+        retry:
+          automatic:
+            - exit_status: 134
+              limit: 1
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"

.buildkite/amd/test-amd-ready.yml

@@ -328,7 +328,10 @@ steps:
         - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
         - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
 
-    - label: "CosyVoice3-TTS E2E Test"
+    # Keep a fast model/server integration check in the blocking lane. The
+    # full default-backend suite remains in the nightly pipeline because a
+    # cold ROCm AITER MHA build can add more than 25 minutes to this job.
+    - label: "CosyVoice3-TTS E2E Smoke (SDPA)"
       agent_pool: mi300_1
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
@@ -340,7 +343,8 @@ steps:
           - exit_status: 134
             limit: 1
       commands:
-        - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
+        - export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA
+        - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002 -m "slow" --run-level "core_model"
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs

.buildkite/common/ci_source_file_dependencies.yml

@@ -255,6 +255,8 @@ source_file_dependencies:
     - *minicpmo_4_5_duplex
     - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
     - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+    - tests/e2e/online_serving/test_minicpmo_4_5_window.py
+    - tests/e2e/online_serving/helpers/minicpmo_window_e2e.py
     - tests/e2e/online_serving/test_duplex_client_live.py
     - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
     - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
@@ -453,6 +455,7 @@ source_file_dependencies:
     - *lingbot_video
     - *lingbot_world
     - tests/e2e/online_serving/test_lingbot_video.py
+    - tests/e2e/online_serving/test_lingbot_video_expansion.py
     - tests/e2e/online_serving/test_lingbot_video_moe.py
     - tests/e2e/offline_inference/test_lingbot_world_v2.py
 

.buildkite/cuda/test-merge.yml

@@ -130,7 +130,19 @@ steps:
         source_file_dependencies: omni_minicpmo_4_5_duplex_function
         timeout_in_minutes: 50
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py tests/e2e/online_serving/test_minicpmo_4_5_window.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_1
+
+      - label: "Diffusion · Qwen Image 2.1 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_21.py
+          - tests/helpers/
+          - vllm_omni/diffusion/models/qwen_image_21/
+          - vllm_omni/diffusion/distributed/autoencoders/
+          - vllm_omni/diffusion/sched/
+          - vllm_omni/entrypoints/openai/
+        commands:
+          - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen_image_21.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_1
 
       - label: "Diffusion · Qwen Image Test"
@@ -163,6 +175,16 @@ steps:
             "
         mirror_hardwares: h100_1
 
+      - label: "Diffusion · LingBot-Video Test"
+        source_file_dependencies: diffusion_lingbot_function
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_lingbot_video.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            "
+        mirror_hardwares: h100_1
+
       - label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3"
         source_file_dependencies: diffusion_minimax_h3_function
         # Keep four server lifecycles in one Buildkite job, but never in one

.buildkite/cuda/test-nightly.yml

@@ -349,7 +349,7 @@ steps:
         commands:
           - >-
             pytest -s -v
-            tests/e2e/online_serving/test_lingbot_video.py
+            tests/e2e/online_serving/test_lingbot_video_expansion.py
             tests/e2e/online_serving/test_lingbot_video_moe.py
             tests/e2e/offline_inference/test_lingbot_world_v2.py
             -m "full_model and diffusion and H100 and B200 and cards_1"

.buildkite/cuda/test-ready.yml

@@ -229,6 +229,18 @@ steps:
             "
         mirror_hardwares: h100_2
 
+      - label: "Diffusion · Qwen Image 2.1 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_21.py
+          - tests/helpers/
+          - vllm_omni/diffusion/models/qwen_image_21/
+          - vllm_omni/diffusion/distributed/autoencoders/
+          - vllm_omni/diffusion/sched/
+          - vllm_omni/entrypoints/openai/
+        commands:
+          - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen_image_21.py -m 'core_model and cuda' --run-level 'core_model'
+        mirror_hardwares: h100_1
+
       - label: "Diffusion · Qwen Image Test"
         source_file_dependencies: diffusion_qwen_image_function
         commands:
@@ -248,3 +260,13 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion' --run-level 'core_model'
             "
         mirror_hardwares: h100_1
+
+      - label: "Diffusion · LingBot-Video Test"
+        source_file_dependencies: diffusion_lingbot_function
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_lingbot_video.py -m 'core_model and cuda' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_1

docs/configuration/environment_variables.md

@@ -85,9 +85,10 @@ depends on the installed kernels and model path.
 | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
 | `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
 | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
-| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
+| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction and separately when the serving-side final-output drain starts | An explicit env value wins; otherwise both consumers use the engine's pipeline default. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Set before server startup. | Experimental |
 | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
 | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
+| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` or `immutable`; default `overwrite` | Speech server; read when speaker storage initializes | Environment-only setting. `immutable` rejects re-registering an existing uploaded voice name until it is deleted; any other value raises `ValueError` at startup. | Experimental |
 | `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
 | `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
 | `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |

docs/configuration/stage_configs.md

@@ -293,19 +293,25 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/depl
 
 ## Qwen3-TTS with Model Runner V2
 
-Qwen3-TTS can opt into the native CUDA Model Runner V2 pipeline on vLLM
-0.29.0. Select one of these deployment profiles:
+Qwen3-TTS runs the native CUDA Model Runner V2 pipeline by default on vLLM
+0.29.0. MRV2 is an experimental feature for this model: the bundled default
+profile selects it on CUDA only, and its scheduler and delivery paths are
+still being qualified. Set `model_runner: v1` in a copy of the deploy config
+to opt out; the `platforms:` sections of `qwen3_tts.yaml` keep V1 on NPU, XPU,
+ROCm and MUSA. Select one of these deployment profiles:
 
 | Profile | Runner | Code2Wav graph batches | Intended use |
 | --- | --- | --- | --- |
-| `qwen3_tts.yaml` | V1 | Existing defaults | Existing deployment / regression control |
-| `qwen3_tts_mrv2.yaml` | V2 | B1 | Native runner validation |
+| `qwen3_tts.yaml` | V2 (default) | Existing defaults | Shipped default; experimental |
+| `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
 
 ```bash
+# MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
-  --deploy-config vllm_omni/deploy/qwen3_tts_mrv2.yaml
+  --deploy-config /path/to/qwen3_tts_v1.yaml
 ```
 
 The V2 profiles bound Talker prefill to 512 tokens per step and select the

docs/contributing/profiling.md

@@ -22,7 +22,7 @@ profiler_config:
 Supported fields:
 
 | Field | Description |
-|---|---|
+| --- | --- |
 | `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
 | `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
 | `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
@@ -287,6 +287,52 @@ Recommended viewers:
 
 For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
 
+## 6. Profiling Omni AR/TTS stages
+
+For Omni AR and TTS stages, put `profiler_config` under each stage that should
+be captured. The serving API then controls a narrow window with
+`POST /start_profile` and `POST /stop_profile`; model initialization and
+warmup stay outside the trace.
+
+Example stage override for a host-stack torch trace:
+
+```bash
+PROFILE='{"profiler":"torch","torch_profiler_dir":"/tmp/qwen3_tts_profile",'\
+'"torch_profiler_with_stack":true,"torch_profiler_record_shapes":false,'\
+'"torch_profiler_use_gzip":true,"torch_profiler_dump_cuda_time_total":true}'
+
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency.yaml \
+  --stage-overrides "{\"0\":{\"profiler_config\":$PROFILE},\"1\":{\"profiler_config\":$PROFILE}}"
+```
+
+After one or more warmup requests, capture one short measured request:
+
+```bash
+curl -sS -X POST http://127.0.0.1:8000/start_profile \
+  -H 'Content-Type: application/json' -d '{"stages":[0,1]}'
+# send one fixed /v1/audio/speech or /v1/chat/completions request here
+curl -sS -X POST http://127.0.0.1:8000/stop_profile \
+  -H 'Content-Type: application/json' -d '{"stages":[0,1]}'
+```
+
+For an Nsight Systems CUDA timeline, use the same stage override with
+`{"profiler":"cuda"}` and wrap the server:
+
+```bash
+nsys profile --trace=cuda,nvtx,osrt --sample=process-tree \
+  --cpuctxsw=process-tree --trace-fork-before-exec=true \
+  --capture-range=cudaProfilerApi --capture-range-end=stop-shutdown \
+  --cuda-graph-trace=node -o /tmp/qwen3_tts_nsys \
+  vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni ...
+```
+
+For Omni AR workers, `profiler="cuda"` requires the native CUDA profiler
+wrapper and uses `cudaProfilerStart/Stop`; the report is only valid if the
+`/start_profile` and `/stop_profile` calls both succeed. Use `nsys stats` for
+CUDA API, GPU kernel, OS runtime, and NVTX summaries. Treat profiler latency as
+diagnostic only; use a separate no-profiler run for performance numbers.
+
 ## 7. Orchestrator Monitor
 
 > **Warning:** Diagnostic only. Adds lightweight counters on the orchestrator poll loop and per-replica queue depth. Disable in production unless you are actively debugging orchestrator saturation.
@@ -314,9 +360,11 @@ If unset, the monitor writes `./vllm_omni_orch_monitor_<MMDDHHMM>.json` at shutd
 Each 1-second window records:
 
 | Series | Description |
-|---|---|
+| --- | --- |
 | `windows.duration_s` | Wall time covered by the window |
 | `windows.loop_idle` / `windows.loop_active` | Orchestrator poll-loop iterations with no work vs. active forwarding |
+| `windows.dispatch_queue_size` | Ready-queue size sampled at each monitor window boundary |
+| `windows.dispatch_queue_high_water` | Maximum Event Driven ready-queue size observed in each window |
 | `replicas.<stage,replica>.outputs_queue_size` | MP client `outputs_queue` backlog for that replica |
 | `replicas.<stage,replica>.inflight` | Requests currently bound/routed to the replica |
 
@@ -331,12 +379,16 @@ orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
 records one iteration per routed output. Window counts and the `loop_active_pct`
 summary are therefore not comparable across the two modes.
 
+For Event Driven runs, use `dispatch_queue_high_water` rather than only the
+window-boundary `dispatch_queue_size`: a serial dispatcher can drain a burst
+before the window closes, hiding a transient ready-queue backlog in the latter.
+
 ### Relationship to other diagnostics
 
 This monitor is intentionally separate from the existing profiling tools:
 
 | Tool | Scope | Output |
-|---|---|---|
+| --- | --- | --- |
 | `--enable-diffusion-pipeline-profiler` | Diffusion stage functions (`vae.decode`, `diffuse`, …) | Per-request `stage_durations` / logs |
 | `--enable-ar-profiler` | AR stage generation time | Per-request `stage_durations` |
 | `profiler_config` (`torch` / `cuda`) | GPU/CPU kernels inside a stage worker | `trace.json`, Nsight, operator tables |

docs/design/feature/cache_dit.md

@@ -180,17 +180,52 @@ cache_dit.enable_cache(
 
 > **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models.
 
+### Example 3: Sequential-CFG Diffusion Pipeline (MammothModa2)
+
+MammothModa2's DiT stage runs on the diffusion runner, which enables the
+configured backend at startup and transfers ownership to the pipeline through
+the request-scoped protocol (`adopt_cache_dit_backend` / `is_cache_dit_enabled`);
+`forward()` then reconciles per-request state through
+`RequestScopedCacheDiTRuntime` (MiniMax H3 precedent).
+
+**Key difference:** The enabler caches only the repeated main-layer stack and
+marks the pipeline as requiring paired CFG forwards, because cache-dit's
+separate-CFG accounting distinguishes cond/uncond passes purely by forward
+parity:
+
+```python
+pipeline._cache_dit_requires_paired_cfg = True
+block_adapter = BlockAdapter(
+    transformer=pipeline.gen_transformer,
+    blocks=[pipeline.gen_transformer.layers],  # Refiners stay outside
+    forward_pattern=[ForwardPattern.Pattern_3],
+    has_separate_cfg=True,
+    check_forward_pattern=True,
+)
+```
+
+Per request, the pipeline calls `runtime.prepare(spec)` for CFG requests
+(refreshing the step context) and `runtime.prepare(None)` for `guidance=1.0`
+requests (disabling hooks — their single forward per step cannot be expressed
+by the parity accounting).
+
 ### Registering Custom Implementations
 
-After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in
+After writing your custom enabler, register it in
+`CUSTOM_DIT_ENABLERS` via `register_custom_dit_enablers()` in
 `vllm_omni/diffusion/cache/cachedit/model_specific.py`:
 
 ```python
-CUSTOM_DIT_ENABLERS = {
-    "Wan22Pipeline": enable_cache_for_wan22,
-    "LongCatImagePipeline": enable_cache_for_longcat_image,
-    "YourCustomPipeline": enable_cache_for_your_model,  # Add here
-}
+CUSTOM_DIT_ENABLERS.update(
+    {
+        "Wan22Pipeline": enable_cache_for_wan22,
+        "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
+        "Krea2Pipeline": enable_cache_for_krea2,
+        "Magi2Pipeline": enable_cache_for_magi2,
+        "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2,
+        "YourCustomPipeline": enable_cache_for_your_model,  # Add here
+    }
+)
 ```
 
 ---
@@ -274,6 +309,7 @@ Complete examples in the codebase:
 | **Wan2.2** | `cachedit.model_specific::enable_cache_for_wan22` | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
 | **LongCat** | `cachedit.config::CacheDiTAdapterConfig` | Declarative block adapter | Two block lists in one transformer |
 | **BAGEL** | `cachedit.model_specific::BagelCachedAdapter` | Custom cached adapter | Complex architecture |
+| **MammothModa2** | `cachedit.model_specific::enable_cache_for_mammothmoda2` | Sub-stack + paired CFG | Generation-runner pipeline; lifecycle owned by `pipeline_mammothmoda2_dit.py` |
 
 ---
 

docs/design/module/engine_orchestration.md

@@ -79,7 +79,7 @@ the boundary affected by the in-flight stage client/process refactor in
 responsibilities proposed only by that PR are not current contracts.
 
 The orchestration loop also has an opt-in event-driven mode
-(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
+(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off except for Qwen3-TTS) proposed in
 [#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
 cadence only: the routing, ordering, and terminal-state contracts below hold
 identically on both loops.

docs/features/sleep_mode.md

@@ -9,9 +9,11 @@ This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/1
 ## 1. Feature Documentation
 
 ### Overview
+
 Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization.
 
 ### Sleep Levels
+
 We support two levels of hibernation to balance recovery speed and memory efficiency:
 
 | Level | Name | Mechanism | Recovery Speed | Memory Freed |
@@ -29,15 +31,67 @@ Omni Sleep Mode is optimized for high-performance computing backends:
 * **Huawei NPU**: Supported via Ascend memory scavenging.
 
 ### Hardware Requirements
+
 * **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep.
 * **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers.
 
----
+### Quiescing a busy diffusion stage: `pause_generation(mode="keep")`
 
+`sleep()` on a diffusion stage is a worker-level RPC: it offloads memory but does not stop the
+stage's scheduler, so a stage that still has queued requests must be made quiet first. Aborting
+discards the batch that is running. `pause_generation(mode="keep")` keeps it instead:
+
+* The batch that is already running finishes on the weights it started with and its outputs are
+  delivered normally.
+* Requests that have not started stay queued and run after `resume_generation()`, on whatever
+  weights are loaded by then. A request never continues denoising across a weight change.
+* The call returns only after every replica has acknowledged: the batch that was running has
+  finished executing, its asynchronous output copies are done, and the device is synchronized; its
+  outputs are then delivered on the normal path. On a non-preemptible backend that acknowledgement
+  arrives at the next batch boundary, so expect it to take as long as the remaining batch.
+* Control RPCs (`list_loras`, `sleep`, `wake_up`, `resume_generation`) keep working while paused.
+
+Recommended order for a weight update on a diffusion stage:
+
+```python
+await engine.pause_generation(mode="keep", clear_cache=False)
+await engine.sleep(level=1)          # optional: free memory for the trainer
+await engine.wake_up()               # admission stays closed until resume
+# install new weights here
+await engine.resume_generation()
+```
+
+Support matrix and limits:
+
+| Configuration | `mode="keep"` |
+| :--- | :--- |
+| Request-level execution (`step_execution=False`) with the in-process (`uni`) executor | Supported |
+| Request-level execution with the multi-process (`mp`) executor, including asynchronous output | Supported; the acknowledgement waits for the output copies |
+| Request-level execution with asynchronous KV prefetch | Supported; the acknowledgement joins any prefetch still in flight |
+| Step-level execution (`step_execution=True` / streaming output) | Not supported: the call raises `NotImplementedError` |
+| `mode="abort"` / `mode="wait"` / `wait_for_inflight_requests=True` on a diffusion stage | Unchanged: frontend admission is paused only |
+
+Behavior notes:
+
+* Callers that already passed `mode="keep"` for a diffusion stage used to get an immediate return
+  that only closed frontend admission. They now also stop the backend and wait for the ACK.
+* A timeout or error from `pause_generation` means the pause did **not** complete; the scheduler
+  stays closed and the caller must retry the pause or call `resume_generation()` explicitly. Do not
+  sleep or mutate weights after a failed pause.
+* `resume_generation(stage_ids=...)` that leaves another paused stage closed (an AR stage, or a
+  keep-paused diffusion stage) keeps frontend admission closed as well, so new requests cannot
+  queue on a stage that will not schedule them. Admission reopens once every stage paused by
+  `pause_generation` has been resumed.
+* `clear_cache` has no diffusion-specific effect.
+* Without a pause call nothing changes: no barrier runs and batching, abort and output ordering are
+  as before.
+
+---
 
 ## 2. Usage Examples
 
 ### Python API Example
+
 You can programmatically control the lifecycle of stages using the `AsyncOmni` engine.
 
 ```python
@@ -65,6 +119,7 @@ if __name__ == "__main__":
 ```
 
 ### server command Example
+
 Start the server with sleep mode enabled:
 
 The first method
@@ -92,9 +147,6 @@ python3 -m vllm_omni.entrypoints.openai.api_server \
 
 ```
 
-
-
-
 ### Test Scenarios & Commands
 
 #### Scenario 1: LLM Engine Sleep
@@ -113,9 +165,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage.
 
-
-
 #### Scenario 2: Diffusion Sleep
+
 Objective: Verify VRAM reclamation for Stage 1 (Diffusion).
 
 Trigger sleep (Level 1 or Level 2) via client:
@@ -128,9 +179,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 ```
 
-
-
 #### Scenario 3: Multi-Stage Coordinated Stress Test
+
 Objective: Test concurrent sleep and rapid wake-up across multiple stages.
 
 Concurrent Sleep (Stage 0 & 1):
@@ -143,7 +193,6 @@ curl -X POST http://localhost:8000/v1/omni/sleep \
 
 ```
 
-
 Rapid Wake-up:
 
 ```
@@ -154,8 +203,8 @@ curl -X POST http://localhost:8000/v1/omni/wakeup \
 
 ```
 
-
 #### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity
+
 Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation.
 
 Check Initial State: Observe baseline VRAM usage.
@@ -196,20 +245,16 @@ curl -X POST http://localhost:8000/v1/images/generations \
 
 ```
 
-
-
-
 ## 3. API Reference
 
-
 ### Methods
 
 | Method | Arguments | Return Type | Description |
 | :--- | :--- | :--- | :--- |
 | **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
 | **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
-
-
+| **pause_generation** | `mode: str, stage_ids: List[int]` | `None` | Stops admission; with `mode="keep"` also stops diffusion schedulers and returns after their ACK. |
+| **resume_generation** | `stage_ids: List[int]` | `None` | Reopens paused schedulers, then admission. |
 
 ### OmniACK Dataclass Fields
 
@@ -233,7 +278,8 @@ The metadata field is a dynamic dictionary containing hardware-specific telemetr
 }
 ```
 
-#### Core Utility:
+#### Core Utility
+
 **VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM.
 
 **Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines.

docs/models/supported_models.md

@@ -32,6 +32,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImage21Pipeline` | Qwen-Image-2.1 (T2I + image-conditioned generation, up to 4 condition images) | `Qwen/Qwen-Image-2.1` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2.1.md) |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |

docs/serving/speech_api.md

@@ -289,6 +289,19 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 
 Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
 
+**Naming rules:**
+
+- Names that collide with one of the model's built-in or precomputed voices are rejected (400). Voice
+  files already on disk under such a name are ignored at startup with a warning, so an upload can never
+  shadow a built-in voice.
+- Re-uploading an existing uploaded name overwrites it in place (the previous audio file is
+  deleted). Set `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY=immutable` on the server to reject duplicates instead,
+  requiring an explicit `DELETE /v1/audio/voices/{name}` before re-registering — useful when the
+  endpoint is reachable by multiple writers and silent overwrites are a risk.
+- The voice registry has no per-user ownership: any client that can reach the endpoint can
+  overwrite (default policy) or delete any uploaded voice. For multi-tenant deployments, add
+  authentication at a proxy and namespace voice names per user (e.g. `{user}.{name}`).
+
 **Usage Example:**
 
 ```bash
@@ -577,6 +590,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in
 | ---------- | --------- | ------------- |
 | `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
 | `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |
+| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` | `immutable` rejects re-uploading an existing uploaded name (400) until it is deleted; any other value fails startup. |
 
 The in-memory LRU has a fixed 512 MiB byte budget.
 
@@ -957,22 +971,24 @@ Use `/v1/audio/voices` to list available voices for the loaded model.
 ## Orchestration Loop (experimental)
 
 Multi-stage omni deployments route stage outputs through a single orchestrator
-loop. By default that loop polls every stage replica on a 1 ms cadence. An
-opt-in event-driven mode replaces the poll with one reader task per live stage
+loop. The legacy loop polls every stage replica on a 1 ms cadence. The
+event-driven mode replaces the poll with one reader task per live stage
 replica awaiting its client directly, and switches the serving-side
-final-output drain to a condition-variable wakeup at the same time.
+final-output drain to a condition-variable wakeup at the same time. Qwen3-TTS
+uses this mode by default; other pipelines keep the legacy poll unless enabled.
 
 **Configuration (environment variables):**
 
 | Variable | Default | Description |
 | --- | --- | --- |
-| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
+| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | On for Qwen3-TTS; off for other pipelines | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. An explicit value wins; otherwise the pipeline default computed at engine initialization is used. The override is resolved at orchestrator construction and separately when the final-output drain starts. `1`, `true`, `yes`, or `on` enables it, ignoring case and surrounding whitespace; other values select the legacy poll loop. |
 
-Set it on the process that runs the orchestrator (stage 0 of an omni
-deployment) before starting the server:
+Set it on the process that runs the orchestrator (stage 0 of an omni deployment)
+before starting the server when overriding the pipeline default. For example,
+explicitly disable event-driven orchestration for Qwen3-TTS:
 
 ```bash
-export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
+VLLM_OMNI_EVENT_DRIVEN_ORCH=0 \
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
     --omni \
     --port 8091
@@ -982,8 +998,9 @@ The server logs the selected loop mode and its reader/poller counts once at
 startup, so you can confirm which loop is live.
 
 Routing, output ordering, and terminal-state behavior are identical on both
-loops; only the poll cadence changes. Leaving the variable unset keeps the
-legacy poll loop, which is the supported default.
+loops; only the poll cadence changes. Leaving the variable unset selects the
+pipeline default: event-driven for Qwen3-TTS, and legacy polling for other
+pipelines, including Qwen3-Omni.
 
 **Known limitations:**
 

docs/user_guide/diffusion/cache_acceleration/cache_dit.md

@@ -72,6 +72,39 @@ omni = Omni(
 
 ## Example Script
 
+### Multi-Stage Models (MammothModa2)
+
+MammothModa2 is a multi-stage AR→DiT pipeline; the diffusion runner owns the
+Cache-DiT lifecycle of its DiT stage. Configure it on the **DiT stage entry**
+of the deploy YAML so the acceleration only applies to the denoising stage:
+
+```yaml
+# deploy YAML (see vllm_omni/deploy/mammoth_moda2.yaml)
+stages:
+  - stage_id: 1
+    # ... other DiT stage settings ...
+    cache_backend: cache_dit
+    cache_config:
+      Fn_compute_blocks: 1
+      Bn_compute_blocks: 0
+      max_warmup_steps: 4
+      residual_diff_threshold: 0.24
+      max_continuous_cached_steps: 3
+    enable_cache_dit_summary: true   # log skip-ratio stats per request
+```
+
+Behavior notes:
+
+- Only the repeated main-layer stack is cached; the Q-Former refiners always run.
+- Sequential-CFG parity requires the unconditional pass on every denoise step,
+  so the `cfg_range` skip optimization is not used: outside the interval CFG is
+  neutralized with `scale=1.0` instead of skipping the uncond forward.
+- Requests with `text_guidance_scale = 1.0` run with cache hooks disabled.
+- `residual_diff_threshold` defaults to `0.24` (measured ~1.7-2.2x at 50 steps,
+  PSNR 27-34 dB vs uncached). For 20-step generation — especially with a
+  partial `cfg_range` — use `0.12` for noticeably closer parity at a modest
+  speed cost.
+
 ### Offline Inference
 
 Use the example script under `examples/offline_inference/text_to_image`:

docs/user_guide/diffusion_features.md

@@ -130,7 +130,7 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -139,6 +139,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Qwen-Image-2.1**       |     ❌     |     ❌      |   ✅ (Ulysses only)    |       ✅        |         ❓         |          ❌          |   ❓    |             ❓             |      ✅ (decode)      |       ✅        |        ✅         |
 | **SenseNova-U1 / U1.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
@@ -154,6 +155,8 @@ The following tables show which models support each feature:
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
+> 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks.
+> 8. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends.
 
 ### VideoGen
 

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -2,7 +2,6 @@
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech>.
 
-
 vLLM-Omni supports several autoregressive TTS models. They share a common
 CLI shape (`--text`, `--ref-audio`, `--ref-text`, `--output-dir`) and live
 together in this hub. Each model has its own subdirectory containing a
@@ -15,7 +14,7 @@ list of supported architectures across all modalities, see
 ## Supported Models
 
 | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
-|---|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- | --- |
 | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
@@ -45,13 +44,15 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
 
 ### Prerequisites
+
 ```bash
 pip install voxcpm
 # or, for a local source checkout:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
+pip install -e /path/to/voxcpm
 ```
 
 ### Quick start
+
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
     --model openbmb/VoxCPM2 \
@@ -59,7 +60,9 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
 ```
 
 ### Voice cloning
+
 Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation:
+
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
     --text "Hello, this is a voice clone demo." \
@@ -68,6 +71,7 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
 ```
 
 ### Notes
+
 - Output: 48 kHz mono WAV.
 - Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (default).
 
@@ -78,12 +82,14 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
 2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
 
 ### Prerequisites
+
 ```bash
 uv pip install -e .
 # Includes soundfile, onnxruntime, x-transformers, einops via requirements.
 ```
 
 Download the model snapshot:
+
 ```python
 from huggingface_hub import snapshot_download
 snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
@@ -91,15 +97,18 @@ snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
 ```
 
 If your downloaded checkpoint lacks `config.json`, add it:
+
 ```json
 {
     "model_type": "cosyvoice3",
     "architectures": ["CosyVoice3Model"]
 }
 ```
+
 This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class.
 
 ### Quick start
+
 ```bash
 python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
     --model pretrained_models/Fun-CosyVoice3-0.5B \
@@ -107,7 +116,9 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ```
 
 ### Voice cloning
+
 Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
+
 ```bash
 python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
     --model pretrained_models/Fun-CosyVoice3-0.5B \
@@ -117,6 +128,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ```
 
 ### Notes
+
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
 - Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
@@ -129,12 +141,14 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 No extra packages are required; the DAC codec is vendored in vLLM-Omni.
 
 ### Quick start
+
 ```bash
 python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
     --text "Hello, this is a test of the Fish Speech text to speech system."
 ```
 
 ### Voice cloning
+
 ```bash
 python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
     --text "Hello, this is a cloned voice." \
@@ -143,14 +157,17 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
 ```
 
 ### Streaming
+
 ```bash
 python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
     --text "Hello, this is a streaming test." \
     --streaming
 ```
+
 Streaming requires `async_chunk: true` in the deploy config.
 
 ### Notes
+
 - Output: 44.1 kHz mono WAV.
 - DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
 
@@ -161,6 +178,7 @@ Streaming requires `async_chunk: true` in the deploy config.
 2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
 
 ### Quick start
+
 ```bash
 python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
     --model zai-org/GLM-TTS \
@@ -171,12 +189,14 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
 ```
 
 ### Architecture
-```
+
+```text
 Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
         (Llama-based)    (32k vocab)      (Flow Matching)
 ```
 
 ### Notes
+
 - `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
 - Reference audio should be 3-10 seconds.
 - First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
@@ -190,12 +210,15 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24
 Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
 
 ### Prerequisites
+
 ```bash
 huggingface-cli download k2-fsa/OmniVoice
 ```
+
 Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`.
 
 ### Quick start (auto voice)
+
 ```bash
 python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --model k2-fsa/OmniVoice \
@@ -203,6 +226,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
 ```
 
 ### Voice cloning
+
 ```bash
 python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --model k2-fsa/OmniVoice \
@@ -212,6 +236,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
 ```
 
 ### Voice design
+
 ```bash
 python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --model k2-fsa/OmniVoice \
@@ -220,6 +245,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
 ```
 
 ### Language hint
+
 ```bash
 python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --model k2-fsa/OmniVoice \
@@ -228,6 +254,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
 ```
 
 ### Notes
+
 - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
 - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
 
@@ -238,13 +265,16 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
 3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape).
 
 ### Prerequisites
+
 For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`:
+
 ```bash
 pip uninstall onnxruntime
 pip install onnxruntime-rocm
 ```
 
 ### Task variants
+
 - `CustomVoice`: predefined speaker (speaker ID) with optional style instruction.
 - `VoiceDesign`: text + descriptive instruction designs a new voice.
 - `Base`: voice cloning from reference audio + transcript.
@@ -264,16 +294,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
 ```
 
 ### Streaming
+
 ```bash
 python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
     --query-type CustomVoice \
     --streaming \
     --output-dir /tmp/out_stream
 ```
+
 Streaming requires `async_chunk: true` in the deploy config.
 
 ### Batched decoding
+
 The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`:
+
 ```bash
 python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
     --query-type CustomVoice \
@@ -281,9 +315,11 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
     --batch-size 4 \
     --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
 ```
+
 `--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…).
 
 ### Notes
+
 - Run `--help` for the full argument surface.
 - See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses.
 
@@ -294,12 +330,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
 Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
 
 ### Prerequisites
+
 Latest `mistral_common` with `SpeechRequest` support:
+
 ```bash
 pip install -e /path/to/mistral-common  # or upgrade from PyPI when available
 ```
 
 ### Quick start (voice preset)
+
 ```bash
 python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
     --write-audio --voice cheerful_female \
@@ -308,6 +347,7 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
 ```
 
 ### Voice cloning (capability gated upstream)
+
 ```bash
 python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
     --write-audio \
@@ -317,15 +357,18 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
 ```
 
 ### Streaming + concurrency
+
 ```bash
 python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
     --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
     --model mistralai/Voxtral-4B-TTS-2603 \
     --text "..."
 ```
+
 Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
 
 ### Notes
+
 - `--num-prompts N` replicates the prompt for performance measurement.
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.

docs/user_guide/examples/online_serving/aura_omni.md

@@ -6,29 +6,29 @@
 ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav
 ```
 
-Qwen3-TTS remains two engine stages so the pipeline reuses the existing native
-Talker and Code2Wav implementation.
+**Supported online path:** Realtime duplex at `/v1/realtime?duplex=1`.
+
+The default deploy profile (`vllm_omni/deploy/aura_omni.yaml`) sets
+`session_mode: duplex`. The `aura_omni` pipeline declares `duplex_plugin`, and
+`DuplexOmniEngine` requires that mode (same pattern as MiniCPM-o duplex
+deploys). Turn-based `chat/completions` is not a supported AURA online mode
+against this profile.
+
+`async_chunk` defaults to `false` in that file (safer if a turn-based
+`OmniOrchestrator` ever loads it). Duplex Realtime Stage2→3 still needs
+async chunks — use the smoke deploy below, or set `async_chunk: true`
+yourself.
 
 ```bash
-vllm serve aurateam/AURA \
-  --omni \
-  --deploy-config vllm_omni/deploy/aura_omni.yaml \
-  --served-model-name aurateam/AURA \
-  --trust-remote-code
+bash examples/online_serving/aura_omni/run_duplex_smoke_serve.sh
+python examples/online_serving/aura_omni/smoke_duplex_realtime_client.py
 ```
 
-Configure local checkpoints by editing per-stage `model` values in
-`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets
-`pipeline: aura_omni`, so the four-stage topology is used even if the
-command-line `--model` points at one of the component checkpoints.
+Configure local checkpoints by editing per-stage `model` values in the deploy
+or smoke YAML. The file sets `pipeline: aura_omni`, so the four-stage topology
+is used even if the command-line `--model` points at one component checkpoint.
 
-Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS
-checkpoint paths are internal stage models from the deploy YAML, not the
-OpenAI-facing served model name.
-
-The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate:
-they produce no Qwen3-TTS Talker input, so no audio is synthesized for that
-turn.
+Silent Stage1 outputs (`<|silent|>` / id `151669`) skip TTS for that turn.
 
 ## GPU Utilization Recommendation
 
@@ -40,26 +40,21 @@ VRAM each stage can reserve. Start with this split for a single GPU:
 - Stage 2 (Qwen3-TTS Talker): `0.20`
 - Stage 3 (Qwen3-TTS Code2Wav): `0.20`
 
-## TTS Modes
+## TTS modes (stage extras)
 
-`aura_omni` can pass AURA text to Qwen3-TTS in two task modes:
+When the duplex session supplies TTS extras, AURA text can feed Qwen3-TTS as:
 
-- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA
-  pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set
-  `tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only.
-- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice
-  checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass
-  `tts_task_type=CustomVoice` and `tts_speaker`.
+- `Base`: voice clone from `tts_ref_audio` (optional x-vector-only mode)
+- `CustomVoice`: predefined speaker (`tts_speaker`) with a CustomVoice checkpoint
+  on stages 2 and 3
 
-By default, AURA responses are passed to Qwen3-TTS as text. Set
-`additional_information.tts_pass_token_ids=true` to pass AURA-generated
-assistant token ids directly instead. Even when token passthrough is disabled,
-the stage processor uses AURA token ids when available to estimate the Talker
-prefill length, so scheduling does not rely on raw character length.
+Optional `tts_pass_token_ids` passes AURA assistant token ids into Talker
+instead of detokenized text.
 
-The example client exposes this as:
+## Old chat / Gradio / curl scripts (unsupported)
 
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --tts-pass-token-ids
-```
+`examples/online_serving/aura_omni/` still contains chat-completions, curl, and
+Gradio helpers from an older turn-based path. They do **not** work against the
+shipped duplex profile. That path was an incomplete stand-in before streaming
+I/O and session history were available; it is not a supported AURA mode. Use
+the Realtime duplex smoke client for online verification.

docs/user_guide/quantization/fp8.md

@@ -1,4 +1,4 @@
-# FP8 Quantization
+# FP8 Quantization
 
 ## Overview
 
@@ -85,6 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
 |-------|-----------|:-------:|:------:|----------------|------------------|------------------|
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
+| Qwen-Image-2.1 | `Qwen/Qwen-Image-2.1` | Yes | Not validated | Skip image-stream MLPs for best fidelity; all-layer FP8 is loadable but composition drifts on some prompts. On Blackwell GB200 FP8 saves ~17% peak memory but is not faster than BF16 at 1024x1024. Text encoder FP8 covers the language model only — the vision tower and LM head stay BF16 | `img_mlp` | ✅︎ (language model only) |
 | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
 | LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
 | LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |

examples/offline_inference/image_to_image/image_edit.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
 
 """
 Example script for image editing with OmniGen2.
@@ -209,14 +209,14 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
+    parser = argparse.ArgumentParser(description="Edit images with vLLM-Omni.")
     parser.add_argument(
         "--model",
         default="Qwen/Qwen-Image-Edit",
         help=(
             "Diffusion model name or local path. "
-            "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511"
-            "which supports QwenImageEditPlusPipeline."
+            "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509, "
+            "Qwen/Qwen-Image-Edit-2511, or Qwen/Qwen-Image-2.1."
         ),
     )
     parser.add_argument(
@@ -368,7 +368,7 @@ def parse_args() -> argparse.Namespace:
         "--color-format",
         type=str,
         default="RGB",
-        help="For Qwen-Image-Layered, set to RGBA.",
+        help="For Qwen-Image-Layered or Qwen-Image-2.1, set to RGBA to preserve transparency.",
     )
 
     # Cache-DiT specific parameters
@@ -647,6 +647,7 @@ def main():
 
     diffusion_params = OmniDiffusionSamplingParams(
         generator=generator,
+        seed=args.seed,
         true_cfg_scale=args.cfg_scale,
         guidance_scale=args.guidance_scale,
         guidance_scale_2=args.guidance_scale_2,

examples/offline_inference/image_to_image/run_qwen_image_21.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+set -euo pipefail
+
+python "$(dirname "${BASH_SOURCE[0]}")/image_edit.py" \
+    --model Qwen/Qwen-Image-2.1 \
+    --color-format RGBA \
+    --seed 42 \
+    --image qwen_bear.png \
+    --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
+    --output qwen_image_21_edit.png \
+    --num-inference-steps 50 \
+    --cfg-scale 1.0 \
+    "$@"

examples/offline_inference/text_to_image/text_to_image.py

@@ -89,7 +89,8 @@ def parse_args() -> argparse.Namespace:
         "--model",
         default="Qwen/Qwen-Image",
         help="Diffusion model name or local path. Supported models: "
-        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
+        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-2.1, "
+        "stepfun-ai/NextStep-1.1, "
         "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
         "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
         "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
@@ -157,7 +158,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enable-cache-dit-summary",
         action="store_true",
-        help="Enable cache-dit summary logging after diffusion forward passes.",
+        default=None,
+        help=(
+            "Enable cache-dit summary logging after diffusion forward passes. "
+            "Default: unset (defer to the deploy YAML's enable_cache_dit_summary)."
+        ),
     )
     parser.add_argument(
         "--ulysses-degree",

examples/offline_inference/text_to_speech/README.md

@@ -563,7 +563,7 @@ Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → F
 ```bash
 pip install voxcpm
 # or, for a local source checkout:
-export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
+pip install -e /path/to/voxcpm
 ```
 
 ### Quick start

examples/offline_inference/text_to_speech/voxcpm2/end2end.py

@@ -1,7 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Offline VoxCPM2 inference example (native AR pipeline).
 
 Uses the single-stage native AR config (voxcpm2.yaml).
-Requires the `voxcpm` package or VLLM_OMNI_VOXCPM_CODE_PATH env var.
+Requires the `voxcpm` package (pip install voxcpm).
 """
 
 from __future__ import annotations

examples/online_serving/aura_omni/README.md

@@ -6,139 +6,89 @@
 Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
 ```
 
-The pipeline has three semantic modules, but four engine stages because the
-existing Qwen3-TTS implementation is natively split into Talker and Code2Wav.
+**AURA's supported online path is Realtime duplex** (`/v1/realtime?duplex=1`).
+The default deploy profile sets `session_mode: duplex` because the pipeline
+declares `duplex_plugin` and `DuplexOmniEngine` requires that mode.
 
-Start the server with the deploy profile:
+This is AURA v1 / Qwen3-VL (silent / ChatML ids `151669` / `151645`), not
+AURA v2 / Qwen3.5-VL.
+
+## Duplex Realtime (only supported online path)
+
+Stage2→3 codec handoff needs `async_chunk: true`. The default
+`vllm_omni/deploy/aura_omni.yaml` keeps `async_chunk: false` as a safer
+baseline (turn-based `OmniOrchestrator` never sees the duplex async-chunk
+gate). For Realtime, use the smoke deploy (already `async_chunk: true`) or
+set that flag yourself:
 
 ```bash
+bash examples/online_serving/aura_omni/run_duplex_smoke_serve.sh
+python examples/online_serving/aura_omni/smoke_duplex_realtime_client.py
+```
+
+Smoke deploy file: `examples/online_serving/aura_omni/aura_omni_duplex_smoke.yaml`.
+
+Equivalent serve with the default topology and async chunks enabled:
+
+```bash
+# Copy or edit so async_chunk: true, then:
 vllm serve aurateam/AURA \
   --omni \
   --port 8091 \
-  --deploy-config vllm_omni/deploy/aura_omni.yaml \
+  --deploy-config examples/online_serving/aura_omni/aura_omni_duplex_smoke.yaml \
   --served-model-name aurateam/AURA \
   --trust-remote-code
 ```
 
-The deploy file sets per-stage model repos:
+Connect clients to `/v1/realtime?duplex=1`. Silent Stage1 outputs gate TTS (no
+audio for that turn). Overlapped input and vision-follow are AURA duplex
+capabilities.
+
+### Browser UI
+
+AURA has no client VAD: control is **push-to-talk** on the shared shell from
+[#7585](https://github.com/vllm-project/vllm-omni/pull/7585)
+([`examples/online_serving/realtime_web/`](../realtime_web/README.md)), profile
+`aura-ptt` (hold = `is_speech=true` PCM + sticky frames; release = `commit`;
+silent+frames for vision-follow). Thin wrapper:
+
+```bash
+python -m examples.online_serving.aura_omni.realtime_web \
+    --ws-backend ws://127.0.0.1:8099 --model aurateam/AURA --port 7862
+```
+
+Do **not** fork a separate `aura_omni/realtime_web` asset tree; MiniCPM / Qwen
+profiles are unchanged and hide the PTT control.
+
+### Per-stage models
+
+Edit `model` on each stage in `vllm_omni/deploy/aura_omni.yaml` (or the smoke
+YAML) for local checkpoints:
 
 - Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B`
 - Stage 1 AURA: `aurateam/AURA`
-- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base`
+- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base` (or CustomVoice)
 
-For local weights, edit the `model` value on each stage in
-`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes
-`pipeline: aura_omni`, so the server uses this four-stage topology even when
-the command-line model path points at one component checkpoint.
+### GPU utilization
 
-Expected request shape:
-
-- Send microphone audio as the Stage 0 multimodal audio input.
-- Include video frames in the original request `multi_modal_data`; the
-  `asr2aura` processor carries them forward to AURA.
-- Optional `additional_information` keys:
-  - `aura_system_prompt`
-  - `tts_task_type`
-  - `tts_language`
-  - `tts_speaker`
-  - `tts_instruct`
-  - `tts_ref_audio`
-  - `tts_ref_text`
-  - `tts_x_vector_only_mode`
-  - `tts_pass_token_ids`
-
-If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so
-the TTS stages are skipped for that turn.
-
-## GPU Utilization Recommendation
-
-Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`.
-Recommended baseline on one GPU for H200
+Tune `gpu_memory_utilization` per stage. Baseline on one large GPU:
 
 - Stage 0 (ASR): `0.10`
-- Stage 1 (AURA): `0.4`
-- Stage 2 (Qwen3-TTS Talker): `0.20`
-- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
+- Stage 1 (AURA): `0.40`
+- Stage 2 (Talker): `0.20`
+- Stage 3 (Code2Wav): `0.20`
 
-## Python Client
+## Old chat / Gradio / curl scripts (unsupported)
 
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --host localhost \
-  --port 8091 \
-  --model aurateam/AURA \
-  --modalities text,audio
-```
+`openai_chat_completion_client.py`, `run_curl_multimodal_generation.sh`, and
+`run_gradio_demo.sh` target the older turn-based Omni path. They do **not**
+work against the shipped `session_mode: duplex` profile.
 
-Use local media:
-
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --audio-path /path/to/input.wav \
-  --video-path /path/to/video.mp4 \
-  --output-dir output_aura_omni_online
-```
-
-Base voice clone mode (default, recommended as x-vector while debugging ICL):
-
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --tts-task-type Base \
-  --tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \
-  --tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-```
-
-Enable AURA token-id passthrough explicitly:
-
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --tts-pass-token-ids
-```
-
-CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a
-Qwen3-TTS CustomVoice checkpoint:
-
-```bash
-python examples/online_serving/aura_omni/openai_chat_completion_client.py \
-  --tts-task-type CustomVoice \
-  --tts-speaker Vivian
-```
-
-By default, AURA responses are passed to Qwen3-TTS as text. Set
-`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly
-to Qwen3-TTS instead. The processor still uses AURA token ids, when available,
-to estimate the Talker prompt length in the default text path.
-
-## Curl
-
-```bash
-cd examples/online_serving/aura_omni
-bash run_curl_multimodal_generation.sh
-```
-
-Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults:
-
-```bash
-PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
-TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
-```
-
-## Gradio
-
-Launch the server and Gradio UI together:
-
-```bash
-cd examples/online_serving/aura_omni
-bash run_gradio_demo.sh
-```
-
-If the server is already running:
-
-```bash
-python examples/online_serving/aura_omni/gradio_demo.py \
-  --model aurateam/AURA \
-  --api-base http://localhost:8091/v1
-```
+That turn-based path was an incomplete June-era stand-in: vLLM-Omni did not
+yet support streaming I/O for this stack, and adding session history then would
+have required a much larger orchestrator change. It is not a supported AURA
+mode. Do not run those scripts against the default deploy. Prefer the duplex
+smoke client above.
 
 ## Offline
 

examples/online_serving/aura_omni/aura_omni_duplex_smoke.yaml

@@ -0,0 +1,119 @@
+# AURA duplex Realtime smoke (local weights, session_mode: duplex).
+# Stage1 = /workspace/models/AURA + AuraQwen3VLForConditionalGeneration (AURA v1 / Qwen3-VL).
+pipeline: aura_omni
+# Qwen3-TTS Stage2→3 codec handoff requires async_chunk=true (same as native /
+# duplex explore). async_chunk=false reaches talker/code_predictor then never
+# delivers audio to Stage3/client.
+async_chunk: true
+session_mode: duplex
+duplex_session:
+  max_sessions: 1
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      codec_left_context_frames: 72
+      # Match native/duplex explore. initial=1 can stall Code2Wav on the first
+      # undersized chunk (observed hang after "Code2Wav codec: frames=1").
+      initial_codec_chunk_frames: 8
+      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_model_len: 4096
+    max_num_batched_tokens: 2048
+    devices: "0"
+    model: /workspace/models/Qwen3-ASR-1.7B
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 256
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.55
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_model_len: 32768
+    max_num_batched_tokens: 8192
+    devices: "0"
+    model: /workspace/models/AURA
+    model_arch: AuraQwen3VLForConditionalGeneration
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 256
+      repetition_penalty: 1.0
+      stop_token_ids: [151669, 151645]
+
+  - stage_id: 2
+    # Allow prior TTS to drain while the next Talker starts (not assumed 1).
+    max_num_seqs: 2
+    gpu_memory_utilization: 0.12
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    # Stage1→2 is orchestrator-fed via aura2tts (custom_process_input_func).
+    # A from_stage_1 SHM edge parks Talker in WAITING_FOR_CHUNK and clears
+    # additional_information.text → Missing Qwen3-TTS text conditioning.
+    # Only Stage2→3 codec handoff needs a connector.
+    output_connectors:
+      to_stage_3: connector_of_shared_memory
+    model: /workspace/models/Qwen3-TTS-12Hz-1.7B-CustomVoice
+    default_sampling_params:
+      temperature: 0.5
+      top_k: 50
+      # 12Hz × 240 ≈ 20s. Stay under the smoke 25s audio gate so a missed
+      # 2150 (English/quotes humming) cannot overlap the next vision-follow.
+      max_tokens: 240
+      min_tokens: 8
+      repetition_penalty: 1.05
+      # Qwen3-TTS Talker codec EOS — without this Talker can run to max_tokens
+      # (~26s garbage audio while Stage1 text was correct).
+      stop_token_ids: [2150]
+    subtalker_sampling_params:
+      do_sample: true
+      temperature: 0.5
+      top_k: 50
+      top_p: 1.0
+
+  - stage_id: 3
+    max_num_seqs: 2
+    gpu_memory_utilization: 0.12
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 16384
+    max_model_len: 16384
+    devices: "0"
+    input_connectors:
+      from_stage_2: connector_of_shared_memory
+    model: /workspace/models/Qwen3-TTS-12Hz-1.7B-CustomVoice
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 16384
+      seed: 42
+      repetition_penalty: 1.0

examples/online_serving/aura_omni/realtime_web/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""AURA push-to-talk browser call UI."""

examples/online_serving/aura_omni/realtime_web/__main__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from examples.online_serving.realtime_web.server import main
+
+if __name__ == "__main__":
+    main(default_profile="aura-ptt")

examples/online_serving/aura_omni/run_duplex_smoke_serve.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Minimal AURA duplex Realtime server for smoke.
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+cd "$REPO_ROOT"
+
+PORT="${PORT:-8099}"
+MODEL="${MODEL:-/workspace/models/AURA}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/examples/online_serving/aura_omni/aura_omni_duplex_smoke.yaml}"
+
+# AURA v1 / Qwen3-VL ChatML ids (do not use AURA v2 248070/248046 here).
+export VLLM_AURA_SILENT_TOKEN_ID="${VLLM_AURA_SILENT_TOKEN_ID:-151669}"
+export VLLM_AURA_IM_END_TOKEN_ID="${VLLM_AURA_IM_END_TOKEN_ID:-151645}"
+export VLLM_AURA_IM_START_TOKEN_ID="${VLLM_AURA_IM_START_TOKEN_ID:-151644}"
+export VLLM_AURA_ASSISTANT_TOKEN_ID="${VLLM_AURA_ASSISTANT_TOKEN_ID:-77091}"
+# Multi-stage colocated init can make earlier stages' Process.sentinel spuriously
+# readable; prefer spawn for stage engine subprocesses (override if needed).
+export VLLM_WORKER_MULTIPROC_METHOD="${VLLM_WORKER_MULTIPROC_METHOD:-spawn}"
+# Do NOT set TORCHDYNAMO_DISABLE here: it breaks Stage0/1 AOT when those
+# stages are not enforce_eager. Stage2/3 use enforce_eager; code_predictor
+# honors that and skips its own torch.compile.
+
+exec "$REPO_ROOT/.venv/bin/vllm-omni" serve "$MODEL" \
+  --omni \
+  --deploy-config "$DEPLOY_CONFIG" \
+  --trust-remote-code \
+  --host 0.0.0.0 \
+  --port "$PORT"

examples/online_serving/aura_omni/smoke_duplex_realtime_client.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Minimal AURA duplex Realtime smoke via DuplexClient."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import sys
+import time
+from io import BytesIO
+from pathlib import Path
+
+from PIL import Image
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni.clients.duplex import (  # noqa: E402
+    DuplexClient,
+    EventCollector,
+    SessionConfig,
+    read_pcm16_wav,
+)
+
+PROGRESS_TYPES = {
+    "response.created",
+    "response.done",
+    "response.audio.delta",
+    "response.output_audio.delta",
+    "response.audio_transcript.delta",
+    "response.output_text.delta",
+    "response.text.delta",
+    "response.listen",
+    "response.model_listen",
+    "error",
+}
+
+
+def _tiny_jpeg_b64() -> str:
+    img = Image.new("RGB", (64, 64), color=(200, 40, 40))
+    buf = BytesIO()
+    img.save(buf, format="JPEG", quality=85)
+    return base64.b64encode(buf.getvalue()).decode("ascii")
+
+
+async def run(url: str, model: str, wav: Path, timeout_s: float) -> int:
+    pcm16 = read_pcm16_wav(wav)
+    if not pcm16:
+        print("empty wav", file=sys.stderr)
+        return 2
+
+    frame = _tiny_jpeg_b64()
+    collector = EventCollector()
+    config = SessionConfig(
+        modalities=("text", "audio"),
+        instructions="You are AURA. Reply briefly in Chinese.",
+        extra_body={
+            "aura_system_prompt": "You are AURA. Reply briefly in Chinese.",
+            "tts_task_type": "CustomVoice",
+            "tts_language": "Chinese",
+            "tts_speaker": "Vivian",
+        },
+    )
+
+    started = time.time()
+    client = DuplexClient(
+        url,
+        model=model,
+        config=config,
+        reconnect=None,
+        heartbeat_interval_s=None,
+        handshake_timeout_s=min(60.0, timeout_s),
+    )
+
+    async with client:
+        consume_task = asyncio.create_task(collector.consume(client))
+        collector.add({"type": "session.created", "session": getattr(client, "session_info", {})})
+
+        await client.append_audio(pcm16, is_speech=True, video_frames=[frame])
+        await client.commit(create_response=True)
+
+        progress = False
+        got_done = False
+        while time.time() - started < timeout_s:
+            types = [e.get("type") for e in collector.events]
+            if any(t in PROGRESS_TYPES for t in types):
+                progress = True
+            if "response.done" in types or "response.completed" in types:
+                got_done = True
+                progress = True
+                break
+            for e in collector.events:
+                meta = e.get("metadata") if isinstance(e.get("metadata"), dict) else {}
+                if meta.get("model_listen") or meta.get("listen_source") == "aura_silent":
+                    progress = True
+                    break
+            if collector.count("error") > 0 or collector.count("session.closed") > 0:
+                break
+            # Keep draining until response.done so TTS is not truncated.
+            await asyncio.sleep(0.2)
+
+        try:
+            await client.close(timeout_s=30.0)
+        except Exception as exc:  # noqa: BLE001
+            print(f"close_warn={exc}", file=sys.stderr)
+        try:
+            await asyncio.wait_for(consume_task, timeout=5.0)
+        except (TimeoutError, asyncio.CancelledError):
+            consume_task.cancel()
+
+    event_types = [e.get("type") for e in collector.events]
+    errors = [json.dumps(e, ensure_ascii=False)[:500] for e in collector.errors()]
+    result = {
+        "ok": progress and (got_done or collector.count("response.done") > 0),
+        "got_response_progress": progress,
+        "got_response_done": got_done or collector.count("response.done") > 0,
+        "errors": errors,
+        "event_types": event_types,
+        "n_events": len(event_types),
+        "audio_bytes": len(collector.audio_bytes()),
+        "llm_text": collector.response_text(collector.response_ids[0]) if collector.response_ids else "",
+        "elapsed_s": round(time.time() - started, 2),
+    }
+    print(json.dumps(result, ensure_ascii=False, indent=2))
+    log_path = Path("/tmp/aura_duplex_smoke/client_events.jsonl")
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    log_path.write_text(
+        "".join(json.dumps(e, ensure_ascii=False) + "\n" for e in collector.events),
+        encoding="utf-8",
+    )
+    print(f"events_log={log_path}", file=sys.stderr)
+    return 0 if progress else 1
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime")
+    p.add_argument("--model", default="/workspace/models/AURA")
+    p.add_argument(
+        "--wav",
+        default=str(REPO_ROOT / "tests/assets/minicpmo_4_5/response_required_16k.wav"),
+    )
+    p.add_argument("--timeout-s", type=float, default=300.0)
+    args = p.parse_args()
+    return asyncio.run(run(args.url, args.model, Path(args.wav), args.timeout_s))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

examples/online_serving/realtime_web/README.md

@@ -9,6 +9,7 @@ This implements [RFC #7222](https://github.com/vllm-project/vllm-omni/issues/722
 | `minicpm-native` | Model-controlled listen/speak, continuous audio input | Frames accompany audio | Yes |
 | `qwen3-turn --stt` (default) | User presses **Send turn** | No | No |
 | `qwen3-turn --vad` | Server detects trailing silence; speech interrupts replies | Sampled frames with each spoken turn | Yes |
+| `aura-ptt` | Hold **Hold to talk**; release commits. Hold also stops local playback | Sticky frames with speech; vision-follow every 2 frames only while unlocked (locked from PTT release until text-final / listen) | Yes |
 
 Qwen3 VAD uses an engine-owned duplex plugin: microphone upload continues
 while replies stream, and speech can interrupt generation and playback. The
@@ -29,6 +30,38 @@ The compatibility wrapper serves the shared assets with `minicpm-native`.
 The current native query, `extra_body`, reference voice, continuous microphone
 upload, camera frames and playback acknowledgements are retained.
 
+## AURA: push-to-talk duplex
+
+AURA has no client VAD. Start the duplex serve, then the thin wrapper:
+
+```bash
+bash examples/online_serving/aura_omni/run_duplex_smoke_serve.sh
+
+python -m examples.online_serving.aura_omni.realtime_web \
+    --ws-backend ws://127.0.0.1:8099 --model aurateam/AURA --port 7862
+```
+
+Or select the profile on the shared host:
+
+```bash
+python -m examples.online_serving.realtime_web --profile aura-ptt \
+    --ws-backend ws://127.0.0.1:8099 --model aurateam/AURA --port 7862
+```
+
+Open the UI, start a session, optionally enable **Camera**, then **hold**
+**Hold to talk** to stream `is_speech=true` PCM (with sticky camera frames).
+**Release** to `input_audio_buffer.commit`. Holding the button stops playback
+in the browser only; the in-flight response is not cancelled on the server.
+With the camera on, frames are still captured at 2 fps. While the button is
+up they are committed once per two frames (`is_speech=false` + silent PCM).
+Both frames of that commit are sent and packed as one video (`<|video_pad|>`).
+Turns still open from this client commit, not from a server-side auto trigger.
+but only after the prior spoken turn unlocks (text-final or `response.listen`).
+Release itself locks vision immediately so a follow-up frame cannot abort the
+just-committed speech turn on Stage0/1 (`max_num_seqs=1`). Audio still playing
+does not hold the next vision turn once text is done.
+MiniCPM / Qwen profiles do not set `pushToTalk`, so the PTT control stays hidden.
+
 ## Qwen3: explicit-turn STT adapter
 
 Start the backend, then the UI in a second terminal:
@@ -49,6 +82,7 @@ accepts only model selection, so system-prompt controls are hidden in this mode.
 The adapter follows the shipped `qwen3_omni/openai_realtime_client.py` and
 `vllm_omni/entrypoints/openai/realtime_connection.py`:
 
+
 1. Explicit `duplex=0` selects the legacy STT handler on a turn deployment.
    It does not enable STT on a duplex deployment.
 2. Send `{type: "session.update", model: ...}` and `commit(final=false)`.
@@ -181,7 +215,7 @@ conversation API. This demo does not add WebRTC, semantic VAD, or tool calling.
 
 ## Shared host options
 
-The shared entry point can select either profile explicitly:
+The shared entry point can select a profile explicitly:
 
 ```bash
 python -m examples.online_serving.realtime_web --profile qwen3-turn \
@@ -193,7 +227,7 @@ python -m examples.online_serving.realtime_web --profile qwen3-turn \
 - `--public-realtime-url`: optional browser-visible WebSocket URL; otherwise the
   static host proxies `/v1/realtime` on the same origin.
 - `--host`, `--port`: UI bind address and port (default port 7862).
-- `--ref-audio`: required for MiniCPM; rejected for Qwen3.
+- `--ref-audio`: required for MiniCPM; rejected for Qwen3 and AURA.
 
 Microphone and camera access require `localhost` or HTTPS. For a remote backend,
 use an SSH tunnel to the UI host or an HTTPS reverse proxy with WebSocket support.

examples/online_serving/realtime_web/app/index.html

@@ -22,6 +22,7 @@
         <div class="call-actions">
           <button id="callButton" class="primary" type="button">Start session</button>
           <button id="sendTurnButton" class="secondary" type="button" hidden disabled>Send turn</button>
+          <button id="pttButton" class="secondary" type="button" hidden disabled>Hold to talk</button>
           <button id="muteButton" class="secondary" type="button" disabled>Mute</button>
           <button id="cameraButton" class="secondary" type="button" disabled>Camera</button>
           <video id="cameraPreview" muted playsinline style="display:none;width:96px;height:72px;border-radius:8px;object-fit:cover;background:#000"></video>
@@ -75,6 +76,7 @@
     <script src="static/profiles/common.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
     <script src="static/profiles/minicpm_native.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
     <script src="static/profiles/qwen3_turn.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
+    <script src="static/profiles/aura_ptt.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
     <script src="static/app.js?v=__FULL_DUPLEX_APP_VERSION__"></script>
   </body>
 </html>

examples/online_serving/realtime_web/app/static/app.js

@@ -5,6 +5,7 @@
   const profile = window.OmniRealtimeProfiles[config.profile || 'minicpm-native'](config);
   const callButton = document.getElementById('callButton');
   const sendTurnButton = document.getElementById('sendTurnButton');
+  const pttButton = document.getElementById('pttButton');
   const muteButton = document.getElementById('muteButton');
   const cameraButton = document.getElementById('cameraButton');
   const cameraPreview = document.getElementById('cameraPreview');
@@ -28,6 +29,7 @@
   const ECHO_GUARD_MS = 300;
   const INITIAL_PLAYBACK_BUFFER_MS = 400;
   const SESSION_CLOSE_TIMEOUT_MS = 1000;
+  const SILENT_PCM_SAMPLES = 160;
 
   const PROMPT_PRESETS = profile.presets;
   document.title = profile.title;
@@ -36,7 +38,21 @@
   document.getElementById('profileDescription').textContent = profile.description;
   document.getElementById('policyLabel').textContent = profile.policy;
   cameraButton.hidden = !profile.camera;
+  if (profile.cameraPreviewLarge) {
+    cameraPreview.classList.add('camera-preview-large');
+    // Inline, so a cached stylesheet cannot leave the 96×72 HTML default.
+    cameraPreview.style.width = '100%';
+    cameraPreview.style.height = 'auto';
+    cameraPreview.style.maxHeight = '480px';
+    cameraPreview.style.flex = '0 0 100%';
+    const cameraActions = cameraPreview.closest('.call-actions');
+    if (cameraActions) {
+      cameraActions.style.flexWrap = 'wrap';
+      cameraActions.style.gridColumn = '1 / -1';
+    }
+  }
   sendTurnButton.hidden = !profile.clientCommit;
+  pttButton.hidden = !profile.pushToTalk;
   promptPreset.replaceChildren();
   for (const name of [...Object.keys(PROMPT_PRESETS), 'custom']) {
     const option = document.createElement('option');
@@ -64,6 +80,12 @@
   let cameraStream = null;
   let cameraTimer = null;
   let cameraPendingFrame = null;
+  let cameraLastFrame = null;
+  // AURA vision clock only. Other profiles leave this empty and never read it.
+  let visionFollowQueue = [];
+  let visionTurnLocked = false;
+  let suppressedPlaybackId = null;
+  let pttHeld = false;
   const cameraCanvas = document.createElement('canvas');
   let playbackRate = OUTPUT_RATE;
   let pendingCapture = [];
@@ -87,6 +109,8 @@
   const interruptedResponses = new Set();
   const pendingEvents = new Set();
   let assistantTextChannel = null;
+  const assistantRowByResponse = new Map();
+  let lastClosedAssistantText = '';
   let connectionReady = false;
   let turnCounter = 0;
   let stopping = null;
@@ -134,7 +158,7 @@
       assistantTextChannel = null;
       if (profile.halfDuplex) pendingCapture = [];
       sendTurnButton.disabled = !profile.clientCommit;
-      setModel(profile.waiting);
+      if (!pttHeld) setModel(profile.waiting);
     }, ECHO_GUARD_MS);
   }
 
@@ -240,6 +264,7 @@
       current.value = finalText;
       current.text.textContent = finalText;
     }
+    if (role === 'assistant' && current.value) lastClosedAssistantText = current.value;
     current.row.classList.remove('turn-live');
     if (role === 'user') liveUserTurn = null;
     else liveAssistantTurn = null;
@@ -317,11 +342,67 @@
   }
 
   function microphoneUploadEnabled() {
-    return running && connectionReady && !muted && (!profile.halfDuplex || !assistantActive);
+    if (!running || !connectionReady || muted) return false;
+    if (profile.pushToTalk) return pttHeld;
+    return !profile.halfDuplex || !assistantActive;
+  }
+
+  function queueVisionFrame(frame) {
+    visionFollowQueue.push(frame);
+    if (visionFollowQueue.length > 2) visionFollowQueue.shift();
+  }
+
+  function commitVisionFollow() {
+    // Two frames still gate the turn (client commit). Both frames are one clip;
+    // the server packs them as a single <|video_pad|>, same as Native.
+    const frames = visionFollowQueue.slice();
+    visionFollowQueue = [];
+    const silent = new Int16Array(SILENT_PCM_SAMPLES);
+    const event = profile.append(int16ToBase64(silent), frames, { isSpeech: false });
+    socket.send(JSON.stringify(event));
+    for (const message of profile.commitMessages()) socket.send(JSON.stringify(message));
+  }
+
+  function releaseVisionTurn(responseId) {
+    if (!profile.visionFollowWhileSpeaking || !visionTurnLocked) return;
+    if (responseId && currentResponseId && responseId !== currentResponseId) return;
+    visionTurnLocked = false;
+  }
+
+  function suppressLocalPlayback() {
+    // Hold-to-talk cuts the speaker only. The in-flight response stays on the server.
+    if (currentResponseId) suppressedPlaybackId = currentResponseId;
+    if (playbackNode) playbackNode.port.postMessage({ type: 'clear' });
+    playbackComplete = true;
+    responseHasAudio = false;
+    setPlayback('Idle');
+  }
+
+  function playbackSuppressed(responseId) {
+    return Boolean(suppressedPlaybackId) && responseId === suppressedPlaybackId;
   }
 
   function flushCapture() {
-    if (!socket || socket.readyState !== WebSocket.OPEN || pendingCapture.length === 0) return;
+    if (!socket || socket.readyState !== WebSocket.OPEN) return;
+
+    // AURA: capture stays 2 fps, but a turn opens only every 2 frames, and
+    // not until the previous response's text is done (or it listened).
+    // Playback still running is not the lock. PTT speech is not batched here.
+    if (profile.pushToTalk && profile.visionFollowWhileSpeaking && !pttHeld) {
+      if (cameraPendingFrame) {
+        queueVisionFrame(cameraPendingFrame);
+        cameraPendingFrame = null;
+      }
+      if (!visionTurnLocked && visionFollowQueue.length >= 2) commitVisionFollow();
+      return;
+    }
+
+    if (profile.pushToTalk && !pttHeld) {
+      pendingCapture = [];
+      return;
+    }
+
+    if (pendingCapture.length === 0) return;
     if (!microphoneUploadEnabled()) {
       pendingCapture = [];
       return;
@@ -335,11 +416,44 @@
     }
     pendingCapture = [];
     const pcm = resampleInt16(merged, captureRate, profile.inputSampleRate || INPUT_RATE);
-    const appendEvent = profile.append(int16ToBase64(pcm), cameraPendingFrame);
+    const frame = cameraPendingFrame || (profile.stickyCamera ? cameraLastFrame : null);
     cameraPendingFrame = null;
+    const appendEvent = profile.pushToTalk
+      ? profile.append(int16ToBase64(pcm), frame, { isSpeech: true })
+      : profile.append(int16ToBase64(pcm), frame);
     socket.send(JSON.stringify(appendEvent));
   }
 
+  function setPttHeld(held) {
+    if (!profile.pushToTalk || !running || pttHeld === held) return;
+    if (held) {
+      pttHeld = true;
+      suppressLocalPlayback();
+      pttButton.classList.toggle('is-active', true);
+      pttButton.textContent = 'Release to send';
+      setModel('Talking');
+      appendLog('PTT down');
+      return;
+    }
+    // Flush remaining speech while still held, then commit.
+    flushCapture();
+    if (socket && socket.readyState === WebSocket.OPEN) {
+      for (const event of profile.commitMessages()) socket.send(JSON.stringify(event));
+    }
+    // Lock vision before clearing pttHeld. Stage0/1 are max_num_seqs=1: a
+    // vision-follow commit in the gap before response.begin aborts the speech
+    // turn (server sees a new ephemeral and kills the prior ASR request).
+    if (profile.visionFollowWhileSpeaking) {
+      visionTurnLocked = true;
+      visionFollowQueue = [];
+    }
+    pttHeld = false;
+    pttButton.classList.toggle('is-active', false);
+    pttButton.textContent = 'Hold to talk';
+    setModel(profile.waiting);
+    appendLog('PTT up · commit');
+  }
+
   function beginAssistant(responseId) {
     currentResponseId = responseId || currentResponseId;
     responseHasAudio = false;
@@ -351,6 +465,7 @@
   }
 
   function feedPlayback(decoded, responseId) {
+    if (playbackSuppressed(responseId || currentResponseId)) return;
     if (!decoded || !decoded.pcm || decoded.pcm.length === 0 || !playbackNode) return;
     const pcm = resampleInt16(decoded.pcm, decoded.sourceRate, playbackRate);
     responseHasAudio = true;
@@ -392,13 +507,18 @@
   }
 
   async function handleAudioEvent(action) {
+    const responseId = action.responseId || currentResponseId || `turn-${turnCounter}`;
+    if (playbackSuppressed(responseId)) return;
     markBusy();
-    currentResponseId = action.responseId || currentResponseId || `turn-${turnCounter}`;
+    currentResponseId = responseId;
     setModel('Speaking');
     const generation = sessionGeneration;
-    const responseId = currentResponseId;
     const decoded = await decodeAudioDelta(action.event);
-    if (generation === sessionGeneration && !interruptedResponses.has(responseId)) feedPlayback(decoded, responseId);
+    if (
+      generation === sessionGeneration
+      && !interruptedResponses.has(responseId)
+      && !playbackSuppressed(responseId)
+    ) feedPlayback(decoded, responseId);
   }
 
   function handleTranscriptEvent(action) {
@@ -408,12 +528,43 @@
       if (profile.deduplicateTranscript && assistantTextChannel && action.channel !== assistantTextChannel) return;
       if (action.text) assistantTextChannel = action.channel || assistantTextChannel;
     }
-    if (action.kind === 'text') addTranscript(action.role, action.text);
-    else if (profile.deduplicateTranscript && action.role === 'assistant') {
+    const responseId = action.responseId || null;
+    const existing = responseId && action.role === 'assistant' ? assistantRowByResponse.get(responseId) : null;
+    if (existing) {
+      if (action.kind === 'text') {
+        if (!action.text) return;
+        existing.value += action.text;
+      } else if (action.text) {
+        existing.value = action.text;
+      }
+      existing.text.textContent = existing.value;
+      liveAssistantTurn = existing;
+      return;
+    }
+    if (
+      action.role === 'assistant'
+      && !liveAssistantTurn
+      && action.text
+      && action.text === lastClosedAssistantText
+    ) {
+      // response.done already closed this sentence. A later full transcript,
+      // including one stamped with a newer vision-follow response id, must
+      // not open a second bubble with the same words.
+      return;
+    }
+    if (action.kind === 'text') {
+      addTranscript(action.role, action.text);
+      if (action.role === 'assistant' && responseId && liveAssistantTurn) {
+        assistantRowByResponse.set(responseId, liveAssistantTurn);
+      }
+      return;
+    }
+    if (profile.deduplicateTranscript && action.role === 'assistant') {
       if (action.text) {
         const turn = ensureTurn('assistant');
         turn.value = action.text;
         turn.text.textContent = action.text;
+        if (responseId) assistantRowByResponse.set(responseId, turn);
       }
     } else finishTranscript(action.role, action.text);
   }
@@ -459,12 +610,14 @@
       }
       case 'listen':
         assistantActive = false;
+        releaseVisionTurn(responseId);
         setModel(profile.waiting);
         break;
       case 'begin':
         if (profile.halfDuplex) armTurnTimeout();
         if (echoTimer !== null) { clearTimeout(echoTimer); echoTimer = null; }
         beginAssistant(responseId);
+        if (profile.visionFollowWhileSpeaking) visionTurnLocked = true;
         break;
       case 'audio':
         await handleAudioEvent(action);
@@ -474,11 +627,13 @@
         break;
       case 'text': case 'text-final':
         handleTranscriptEvent(action);
+        if (action.kind === 'text-final' && action.role === 'assistant') releaseVisionTurn(responseId);
         break;
       case 'done':
         clearTimeout(turnTimeout);
         responseComplete = true;
         finishTranscript('assistant');
+        releaseVisionTurn(responseId);
         requestPlaybackDrain(responseId);
         finishResponseIfReady();
         break;
@@ -671,7 +826,11 @@
       await openSocket();
       running = true;
       muted = false;
+      pttHeld = false;
       assistantActive = false;
+      visionFollowQueue = [];
+      visionTurnLocked = false;
+      suppressedPlaybackId = null;
       sendTimer = window.setInterval(flushCapture, profile.sendIntervalMs || SEND_INTERVAL_MS);
       startClock();
       callButton.textContent = 'End session';
@@ -679,6 +838,7 @@
       muteButton.disabled = false;
       cameraButton.disabled = !profile.camera;
       sendTurnButton.disabled = !profile.clientCommit;
+      pttButton.disabled = !profile.pushToTalk;
       setConnection('Connected', 'online');
       setModel(profile.waiting);
       appendLog('session started');
@@ -698,8 +858,9 @@
     cameraPreview.srcObject = cameraStream;
     cameraPreview.style.display = '';
     await cameraPreview.play().catch(() => {});
-    // Official omni-duplex cadence: one JPEG (quality 0.7) per ~1 s chunk,
-    // no client-side resize (the server normalizes at scale_resolution=448).
+    // JPEG quality 0.7; no client-side resize (server scale_resolution=448).
+    // Interval is per-profile: AURA 2 fps, MiniCPM/Qwen stay at 1 s.
+    const cameraIntervalMs = profile.cameraIntervalMs || 1000;
     const captureCameraFrame = () => {
       if (!cameraStream || cameraPreview.videoWidth === 0) return;
       const scale = profile.cameraMaxDimension
@@ -712,14 +873,17 @@
         if (connectionReady && socket?.readyState === WebSocket.OPEN) {
           for (const event of profile.imageMessages(frame)) socket.send(JSON.stringify(event));
         }
-      } else cameraPendingFrame = frame;
+      } else {
+        cameraPendingFrame = frame;
+        if (profile.stickyCamera) cameraLastFrame = frame;
+      }
     };
     // Do not make the first spoken turn race a one-second timer.
     captureCameraFrame();
-    cameraTimer = window.setInterval(captureCameraFrame, 1000);
+    cameraTimer = window.setInterval(captureCameraFrame, cameraIntervalMs);
     cameraButton.textContent = 'Camera off';
     cameraButton.classList.add('is-active');
-    appendLog('camera on (1 fps omni frames)');
+    appendLog(`camera on (${(1000 / cameraIntervalMs).toFixed(0)} fps omni frames)`);
   }
 
   function stopCamera() {
@@ -730,6 +894,7 @@
     }
     cameraStream = null;
     cameraPendingFrame = null;
+    cameraLastFrame = null;
     cameraPreview.srcObject = null;
     cameraPreview.style.display = 'none';
     cameraButton.textContent = 'Camera';
@@ -778,9 +943,18 @@
     playbackComplete = true;
     turnSubmitted = false;
     assistantTextChannel = null;
+    assistantRowByResponse.clear();
+    lastClosedAssistantText = '';
     sendTurnButton.disabled = true;
+    pttHeld = false;
+    pttButton.disabled = true;
+    pttButton.classList.remove('is-active');
+    pttButton.textContent = 'Hold to talk';
     running = false;
     assistantActive = false;
+    visionFollowQueue = [];
+    visionTurnLocked = false;
+    suppressedPlaybackId = null;
     pendingCapture = [];
     if (sendTimer !== null) clearInterval(sendTimer);
     if (clockTimer !== null) clearInterval(clockTimer);
@@ -854,6 +1028,23 @@
     setModel('Thinking');
     appendLog('turn submitted');
   });
+  const pttDown = (event) => {
+    event.preventDefault();
+    setPttHeld(true);
+  };
+  const pttUp = (event) => {
+    event.preventDefault();
+    setPttHeld(false);
+  };
+  pttButton.addEventListener('pointerdown', pttDown);
+  pttButton.addEventListener('pointerup', pttUp);
+  pttButton.addEventListener('pointercancel', pttUp);
+  pttButton.addEventListener('pointerleave', () => {
+    if (pttHeld) setPttHeld(false);
+  });
+  window.addEventListener('pointerup', () => {
+    if (pttHeld) setPttHeld(false);
+  });
   muteButton.addEventListener('click', toggleMute);
   clearLogButton.addEventListener('click', () => {
     eventLog.textContent = '';

examples/online_serving/realtime_web/app/static/profiles/aura_ptt.js

@@ -0,0 +1,79 @@
+(function (root) {
+  'use strict';
+  const profiles = root.OmniRealtimeProfiles;
+  profiles['aura-ptt'] = () => ({
+    title: 'AURA Realtime',
+    eyebrow: 'Push-to-talk duplex',
+    policy: 'Hold to talk • release to commit',
+    description:
+      'Hold to talk, release to send. Camera stays on whether you speak or not. '
+    + 'Holding the button stops local playback.',
+    waiting: 'Ready',
+    camera: true,
+    // Display only. Shared CSS puts .camera-preview-large on its own row.
+    // Other profiles keep the 96×72 preview in index.html.
+    cameraPreviewLarge: true,
+    playbackAck: true,
+    clientCommit: false,
+    pushToTalk: true,
+    stickyCamera: true,
+    // OmniInteract / Native AURA clock: 2 fps. Other profiles omit this and stay at 1 s.
+    cameraIntervalMs: 500,
+    visionFollowWhileSpeaking: true,
+    deduplicateTranscript: true,
+    halfDuplex: false,
+    closeSession: true,
+    reconnectEachTurn: false,
+    readyEvent: 'session.updated',
+    instructions: true,
+    sendIntervalMs: 200,
+    presets: {
+      omni: 'You are receiving a live video stream where the final frame is the present moment. Respond only when a response is needed based on the user\'s message or the visual context. Otherwise, output `<|silent|>` to signify silence.',
+    },
+    url(config, location) {
+      const url = profiles.url(config, location);
+      url.searchParams.set('duplex', '1');
+      url.searchParams.set('model', config.model || 'aurateam/AURA');
+      return url.toString();
+    },
+    initialMessages(config, instructions) {
+      const session = {
+        modalities: ['audio', 'text'],
+        extra_body: { auto_response: true },
+      };
+      if (instructions) session.instructions = instructions;
+      return [{ type: 'session.update', session }];
+    },
+    append(audio, frame, opts) {
+      const isSpeech = !opts || opts.isSpeech !== false;
+      const event = {
+        type: 'input_audio_buffer.append',
+        audio,
+        format: 'pcm16',
+        sample_rate_hz: 16000,
+        is_speech: isSpeech,
+      };
+      const frames = Array.isArray(frame) ? frame.filter(Boolean) : (frame ? [frame] : []);
+      if (frames.length) event.video_frames = frames;
+      return event;
+    },
+    commitMessages: () => [{ type: 'input_audio_buffer.commit', create_response: true }],
+    ack(responseId, playedMs) {
+      return {
+        type: 'playback.ack',
+        response_id: responseId,
+        item_id: `item_${responseId}`,
+        played_ms: playedMs,
+        committed_ms: playedMs,
+      };
+    },
+    mapEvent(event) {
+      if (event.type === 'response.listen') return { kind: 'listen' };
+      if (event.type === 'playback.acknowledged') {
+        return { kind: 'ack', committedMs: (event.event || event).committed_ms || 0 };
+      }
+      return profiles.event(event);
+    },
+    connectionHint: 'Check the AURA duplex serve (run_duplex_smoke_serve.sh) and --ws-backend port.',
+  });
+})(globalThis);

examples/online_serving/realtime_web/app/static/styles.css

@@ -71,6 +71,18 @@ main { width: min(1040px, calc(100% - 40px)); margin: 28px auto 48px; }
 }
 
 .call-actions { display: flex; gap: 8px; }
+/* AURA only: class is added solely by the aura-ptt profile. */
+.call-bar:has(.camera-preview-large) .call-actions { grid-column: 1 / -1; }
+.call-actions:has(.camera-preview-large) { flex-wrap: wrap; align-items: center; }
+.call-actions:has(.camera-preview-large) .camera-preview-large {
+  order: -1;
+  flex: 0 0 100%;
+  width: 100% !important;
+  height: auto !important;
+  max-height: 480px;
+  aspect-ratio: 4 / 3;
+}
+.call-actions:has(.camera-preview-large) button { flex: 0 0 auto; }
 button { border-radius: 5px; border: 1px solid transparent; cursor: pointer; }
 button:disabled { opacity: .46; cursor: default; }
 .primary { padding: 10px 15px; color: #ffffff; background: #16834b; font-weight: 700; }

examples/online_serving/realtime_web/server.py

@@ -76,15 +76,20 @@ def build_app(
     profile: str = "minicpm-native",
     adapter: str = "stt",
 ) -> FastAPI:
-    if profile not in {"minicpm-native", "qwen3-turn"}:
+    if profile not in {"minicpm-native", "qwen3-turn", "aura-ptt"}:
         raise ValueError(f"Unknown profile: {profile}")
     if adapter not in {"stt", "vad"}:
         raise ValueError(f"Unknown adapter: {adapter}")
     if profile == "qwen3-turn" and ref_audio:
         raise ValueError("Qwen3 profile does not accept --ref-audio")
+    if profile == "aura-ptt" and ref_audio:
+        raise ValueError("AURA profile does not accept --ref-audio")
     native = profile == "minicpm-native"
-    ws_backend = ws_backend or ("ws://127.0.0.1:8099" if native else "ws://127.0.0.1:8091")
-    model = model or ("openbmb/MiniCPM-o-4_5" if native else "Qwen/Qwen3-Omni-30B-A3B-Instruct")
+    aura = profile == "aura-ptt"
+    ws_backend = ws_backend or ("ws://127.0.0.1:8099" if native or aura else "ws://127.0.0.1:8091")
+    model = model or (
+        "openbmb/MiniCPM-o-4_5" if native else "aurateam/AURA" if aura else "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    )
     app = FastAPI(title="Omni Realtime Web UI")
     index_path = APP_DIR / "index.html"
     app_version_hash = hashlib.sha256()
@@ -182,7 +187,7 @@ def main(default_profile: str = "minicpm-native") -> None:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--host", default="0.0.0.0")
     parser.add_argument("--port", type=int, default=7862)
-    parser.add_argument("--profile", choices=("minicpm-native", "qwen3-turn"), default=default_profile)
+    parser.add_argument("--profile", choices=("minicpm-native", "qwen3-turn", "aura-ptt"), default=default_profile)
     parser.add_argument("--ws-backend", "--backend", dest="ws_backend")
     mode = parser.add_mutually_exclusive_group()
     mode.add_argument("--stt", dest="adapter", action="store_const", const="stt")
@@ -202,12 +207,15 @@ def main(default_profile: str = "minicpm-native") -> None:
     )
     args = parser.parse_args()
     native = args.profile == "minicpm-native"
+    aura = args.profile == "aura-ptt"
     if native and not args.ref_audio:
         parser.error("--ref-audio is required for the minicpm-native profile")
     if not native and args.ref_audio:
         parser.error("--ref-audio is only supported by minicpm-native")
-    args.model = args.model or ("openbmb/MiniCPM-o-4_5" if native else "Qwen/Qwen3-Omni-30B-A3B-Instruct")
-    args.ws_backend = args.ws_backend or ("ws://127.0.0.1:8099" if native else "ws://127.0.0.1:8091")
+    args.model = args.model or (
+        "openbmb/MiniCPM-o-4_5" if native else "aurateam/AURA" if aura else "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    )
+    args.ws_backend = args.ws_backend or ("ws://127.0.0.1:8099" if native or aura else "ws://127.0.0.1:8091")
 
     logging.basicConfig(level=logging.INFO)
     uvicorn.run(

recipes/Qwen/Qwen-Image-2.1.md

@@ -0,0 +1,398 @@
+# Qwen-Image-2.1
+
+> Text-to-image and image-conditioned generation with Qwen-Image 2.1
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen-Image-2.1`
+- Task: Text-to-image generation and image-conditioned generation (editing)
+- Mode: Offline inference and online serving with optional step-wise execution
+- Maintainer: Community
+
+> **Model id:** `Qwen/Qwen-Image-2.1` follows the reference diffusers
+> documentation and may change before the official release — the final
+> checkpoint name is subject to the official release.
+
+## When to use this recipe
+
+Use this recipe as a starting point for running `Qwen/Qwen-Image-2.1`. A single
+pipeline class, `QwenImage21Pipeline`, handles both pure text-to-image requests
+(no condition image) and image-conditioned requests with up to 4 condition
+images: prompt and condition images are encoded together by a Qwen3-VL text
+encoder, so editing-style requests go through the same serving path as
+text-to-image.
+
+## References
+
+- Related offline examples:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py),
+  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
+- Supported models table:
+  [`docs/models/supported_models.md`](../../docs/models/supported_models.md)
+- Feature compatibility matrix:
+  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
+- Sibling recipes: [`Qwen-Image.md`](./Qwen-Image.md),
+  [`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md)
+
+## Offline Inference
+
+### Text-to-image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --output qwen_image_21_t2i.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+### Image-conditioned generation (editing)
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --color-format RGBA \
+  --seed 42 \
+  --image qwen_bear.png \
+  --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
+  --output qwen_image_21_edit.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+or via the bundled launcher:
+
+```bash
+bash examples/offline_inference/image_to_image/run_qwen_image_21.sh
+```
+
+Up to 4 condition images can be passed via `--image`:
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --color-format RGBA \
+  --seed 42 \
+  --image input1.png input2.png \
+  --prompt "Combine these images into a single scene" \
+  --output qwen_image_21_multi.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+Use `--color-format RGBA` to preserve transparency in condition images. The
+shared example otherwise loads inputs as RGB. Add `--width` and `--height`
+(multiples of 32) to set the output size; when omitted, Qwen-Image 2.1 derives
+it from the last condition image's aspect ratio at approximately 1024×1024.
+Use `--num-outputs-per-prompt` to generate multiple images and `--vae-use-tiling`
+to reduce VAE memory usage. The launcher accepts these options as well.
+
+## Online Serving
+
+Start the server:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --port 8091
+```
+
+To enable the step-wise runtime:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni \
+  --port 8091 \
+  --step-execution \
+  --max-num-seqs 8
+```
+
+Step execution batches compatible requests at the same KV-cache phase. The
+pipeline tags every request with a `batch_compatibility_key` (conditioning
+mode plus condition-image layout), so a newly arriving request that would
+otherwise join a batch already past its first denoising step is isolated by
+the scheduler and deferred until the in-flight batch drains — it no longer
+fails the whole batch. `--max-num-seqs` controls batch capacity.
+
+### CUDA Graph decode
+
+Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising
+decode steps when `enable_cuda_graph_decode=True` (the default) and
+`enforce_eager=False`. Prefill remains eager. This works with both request
+execution and `--step-execution`.
+
+To disable graph capture while keeping the rest of the compilation setup:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode
+```
+
+The same control is available as `Omni(..., enable_cuda_graph_decode=False)`
+and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML.
+`--enforce-eager` / `enforce_eager: true` remains the broader switch: it
+disables both graph capture and automatic `torch.compile`. The two
+optimizations stack: the DiT blocks are regionally `torch.compile`'d and
+graph capture records the compiled (fused) kernels, while inductor's own
+cudagraphs stay disabled so the two graph layers never nest. Measured on a
+single GB200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`,
+warmup 1 + median of 3): compile+graph 3.04 s, compile-only 3.2 s,
+graph-only 4.0 s, eager 4.3 s end-to-end per image; combo output vs eager is
+44.5–45.6 dB PSNR (the same magnitude as pure `torch.compile` fusion
+divergence), and graph-only vs eager is bit-identical.
+
+The autoregressive engine's `compilation_config.cudagraph_mode` does not
+control this diffusion path; use `enable_cuda_graph_decode` (or
+`enforce_eager` for everything) to force eager decode execution.
+
+TP/SP/ring parallelism, HSDP, offload/cache hooks, quantized KV caches,
+dynamic LoRA, padded text masks, and a second in-flight request whose cache
+aliases an already-owned graph key fall back to eager decode.
+Graphs keep separate entries for different image layouts and copy
+request-owned prefix K/V into static buffers before replay; each graph key
+has a single live owner at a time, so concurrent same-shape requests never
+overwrite each other's prefix. Those buffers require additional memory; graph
+entries are bounded by the transformer's cache limit.
+
+### Verification
+
+For a direct API smoke test:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen-Image-2.1",
+    "prompt": "A ceramic teapot on a wooden table",
+    "size": "1024x1024",
+    "num_inference_steps": 50,
+    "true_cfg_scale": 1.0,
+    "seed": 42
+  }'
+```
+
+For image-conditioned generation, pass condition images through the
+OpenAI-compatible multimodal endpoint (`/v1/chat/completions` with
+`image_url` content parts, or `/v1/images/edits`), same as
+[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md).
+
+## Parameters
+
+| Parameter | Default | Notes |
+| --- | --- | --- |
+| `true_cfg_scale` | 1.0 | Disables the negative-prompt branch by default. Set above 1 with a `negative_prompt` to enable true CFG. Qwen-Image 2.1 supports true CFG only — there is no `guidance_scale` parameter. |
+| `num_inference_steps` | 50 | Number of denoising steps. |
+| `height` / `width` | 1024×1024 | Must be multiples of 32. For image-conditioned requests, leaving them unset derives the output size from the last condition image's aspect ratio at ~1024×1024. |
+| `negative_prompt` | None | Ignored at the default CFG scale of 1.0; used only when `true_cfg_scale > 1`. |
+| condition images | — | Up to 4 input images per request; more than 4 raises an error. |
+| `seed` | — | Fix for reproducible outputs. |
+
+The default CFG scale of 1.0 runs only the positive-prompt branch at each
+denoising step. To enable negative-prompt guidance, explicitly set
+`--cfg-scale` (offline) or `true_cfg_scale` (API) above 1 and provide a
+negative prompt. This changes the generated output as well as compute cost.
+
+### Prefix KV cache
+
+The transformer caches the text/condition-image prefix across denoising steps
+(a prefill/decode split keyed by CFG branch). The cache is enabled
+automatically when the checkpoint declares `causal_condition: true` and is
+re-created per generation, so no state leaks between requests. It changes the
+compute path — cached prefix vs. full recompute at every step — so exact
+numerical reproduction against a no-cache reference (e.g. the diffusers
+pipeline) requires matching the cache setting on both sides. The prefix KV
+cache is orthogonal to weight quantization and works unchanged under FP8.
+
+### FP8 quantization
+
+Online FP8 quantization of the transformer is supported; the fused
+`to_qkv` projections load correctly, and the prefix KV cache path is
+unaffected. Only block-internal linears are eligible — the boundary
+projections (`img_in`/`txt_in`/`modulation`/`proj_out`/`time_text_embed`) stay
+in BF16 by construction:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --output qwen_image_21_fp8.png \
+  --num-inference-steps 50 --cfg-scale 1.0 \
+  --quantization fp8 --ignored-layers "img_mlp"
+```
+
+`ignored_layers` entries are matched as name patterns for this model (e.g.
+`img_mlp` matches every block's `img_mlp.proj`/`img_mlp.gate_layer`/
+`img_mlp.out`); exact full prefixes like
+`transformer.transformer_blocks.3.attn.to_qkv` also work.
+
+A per-module sensitivity sweep (fixed prompt embeds / latent / timestep, 5
+timesteps across the trajectory, relative L2 of the noise prediction vs BF16)
+ranks the groups: `img_mlp` is the most sensitive (3.8% max error on its own,
+driven by `img_mlp.proj` 2.8% and `img_mlp.out` 2.5%), then `attn.to_out`
+(2.7%), then `attn.to_qkv` (1.6%); errors are largest at the late (low-sigma)
+steps. End-to-end on 4 fixed prompts at 1024x1024, 50 steps, seed 42
+(GB200, single GPU):
+
+| Config | Quantized block linears | Avg PSNR vs BF16 | Min PSNR | Steady-state time/image | Peak memory |
+| --- | --- | --- | --- | --- | --- |
+| BF16 | 0 / 160 | — | — | 7.2 s | 40.0 GB |
+| FP8 all layers | 160 / 160 | 26.1 dB | 19.1 dB | 8.4 s | 33.4 GB |
+| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | 29.7 dB | 21.6 dB | 8.8 s | 38.3 GB |
+
+All-layer FP8 is usable but its composition can drift on some prompts
+(fine detail and layout shift, worst case ~19 dB); keeping `img_mlp` in BF16
+preserves composition noticeably better and is the recommended setting. On
+Blackwell (GB200) FP8 here is a memory optimization, not a speedup: dynamic
+per-token activation quantization costs more than the FP8 GEMM saves at these
+shapes, so BF16 remains the fastest option.
+
+## Quantization
+
+Online FP8 is supported per component via `quantization_config`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image-2.1",
+    quantization_config={
+        # DiT: skip the sensitive image-stream MLPs if quality regresses.
+        "transformer": {"method": "fp8", "ignored_layers": ["img_mlp"]},
+        # Qwen3-VL text encoder: only the language-model linear layers are
+        # quantized; the vision tower and the (unused) LM head stay BF16.
+        "text_encoder": {"method": "fp8"},
+    },
+)
+```
+
+Either component can be quantized on its own. Measured on GB200 at 1024×1024
+(seed 42, 50 steps): text-encoder FP8 lowers peak GPU memory from ~41.0 GiB to
+~34.4 GiB (−6.6 GiB); adding DiT FP8 reaches ~28.1 GiB. Output quality stays
+close to BF16 (T2I PSNR vs BF16 ≈ 28–32 dB with text-encoder FP8; the
+image-conditioned edit path ≈ 36 dB). FP8 saves memory but does not speed up
+generation on this hardware. The edit path routes condition images through the
+BF16 vision tower, so it is unaffected by text-encoder FP8. See
+[`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md)
+for the scope rules.
+
+### FP8 prefix KV storage
+
+The cached prefix K/V can be stored in FP8 E4M3 (with per-token-per-head fp32
+scales), halving the prefix-cache memory — relevant for long prompts (up to
+8192 text tokens) and up to 4 condition images (~4096 latent tokens each).
+Prefill quantizes once; decode dequantizes per step, so attention compute is
+unchanged. Enable it through the model-specific `extras` config (e.g. in the
+deploy YAML's stage `extras` or as an `Omni(...)` engine kwarg):
+
+```yaml
+extras:
+  prefix_kv_cache_dtype: "fp8"
+```
+
+Accepted values: `"fp8"` / `"fp8_e4m3"` (K and V quantized), `"fp8_v"` (V
+only — K stays in the native dtype), and `None` / `"auto"` (default, native
+dtype — behavior unchanged). This is independent of
+`diffusion_kv_cache_dtype`, which quantizes attention Q/K/V *compute* per
+forward pass on supported backends.
+
+Measured on T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0, PSNR vs. the
+bf16 baseline): **`"fp8"` 34.9 dB**, **`"fp8_v"` 40.9 dB** — same composition
+and semantics, with texture-level drift in fine detail for `"fp8"`, and
+visually indistinguishable output for `"fp8_v"`. The error is dominated by
+K quantization: post-RoPE keys are the precision-sensitive half of the
+cache, so `"fp8_v"` buys back ~6 dB at 75% (instead of 50%) of the original
+cache size. The residual error is inherent e4m3 precision accumulated
+coherently over the denoising trajectory; finer scale granularity
+(per-tensor/per-head/per-token were compared) or Hadamard-rotated V did not
+improve it. The saving scales with prefix length — at the limit
+(8192 text tokens + 4 condition images, ~24.6k prefix tokens) the prefix
+cache is ~12.9 GB per CFG branch in bf16, ~6.6 GB in `"fp8"` and ~9.7 GB in
+`"fp8_v"`, which is where this option matters. Treat it as an opt-in for
+memory-bound long-prompt / multi-image workloads, not a free lunch;
+`"fp8_v"` is the better default trade-off when quality matters.
+
+Note: a quantized prefix cache is not CUDA-graph capturable — requests using
+this option stay on the eager decode path (a warning is logged once).
+
+### Why not the scheduler-managed paged KV (`DiffusionKVCacheMode.PAGED_SCHEDULER`)?
+
+Evaluated and deliberately not adopted (see the `qwen21-p0-kv` commit message
+for the full analysis): the paged scheduler targets *cross-request* KV
+management — a statically sized page pool, block sharing/dedup, and eviction —
+whereas this prefix cache lives exactly as long as one generation and is
+sized once per request. Migrating would require routing decode attention
+through the worker paged adapter (writing fresh target K/V into pages every
+step), per-request page tables, and re-validating the block-causal
+piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over
+the per-request dense tensors. The FP8 storage above captures the actual
+memory win at ~1% of the integration cost.
+
+## Quality Comparison (head vs diffusers)
+
+The reference is the diffusers Qwen-Image 2.1 pipeline
+([huggingface/diffusers#14804](https://github.com/huggingface/diffusers/pull/14804)).
+Both sides must pin the same checkpoint, inputs, seed, precision (BF16), and
+hardware.
+
+vLLM-Omni (this recipe):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --seed 42 \
+  --output qwen_image_21_vllm_omni.png \
+  --num-inference-steps 50 \
+  --cfg-scale 1.0
+```
+
+diffusers reference:
+
+```python
+import torch
+from diffusers import QwenImage21Pipeline  # diffusers PR #14804
+
+pipe = QwenImage21Pipeline.from_pretrained(
+    "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
+).to("cuda")
+image = pipe(
+    prompt="A ceramic teapot on a wooden table",
+    num_inference_steps=50,
+    true_cfg_scale=1.0,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+    height=1024,
+    width=1024,
+).images[0]
+image.save("qwen_image_21_diffusers.png")
+```
+
+| Metric | vLLM-Omni | diffusers (#14804) | Status |
+| --- | --- | --- | --- |
+| Output image (1024×1024, 50 steps, seed 42, BF16) | pending measurement | pending measurement | not yet run |
+| PSNR / LPIPS vs reference | pending measurement | — | not yet run |
+| Latency (time/image, same GPU) | pending measurement | pending measurement | not yet run |
+
+The prefix KV cache changes the compute path (cached prefix vs. full recompute
+at every step), so exact numerical parity is not expected; match the cache
+setting on both sides or compare with PSNR/LPIPS rather than exact pixels.
+
+The following capabilities are implemented but the PR-body checkboxes predate
+any attached measurements; treat them as unverified until the A/B above is
+filled in:
+
+- [ ] TP=2/4 output parity vs single-card — pending measurement.
+- [ ] VAE tiling vs non-tiled decode — pending measurement.
+- [x] CUDA graph decode vs eager — measured on GB200 (1024×1024, 50 steps,
+  seed 42, BF16): compile+graph 3.04 s/image vs eager 4.3 s/image (~1.4×),
+  44.5–45.6 dB PSNR vs eager; graph-only is bit-identical to eager. See
+  "CUDA Graph decode" above.
+
+## Known Limitations
+
+- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;
+  `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`.
+- Sequence parallelism supports Ulysses only; ring attention is not supported.
+- Only true CFG is exposed; there is no `guidance_scale` knob.
+- VAE tiling (`--vae-use-tiling`, also implied by `--vae-patch-parallel-size > 1`)
+  decodes in 512px tiles with 384px stride (32 latent pixels per tile at 16x
+  spatial compression). On CUDA OOM the tile size is automatically halved
+  (keeping 25% overlap) and decoding retried, down to 128px tiles. Tile sizes
+  are not yet exposed as engine arguments.

recipes/Qwen/Qwen3-TTS.md

@@ -71,8 +71,10 @@ Alternatively, use the convenience script:
 ```
 
 The bundled deploy config (`vllm_omni/deploy/qwen3_tts.yaml`) enables async
-chunking for low first-audio latency. For advanced deployment tuning, pass a
-custom deploy config:
+chunking for low first-audio latency, and selects the experimental Model
+Runner V2 on CUDA. Copy it and set `model_runner: v1` to run V1; NPU, XPU,
+ROCm and MUSA keep V1 through the `platforms:` sections. For advanced
+deployment tuning, pass a custom deploy config:
 
 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \

recipes/README.md

@@ -83,6 +83,7 @@ recipes/
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
 | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
+| [`Qwen/Qwen-Image-2.1.md`](./Qwen/Qwen-Image-2.1.md) | Text-to-image and image-conditioned generation (editing) with prefix KV cache and CUDA graph decode | 1x GPU with CPU offload / 1x GB200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |

tests/benchmarks/test_aura_omniinteract_clock.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.benchmarks.omniinteract import _session_capabilities, _wants_annotation_video_clock
+from vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock import (
+    stream_annotation_clock,
+    wants_annotation_video_clock,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Events:
+    def __init__(self, events):
+        self.events = events
+
+
+class _Client:
+    def __init__(self, capabilities):
+        self.events = _Events([{"type": "session.created", "session": {"capabilities": capabilities}}])
+        self._client = None
+
+
+def test_annotation_clock_follows_capabilities_not_model_name():
+    aura = {"required_input_modalities": ["video"], "supports_turn_commit_only": True}
+    minicpm = {"required_input_modalities": ["audio"], "supports_turn_commit_only": False}
+    audio_commit = {"required_input_modalities": ["audio"], "supports_turn_commit_only": True}
+    assert wants_annotation_video_clock(aura)
+    assert not wants_annotation_video_clock(minicpm)
+    assert not wants_annotation_video_clock(audio_commit)
+    assert not wants_annotation_video_clock(None)
+    assert _wants_annotation_video_clock(_Client(minicpm)) is False
+    assert _wants_annotation_video_clock(_Client(aura)) is True
+    assert _session_capabilities(_Client(aura))["required_input_modalities"] == ["video"]
+
+
+def test_annotation_clock_commits_only_the_due_wav(monkeypatch, tmp_path: Path):
+    video = tmp_path / "videos" / "0002.mp4"
+    video.parent.mkdir()
+    video.write_bytes(b"")
+    sent = []
+
+    class Client:
+        async def send(self, event):
+            sent.append(event)
+
+        async def commit(self):
+            sent.append({"type": "commit"})
+
+    class Playback:
+        async def acknowledge(self, client):
+            del client
+
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock.load_annotation_utterances",
+        lambda path: [{"at_sec": 0.0, "pcm": b"\x01\x00" * 160, "sent": False, "qa_index": 0}],
+    )
+    chunks, frames, _, _ = asyncio.run(
+        stream_annotation_clock(
+            Client(),
+            ["frame-a", "frame-b"],
+            Playback(),
+            video_path=video,
+            fps=2.0,
+        )
+    )
+    assert (chunks, frames) == (2, 2)
+    assert sent[0]["is_speech"] is True
+    assert sent[0]["video_frames"] == ["frame-a"]
+    assert sent[1] == {"type": "commit"}
+    # The trailing silent frame is one vision turn, committed after the clip.
+    assert sent[2]["is_speech"] is False
+    assert sent[2]["video_frames"] == ["frame-b"]
+    assert sent[3] == {"type": "commit"}
+    assert sent.count({"type": "commit"}) == 2
+
+
+def test_annotation_clock_commits_silent_frames_in_pairs(monkeypatch, tmp_path: Path):
+    video = tmp_path / "videos" / "0002.mp4"
+    video.parent.mkdir()
+    video.write_bytes(b"")
+    sent = []
+
+    class Client:
+        async def send(self, event):
+            sent.append(event)
+
+        async def commit(self):
+            sent.append({"type": "commit"})
+
+    class Playback:
+        async def acknowledge(self, client):
+            del client
+
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock.load_annotation_utterances",
+        lambda path: [{"at_sec": 100.0, "pcm": b"\x01\x00" * 160, "sent": False, "qa_index": 0}],
+    )
+    asyncio.run(
+        stream_annotation_clock(
+            Client(),
+            ["f1", "f2", "f3", "f4"],
+            Playback(),
+            video_path=video,
+            fps=2.0,
+        )
+    )
+    vision = [event for event in sent if event.get("is_speech") is False]
+    assert [event["video_frames"] for event in vision] == [["f1", "f2"], ["f3", "f4"]]
+    assert sent.count({"type": "commit"}) == 2
+
+
+def test_annotation_clock_holds_vision_while_a_response_is_open(monkeypatch, tmp_path: Path):
+    video = tmp_path / "videos" / "0002.mp4"
+    video.parent.mkdir()
+    video.write_bytes(b"")
+    sent = []
+
+    class Events:
+        def __init__(self):
+            self.events = [{"type": "response.created", "response": {"id": "resp-open"}}]
+
+    class Client:
+        def __init__(self):
+            self.events = Events()
+
+        async def send(self, event):
+            sent.append(event)
+
+        async def commit(self):
+            sent.append({"type": "commit"})
+
+    class Playback:
+        async def acknowledge(self, client):
+            del client
+
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.aura_omni.benchmarks.omniinteract_clock.load_annotation_utterances",
+        lambda path: [{"at_sec": 100.0, "pcm": b"\x01\x00" * 160, "sent": False, "qa_index": 0}],
+    )
+    asyncio.run(
+        stream_annotation_clock(
+            Client(),
+            ["f1", "f2", "f3", "f4"],
+            Playback(),
+            video_path=video,
+            fps=2.0,
+        )
+    )
+    assert sent == []

tests/buildkite/test_amd_pipeline.py

@@ -67,11 +67,31 @@ def test_ready_diffusion_cpu_suite_is_sharded() -> None:
     assert "--shard-id=$$BUILDKITE_PARALLEL_JOB" in pytest_command
 
 
-def test_cosyvoice_gpu_abort_gets_one_fresh_job_retry() -> None:
-    step = _find_step("CosyVoice3-TTS E2E Test", AMD_READY_PIPELINE)
+def test_cosyvoice_ready_smoke_uses_sdpa() -> None:
+    step = _find_step("CosyVoice3-TTS E2E Smoke (SDPA)", AMD_READY_PIPELINE)
 
+    assert step["grade"] == "Blocking"
     assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]}
+    assert "export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" in step["commands"]
 
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002" in pytest_command
+
+
+def test_cosyvoice_full_default_backend_suite_runs_nightly() -> None:
+    step = _find_step("CosyVoice3-TTS E2E Test", AMD_NIGHTLY_PIPELINE)
+
+    assert step["grade"] == "NonBlocking"
+    assert step["timeout_in_minutes"] == 90
+    assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]}
+    assert all("DIFFUSION_ATTENTION_BACKEND" not in command for command in step["commands"])
+
+    pytest_command = next(command for command in step["commands"] if "pytest" in command)
+    assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py" in pytest_command
+    assert "::" not in pytest_command
+
+
+def test_amd_template_preserves_step_retry_policy() -> None:
     template = AMD_TEMPLATE.read_text(encoding="utf-8")
     # Both grouped and top-level AMD steps must preserve an explicit retry
     # policy when the source suite is rendered into the uploaded pipeline.

tests/clients/test_duplex_client.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 import asyncio
 import base64
 import json
+import struct
 import wave
 from collections.abc import Callable
 from typing import Any
@@ -33,6 +34,7 @@ from vllm_omni.clients.duplex import (
     ReconnectPolicy,
     SessionConfig,
     SessionResumed,
+    SessionUpdated,
     build_realtime_url,
     chunk_period_ms,
     duplex_unit_boundary_ms,
@@ -797,6 +799,52 @@ async def test_inline_client_submits_typed_commands_in_order():
         assert ack.played_ms == 120 and ack.response_id == "resp-1"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("accepted", [True, False])
+@pytest.mark.parametrize(
+    "patch",
+    [
+        {"input_audio_format": "pcm_f32le", "sample_rate_hz": 24_000},
+        {"audio": {"input": {"format": "pcm_f32le", "sample_rate_hz": 24_000}}},
+    ],
+)
+async def test_inline_input_defaults_change_only_after_session_updated(accepted, patch):
+    client, omni = make_inline_client()
+    patch = {**patch, "overlap_silence_rms": 0.1}
+    async with client:
+        handle = omni.handles[0]
+        previous_session = dict(client.session_info)
+        await client.send({"type": "session.update", "session": patch})
+        # Sending a patch is not an acknowledgement of its acceptance.
+        assert client.session_info == previous_session
+        assert client._input_defaults().input_sample_rate_hz == 16_000
+        updated_session = {**previous_session, **patch}
+        event: dict[str, object] = (
+            {"type": "session.updated", "session": updated_session}
+            if accepted
+            else {"type": "error", "error": {"code": "invalid_session_update", "message": "rejected"}}
+        )
+        handle.feed(event)
+        received = await client.wait_for("session.updated" if accepted else "error", timeout_s=2.0)
+        assert isinstance(received, SessionUpdated if accepted else ErrorEvent)
+
+        pcm = struct.pack("<f", 0.05) * 16
+        await client.send({"type": "input_audio_buffer.append", "audio": _b64(pcm)})
+        command = handle.commands[-1]
+        assert isinstance(command, duplex_commands.AppendAudio)
+        assert command.format == "pcm_f32le"
+        if accepted:
+            assert client.session_info == updated_session
+            assert command.audio == pcm
+            assert command.sample_rate_hz == 24_000
+            assert command.is_speech is False
+        else:
+            assert client.session_info == previous_session
+            assert len(command.audio) == 2 * len(pcm)  # Still decoded as pcm16.
+            assert command.sample_rate_hz == 16_000
+            assert command.is_speech is True
+
+
 @pytest.mark.asyncio
 async def test_inline_session_close_maps_to_handle_close():
     client, omni = make_inline_client()

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 29,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 59,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         ModelEnvironmentVariableDisposition.PROMOTE: 31,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
-        ModelEnvironmentVariableDisposition.INTERNALIZE: 14,
+        ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
         ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
     }
 

tests/config/test_omni_config.py

@@ -816,19 +816,25 @@ def test_mrv2_fails_fast_on_platforms_without_native_workers(platform: str):
         _apply_platform_overrides(DeployConfig(model_runner="v2"), platform=platform)
 
 
-@pytest.mark.parametrize(
-    ("default_name", "mrv2_name"),
-    [
-        ("qwen3_tts.yaml", "qwen3_tts_mrv2.yaml"),
-        (
-            "qwen3_tts_high_concurrency.yaml",
-            "qwen3_tts_high_concurrency_mrv2.yaml",
-        ),
-    ],
-)
-def test_qwen3_mrv2_profiles_are_explicit_opt_in(default_name: str, mrv2_name: str):
-    assert load_deploy_config(_DEPLOY_DIR / default_name).model_runner == "v1"
-    assert load_deploy_config(_DEPLOY_DIR / mrv2_name).model_runner == "v2"
+def test_qwen3_tts_high_concurrency_mrv2_profile_is_explicit_opt_in():
+    assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency.yaml").model_runner == "v1"
+    assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2.yaml").model_runner == "v2"
+
+
+def test_qwen3_tts_default_profile_is_experimental_mrv2_with_v1_platform_fallback():
+    deploy_path = _DEPLOY_DIR / "qwen3_tts.yaml"
+    assert load_deploy_config(deploy_path).model_runner == "v2"
+
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
+    assert all(stage.yaml_engine_args["use_v2_model_runner"] is True for stage in stages)
+
+    assert _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda").model_runner == "v2"
+    for platform in ("npu", "xpu", "rocm", "musa"):
+        assert _apply_platform_overrides(load_deploy_config(deploy_path), platform=platform).model_runner == "v1"
+
+    # The explicit MRV2 profile resolves to the same runner selection.
+    assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_mrv2.yaml").model_runner == "v2"
 
 
 def test_qwen3_tts_mrv2_retunes_do_not_change_default_mrv1_profile():

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -69,6 +69,39 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
     )
 
 
+def _make_minicpm_window_update(
+    *,
+    seq: int,
+    mode: str,
+    high: int = 15,
+    low: int = 10,
+    context_max_units: int = 1,
+) -> StreamingUpdate:
+    update = _make_update([0] * 8)
+    update.model_intermediate_buffer = {
+        "duplex": {
+            "data_plane": True,
+            "seq": seq,
+            "runtime_config": {
+                "duplex_scheduler_token_id": 0,
+                "duplex_first_append_context_tokens": 3,
+                "duplex_window_prefix_tokens": 2,
+                "duplex_window_suffix_token_ids": [3],
+                "duplex_window_previous_marker_token_ids": [70, 71],
+                "duplex_window_special_token_ids": [99],
+                "duplex_window_config": {
+                    "sliding_window_mode": mode,
+                    "basic_window_high_tokens": high,
+                    "basic_window_low_tokens": low,
+                    "context_previous_max_tokens": 4,
+                    "context_max_units": context_max_units,
+                },
+            },
+        }
+    }
+    return update
+
+
 def _make_talker_adapter(
     *,
     max_model_len: int = 100,
@@ -510,6 +543,225 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
     assert sched._new_prompt_len_snapshot[session.request_id] == 2
 
 
+def test_stage0_basic_window_rebuilds_below_low_watermark() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids.clear()
+    session._all_token_ids.extend(session.prompt_token_ids)
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 9
+    assert session.num_computed_tokens == 0
+    assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan == {
+        "completed_token_ids": [40],
+        "replace": True,
+        "mode": "basic",
+        "drop_units": 1,
+        "dropped_tokens": 9,
+        "previous_token_ids": [],
+        "previous_marker_token_ids": [70, 71],
+        "replacement_prompt_len": 9,
+    }
+    sched._free_request_blocks.assert_called_once_with(session)
+
+
+def test_stage0_context_window_compacts_dropped_speech() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids.clear()
+    session._all_token_ids.extend(session.prompt_token_ids)
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+    first = _make_minicpm_window_update(seq=2, mode="context")
+    sched._update_request_as_session(session, first)
+    assert session.num_prompt_tokens == 18
+    assert first.model_intermediate_buffer["duplex"]["stage0_window"] == {"completed_token_ids": [40]}
+
+    session.append_output_token_ids([50])
+    session.num_computed_tokens = 19
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    second = _make_minicpm_window_update(seq=3, mode="context")
+    sched._update_request_as_session(session, second)
+
+    assert session.prompt_token_ids == [0] * 21
+    plan = second.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["completed_token_ids"] == [50]
+    assert plan["drop_units"] == 1
+    assert plan["previous_token_ids"] == [40]
+    # The worker embeds this exact marker, not a re-tokenized one, so the
+    # previous region is len(marker) + len(previous).
+    assert plan["previous_marker_token_ids"] == [70, 71]
+    assert plan["replacement_prompt_len"] == 21
+    assert getattr(session, "_minicpmo45_window_previous_len") == 3
+
+
+@pytest.mark.parametrize("cleared_outputs", [False, True])
+@pytest.mark.parametrize("in_flight", [0, 1])
+def test_stage0_window_uses_confirmed_span_and_terminator(cleared_outputs, in_flight) -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = session.prompt_token_ids
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40, 99])
+    session.num_computed_tokens = 10 + in_flight
+    session.num_output_placeholders = in_flight
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    if cleared_outputs:
+        # A stop without a queued append clears this before the next update.
+        session._output_token_ids.clear()
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+
+    sched._update_request_as_session(session, update)
+
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["completed_token_ids"] == [40]
+    assert plan["completed_terminator_token_id"] == 99
+    assert plan["replacement_prompt_len"] == 9
+    assert plan["dropped_tokens"] == 9
+
+
+def test_stage0_window_rebuild_that_overflows_max_model_len_finishes_the_session() -> None:
+    """A rebuilt window prompt is bounded by the client's window settings, not
+    by the model: a unit carrying camera frames is hundreds of tokens, and a
+    ``context`` window adds the ``previous`` region on top. The replacement
+    branch must therefore check ``replacement_prompt_len`` against
+    ``max_model_len - sample_room`` and finish the session with
+    ``context_length_exceeded``, exactly like a plain extension."""
+    sched = _make_live_session_scheduler(max_model_len=21)
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = [0] * 9
+    session.num_prompt_tokens = 9
+    session._output_token_ids[:] = [40]
+    session._all_token_ids.append(40)
+    session.num_computed_tokens = 10
+    session.num_output_placeholders = 0
+    _park_session(sched, session)
+    # The parked append state above is the frame before this update: one
+    # confirmed output token, fully computed.
+    session._output_token_ids[:] = [40]
+    session._all_token_ids[:] = [*session.prompt_token_ids, 40]
+    session.num_computed_tokens = 10
+    session.num_in_flight_tokens = 0
+    sched.num_sampled_tokens_per_step = 1
+    # The replacement plan rebuilds 21 tokens (retained context plus the
+    # append) and one step samples one more, so 21 leaves no room.
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 20
+    original_prompt = list(session.prompt_token_ids)
+
+    sched._update_request_as_session(session, update)
+
+    assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
+    assert session.prompt_token_ids == original_prompt
+    assert session.num_prompt_tokens == 9
+    assert session.status == RequestStatus.FINISHED_ERROR
+    assert session.request_id not in sched.requests
+    assert sched.finished_req_ids == {session.request_id}
+    sched._free_request_blocks.assert_called_once_with(session)
+    client_index, reason = sched._streaming_context_overflow[session.request_id]
+    assert client_index == session.client_index
+    assert reason.startswith("context_length_exceeded: ")
+    assert "21 tokens" in reason and "max_model_len 21" in reason
+
+    engine_core_outputs = _run_idle_step(sched)
+
+    (output,) = engine_core_outputs[session.client_index].outputs
+    assert output.request_id == session.request_id
+    assert output.finish_reason == FinishReason.ERROR
+    assert output.stop_reason == reason
+
+
+def test_stage0_window_rebuild_that_leaves_room_to_sample_replaces_the_prompt() -> None:
+    """One slot above the plan fits: the replacement applies through the normal
+    replacement path."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 25
+    sched.num_sampled_tokens_per_step = 1
+    session = _make_request()
+    session.prompt_token_ids = [0] * 9
+    session._all_token_ids[:] = session.prompt_token_ids
+    session.num_prompt_tokens = 9
+    session.append_output_token_ids([40])
+    session.num_computed_tokens = 10
+    session.num_output_placeholders = 0
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 20
+
+    sched._update_request_as_session(session, update)
+
+    assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
+    assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
+    assert session.prompt_token_ids == [0] * 21
+    assert session.num_prompt_tokens == 21
+    assert session.num_computed_tokens == 0
+    assert session.status == RequestStatus.WAITING
+    sched._free_request_blocks.assert_called_once_with(session)
+    assert not getattr(sched, "_streaming_context_overflow", {})
+
+
+def test_stage0_window_open_start_zero_is_not_replaced_by_the_context_reserve() -> None:
+    """A recorded ``open_start`` of 0 is a legitimate empty context prefix. The
+    fallback must be an explicit ``is None`` check: treating 0 as missing
+    substitutes the context reserve, counts the suffix twice, and the worker's
+    rebuild-length check raises on the first replacement."""
+    session = SimpleNamespace(
+        num_prompt_tokens=9,
+        _all_token_ids=[0] * 9 + [40],
+        num_computed_tokens=10,
+        num_output_placeholders=0,
+        _minicpmo45_window_open_start=0,
+    )
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 21
+
+    assert OmniARScheduler._prepare_minicpmo45_stage0_window(
+        session, update, segment_output_ids=[40], completed_terminator=3
+    )
+
+    # open_start stays 0, so the reported unit spans the whole 12-token
+    # boundary instead of the 9 the context reserve would fold in.
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["dropped_tokens"] == 12
+    assert plan["replacement_prompt_len"] == 19
+
+
+def test_stage0_window_open_start_falls_back_to_the_context_reserve_when_unset() -> None:
+    session = SimpleNamespace(
+        num_prompt_tokens=9,
+        _all_token_ids=[0] * 9 + [40],
+        num_computed_tokens=10,
+        num_output_placeholders=0,
+    )
+    update = _make_minicpm_window_update(seq=2, mode="basic")
+    update.prompt_token_ids = [0] * 21
+
+    assert OmniARScheduler._prepare_minicpmo45_stage0_window(
+        session, update, segment_output_ids=[40], completed_terminator=3
+    )
+
+    # fallback open_start is preserve_len 3 from duplex_first_append_context_tokens,
+    # so the unit is three tokens shorter.
+    plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
+    assert plan["dropped_tokens"] == 9
+    assert plan["replacement_prompt_len"] == 22
+
+
 def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
     sched = _make_scheduler(stage_id=1)
     sched.chunk_transfer_adapter = SimpleNamespace(

tests/diffusion/attention/test_attention_config.py

@@ -42,6 +42,14 @@ from vllm_omni.diffusion.data import (
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def _clear_inherited_attention_backend(monkeypatch):
+    # Weekly CPU exports DIFFUSION_ATTENTION_BACKEND; OmniDiffusionConfig
+    # treats that as an explicit default so auto-backend cases break.
+    # Tests that setenv the same variable still work: this fixture runs first.
+    monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
+
+
 class TestAttentionSpec:
     def test_construct_no_skip_softmax(self):
         spec = AttentionSpec(backend="FLASH_ATTN")
@@ -416,13 +424,6 @@ class TestBuildAttentionConfig:
 class TestOmniDiffusionConfigAttentionParsing:
     """Test OmniDiffusionConfig attention shorthand and structured config."""
 
-    @pytest.fixture(autouse=True)
-    def _clear_diffusion_attention_backend_env(self, monkeypatch):
-        # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
-        # build_attention_config(); clear it so these tests assert config defaults,
-        # not whatever the process inherited from CI / sibling tests.
-        monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
-
     def test_diffusion_attention_backend_sets_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
         assert isinstance(config.diffusion_attention_config, AttentionConfig)

tests/diffusion/attention/test_cudnn_attn.py

@@ -12,9 +12,10 @@ import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
+@pytest.mark.cpu
 def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
     """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
     selected_backends = []
@@ -40,6 +41,7 @@ def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
     assert selected_backends == [(SDPBackend.MATH,)]
 
 
+@pytest.mark.cpu
 def test_explicit_cudnn_rejects_kv_seq_len_one():
     impl = CuDNNAttentionImpl(
         num_heads=2,
@@ -54,6 +56,7 @@ def test_explicit_cudnn_rejects_kv_seq_len_one():
         impl.forward_cuda(query, singleton_kv, singleton_kv)
 
 
+@pytest.mark.cpu
 def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
     selected_backends = []
 
@@ -77,6 +80,7 @@ def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
     assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]
 
 
+@pytest.mark.cpu
 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
     observed: dict[str, Any] = {}
 
@@ -112,6 +116,7 @@ def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
     assert observed["kwargs"]["attn_mask"] is None
 
 
+@pytest.mark.cpu
 def test_cudnn_rejects_invalid_valid_kv_length():
     impl = CuDNNAttentionImpl(
         num_heads=2,
@@ -130,12 +135,97 @@ def test_cudnn_rejects_invalid_valid_kv_length():
 
 
 @pytest.mark.parametrize("head_size", [8, 64, 128, 256])
+@pytest.mark.cpu
 def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
     assert CuDNNAttentionBackend.supports_head_size(head_size)
     assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
 
 
 @pytest.mark.parametrize("head_size", [0, 7, 12, 320])
+@pytest.mark.cpu
 def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
     assert not CuDNNAttentionBackend.supports_head_size(head_size)
     assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("kv_heads", [1, 2])
+@pytest.mark.parametrize("query_offset", [0, 4])
+@pytest.mark.cpu
+def test_piecewise_matches_dense_block_causal_mask(monkeypatch, batch_size, kv_heads, query_offset):
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    torch.manual_seed(42)
+    query = torch.randn(batch_size, 12, 2, 8)
+    key = torch.randn(batch_size, 12, kv_heads, 8)
+    value = torch.randn_like(key)
+    spans = [(2, 5), (7, 10)]
+    positions = torch.arange(12)
+    mask = positions[:, None] >= positions[None, :]
+    for start, end in spans:
+        mask[start:end, start:end] = True
+    expected = torch.nn.functional.scaled_dot_product_attention(
+        query[:, query_offset:].transpose(1, 2),
+        key.transpose(1, 2),
+        value.transpose(1, 2),
+        attn_mask=mask[None, None, query_offset:],
+        scale=0.5,
+        enable_gqa=kv_heads == 1,
+    ).transpose(1, 2)
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    actual = impl.forward_cuda(
+        query[:, query_offset:], key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size)
+    )
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.cpu
+def test_piecewise_preserves_discontinuous_query_ranges(monkeypatch):
+    from vllm_omni.diffusion.attention.backends.abstract import QueryRange
+
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    query, key, value = [torch.randn(2, 12, 2, 8) for _ in range(3)]
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (7, 12)]] * 2)
+    expected = impl.forward_cuda(query, key, value, metadata)
+    metadata.query_ranges = (QueryRange(0, 3, 0), QueryRange(3, 6, 8))
+    actual = impl.forward_cuda(torch.cat([query[:, :3], query[:, 8:11]], dim=1), key, value, metadata)
+    torch.testing.assert_close(actual, torch.cat([expected[:, :3], expected[:, 8:11]], dim=1))
+
+
+@pytest.mark.cpu
+def test_piecewise_rejects_padding_mask(monkeypatch):
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    query = torch.randn(1, 8, 2, 8)
+    metadata = AttentionMetadata(full_attn_spans=[[(2, 8)]], attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    with pytest.raises(ValueError, match="padding mask"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
+@pytest.mark.cpu
+def test_dense_mask_takes_precedence_over_spans(monkeypatch):
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
+    query, key, value = [torch.randn(2, 8, 2, 8) for _ in range(3)]
+    mask = torch.ones(2, 1, 8, 8, dtype=torch.bool).tril()
+    mask[1, :, :, 3] = False
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask))
+    actual = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask, full_attn_spans=[[(2, 8)]] * 2))
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_cudnn_piecewise_matches_dense_on_device(batch_size):
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    query, key, value = [
+        torch.randn(batch_size, 96, 2, 128, generator=generator, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+    ]
+    spans = [(8, 40), (48, 96)]
+    mask = torch.ones(96, 96, device="cuda", dtype=torch.bool).tril()
+    for start, end in spans:
+        mask[start:end, start:end] = True
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=128, softmax_scale=128**-0.5)
+    actual = impl.forward_cuda(query, key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size))
+    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask[None, None]))
+    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.004)

tests/diffusion/cache/test_cache_dit.py

@@ -60,6 +60,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
         "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
         "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
         "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2,
+        "MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2,
     }
 
     with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
@@ -116,6 +117,33 @@ def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter,
     assert get_transformer(pipeline) is transformer_block
 
 
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+def test_mammothmoda2_cache_dit_targets_only_main_layers(mock_block_adapter, mock_enable_cache):
+    pipeline = Mock()
+    transformer = pipeline.gen_transformer
+    layers = torch.nn.ModuleList([torch.nn.Identity()])
+    transformer.layers = layers
+    adapter = mock_block_adapter.return_value
+    refresh = Mock()
+    mock_enable_cache.return_value = refresh
+
+    result = cd_model_specific.enable_cache_for_mammothmoda2(pipeline, SAMPLE_CACHE_CONFIG)
+
+    mock_block_adapter.assert_called_once()
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    assert adapter_kwargs["transformer"] is transformer
+    assert adapter_kwargs["blocks"] == [layers]
+    assert adapter_kwargs["has_separate_cfg"] is True
+    assert adapter_kwargs["check_forward_pattern"] is True
+    assert result.refresh is refresh
+    assert result.targets == (adapter,)
+    get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"]
+    assert get_transformer(pipeline) is transformer
+    # Sequential-CFG parity is mandatory while cache-dit is installed.
+    assert pipeline._cache_dit_requires_paired_cfg is True
+
+
 @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary")
 @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True)
 def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary):
@@ -338,3 +366,83 @@ def test_summary_with_no_transformer_is_nonfatal():
         pass
 
     cache_summary(pipeline=FakePipeline())
+
+
+# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM ROCm custom ops lack CPU fallback",
+)
+def test_mammothmoda2_cache_dit_runs_end_to_end_on_tiny_model(request: pytest.FixtureRequest):
+    """A tiny MammothModa2 DiT runs through Cache-DiT with paired CFG forwards.
+
+    Exercises the Pattern_3 contract end to end: after enabling the custom
+    enabler, blocks must be callable as ``block(hidden_states, **kwargs)``
+    and the pipeline's two-forwards-per-step CFG cadence must drive the
+    cache context without shape or parity errors.
+    """
+    from cache_dit import BlockAdapter
+
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
+    from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
+
+    # MammothModa2 attention resolves through the shared Omni attention layer,
+    # which picks the platform-default backend at construction time. Pin
+    # TORCH_SDPA so a CUDA-visible host does not select FA3 for CPU tensors.
+    with _force_torch_sdpa():
+        model = Transformer2DModel(
+            patch_size=2,
+            in_channels=4,
+            hidden_size=96,
+            num_layers=4,
+            num_refiner_layers=1,
+            num_attention_heads=2,
+            num_kv_heads=2,
+            multiple_of=8,
+            axes_dim_rope=(16, 16, 16),
+            axes_lens=(300, 512, 512),
+            text_feat_dim=16,
+        )
+    model.eval()
+
+    MammothModa2Pipeline = type("MammothModa2DiTPipeline", (), {})
+    pipeline = MammothModa2Pipeline()
+    pipeline.gen_transformer = model
+    backend = CacheDiTBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    request.addfinalizer(lambda: backend.disable(pipeline))
+    assert BlockAdapter.is_cached(model)
+
+    num_inference_steps = 8
+    backend.refresh(pipeline, num_inference_steps=num_inference_steps)
+
+    freqs_cis = RotaryPosEmbedReal.get_freqs_real((16, 16, 16), (300, 512, 512), theta=10000)
+    latents = torch.randn(1, 4, 8, 8)
+    text = torch.randn(1, 5, 16)
+    negative = torch.randn(1, 3, 16)
+    text_mask = torch.ones(1, 5, dtype=torch.bool)
+    negative_mask = torch.ones(1, 3, dtype=torch.bool)
+
+    with torch.no_grad():
+        for step in range(num_inference_steps):
+            timestep = torch.full((1,), 1.0 - step / num_inference_steps)
+            # Conditional pass ...
+            cond = model(
+                hidden_states=latents,
+                timestep=timestep,
+                text_hidden_states=text,
+                freqs_cis=freqs_cis,
+                text_attention_mask=text_mask,
+            )
+            # ... always followed by the unconditional pass (parity).
+            uncond = model(
+                hidden_states=latents,
+                timestep=timestep,
+                text_hidden_states=negative,
+                freqs_cis=freqs_cis,
+                text_attention_mask=negative_mask,
+            )
+            latents = uncond + 4.0 * (cond - uncond)
+
+    assert latents.shape == (1, 4, 8, 8)
+    assert torch.isfinite(latents).all()

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass
 from types import SimpleNamespace
 
@@ -62,6 +65,7 @@ class E2EOperator:
 @dataclass
 class FakeWorldGroup:
     device_group: object
+    cpu_group: object = None
 
 
 class DummyMixin(DistributedVaeMixin):

tests/diffusion/distributed/test_pause_keep_two_ranks.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Two-rank ``pause_generation(mode="keep")``: every rank runs the pause barrier
+before the ACK, and a barrier failure on a non-zero rank fails the pause
+without aborting or dropping the running batch.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from contextlib import ExitStack
+
+import pytest
+
+from tests.e2e.features.custom_pipeline.test_diffusion_pause_keep import (
+    CUSTOM_PIPELINE_CLASS,
+    MODEL,
+    WORKER_EXTENSION_CLASS,
+    _generate,
+    _pause_round,
+    _wait_for,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+
+def _all_true(results) -> bool:
+    flat = []
+    stack = list(results)
+    while stack:
+        item = stack.pop()
+        if isinstance(item, list):
+            stack.extend(item)
+        else:
+            flat.append(item)
+    return bool(flat) and all(item is True for item in flat)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_keep_pause_two_ranks_ack_and_rank_failure():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            num_gpus=2,
+            parallel_config={"ulysses_degree": 2, "ulysses_mode": "advanced_uaa"},
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+        assert _all_true(await engine.collective_rpc(method="start_counting_synchronize_device"))
+
+        pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(3)]
+        # The result is the all-rank AND, so True means both ranks ran the barrier.
+        assert _all_true(await engine.collective_rpc(method="synchronize_device_seen"))
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, True)))
+        a = asyncio.create_task(_generate(engine, "a-fail", events))
+        await _wait_for(lambda: "started" in events["a-fail"])
+        with pytest.raises(RuntimeError, match="pause_scheduler failed"):
+            await engine.pause_generation(mode="keep", clear_cache=False)
+        assert await engine.is_paused()
+        a_output = await asyncio.wait_for(a, 60.0)
+        assert a_output.images, "a failed barrier must not abort the running batch"
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, False)))
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        recovered = await asyncio.wait_for(_generate(engine, "after-recovery", events), 60.0)
+        assert recovered.images
+
+        assert not engine.request_states
+        print(f"[sp2] rounds=3 pause_latency_s min/max={min(pause_latencies):.3f}/{max(pause_latencies):.3f}")

tests/diffusion/distributed/test_vae_error_recovery.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as vae_executor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TileOperator:
+    def split(self, tensor):
+        return [vae_executor.TileTask(i, (i,), row[None]) for i, row in enumerate(tensor)], vae_executor.GridSpec(
+            split_dims=(0,), grid_shape=(tensor.shape[0],)
+        )
+
+    def exec(self, task):
+        return task.tensor * 2
+
+    def merge(self, tiles, grid):
+        return torch.cat([tiles[(i,)] for i in range(grid.grid_shape[0])])
+
+
+def run_recovery_worker(rank, init_file):
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo", init_method=f"file://{init_file}", rank=rank, world_size=2, timeout=timedelta(seconds=30)
+    )
+    try:
+        group = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(vae_executor, "get_world_group", lambda: group)
+            executor = vae_executor.DistributedVaeExecutor()
+        executor.set_parallel_size(2)
+        tensor = torch.arange(12, dtype=torch.float32).reshape(4, 3)
+        operator = TileOperator()
+        for stage, failing_rank in [("decode", 1), ("packing", 1), ("merge", 0), ("gather", 1)]:
+            with pytest.MonkeyPatch.context() as patch:
+                if rank == failing_rank:
+
+                    def fail(*args, **kwargs):
+                        raise torch.OutOfMemoryError("injected OOM")
+
+                    target, attribute = {
+                        "decode": (operator, "exec"),
+                        "packing": (executor, "_pack_local_tiles"),
+                        "merge": (operator, "merge"),
+                        "gather": (torch, "empty_like"),
+                    }[stage]
+                    patch.setattr(target, attribute, fail)
+                with pytest.raises(torch.OutOfMemoryError, match="Distributed VAE"):
+                    executor.execute(tensor, operator)
+            torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
+
+        with pytest.raises(RuntimeError, match="Distributed VAE mixed failures failed"):
+            with executor._sync_errors("mixed failures"):
+                if rank == 0:
+                    raise torch.OutOfMemoryError("retryable")
+                raise ValueError("not retryable")
+        torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2)
+    finally:
+        dist.destroy_process_group()
+
+
+def test_all_ranks_observe_failures_and_can_retry(tmp_path):
+    torch.multiprocessing.spawn(run_recovery_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)

tests/diffusion/model_loader/test_hub_prefetch.py

@@ -3,21 +3,21 @@
 
 import contextlib
 
-import huggingface_hub
 import pytest
 
+from tests.helpers.mock import patch_hf_snapshot_download
 from vllm_omni.diffusion.model_loader import hub_prefetch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
-def test_prefetch_subfolders_propagates_revision(monkeypatch):
+def test_prefetch_subfolders_propagates_revision(monkeypatch, tmp_path):
     calls = []
 
-    def fake_snapshot_download(self, **kwargs):
+    def fake_snapshot_download(**kwargs):
         calls.append(kwargs)
 
-    monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
     monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
 
     hub_prefetch.prefetch_subfolders(

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_it2i_multi_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
 """Multi-image input regression for HunyuanImage3 IT2I prompt construction.
 
 The official HunyuanImage-3.0-Instruct supports up to 3 reference images
@@ -41,7 +41,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class FakeTokenizer:
-    """Recording fake tokenizer mirroring the one in test_prompt_utils.
+    """Recording fake tokenizer mirroring the one in test_hunyuan_image3_prompt_utils.
 
     Special token ids: `<|startoftext|>`=1, `<img>`=2, `<think>`=3,
     `<recaption>`=4. encode() returns one id per character starting at

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_prompt_utils.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
 
 import ast

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -171,7 +171,7 @@ def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
         else pytest.fail("unexpected snapshot request")
     )
     monkeypatch.setattr(
-        "vllm.transformers_utils.repo_utils.hf_api",
+        "vllm_omni.transformers_utils.repo_utils.hf_api",
         lambda: api,
     )
     assert _resolve_checkpoint_root(

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -1,13 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from contextlib import contextmanager
 from dataclasses import dataclass
+from types import SimpleNamespace
 
+import cache_dit
 import pytest
 import torch
+from cache_dit import BlockAdapter
 from torch import nn
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, RequestScopedCacheDiTRuntime
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import (
+    AttentionConfig,
+    DiffusionCacheConfig,
+    DiffusionOutput,
+    OmniDiffusionConfig,
+    TransformerConfig,
+)
+from vllm_omni.diffusion.models.interface import adopt_request_scoped_cache_dit
 from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
 from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     MammothModa2DiTPipeline,
@@ -17,6 +31,7 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -110,6 +125,10 @@ def _pipeline_shell() -> MammothModa2DiTPipeline:
     pipeline.device = torch.device("cpu")
     pipeline.config = _build_mammoth_config(_od_config())
     pipeline._llm_hidden_size = 8
+    # Cache-DiT runtime state that __init__ normally creates (disabled here,
+    # matching a pipeline built without a cache_backend).
+    pipeline._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipeline)
+    pipeline._cache_dit_config = None
     return pipeline
 
 
@@ -118,12 +137,14 @@ def _batch(
     request_id: str = "req-7",
     prompt: object | None = None,
     sampling: OmniDiffusionSamplingParams | None = None,
+    height: int = 32,
+    width: int = 48,
 ) -> DiffusionRequestBatch:
     if prompt is None:
         prompt = {
             "prompt": "",
-            "height": 32,
-            "width": 48,
+            "height": height,
+            "width": width,
             "additional_information": {
                 "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8),
                 "full_token_ids": [10, 11, 100, 101],
@@ -390,3 +411,162 @@ def test_forward_rejects_missing_visual_tokens_before_model_access() -> None:
     }
     with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"):
         _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt))
+
+
+@contextmanager
+def _force_torch_sdpa():
+    """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
+    od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with set_current_diffusion_config(od_config):
+        yield
+
+
+def _constructible_raw_config() -> dict:
+    """Checkpoint config whose DiT/VAE are small enough to actually run on CPU."""
+    return {
+        "model_type": "mammothmoda2",
+        "llm_config": {
+            "model_type": "mammothmoda2_qwen2_5_vl",
+            "text_config": {
+                "model_type": "mammothmoda2_qwen2_5_vl_text",
+                "hidden_size": 8,
+                "gen_vocab_start_index": 100,
+            },
+        },
+        "gen_vae_config": {"block_out_channels": [8, 8], "norm_num_groups": 8},
+        "gen_dit_config": {
+            "patch_size": 2,
+            "in_channels": 4,
+            "hidden_size": 96,
+            "num_layers": 4,
+            "num_refiner_layers": 1,
+            "num_attention_heads": 2,
+            "num_kv_heads": 2,
+            "multiple_of": 8,
+            "axes_dim_rope": [16, 16, 16],
+            "axes_lens": [300, 512, 512],
+            "text_feat_dim": 16,
+        },
+        # Top-level axes feed the pipeline-owned rotary table (see pipeline __init__).
+        "gen_axes_dim_rope": [16, 16, 16],
+        "gen_axes_lens": [300, 512, 512],
+    }
+
+
+def _cache_dit_od_config() -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(_constructible_raw_config()),
+        cache_backend="cache_dit",
+        cache_config=DiffusionCacheConfig(),
+    )
+
+
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM ROCm custom ops lack CPU fallback",
+)
+def test_forward_transitions_request_scoped_cache_dit_across_requests(request: pytest.FixtureRequest) -> None:
+    """Cache-DiT request lifecycle on one production-built pipeline instance.
+
+    Mirrors the diffusion runner startup (``get_cache_backend`` → ``enable`` →
+    ``adopt_request_scoped_cache_dit``) on a tiny pipeline built through the
+    real ``__init__``, then drives three requests through the real ``forward()``:
+
+    1. a partial-CFG request installs the request-scoped cache context:
+       cond/uncond forwards stay paired on every denoise step (parity
+       accounting), even outside ``cfg_range``, and blocks record cached steps;
+    2. a guidance=1.0 request runs with hooks disabled (``prepare(None)``) and
+       its output is bit-identical to a reference captured before the backend
+       was ever enabled, proving a clean uninstall with no stale module state;
+    3. a CFG request with a different step count and image size re-installs the
+       hooks through a fresh context, so no stale shape or step count survives.
+    """
+    torch.manual_seed(0)  # Deterministic random init → deterministic cache hits.
+    with _force_torch_sdpa():
+        pipeline = MammothModa2DiTPipeline(od_config=_cache_dit_od_config())
+    pipeline.eval()
+    assert pipeline._cache_dit_config is not None
+    request.addfinalizer(pipeline._cache_dit_runtime.disable)
+
+    # The default prompt carries 2 text + 2 visual condition tokens, so the
+    # conditional pass sees 4 tokens and the uncond pass (empty prompt) sees 0.
+    cond_tokens, uncond_tokens = 4, 0
+    seq_lens: list[int] = []
+
+    def record_call(_module, _args, kwargs):
+        text = kwargs.get("text_hidden_states")
+        seq_lens.append(int(text.shape[1]) if text is not None else -1)
+
+    hook = pipeline.gen_transformer.register_forward_pre_hook(record_call, with_kwargs=True)
+
+    unguided_sampling = OmniDiffusionSamplingParams(
+        height=32, width=48, seed=4321, guidance_scale=1.0, num_inference_steps=8
+    )
+    partial_cfg_sampling = OmniDiffusionSamplingParams(
+        height=32,
+        width=48,
+        seed=1234,
+        guidance_scale=4.0,
+        num_inference_steps=8,
+        extra_args={"cfg_range": [0.2, 0.8]},
+    )
+    changed_cfg_sampling = OmniDiffusionSamplingParams(
+        height=48,
+        width=64,
+        seed=777,
+        guidance_scale=6.0,
+        num_inference_steps=10,
+        extra_args={"cfg_range": [0.1, 0.9]},
+    )
+
+    # Reference run for phase 2: identical request while hooks were never
+    # installed on this pipeline.
+    assert not pipeline.is_cache_dit_enabled()
+    reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    assert seq_lens == [cond_tokens] * 8
+    assert torch.isfinite(reference.float()).all()
+
+    # Runner startup: enable once, then the pipeline owns all transitions.
+    backend = get_cache_backend("cache_dit", DiffusionCacheConfig())
+    assert isinstance(backend, CacheDiTBackend)
+    backend.enable(pipeline)
+    assert adopt_request_scoped_cache_dit(pipeline, backend)
+    assert pipeline.is_cache_dit_enabled()
+
+    # Phase 1: partial CFG keeps cond/uncond paired on every step, including
+    # the steps whose normalized index falls outside cfg_range [0.2, 0.8].
+    seq_lens.clear()
+    pipeline.forward(_batch(sampling=partial_cfg_sampling))
+    assert seq_lens == [cond_tokens, uncond_tokens] * 8
+    assert pipeline.is_cache_dit_enabled()
+    assert BlockAdapter.is_cached(pipeline.gen_transformer)
+    hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
+    assert hits.cached_steps or hits.cfg_cached_steps
+
+    # Phase 2: no-CFG request disables the hooks and must reproduce the
+    # never-enabled reference bit-for-bit.
+    seq_lens.clear()
+    unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    assert seq_lens == [cond_tokens] * 8
+    assert not pipeline.is_cache_dit_enabled()
+    assert not BlockAdapter.is_cached(pipeline.gen_transformer)
+    assert torch.equal(unguided, reference)
+
+    # Phase 3: CFG with a new step count and image size re-installs the hooks
+    # on a fresh context; output shape follows the new request geometry.
+    seq_lens.clear()
+    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
+    assert seq_lens == [cond_tokens, uncond_tokens] * 10
+    assert pipeline.is_cache_dit_enabled()
+    assert BlockAdapter.is_cached(pipeline.gen_transformer)
+    hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
+    assert hits.cached_steps or hits.cfg_cached_steps
+    assert regenerated.shape == (1, 3, 12, 16)
+    assert torch.isfinite(regenerated.float()).all()
+
+    hook.remove()

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -16,6 +16,7 @@ from tests.diffusion.quantization.test_quantization_quality import (
     _maybe_save_output,
 )
 from tests.helpers.mark import hardware_marks
+from tests.helpers.mock import patch_hf_snapshot_download
 
 _MINIMAX_H3_REPO = "MiniMaxAI/MiniMax-H3"
 _MINIMAX_H3_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc"
@@ -258,13 +259,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
     fl2va_root.mkdir()
     (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
 
-    def fake_snapshot_download(self, *, repo_id, revision, allow_patterns):
+    def fake_snapshot_download(*, repo_id, revision, allow_patterns):
         assert repo_id == _MINIMAX_H3_REPO
         assert revision == _MINIMAX_H3_REVISION
         assert allow_patterns == ["FL2VA/**"]
         return str(tmp_path)
 
-    monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
     assert _resolve_fl2va_model_ref() == str(fl2va_root)
 
 

tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import MethodType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.nextstep_1_1 import pipeline_nextstep_1_1 as pipeline_module
+from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import NextStepModel
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("height,width", [(32, 32), (32, 48), (48, 32)])
+@pytest.mark.parametrize("batch", [1, 2])
+def test_forward_preserves_latent_grid(monkeypatch, height, width, batch):
+    pipeline = object.__new__(pipeline_module.NextStep11Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._NextStep11Pipeline__device = torch.device("cpu")
+    pipeline.config = SimpleNamespace(use_gen_pos_embed=False)
+    pipeline.down_factor = 16
+    pipeline.scaling_factor = 2.0
+    pipeline.shift_factor = 0.25
+    grid = SimpleNamespace(config=SimpleNamespace(latent_channels=2, latent_patch_size=2))
+    expected = torch.arange(batch * 2 * (height // 8) * (width // 8), dtype=torch.float32)
+    expected = expected.reshape(batch, 2, height // 8, width // 8)
+    tokens = NextStepModel.patchify(grid, expected)
+    pipeline.model = Mock(
+        return_value=SimpleNamespace(last_hidden_state=torch.zeros(batch, 1, 4), past_key_values=None)
+    )
+    pipeline.model.forward_model.return_value = pipeline.model.return_value
+    pipeline.model.unpatchify = MethodType(NextStepModel.unpatchify, grid)
+    pipeline.tokenizer = Mock(
+        bos_token=None,
+        return_value=SimpleNamespace(
+            input_ids=torch.ones(batch, 2, dtype=torch.long), attention_mask=torch.ones(batch, 2, dtype=torch.long)
+        ),
+    )
+    pipeline._check_input = Mock(return_value=(["teapot"], None))
+    pipeline._build_captions = Mock(return_value=(["teapot"] * batch, None, 1, 1.0))
+    pipeline._add_prefix_ids = lambda _hw, ids, mask: (ids, mask)
+    pipeline.decoding = Mock(return_value=tokens)
+    pipeline.vae = SimpleNamespace(dtype=torch.float32, decode=lambda latent: SimpleNamespace(sample=latent))
+    monkeypatch.setattr(pipeline_module, "StaticCache", lambda **_kwargs: None)
+    params = OmniDiffusionSamplingParams(
+        height=height, width=width, guidance_scale=1.0, num_outputs_per_prompt=batch, seed=None
+    )
+    output = pipeline.forward(SimpleNamespace(prompts=["teapot"], sampling_params=params))
+    torch.testing.assert_close(output.output, expected / 2 + 0.25, rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_batch_compatibility.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from types import SimpleNamespace
+
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+    get_qwen_image_21_pre_process_func,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler, StepScheduler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def preprocess(tmp_path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
+    return get_qwen_image_21_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+
+
+def make_request(request_id, sizes, color="red"):
+    return OmniDiffusionRequest(
+        request_id=request_id,
+        prompt={
+            "prompt": "A ceramic teapot",
+            "multi_modal_data": {"image": [Image.new("RGB", size, color) for size in sizes]},
+        },
+        sampling_params=OmniDiffusionSamplingParams(
+            height=1024, width=1024, num_inference_steps=2, true_cfg_scale=1.0, seed=42
+        ),
+    )
+
+
+@pytest.mark.parametrize("scheduler_class", [RequestScheduler, StepScheduler])
+@pytest.mark.parametrize(
+    "first_sizes,second_sizes,compatible",
+    [
+        pytest.param([], [(64, 64)], False, id="text-and-edit"),
+        pytest.param([(64, 64)], [], False, id="edit-and-text"),
+        pytest.param([(64, 64)], [(64, 64), (64, 64)], False, id="image-count"),
+        pytest.param([(64, 64)], [(128, 64)], False, id="image-aspect-ratio"),
+        pytest.param([(64, 64), (128, 64)], [(128, 64), (64, 64)], False, id="slot-order"),
+        pytest.param([], [], True, id="two-text-requests"),
+        pytest.param([(64, 64)], [(64, 64)], True, id="different-image-content"),
+        pytest.param([(64, 64)], [(128, 128)], True, id="same-resized-layout"),
+        pytest.param([(64, 64), (128, 64)], [(128, 128), (256, 128)], True, id="same-multi-image-layout"),
+    ],
+)
+def test_scheduler_groups_by_condition_layout(preprocess, scheduler_class, first_sizes, second_sizes, compatible):
+    first = preprocess(make_request("first", first_sizes))
+    second = preprocess(make_request("second", second_sizes, color="blue"))
+    scheduler = scheduler_class()
+    scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+    scheduler.add_request(first)
+    scheduler.add_request(second)
+
+    scheduled = scheduler.schedule()
+    expected_ids = ["first", "second"] if compatible else ["first"]
+    assert scheduled.scheduled_request_ids == expected_ids
+
+    scheduler.finish_requests(expected_ids, DiffusionRequestStatus.FINISHED_COMPLETED)
+    if not compatible:
+        assert scheduler.schedule().scheduled_request_ids == ["second"]
+        scheduler.finish_requests("second", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert not scheduler.has_requests()
+
+
+@pytest.mark.parametrize(
+    "prompt", ["A teapot", {"prompt": "A teapot"}, {"prompt": "A teapot", "multi_modal_data": {"image": None}}]
+)
+def test_text_input_forms_share_compatibility_key(preprocess, prompt):
+    request = make_request("text", [])
+    request.prompt = prompt
+    processed = preprocess(request)
+    empty_list_request = preprocess(make_request("empty-list", []))
+    assert processed.batch_compatibility_key is not None
+    assert processed.batch_compatibility_key == empty_list_request.batch_compatibility_key
+
+
+@pytest.mark.parametrize("image_type", ["numpy", "tensor"])
+def test_preprocess_accepts_array_images_without_truth_testing(preprocess, image_type):
+    import numpy as np
+    import torch
+
+    pixels = np.full((32, 32, 4), 255, dtype=np.uint8)
+    image = torch.from_numpy(pixels).permute(2, 0, 1) if image_type == "tensor" else pixels
+    request = make_request("array", [])
+    request.prompt["multi_modal_data"]["image"] = image
+    processed = preprocess(request)
+    assert processed.batch_compatibility_key is not None
+    assert processed.allow_mixed_step_phases is False
+
+
+def test_late_prefill_waits_until_all_decode_requests_finish(preprocess):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    scheduler = StepScheduler()
+    scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+    for request_id in ("a", "b"):
+        scheduler.add_request(preprocess(make_request(request_id, [])))
+    wave = scheduler.schedule()
+    assert wave.scheduled_request_ids == ["a", "b"]
+    scheduler.update_from_output(
+        wave,
+        BatchRunnerOutput.from_list(
+            [RunnerOutput(request_id=request_id, step_index=1, finished=False) for request_id in ("a", "b")]
+        ),
+    )
+    scheduler.add_request(preprocess(make_request("late", [])))
+    assert scheduler.schedule().scheduled_request_ids == ["a", "b"]
+    scheduler.finish_requests("a", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert scheduler.schedule().scheduled_request_ids == ["b"]
+    scheduler.finish_requests("b", DiffusionRequestStatus.FINISHED_COMPLETED)
+    assert scheduler.schedule().scheduled_request_ids == ["late"]
+
+
+def test_preprocess_normalizes_tuple_images_like_a_list(preprocess):
+    image = Image.new("RGB", (64, 64), "red")
+    tuple_request = make_request("tuple", [])
+    tuple_request.prompt["multi_modal_data"]["image"] = (image,)
+
+    processed = preprocess(tuple_request)
+    reference = preprocess(make_request("list", [(64, 64)]))
+
+    assert processed.batch_compatibility_key == reference.batch_compatibility_key
+    assert processed.batch_compatibility_key[0] == "qwen_image_21"
+    assert len(processed.batch_compatibility_key[1]) == 1
+    info = processed.prompt["additional_information"]
+    assert len(info["prompt_image"]) == 1
+    assert len(info["vae_images"]) == 1
+    assert len(info["input_image_sizes"]) == 1
+
+
+def test_preprocess_rejects_more_images_than_the_model_supports(preprocess):
+    from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+        MAX_QWEN_IMAGE_21_INPUT_IMAGES,
+    )
+
+    request = make_request("too-many", [(64, 64)] * (MAX_QWEN_IMAGE_21_INPUT_IMAGES + 1))
+    with pytest.raises(ValueError, match="At most"):
+        preprocess(request)
+
+
+def test_preprocess_defaults_true_cfg_scale(preprocess):
+    request = make_request("cfg", [])
+    request.sampling_params.true_cfg_scale = None
+    processed = preprocess(request)
+    assert processed.sampling_params.true_cfg_scale == 1.0
+
+    request = make_request("cfg-set", [])
+    request.sampling_params.true_cfg_scale = 4.0
+    assert preprocess(request).sampling_params.true_cfg_scale == 4.0
+
+
+@pytest.mark.parametrize("size,expected_wider", [((128, 64), True), ((64, 128), False)])
+def test_preprocess_derives_output_size_from_condition_aspect(preprocess, size, expected_wider):
+    request = make_request("aspect", [size])
+    request.sampling_params.height = None
+    request.sampling_params.width = None
+
+    processed = preprocess(request)
+
+    height = processed.sampling_params.height
+    width = processed.sampling_params.width
+    assert height % 32 == 0 and width % 32 == 0
+    assert (width > height) is expected_wider
+    info = processed.prompt["additional_information"]
+    assert info["calculated_height"] == height
+    assert info["calculated_width"] == width
+
+
+def test_preprocess_stages_rgba_vae_images_and_resized_prompt_images(preprocess):
+    processed = preprocess(make_request("staged", [(96, 64)]))
+    info = processed.prompt["additional_information"]
+
+    (vae_image,) = info["vae_images"]
+    # RGB input gains an opaque alpha channel for the RGBA VAE; the frame dim is
+    # inserted for the 3D (video-capable) VAE.
+    assert vae_image.shape[1] == 4
+    assert vae_image.shape[2] == 1
+    (prompt_image,) = info["prompt_image"]
+    assert isinstance(prompt_image, Image.Image)
+    (width, height) = info["input_image_sizes"][0]
+    assert prompt_image.size == (width, height)
+    assert processed.batch_compatibility_key == ("qwen_image_21", ((width, height),))

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -0,0 +1,318 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import weakref
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
+from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
+from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class DecodeModel(torch.nn.Module):
+    in_channels = 4
+
+    def _decode_graph_forward(self, entry):
+        return (
+            entry.hidden.sin()
+            + entry.k[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.v[0].mean(dim=(1, 2, 3))[:, None, None]
+            + entry.freqs[None, :, :]
+        )
+
+
+@pytest.fixture
+def graph_case(request, monkeypatch):
+    device = "cuda" if request.node.get_closest_marker("cuda") else "cpu"
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2)
+    monkeypatch.setattr(manager, "eligible", lambda: True)
+    monkeypatch.setattr(manager, "_backend_name", lambda: "test")
+    cleanup = []
+    monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True))
+    return manager, device, cleanup
+
+
+def make_cache(device, value):
+    return [
+        {
+            "cond": {
+                "key": torch.full((1, 2, 1, 4), value, device=device),
+                "value": torch.full((1, 2, 1, 4), value + 1, device=device),
+            }
+        }
+    ]
+
+
+def register(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    manager.register_prefill(
+        kv_cache=cache,
+        cache_branch="cond",
+        prefix_len=2,
+        img_shapes=[layout],
+        target_freqs=torch.full((4, 4), float(layout[1]), device=device),
+        joint_key_valid=None,
+        dtype=torch.float32,
+    )
+
+
+def decode(manager, cache, layout=(1, 2, 2)):
+    device = cache[0]["cond"]["key"].device
+    return manager.try_decode(
+        hidden_states=torch.ones(1, 4, 4, device=device),
+        timestep=torch.ones(1, device=device),
+        kv_cache=cache,
+        cache_branch="cond",
+        img_shapes=[layout],
+        img_mask=torch.tensor([[False, False, True]], device=device),
+        encoder_hidden_states_mask=None,
+    )
+
+
+@pytest.mark.cpu
+def test_model_level_offload_without_staging_disables_decode_graphs():
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True)
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_top_level_offload_hook_without_staging_disables_decode_graphs():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model)
+    registry = HookRegistry.get_or_create(model)
+    # A present-but-empty registry (e.g. after hooks were removed) must not disqualify.
+    assert manager._offload_reason() is None
+    registry.register_hook("sequential_offload", ModelHook())
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_persistent_staging_makes_model_level_offload_graph_eligible():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True)
+    hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu"))
+    HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook)
+    # No staging storage yet: a swap would allocate fresh storage per generation.
+    assert manager._offload_reason() is not None
+    # Staging established: fixed device storage keeps captured pointers valid.
+    hook._stager = object()
+    assert manager._offload_reason() is None
+    # A second, non-offload hook on the top-level module still disqualifies.
+    model._hook_registry.register_hook("teacache", ModelHook())
+    assert manager._offload_reason() is not None
+
+
+@pytest.mark.cpu
+def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+    allocated = []
+
+    def fail_allocation(**kwargs):
+        entry = constructor(**kwargs)
+        allocated.append(weakref.ref(entry.k[0]))
+        raise torch.OutOfMemoryError("injected allocation OOM")
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation)
+    register(manager, cache)
+    register(manager, cache)
+    assert len(allocated) == 1
+    assert allocated[0]() is None
+    assert list(manager.entries.values()) == [None]
+    assert len(cleanup) == 1
+    assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4))
+
+
+@pytest.mark.cpu
+def test_evict_before_allocating_replacement(graph_case, monkeypatch):
+    manager, device, _ = graph_case
+    manager.max_entries = 1
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0])
+    constructor = decode_graph.QwenImage21DecodeGraphEntry
+
+    def allocate_after_release(**kwargs):
+        assert old_buffer() is None
+        return constructor(**kwargs)
+
+    monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release)
+    register(manager, cache, (1, 1, 4))
+    assert len(manager.entries) == 1
+
+
+@pytest.mark.cpu
+def test_registration_preserves_request_cache_and_separates_layouts(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    original = first[0]["cond"]["key"]
+    register(manager, first)
+    register(manager, second)
+    assert first[0]["cond"]["key"] is original
+    assert torch.equal(original, torch.ones_like(original))
+    register(manager, second, (1, 1, 4))
+    assert len(manager.entries) == 2
+    entries = list(manager.entries.values())
+    assert not torch.equal(entries[0].freqs, entries[1].freqs)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+@pytest.mark.parametrize("failure_phase", ["warmup", "capture"])
+def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase):
+    manager, device, cleanup = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0])
+    del entry
+
+    original = manager.model._decode_graph_forward
+
+    def fail_capture(entry):
+        output = original(entry)
+        if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing():
+            raise torch.OutOfMemoryError("injected capture OOM")
+        return output
+
+    monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture)
+    assert decode(manager, cache) is None
+    assert entry_ref() is None and buffer_ref() is None
+    assert list(manager.entries.values()) == [None]
+    register(manager, cache)
+    assert decode(manager, cache) is None
+    assert len(cleanup) == 1
+    assert bool((cache[0]["cond"]["key"] == 1).all())
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_replay_uses_active_request_and_layout(graph_case):
+    manager, device, _ = graph_case
+    first, second = make_cache(device, 1.0), make_cache(device, 5.0)
+    first_key = weakref.ref(first[0]["cond"]["key"])
+    with torch.inference_mode():
+        register(manager, first)
+        # A second live request aliasing the same key is refused: it keeps its
+        # own cache and decodes eagerly while the owner is in flight.
+        register(manager, second)
+        assert decode(manager, second) is None
+        output = decode(manager, first)
+        expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
+        torch.testing.assert_close(output, expected)
+        # Once the owner's cache is released, the key is free to adopt.
+        del first
+        assert first_key() is None
+        register(manager, second)
+        output = decode(manager, second)
+        expected = (
+            torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2
+        )
+        torch.testing.assert_close(output, expected)
+        original = decode(manager, second).clone()
+        register(manager, second, (1, 1, 4))
+        other = decode(manager, second, (1, 1, 4))
+        torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1)
+        torch.testing.assert_close(decode(manager, second), original)
+        assert all(entry.captures == 1 for entry in manager.entries.values())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("inference", [False, True])
+def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference):
+    manager, device, _ = graph_case
+    with torch.inference_mode(inference):
+        cache = make_cache(device, 1.0) + make_cache(device, 2.0)
+        register(manager, cache)
+        entry = next(iter(manager.entries.values()))
+        entry.refresh_prefix(cache)
+        versions = [tensor._version for tensor in entry.k + entry.v]
+        entry.refresh_prefix(cache)
+        assert versions == [tensor._version for tensor in entry.k + entry.v]
+        cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0)
+        entry.refresh_prefix(cache)
+        assert bool((entry.v[1] == 7).all())
+        source = weakref.ref(cache[0]["cond"]["key"])
+        del cache
+        assert source() is None
+
+
+@pytest.mark.cpu
+def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    entry.refresh_prefix(cache)
+    cache[0]["cond"]["value"].add_(3)
+    entry.refresh_prefix(cache)
+    assert bool((entry.v[0] == 5).all())
+    register(manager, cache)
+    assert not entry.prefix_sources
+
+
+@pytest.mark.cpu
+def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case):
+    manager, device, _ = graph_case
+    cache = make_cache(device, 1.0)
+    register(manager, cache)
+    entry = next(iter(manager.entries.values()))
+    for valid in (True, False):
+        manager.register_prefill(
+            kv_cache=cache,
+            cache_branch="cond",
+            prefix_len=2,
+            img_shapes=[(1, 2, 2)],
+            target_freqs=torch.ones(4, 4),
+            joint_key_valid=torch.full((1, 6), valid),
+            dtype=torch.float32,
+        )
+        assert list(manager.entries.values()) == [entry]
+        assert entry.attn_metadata is None
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case):
+    manager, device, _ = graph_case
+    first = make_cache(device, 1.0)
+    second = make_cache(device, 4.0)
+    merged = [
+        {"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}}
+    ]
+    with torch.inference_mode():
+        register(manager, first)
+        single = decode(manager, first).clone()
+        register(manager, merged)
+        kwargs = dict(
+            hidden_states=torch.ones(2, 4, 4, device=device),
+            timestep=torch.ones(2, device=device),
+            kv_cache=merged,
+            cache_branch="cond",
+            img_shapes=[(1, 2, 2)],
+            img_mask=torch.tensor([[False, False, True]] * 2, device=device),
+            encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device),
+        )
+        result = manager.try_decode(**kwargs)
+        assert result is not None and result.shape == (2, 4, 4)
+        expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None]
+        torch.testing.assert_close(result, expected)
+        # A rebuilt cache (e.g. reordered rows) is a different live cache
+        # aliasing the same key: decode refuses the shared entry until the
+        # previous owner's tensors are released.
+        flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}]
+        kwargs["kv_cache"] = flipped
+        assert manager.try_decode(**kwargs) is None
+        del merged
+        register(manager, flipped)
+        torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0))
+        kwargs["encoder_hidden_states_mask"][1, 0] = False
+        assert manager.try_decode(**kwargs) is None
+        torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_pipeline.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU unit tests for QwenImage21Pipeline helpers: prepare_latents duplication
+order, KV-cache assembly across step phases, and the post-process func."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+    QwenImage21Pipeline,
+    get_qwen_image_21_post_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _toy_pipeline() -> QwenImage21Pipeline:
+    """A pipeline shell with the attributes prepare_latents reads, no weights."""
+    pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.vae_scale_factor = 16
+    pipeline.latent_channels = 4
+    return pipeline
+
+
+# ---------------------------------------------------------------------------
+# prepare_latents
+# ---------------------------------------------------------------------------
+
+
+def test_prepare_latents_duplicates_condition_latents_request_major():
+    """batch=2 prompts x n=2 outputs must order condition latents A,A,B,B.
+
+    ``encode_prompt`` expands prompt embeddings request-major (A,A,B,B); the
+    condition latents must follow the same order or generations 2 and 3 swap
+    condition images. ``repeat`` would produce the wrong (A,B,A,B) tile-major
+    order; ``repeat_interleave`` is correct.
+    """
+    pipeline = _toy_pipeline()
+    cond_a = torch.full((1, 4, 1, 8, 8), 1.0)
+    cond_b = torch.full((1, 4, 1, 8, 8), 2.0)
+    stacked = torch.cat([cond_a, cond_b], dim=0)  # one row per request
+
+    latents, image_latents = pipeline.prepare_latents(
+        [stacked],
+        batch_size=4,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    assert latents.shape == (4, 16, 4)  # packed: 4x4 latent tokens, 4 channels
+    assert image_latents.shape == (4, 64, 4)  # packed: 8x8 condition tokens
+    row_means = image_latents.mean(dim=(1, 2))
+    torch.testing.assert_close(row_means, torch.tensor([1.0, 1.0, 2.0, 2.0]))
+
+
+def test_prepare_latents_without_duplication_when_rows_match_batch():
+    pipeline = _toy_pipeline()
+    stacked = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 4, 1, 2, 2).contiguous()
+
+    _, image_latents = pipeline.prepare_latents(
+        [stacked],
+        batch_size=4,
+        num_channels_latents=4,
+        height=32,
+        width=32,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    torch.testing.assert_close(image_latents.mean(dim=(1, 2)), torch.arange(4, dtype=torch.float32))
+
+
+def test_prepare_latents_rejects_non_divisible_duplication():
+    pipeline = _toy_pipeline()
+    stacked = torch.zeros(2, 4, 1, 8, 8)
+
+    with pytest.raises(ValueError, match="Cannot duplicate"):
+        pipeline.prepare_latents(
+            [stacked],
+            batch_size=3,
+            num_channels_latents=4,
+            height=64,
+            width=64,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            generator=torch.Generator().manual_seed(0),
+        )
+
+
+def test_prepare_latents_rejects_generator_list_length_mismatch():
+    pipeline = _toy_pipeline()
+
+    with pytest.raises(ValueError, match="list of generators"):
+        pipeline.prepare_latents(
+            None,
+            batch_size=2,
+            num_channels_latents=4,
+            height=32,
+            width=32,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            generator=[torch.Generator().manual_seed(0)],
+        )
+
+
+def test_prepare_latents_concatenates_multiple_condition_images():
+    pipeline = _toy_pipeline()
+    first = torch.full((2, 4, 1, 4, 4), 1.0)
+    second = torch.full((2, 4, 1, 2, 2), 3.0)
+
+    _, image_latents = pipeline.prepare_latents(
+        [first, second],
+        batch_size=2,
+        num_channels_latents=4,
+        height=32,
+        width=32,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(0),
+    )
+
+    # Token dim concatenates the packed condition images: 16 + 4 per row.
+    assert image_latents.shape == (2, 20, 4)
+    torch.testing.assert_close(image_latents[:, :16].mean(), torch.tensor(1.0))
+    torch.testing.assert_close(image_latents[:, 16:].mean(), torch.tensor(3.0))
+
+
+# ---------------------------------------------------------------------------
+# KV-cache assembly across step phases
+# ---------------------------------------------------------------------------
+
+
+def _state(cache) -> SimpleNamespace:
+    return SimpleNamespace(extra={"kv_cache": cache})
+
+
+def _decode_cache(rows: int, value: float) -> list[dict]:
+    return [
+        {
+            "cond": {
+                "key": torch.full((rows, 2), value),
+                "value": torch.full((rows, 2), value + 0.5),
+            }
+        }
+    ]
+
+
+def test_kv_cache_phase_classification():
+    assert QwenImage21Pipeline._kv_cache_phase(None) == "none"
+    assert QwenImage21Pipeline._kv_cache_phase([{}]) == "prefill"
+    assert QwenImage21Pipeline._kv_cache_phase(_decode_cache(1, 0.0)) == "decode"
+
+
+def test_assemble_kv_cache_all_none():
+    assert QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(None)]) == (None, False)
+
+
+def test_assemble_kv_cache_prefill_returns_fresh_owned_cache():
+    states = [_state([{} for _ in range(3)]), _state([{} for _ in range(3)])]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
+
+    assert take_ownership is True
+    assert kv_cache == [{}, {}, {}]
+    # Fresh blocks, not aliases of either request's cache: the batched prefill
+    # writes into these and _scatter_kv_cache splits them back per request.
+    for block in kv_cache:
+        assert all(block is not request_block for state in states for request_block in state.extra["kv_cache"])
+
+
+def test_assemble_kv_cache_decode_concatenates_in_request_order():
+    first, second = _decode_cache(2, 1.0), _decode_cache(1, 3.0)
+    states = [_state(first), _state(second)]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states)
+
+    assert take_ownership is False
+    assert len(kv_cache) == 1
+    cond = kv_cache[0]["cond"]
+    torch.testing.assert_close(cond["key"], torch.tensor([[1.0, 1.0], [1.0, 1.0], [3.0, 3.0]]))
+    torch.testing.assert_close(cond["value"], torch.tensor([[1.5, 1.5], [1.5, 1.5], [3.5, 3.5]]))
+    # Decode only reads the cache: per-request tensors must stay untouched.
+    assert first[0]["cond"]["key"].shape == (2, 2)
+    assert second[0]["cond"]["key"].shape == (1, 2)
+
+
+def test_assemble_kv_cache_decode_merges_every_part():
+    first = [{"cond": {"key": torch.ones(1, 1), "value": torch.ones(1, 1), "key_scale": torch.ones(1)}}]
+    second = [{"cond": {"key": torch.zeros(1, 1), "value": torch.zeros(1, 1), "key_scale": torch.zeros(1)}}]
+
+    kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache([_state(first), _state(second)])
+
+    assert take_ownership is False
+    assert set(kv_cache[0]["cond"]) == {"key", "value", "key_scale"}
+    torch.testing.assert_close(kv_cache[0]["cond"]["key_scale"], torch.tensor([1.0, 0.0]))
+
+
+def test_assemble_kv_cache_mixed_phases_raises():
+    """Defensive backstop: a late prefill must never join a decoding batch.
+
+    The preprocessor sets ``allow_mixed_step_phases = False`` so the
+    StepScheduler defers late requests until the batch finishes decoding; this
+    raise is the last line of defense if a mixed batch is ever assembled.
+    """
+    states = [_state([{}]), _state(_decode_cache(1, 0.0))]
+
+    with pytest.raises(ValueError, match="mixed KV-cache phases"):
+        QwenImage21Pipeline._assemble_kv_cache(states)
+
+    with pytest.raises(ValueError, match="mixed KV-cache phases"):
+        QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(_decode_cache(1, 0.0))])
+
+
+def test_scatter_kv_cache_splits_rows_back_per_request():
+    states = [_state([{}]), _state([{}])]
+    merged = [
+        {
+            "cond": {
+                "key": torch.tensor([[1.0], [2.0], [3.0]]),
+                "value": torch.tensor([[4.0], [5.0], [6.0]]),
+            }
+        }
+    ]
+
+    QwenImage21Pipeline._scatter_kv_cache(states, [1, 2], merged)
+
+    torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[1.0]]))
+    torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["value"], torch.tensor([[4.0]]))
+    torch.testing.assert_close(states[1].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[2.0], [3.0]]))
+    assert states[1].extra["kv_cache"][0]["cond"]["key"].is_contiguous()
+
+
+# ---------------------------------------------------------------------------
+# post-process func
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def postprocess(tmp_path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
+    return get_qwen_image_21_post_process_func(SimpleNamespace(model=str(tmp_path)))
+
+
+def test_postprocess_tensor_returns_pil_images(postprocess):
+    from PIL import Image
+
+    images = torch.rand(2, 4, 16, 16) * 2 - 1
+
+    result = postprocess(images)
+
+    assert len(result) == 2
+    assert all(isinstance(image, Image.Image) for image in result)
+
+
+def test_postprocess_envelope_unwraps_payload_and_keeps_metadata(postprocess):
+    from PIL import Image
+
+    envelope = {
+        "payload": {"image": torch.rand(1, 4, 16, 16), "seed": 42},
+        "metadata": {"request_id": "abc"},
+    }
+
+    result = postprocess(envelope)
+
+    assert result["metadata"] == {"request_id": "abc"}
+    assert result["payload"]["seed"] == 42
+    assert len(result["payload"]["image"]) == 1
+    assert isinstance(result["payload"]["image"][0], Image.Image)
+
+
+def test_postprocess_envelope_requires_image_payload(postprocess):
+    with pytest.raises(ValueError, match=r"payload\['image'\]"):
+        postprocess({"payload": {"seed": 42}, "metadata": {}})
+
+
+def test_postprocess_envelope_normalizes_non_dict_metadata(postprocess):
+    result = postprocess({"payload": {"image": torch.rand(1, 4, 16, 16)}, "metadata": "bogus"})
+
+    assert result["metadata"] == {}

tests/diffusion/models/qwen_image_21/test_registry.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""L1 tests for checkpoint-to-pipeline routing of Qwen-Image 2.1."""
+
+import json
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_registry_entry_loads_the_pipeline_class():
+    from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    assert DiffusionModelRegistry._try_load_model_cls("QwenImage21Pipeline") is QwenImage21Pipeline
+
+
+def test_pipeline_declares_request_batching_and_step_execution():
+    from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline
+
+    assert QwenImage21Pipeline.supports_request_batch
+    assert QwenImage21Pipeline.supports_step_execution
+
+
+def test_process_funcs_resolve_to_the_pipeline_module():
+    from vllm_omni.diffusion.models.qwen_image_21 import pipeline_qwen_image_21
+    from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
+
+    pre_name = _DIFFUSION_PRE_PROCESS_FUNCS["QwenImage21Pipeline"]
+    post_name = _DIFFUSION_POST_PROCESS_FUNCS["QwenImage21Pipeline"]
+    assert pre_name == "get_qwen_image_21_pre_process_func"
+    assert post_name == "get_qwen_image_21_post_process_func"
+    assert callable(getattr(pipeline_qwen_image_21, pre_name))
+    assert callable(getattr(pipeline_qwen_image_21, post_name))
+
+
+def test_multimodal_metadata_caps_condition_images_at_four():
+    from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES, get_diffusion_model_metadata
+
+    metadata = get_diffusion_model_metadata("QwenImage21Pipeline")
+
+    assert metadata.supports_multimodal_inputs
+    assert metadata.max_multimodal_image_inputs == 4
+    assert metadata.max_multimodal_image_inputs == QWEN_IMAGE_21_MAX_INPUT_IMAGES
+
+
+def test_model_index_resolves_and_enriches(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"_class_name": "QwenImage21Pipeline"}),
+        encoding="utf-8",
+    )
+
+    assert resolve_model_class_name(str(tmp_path)) == "QwenImage21Pipeline"
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+
+    assert config.model_class_name == "QwenImage21Pipeline"
+    assert config.supports_multimodal_inputs
+    assert config.max_multimodal_image_inputs == 4

tests/diffusion/models/qwen_image_21/test_transformer_numerics.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image_21 import qwen_image_21_transformer as transformer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "device",
+    [pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
+)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_qk_norm_matches_diffusers_rounding(monkeypatch, device, dtype):
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state.get_tp_group", lambda: SimpleNamespace(rank_in_group=0, world_size=1)
+    )
+    monkeypatch.setattr(transformer, "Attention", lambda **kwargs: torch.nn.Identity())
+    attention = transformer.QwenImage21Attention(dim=256, heads=2, dim_head=128).to(device=device, dtype=dtype)
+    values = torch.randn(2, 17, 2, 128, generator=torch.Generator().manual_seed(42)).to(device=device, dtype=dtype)
+    for norm in (attention.norm_q, attention.norm_k):
+        with torch.no_grad():
+            norm.weight.copy_(torch.linspace(0.5, 1.5, 128, device=device, dtype=dtype))
+        variance = values.float().square().mean(-1, keepdim=True)
+        expected = (values * torch.rsqrt(variance + norm.eps)).to(dtype) * norm.weight
+        torch.testing.assert_close(norm(values), expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [pytest.param("meta", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])],
+)
+def test_rope_frequencies_do_not_depend_on_loader_device(device):
+    reference = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
+    with torch.device(device):
+        loaded = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56])
+    for expected, actual in zip(reference.freqs, loaded.freqs):
+        assert actual.device.type == "cpu"
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_scheduler_matches_diffusers_on_latent_device():
+    import numpy as np
+    from diffusers import FlowMatchEulerDiscreteScheduler
+
+    from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import QwenImage21Pipeline
+
+    pipeline = object.__new__(QwenImage21Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cuda", 0)
+    pipeline.scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True)
+    timesteps, _ = pipeline.prepare_timesteps(50, None, 4096)
+    reference = FlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
+    reference.set_timesteps(sigmas=np.linspace(1.0, 1 / 50, 50), device=pipeline.device, mu=1.15)
+    sample = torch.randn(1, 16, 64, generator=torch.Generator().manual_seed(42)).to(
+        device=pipeline.device, dtype=torch.bfloat16
+    )
+    noise = sample.cos()
+    actual = pipeline.scheduler.step(noise, timesteps[0], sample, return_dict=False)[0]
+    expected = reference.step(noise, reference.timesteps[0], sample, return_dict=False)[0]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert pipeline.scheduler.sigmas.device == sample.device
+    assert timesteps.device == sample.device
+
+
+@pytest.mark.cuda
+@pytest.mark.gpu
+def test_timestep_frequencies_match_cpu_initialization():
+    reference = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
+    with torch.device("cuda"):
+        loaded = transformer.QwenImage21TemporalTimesteps(timestep_dim=256)
+    assert loaded.freqs.device.type == "cuda"
+    torch.testing.assert_close(loaded.freqs.cpu(), reference.freqs, rtol=0, atol=0)
+    timesteps = torch.tensor([0.0, 0.984375, 1.0], device="cuda")
+    torch.testing.assert_close(loaded(timesteps), reference.to("cuda")(timesteps), rtol=0, atol=0)

tests/diffusion/models/qwen_image_21/test_vae.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import math
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21 import (
+    DistributedAutoencoderKLQwenImage21,
+)
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def single_thread():
+    previous = torch.get_num_threads()
+    torch.set_num_threads(1)
+    yield
+    torch.set_num_threads(previous)
+
+
+def _make_vae(cls=AutoencoderKLQwenImage21, **overrides):
+    """Toy VAE with the 2.1 residual architecture; 4x spatial compression."""
+    kwargs = dict(
+        base_dim=16,
+        decoder_base_dim=16,
+        z_dim=4,
+        dim_mult=[1, 2, 2],
+        num_res_blocks=1,
+        temperal_downsample=[False, False],
+        is_residual=True,
+        in_channels=4,
+        out_channels=4,
+        latents_mean=[0.0] * 4,
+        latents_std=[1.0] * 4,
+        scale_factor_spatial=None,
+        scale_factor_temporal=None,
+    )
+    kwargs.update(overrides)
+    return cls(**kwargs)
+
+
+def test_spatial_compression_ratio_derived_from_architecture():
+    # One 2x downsampling stage per `temperal_downsample` entry.
+    vae = _make_vae()
+    assert vae.spatial_compression_ratio == 4
+
+    # The real Qwen-Image-2.1 checkpoint ships scale_factor_spatial=8, but its
+    # four downsampling stages compress 16x; the architecture wins.
+    with torch.device("meta"):
+        vae = AutoencoderKLQwenImage21(
+            base_dim=96,
+            decoder_base_dim=144,
+            z_dim=64,
+            dim_mult=[1, 2, 4, 8, 8],
+            num_res_blocks=2,
+            temperal_downsample=[False, True, True, True],
+            is_residual=True,
+            in_channels=4,
+            out_channels=4,
+            scale_factor_spatial=8,
+        )
+    assert vae.spatial_compression_ratio == 16
+
+
+def test_tiled_decode_output_shape_and_closeness():
+    torch.manual_seed(0)
+    vae = _make_vae()
+    z = torch.randn(1, 4, 1, 16, 16)
+    ratio = vae.spatial_compression_ratio
+
+    with torch.no_grad():
+        ref = vae._decode(z, return_dict=False)[0]
+        vae.enable_tiling(
+            tile_sample_min_height=32,
+            tile_sample_min_width=32,
+            tile_sample_stride_height=24,
+            tile_sample_stride_width=24,
+        )
+        tiled = vae.decode(z, return_dict=False)[0]
+
+    # Regression: with the wrong (config-derived) ratio the tiled output was
+    # cropped to a fraction of the sample size.
+    assert tiled.shape == (1, 4, 1, 16 * ratio, 16 * ratio)
+    assert tiled.shape == ref.shape
+
+    # Tiled decode blends overlapping tiles, so it is not bitwise-identical to
+    # untiled decode; on this toy model the seam error stays small.
+    mse = ((ref - tiled) ** 2).mean().item()
+    psnr = 10 * math.log10(4.0 / mse)
+    assert psnr > 20.0
+
+    # Tiled output is clamped to [-1, 1] like the untiled path.
+    assert tiled.min() >= -1.0 and tiled.max() <= 1.0
+
+
+def test_distributed_tile_hooks_match_tiled_decode():
+    """tile_split -> tile_exec -> tile_merge must reproduce tiled_decode exactly."""
+    torch.manual_seed(0)
+    vae = _make_vae()
+    vae.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_stride_height=24,
+        tile_sample_stride_width=24,
+    )
+    distributed = _make_vae(DistributedAutoencoderKLQwenImage21)
+    distributed.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_stride_height=24,
+        tile_sample_stride_width=24,
+    )
+    distributed.load_state_dict(vae.state_dict())
+
+    z = torch.randn(1, 4, 1, 16, 16)
+    with torch.no_grad():
+        tiled = vae.tiled_decode(z, return_dict=False)[0]
+
+        tasks, grid_spec = DistributedAutoencoderKLQwenImage21.tile_split(distributed, z)
+        # 16x16 latent, 8x8 latent tiles, stride 6 -> range(0, 16, 6) = 3 stops/axis.
+        assert len(tasks) == 9
+        assert grid_spec.grid_shape == (3, 3)
+        assert grid_spec.split_dims == (3, 4)
+        # Every tile carries all frames (T=1 for images).
+        assert all(len(task.tensor) == z.shape[2] for task in tasks)
+
+        coord_tensor_map = {task.grid_coord: distributed.tile_exec(task) for task in tasks}
+        merged = distributed.tile_merge(coord_tensor_map, grid_spec)
+
+    assert merged.shape == tiled.shape
+    torch.testing.assert_close(merged, tiled)
+
+
+def test_tiled_encode_output_shape_and_closeness():
+    """tiled_encode is engaged for image-conditioned generation when tiling is on."""
+    torch.manual_seed(0)
+    vae = _make_vae()
+    x = torch.randn(1, 4, 1, 64, 64)
+
+    with torch.no_grad():
+        ref = vae._encode(x)
+        vae.enable_tiling(
+            tile_sample_min_height=32,
+            tile_sample_min_width=32,
+            tile_sample_stride_height=24,
+            tile_sample_stride_width=24,
+        )
+        tiled = vae._encode(x)
+
+    ratio = vae.spatial_compression_ratio
+    assert tiled.shape == (1, 8, 1, 64 // ratio, 64 // ratio)
+    assert tiled.shape == ref.shape
+    # Blended tile borders deviate from untiled encode but stay small.
+    assert (ref - tiled).abs().max().item() < 1.0
+
+
+def test_encode_decode_round_trip_shape():
+    torch.manual_seed(0)
+    vae = _make_vae()
+    x = torch.randn(1, 4, 1, 32, 32)
+    with torch.no_grad():
+        latents = vae.encode(x, return_dict=False)[0].mode()
+        decoded = vae.decode(latents, return_dict=False)[0]
+    ratio = vae.spatial_compression_ratio
+    assert latents.shape == (1, 4, 1, 32 // ratio, 32 // ratio)
+    assert decoded.shape == x.shape
+
+
+@pytest.mark.parametrize("persistent", [False, True])
+def test_adaptive_oom_restores_tile_configuration(monkeypatch, persistent):
+    vae = _make_vae()
+    vae.enable_tiling(
+        tile_sample_min_height=512,
+        tile_sample_min_width=512,
+        tile_sample_stride_height=384,
+        tile_sample_stride_width=384,
+    )
+    attempts = []
+    z = torch.zeros(1, 4, 1, 4, 4)
+
+    def decode_tile(z, return_dict):
+        attempts.append(vae.tile_sample_min_height)
+        if persistent or vae.tile_sample_min_height > 128:
+            raise torch.OutOfMemoryError("injected tile OOM")
+        return (z,)
+
+    monkeypatch.setattr(vae, "tiled_decode", decode_tile)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+    if persistent:
+        with pytest.raises(torch.OutOfMemoryError):
+            vae._tiled_decode_adaptive(z, return_dict=False)
+    else:
+        assert vae._tiled_decode_adaptive(z, return_dict=False)[0] is z
+    assert attempts == [512, 256, 128]
+    assert (
+        vae.tile_sample_min_height,
+        vae.tile_sample_min_width,
+        vae.tile_sample_stride_height,
+        vae.tile_sample_stride_width,
+    ) == (512, 512, 384, 384)

tests/diffusion/offloader/test_sequential_backend.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for SequentialOffloadBackend."""
 
@@ -284,6 +284,97 @@ def test_component_selective_model_offload_requires_swap_counterpart(component,
         backend.enable(pipeline)
 
 
+def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None:
+    dit = _create_simple_module().to(accelerator_device)
+    encoder = _create_simple_module().to(accelerator_device)
+    original_weight = dit.linear.weight.detach().clone()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=accelerator_device,
+        pin_memory=False,
+        persistent_dit_staging=True,
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook._stager is not None
+    assert encoder_hook._stager is None  # encoders keep plain move semantics
+
+    # Registration snapshots the CPU master and rebinds to it.
+    assert dit.linear.weight.device.type == "cpu"
+
+    dit_hook._to_gpu(dit)
+    first_ptr = dit.linear.weight.data_ptr()
+    assert dit.linear.weight.device == accelerator_device
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+
+    # Offload via a *different* module's hook (the cross-module swap path),
+    # then reload: the same fixed device storage must be reused so that
+    # CUDA-graph-captured weight pointers stay valid.
+    encoder_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+
+    # A second offload/load cycle through the DiT hook itself stays put too.
+    dit_hook._to_cpu(dit)
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+
+    remove_sequential_offload([dit, encoder])
+
+
+def test_default_sequential_offload_releases_device_storage(accelerator_device) -> None:
+    """Without persistent staging, offloading keeps plain move semantics."""
+    dit = _create_simple_module().to(accelerator_device)
+    encoder = _create_simple_module().to(accelerator_device)
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=accelerator_device,
+        pin_memory=False,
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook._stager is None  # no retained staging: GPU storage is freed on offload
+
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.device == accelerator_device
+    dit_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+
+    remove_sequential_offload([dit, encoder])
+
+
+@pytest.mark.parametrize(
+    ("decode_graphs", "expect_stager"),
+    [(True, True), (False, False)],
+)
+def test_model_level_persistent_staging_gated_on_decode_graph(
+    accelerator_device, decode_graphs: bool, expect_stager: bool
+) -> None:
+    """ModelLevelOffloadBackend only retains DiT staging for decode-graph models."""
+    pipeline = nn.Module()
+    pipeline.transformer = _create_simple_module()
+    if decode_graphs:
+        pipeline.transformer.enable_cuda_graph_decode = True
+    pipeline.text_encoder = _create_simple_module()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        accelerator_device,
+    )
+    backend.enable(pipeline)
+
+    hook = pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert (hook._stager is not None) is expect_stager
+
+    backend.disable()
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()

tests/diffusion/test_async_output_worker.py

@@ -497,3 +497,99 @@ class TestWorkerProcShutdown:
         result_mq.shutdown.assert_called_once_with()
         assert proc.mq is None
         assert proc.result_mq is None
+
+
+class TestSynchronizeDeviceBarrier:
+    """The pause barrier must not report a quiet device while outputs still pack."""
+
+    @staticmethod
+    def _rpc(timeout):
+        return {
+            "method": "synchronize_device",
+            "args": (),
+            "kwargs": {"timeout": timeout},
+            "output_rank": 0,
+            "exec_all_ranks": False,
+            "collect_rank_status": False,
+        }
+
+    def test_rejects_undrained_output(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        mocker.patch.object(proc, "drain_async_outputs", return_value=False)
+
+        with pytest.raises(TimeoutError, match="did not drain"):
+            proc._execute_rpc(self._rpc(0.01))
+
+        proc.worker.execute_method.assert_not_called()
+
+    @pytest.mark.parametrize("timeout", [0.5, None])
+    def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        call_order: list = []
+        mocker.patch.object(
+            proc,
+            "drain_async_outputs",
+            side_effect=lambda **kw: call_order.append(("drain", kw)) or True,
+        )
+        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw))
+
+        _, should_reply = proc._execute_rpc(self._rpc(timeout))
+
+        assert call_order == [
+            ("drain", {"timeout": timeout}),
+            ("sync", ("synchronize_device",), {"timeout": timeout}),
+        ]
+        assert should_reply is True
+
+    def test_completes_only_after_output_ready_is_published(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+
+        mock_platform = mocker.MagicMock()
+        mock_platform.record_device_event.return_value = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+            mock_platform,
+        )
+        mock_torch = mocker.MagicMock()
+        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
+        release = threading.Event()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=lambda *a, **kw: release.wait(timeout=5.0),
+        )
+
+        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
+        loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True)
+        loop_thread.start()
+
+        barrier_done = threading.Event()
+
+        def run_barrier():
+            proc._execute_rpc(self._rpc(5.0))
+            barrier_done.set()
+
+        barrier_thread = threading.Thread(target=run_barrier, daemon=True)
+        barrier_thread.start()
+        try:
+            assert not barrier_done.wait(timeout=0.2)
+            proc.worker.execute_method.assert_not_called()
+
+            release.set()
+            assert barrier_done.wait(timeout=5.0)
+
+            kinds = [
+                call.args[0].kind
+                for call in proc.result_mq.enqueue.call_args_list
+                if isinstance(call.args[0], AsyncDiffusionOutput)
+            ]
+            assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY]
+            proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0)
+        finally:
+            release.set()
+            proc._async_output_queue.put(None)
+            loop_thread.join(timeout=2.0)
+            barrier_thread.join(timeout=2.0)

tests/diffusion/test_diffusion_config_propagation.py

@@ -45,6 +45,21 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
 
+def test_direct_config_extras_are_independent_dicts():
+    first = OmniDiffusionConfig(model="test-model")
+    second = OmniDiffusionConfig(model="test-model")
+
+    assert first.extras.get("prefix_kv_cache_dtype") is None
+    first.extras["prefix_kv_cache_dtype"] = "fp8"
+    assert second.extras == {}
+
+
+def test_direct_config_preserves_explicit_extras():
+    config = OmniDiffusionConfig(model="test-model", extras={"prefix_kv_cache_dtype": "fp8"})
+
+    assert config.extras["prefix_kv_cache_dtype"] == "fp8"
+
+
 class TestParallelConfigPropagation:
     """Core regression tests: parallel_config must survive serialization."""
 

tests/diffusion/test_diffusion_engine_rpc_routing.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
 """Regression tests for the busy-loop-based RPC routing in DiffusionEngine.
 
 These tests guard the invariants established when ``collective_rpc`` was
@@ -60,6 +60,12 @@ class _ConcurrencyTrackingExecutor:
         self.is_failed = False
         self._closed = False
         self.od_config = SimpleNamespace()
+        # Pause tests: a call whose gate is registered (by request id, or by
+        # method name for ``synchronize_device``) blocks inside the executor
+        # until the gate is set; ``version`` tags each call so a test can tell
+        # which "weights" a request ran on.
+        self.gates: dict[str, threading.Event] = {}
+        self.version = "v1"
 
     def collective_rpc(
         self,
@@ -81,16 +87,25 @@ class _ConcurrencyTrackingExecutor:
                     "kwargs": kwargs,
                     "unique_reply_rank": unique_reply_rank,
                     "thread": threading.get_ident(),
+                    "version": self.version,
                 }
             )
         try:
             if self.rpc_delay:
                 time.sleep(self.rpc_delay)
+            if method == "synchronize_device":
+                gate = self.gates.get(method)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
+                return [None]
             # Distinguish per-request execution from raw RPC by method name.
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
                 tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
+                gate = self.gates.get(tag)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
                 return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -133,21 +148,27 @@ def _make_request(tag: str):
 
 
 def _make_engine_with_loop(
-    loop: asyncio.AbstractEventLoop,
+    loop: asyncio.AbstractEventLoop | None,
     rpc_delay: float = 0.0,
+    *,
+    scheduler: RequestScheduler | None = None,
+    start_loop: bool = True,
 ):
     """Construct a ``DiffusionEngine`` skeleton with a real busy loop.
 
     The engine is wired with a fake executor that asserts no concurrent
-    calls and a real ``RequestScheduler``.
+    calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the
+    engine in its bootstrap state (no busy-loop thread).
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
     engine._closed = False
     engine.od_config = SimpleNamespace(streaming_output=False)
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
-    sched = RequestScheduler()
-    sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+    sched = scheduler
+    if sched is None:
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
@@ -160,9 +181,14 @@ def _make_engine_with_loop(
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
+    engine._scheduling_paused = False
 
     engine.main_loop = loop
     engine.stop_event = threading.Event()
+    if not start_loop:
+        engine._loop_started = False
+        engine.worker_thread = None
+        return engine
     engine._loop_started = True
     engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
     engine.worker_thread.start()
@@ -170,6 +196,10 @@ def _make_engine_with_loop(
 
 
 def _stop_engine(engine: DiffusionEngine) -> None:
+    if engine.worker_thread is None:
+        return
+    for gate in engine.executor.gates.values():
+        gate.set()
     with engine._cv:
         engine.stop_event.set()
         engine._cv.notify_all()
@@ -565,3 +595,319 @@ async def test_rpc_and_request_results_do_not_swap():
     assert req_res.error == "result_for_req1"
     # Belt and braces: never overlapped on the executor.
     assert engine.executor.max_concurrent == 1
+
+
+# ───────────────────────── scheduler pause (keep) ─────────────────────────
+
+
+class _GatedAdmissionScheduler(RequestScheduler):
+    """Batching wait that only ends when the test opens admission."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.admission_open = threading.Event()
+
+    def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False):
+        decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent)
+        return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs)
+
+    def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool:
+        return self.admission_open.is_set()
+
+
+def _pause(engine: DiffusionEngine, timeout: float | None = None):
+    return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"})
+
+
+def _resume(engine: DiffusionEngine):
+    return engine.async_collective_rpc("resume_scheduler")
+
+
+def _methods(engine: DiffusionEngine) -> list[str]:
+    return [call["method"] for call in engine.executor.calls]
+
+
+def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]:
+    return [
+        (call["args"][0].request_id, call["version"])
+        for call in engine.executor.calls
+        if call["method"] == "execute_model"
+    ]
+
+
+def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task:
+    return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag))))
+
+
+async def _wait_until(predicate, timeout: float = 3.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() > deadline:
+            raise AssertionError("condition not reached in time")
+        await asyncio.sleep(0.005)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        assert await _pause(engine) is None
+        queued = _submit(engine, "queued")
+        result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0)
+        assert result.error == "rpc_result_for_list_loras"
+        await asyncio.sleep(0.05)
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert _executed(engine) == []
+        assert not queued.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued"
+        assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_acks_after_running_batch_and_freezes_queue():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        pause = asyncio.create_task(_pause(engine))
+        await asyncio.sleep(0.1)
+        assert not pause.done()
+
+        gate.set()
+        await asyncio.wait_for(pause, 3.0)
+        assert _methods(engine) == ["execute_model", "synchronize_device"]
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert not b.done()
+
+        engine.executor.version = "v2"
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert _executed(engine) == [("A", "v1"), ("B", "v2")]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted():
+    """A pause ACK must not open the queue to RPCs that touch device memory
+    (sleep) before the batch that just ran has been delivered.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event()
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        pause = asyncio.create_task(_pause(engine))
+        await _wait_until(lambda: engine._scheduling_paused)
+        gate.set()
+        # The barrier holds the busy loop inside the post-execute RPC drain,
+        # so the sleep is queued before the pause ACK is published.
+        await _wait_until(lambda: "synchronize_device" in _methods(engine))
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: not engine._rpc_queue.empty())
+        barrier_gate.set()
+
+        await asyncio.wait_for(pause, 3.0)
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted():
+    """A pause whose caller already gave up is still a drain barrier: its
+    cancelled task must not let a later RPC run before the batch is emitted.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+
+        # The caller times out while A is still executing: the pause future is
+        # cancelled, but its task stays queued ahead of the sleep.
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: engine._rpc_queue.qsize() == 2)
+        gate.set()
+
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_pause_rejects_unsupported_modes_without_closing_the_gate():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}):
+            with pytest.raises(ValueError):
+                await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs)
+        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
+        with pytest.raises(NotImplementedError):
+            await _pause(engine)
+        engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
+
+        assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs"
+        assert _methods(engine) == ["execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_during_batching_wait_holds_the_batch():
+    loop = asyncio.get_running_loop()
+    sched = _GatedAdmissionScheduler()
+    sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0))
+    engine = _make_engine_with_loop(loop, scheduler=sched)
+    try:
+        waiting = _submit(engine, "W")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        await asyncio.sleep(0.05)
+
+        await asyncio.wait_for(_pause(engine), 3.0)
+        assert _methods(engine) == ["synchronize_device"]
+        assert engine.scheduler.num_waiting_requests() == 1
+
+        await _resume(engine)
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        sched.admission_open.set()
+        assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W"
+    finally:
+        sched.admission_open.set()
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_pause_timeout_leaves_gate_closed_until_explicit_resume():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        gate.set()
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        await asyncio.sleep(0.1)
+        assert _executed(engine) == [("A", "v1")]
+        assert not b.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert "synchronize_device" not in _methods(engine)
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_repeated_pause_and_resume_are_idempotent():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _resume(engine)
+        await _pause(engine)
+        await _pause(engine)
+        queued = _submit(engine, "Q")
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        await _resume(engine)
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q"
+        assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_abort_of_waiting_request_while_paused_delivers_terminal():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _pause(engine)
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        engine.abort("B")
+
+        result = await asyncio.wait_for(b, 3.0)
+        assert result.aborted is True
+        assert engine.scheduler.get_request_state("B") is None
+        assert _executed(engine) == []
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C"
+        assert _executed(engine) == [("C", "v1")]
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
+def test_pause_and_resume_before_loop_starts_use_engine_local_path():
+    engine = _make_engine_with_loop(None, start_loop=False)
+
+    assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None
+    assert _methods(engine) == ["synchronize_device"]
+
+    assert engine.collective_rpc("resume_scheduler") is None
+    assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X"
+    assert _executed(engine) == [("X", "v1")]
+
+
+@pytest.mark.asyncio
+async def test_requests_without_pause_never_trigger_the_barrier():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        first = await asyncio.wait_for(_submit(engine, "r1"), 3.0)
+        await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0)
+        second = await asyncio.wait_for(_submit(engine, "r2"), 3.0)
+    finally:
+        _stop_engine(engine)
+
+    assert (first.error, second.error) == ("result_for_r1", "result_for_r2")
+    assert _methods(engine) == ["execute_model", "ping", "execute_model"]

tests/diffusion/test_diffusion_scheduler.py

@@ -2092,3 +2092,19 @@ class TestStepScheduler:
 
         with pytest.raises(ValueError):
             self.scheduler.add_request(request)
+
+
+class TestPendingFinishedRequestIds:
+    def test_reports_finished_ids_that_still_hold_state_without_clearing_them(self):
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+        sched.add_request(_make_request("a"))
+        sched.add_request(_make_request("b"))
+
+        sched.finish_requests("a", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert sched.pending_finished_request_ids() == {"a"}
+
+        sched.pop_request_state("a")
+        assert sched.pending_finished_request_ids() == set()
+        # The next wave still ships the id to the worker for its own cleanup.
+        assert sched.schedule().finished_req_ids == {"a"}

tests/diffusion/test_diffusion_worker_synchronize_device.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for the device step of the diffusion pause barrier."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_worker(runner: object) -> DiffusionWorker:
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = runner
+    worker.rank = 0
+    return worker
+
+
+def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    order: list = []
+    manager = MagicMock()
+    manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
+    platform.synchronize.side_effect = lambda: order.append("sync")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    worker.synchronize_device(timeout=2.5)
+
+    assert order == [("prefetch", 2.5), "sync"]
+
+
+def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    manager = MagicMock()
+    manager.wait_prefetch.return_value = False
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    with pytest.raises(TimeoutError, match="prefetch"):
+        worker.synchronize_device(timeout=0.01)
+
+    platform.synchronize.assert_not_called()
+
+
+def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
+
+    worker.synchronize_device(timeout=None)
+
+    platform.synchronize.assert_called_once_with()

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -17,7 +17,7 @@ import zmq
 from vllm.v1.engine.exceptions import EngineDeadError
 
 import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
@@ -37,6 +37,38 @@ from vllm_omni.outputs import OmniRequestOutput
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _signal_worker_ready(_rank, _config, writer, _broadcast, wake_event, *_args):
+    assert mp.get_start_method() == "spawn"
+    wake_event.set()
+    writer.send({"status": "ready", "result_handle": "ready"})
+    writer.close()
+
+
+@pytest.mark.skipif("fork" not in mp.get_all_start_methods(), reason="Requires a fork default context")
+def test_executor_spawns_workers_with_compatible_events(mocker):
+    previous_method = mp.get_start_method(allow_none=True)
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = OmniDiffusionConfig(model="test-model", step_execution=True)
+    broadcast = mocker.Mock()
+    broadcast.export_handle.return_value = None
+    mocker.patch.object(executor, "_init_broadcast_queue", return_value=broadcast)
+    mocker.patch.object(executor, "_init_result_queue")
+    mocker.patch.object(executor, "_start_worker_monitor")
+    mocker.patch.object(WorkerProc, "worker_main", new=_signal_worker_ready)
+    try:
+        mp.set_start_method("fork", force=True)
+        executor._init_executor()
+        for process in executor._processes:
+            process.join(timeout=30)
+            assert process.exitcode == 0
+        assert all(event.is_set() for event in executor.wake_events)
+        assert mp.get_start_method() == "fork"
+    finally:
+        if getattr(executor, "_finalizer", None) is not None:
+            executor.shutdown()
+        mp.set_start_method(previous_method, force=True)
+
+
 # ───────────────────────────────────────────── helpers ─────────────────────
 
 

tests/distributed/omni_connectors/test_kv_async_prefetch.py

@@ -5,6 +5,7 @@ mock connector, dedup, per-call sender_info isolation, role classification, and
 miss/abort paths.  GPU H2D and D2D paths need CUDA (integration tests).
 """
 
+import threading
 from types import SimpleNamespace
 
 import pytest
@@ -354,3 +355,41 @@ def test_consume_then_apply_attaches_payload():
     receiver._record_stream_for_prefetched(data)
     receiver.apply_kv_cache_to_request(req, data)
     assert req.past_key_values is not None
+
+
+# --------------------------------------------------------------------------- #
+#  wait_prefetch: the pause barrier's view of outstanding background loads
+# --------------------------------------------------------------------------- #
+
+
+def test_wait_prefetch_is_true_when_nothing_outstanding():
+    _, receiver, _ = _make_sender_receiver()
+    assert receiver.wait_prefetch(timeout=0.01) is True
+
+
+def test_wait_prefetch_is_true_when_prefetch_disabled():
+    _, receiver, _ = _make_sender_receiver(async_prefetch=False)
+    assert receiver.wait_prefetch(timeout=0.01) is True
+
+
+def test_wait_prefetch_reports_a_running_prefetch_and_keeps_its_payload():
+    sender, receiver, connector = _make_sender_receiver()
+    _seed_payload(sender, "rid-w")
+    gate = threading.Event()
+    real_get = connector.get
+
+    def slow_get(*args, **kwargs):
+        gate.wait(timeout=5.0)
+        return real_get(*args, **kwargs)
+
+    connector.get = slow_get
+    receiver.start_prefetch({"request_id": "rid-w", "kv_sender_info": _SENDER_INFO})
+    try:
+        assert receiver.wait_prefetch(timeout=0.05) is False
+        gate.set()
+        assert receiver.wait_prefetch(timeout=5.0) is True
+        data, _ = receiver.consume_prefetched_kv(_req("rid-w"))
+        assert data is not None
+    finally:
+        gate.set()
+        receiver.shutdown_prefetch()

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -32,7 +32,6 @@ from huggingface_hub import hf_hub_download, snapshot_download
 from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
 
 from tests.e2e.accuracy.helpers import reset_artifact_dir
-from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.mark import hardware_test
 
 OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git"
@@ -303,32 +302,6 @@ def _resolve_ltx25_image() -> Path:
     )
 
 
-def _require_ltx25_model_access() -> None:
-    configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
-    if not (configured_omni_model and Path(configured_omni_model).exists()):
-        omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID
-        omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
-        if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID:
-            omni_revision = LTX25_OMNI_MODEL_REVISION
-        skip_if_gated_repo_inaccessible(
-            omni_model,
-            revision=omni_revision,
-            filename="model_index.json",
-        )
-
-    configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
-    if not (configured_official_model and Path(configured_official_model).exists()):
-        official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID
-        official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
-        if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID:
-            official_revision = LTX25_OFFICIAL_MODEL_REVISION
-        skip_if_gated_repo_inaccessible(
-            official_model,
-            revision=official_revision,
-            filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"],
-        )
-
-
 def _ltx25_request(image: Path | None = None) -> dict[str, object]:
     request: dict[str, object] = {
         "prompt": LTX25_PROMPT,
@@ -475,7 +448,6 @@ def _assert_strict_similarity(
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
     """Compare Omni with official canonical DiffVAE and matching connector weights."""
-    _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
     official_root, official_revision = _ltx25_official_source(artifact_parent)
@@ -600,7 +572,6 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
 @pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
 def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
     """Compare raw official Full/SFT and DiffVAE weights with Omni."""
-    _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
     model_class_name = {

tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Lifecycle test for ``AsyncOmni.pause_generation(mode="keep")`` on a diffusion stage.
+
+Per round, request A is started, request B is queued behind it, and the pause
+is issued while A is still running. The pause must return only after A has
+been delivered, B must stay queued through ``sleep`` / ``wake_up``, and B may
+start only after ``resume_generation``.
+
+A second test follows the order the docs recommend for a weight update: sleep
+is issued as soon as the pause returns, without awaiting the request that was
+running, which must still be delivered intact.
+
+Both single-GPU executor backends are covered explicitly: ``uni`` (in-process
+worker) and ``mp`` (worker process with the asynchronous D2H output thread).
+
+Usage:
+    pytest tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py -v -s
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+from collections import defaultdict
+from contextlib import ExitStack
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.data import is_diffusion_request_started_output
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "tiny-random/Qwen-Image"
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+
+# CI runs a couple of rounds; set PAUSE_KEEP_ROUNDS=10 for a soak run.
+ROUNDS = int(os.environ.get("PAUSE_KEEP_ROUNDS", "2"))
+# Enough denoising steps that the pause lands while A is still executing.
+NUM_INFERENCE_STEPS = 40
+
+
+def _sampling_params() -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        num_inference_steps=NUM_INFERENCE_STEPS,
+        guidance_scale=0.0,
+        height=256,
+        width=256,
+        seed=42,
+        emit_request_lifecycle=True,
+    )
+
+
+async def _generate(engine: AsyncOmni, request_id: str, events: dict[str, dict[str, float]]):
+    final = None
+    async for output in engine.generate(
+        prompt={"prompt_ids": list(range(50))},
+        request_id=request_id,
+        sampling_params_list=[_sampling_params()],
+        output_modalities=["image"],
+    ):
+        if is_diffusion_request_started_output(output):
+            events[request_id]["started"] = time.monotonic()
+            continue
+        final = output
+    events[request_id]["finished"] = time.monotonic()
+    return final
+
+
+async def _wait_for(predicate, timeout: float = 30.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        assert time.monotonic() < deadline, "condition not reached in time"
+        await asyncio.sleep(0.005)
+
+
+async def _pause_round(engine: AsyncOmni, events: dict[str, dict[str, float]], round_index: int) -> float:
+    """A runs, B queues, keep-pause lands while A runs; B starts only after resume.
+
+    Returns the pause latency.
+    """
+    a_id, b_id = f"a{round_index}", f"b{round_index}"
+    a = asyncio.create_task(_generate(engine, a_id, events))
+    await _wait_for(lambda: "started" in events[a_id])
+    b = asyncio.create_task(_generate(engine, b_id, events))
+    await asyncio.sleep(0.05)
+
+    assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS"
+    pause_start = time.monotonic()
+    await engine.pause_generation(mode="keep", clear_cache=False)
+    pause_latency = time.monotonic() - pause_start
+    assert await engine.is_paused()
+
+    a_output = await asyncio.wait_for(a, 60.0)
+    assert a_output.images, "A must be delivered, not aborted, by a keep pause"
+    assert "started" not in events[b_id], "B must not start while paused"
+
+    # Control RPCs and the sleep cycle stay available on the paused stage.
+    assert await engine.list_loras() == []
+    await engine.sleep(level=1)
+    assert await engine.is_sleeping()
+    await engine.wake_up()
+    assert not await engine.is_sleeping()
+    assert await engine.is_paused()
+    assert not b.done() and "started" not in events[b_id]
+
+    resume_at = time.monotonic()
+    await engine.resume_generation()
+    b_output = await asyncio.wait_for(b, 60.0)
+    assert b_output.images, "B must run after resume"
+    assert events[b_id]["started"] >= resume_at
+    return pause_latency
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["uni", "mp"])
+async def test_keep_pause_sleep_wake_resume_lifecycle(backend: str):
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            distributed_executor_backend=backend,
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+
+        pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(ROUNDS)]
+
+        assert not engine.request_states, "every request must be finalized"
+        print(
+            f"[{backend}] rounds={ROUNDS} steps={NUM_INFERENCE_STEPS} "
+            f"pause_latency_s min/median/max="
+            f"{min(pause_latencies):.3f}/{sorted(pause_latencies)[ROUNDS // 2]:.3f}/{max(pause_latencies):.3f}"
+        )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", ["uni", "mp"])
+async def test_sleep_right_after_keep_pause_keeps_the_running_batch(backend: str):
+    """The documented weight-update order: sleep is issued as soon as the pause
+    returns, with the request that was running still un-awaited.
+    """
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            distributed_executor_backend=backend,
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+
+        a = asyncio.create_task(_generate(engine, "a", events))
+        await _wait_for(lambda: "started" in events["a"])
+        await asyncio.sleep(0.05)
+        assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS"
+
+        await engine.pause_generation(mode="keep", clear_cache=False)
+        await engine.sleep(level=1)
+
+        a_output = await asyncio.wait_for(a, 60.0)
+        assert a_output.images, "A must survive a sleep issued right after the pause ACK"
+
+        await engine.wake_up()
+        await engine.resume_generation()
+        b_output = await asyncio.wait_for(_generate(engine, "b", events), 120.0)
+        assert b_output.images, "the stage must generate again after wake and resume"

tests/e2e/features/helpers/custom_pipeline.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
 
 """Test doubles for custom-pipeline E2E and RLHF integration.
 
@@ -290,6 +290,36 @@ class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
         """Return a stable identifier for assertions in unit tests."""
         return "vllm-omni-colocate-worker-extension-for-test"
 
+    def start_counting_synchronize_device(self) -> bool:
+        """Count pause-barrier runs on this rank (test-only)."""
+        self._synchronize_device_calls = 0
+        inner = type(self).synchronize_device
+
+        def _counted(timeout=None):
+            self._synchronize_device_calls += 1
+            return inner(self, timeout=timeout)
+
+        self.synchronize_device = _counted
+        return True
+
+    def synchronize_device_seen(self) -> bool:
+        """True from every rank only when each rank ran the barrier at least once."""
+        return getattr(self, "_synchronize_device_calls", 0) > 0
+
+    def set_synchronize_device_failure(self, rank: int, fail: bool) -> bool:
+        """Make (or stop making) the pause barrier fail on ``rank`` (test-only)."""
+        if self.rank != rank:
+            return True
+        if fail:
+
+            def _fail(timeout=None):
+                raise RuntimeError(f"injected barrier failure on rank {rank}")
+
+            self.synchronize_device = _fail
+        else:
+            self.__dict__.pop("synchronize_device", None)
+        return True
+
 
 # ---------------------------------------------------------------------------
 # Pipeline layer

tests/e2e/offline_inference/test_hunyuanimage3.py

@@ -558,7 +558,7 @@ def test_shared_script_ar_path_reaches_generation(
         # safe against MODEL_NAME's base checkpoint missing the Instruct-only
         # <img_ratio_33>/<img_ratio_36> tokens. Wiring it here exercises the
         # real production path (get_ar_tokenizer_validator via the registry),
-        # not just the isolated fake-tokenizer coverage in test_prompt_utils.py
+        # not just the isolated fake-tokenizer coverage in test_hunyuan_image3_prompt_utils.py
         # / test_shared_script_ar_integration.py.
         validate_tokenizer=get_ar_tokenizer_validator(model_class_name),
     )

tests/e2e/online_serving/helpers/minicpmo_window_e2e.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Drive complete native duplex turns while Stage-0 repeatedly rebuilds KV."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+from pathlib import Path
+
+from vllm_omni.clients.duplex import DuplexClient, EventCollector, read_pcm16_wav, wait_for_condition
+from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
+
+
+async def run_window_turn(
+    *,
+    url: str,
+    model: str,
+    input_wav: Path,
+    mode: str,
+    ref_audio: Path | None = None,
+    repeats: int = 2,
+    video_frames: list[str] | None = None,
+    buffered_flush: bool = False,
+    timeout_s: float = 240.0,
+) -> dict[str, object]:
+    """Use aggressive limits to exercise multiple windows in one audio turn."""
+    pcm = read_pcm16_wav(input_wav) * repeats
+    collector = EventCollector()
+    reference = (
+        "data:audio/wav;base64," + base64.b64encode(ref_audio.read_bytes()).decode("ascii")
+        if ref_audio is not None
+        else None
+    )
+    config = create_duplex_session_config(
+        modalities=("text", "audio") if ref_audio is not None else ("text",),
+        ref_audio=reference,
+        temperature=0.0,
+        extra_body={
+            "sliding_window_mode": mode,
+            "context_max_units": 1,
+            "context_previous_max_tokens": 16,
+            "basic_window_high_tokens": 64,
+            "basic_window_low_tokens": 32,
+        },
+    )
+    client = DuplexClient(
+        url, model=model, config=config, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=timeout_s
+    )
+    async with client:
+        reader = asyncio.create_task(collector.consume(client))
+        collector.add({"type": "session.created", "session": client.session_info})
+        try:
+            frames_sent = await client.stream_pcm(
+                pcm,
+                chunk_ms=2500 if buffered_flush else 200,
+                realtime=not buffered_flush,
+                is_speech=True,
+                video_frames=video_frames,
+            )
+            await client.commit(final=True)
+            await wait_for_condition(
+                lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()),
+                timeout_s=timeout_s,
+                label="final input commit",
+            )
+            await wait_for_condition(
+                lambda: collector.count("response.done") > 0 or bool(collector.errors()),
+                timeout_s=timeout_s,
+                label="completed windowed response",
+            )
+            # Final execution follows the commit acknowledgement asynchronously.
+            await asyncio.sleep(3.0)
+            await client.close(timeout_s=timeout_s)
+            await asyncio.wait_for(reader, timeout=5.0)
+        finally:
+            if not reader.done():
+                reader.cancel()
+                await asyncio.gather(reader, return_exceptions=True)
+    return {
+        "created": collector.count("session.created"),
+        "committed": collector.count("input_audio_buffer.committed"),
+        "done": collector.count("response.done"),
+        "closed": collector.count("session.closed"),
+        "errors": collector.errors(),
+        "audio_bytes": len(collector.audio_bytes()),
+        "frames_sent": frames_sent,
+        "input_seconds": len(pcm) / 32000,
+    }

tests/e2e/online_serving/test_lingbot_video.py

@@ -1,7 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Online serving smoke coverage for the dense LingBot-Video checkpoint."""
+"""
+Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
+
+``test_text_to_image_001`` and ``test_video_generation_modes_001`` are the
+baseline T2I / T2V+TI2V cases (``core_model`` + ``advanced_model``) shared by
+ready (L2) and merge (L3). Heavier MoE and World coverage stays on
+``full_model`` in ``test_lingbot_video_moe.py`` / ``test_lingbot_world_v2.py``.
+
+From ``tests/``::
+
+    pytest -s -v e2e/online_serving/test_lingbot_video.py -m "core_model and diffusion" --run-level=core_model
+    pytest -s -v e2e/online_serving/test_lingbot_video.py -m "advanced_model and diffusion" --run-level=advanced_model
+"""
 
 import json
 import os
@@ -41,7 +53,8 @@ def _get_diffusion_feature_cases(model: str):
     ]
 
 
-@pytest.mark.full_model
+@pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
 def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
@@ -63,7 +76,8 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH
     )
 
 
-@pytest.mark.full_model
+@pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
 def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:

tests/e2e/online_serving/test_lingbot_video_expansion.py

@@ -4,9 +4,9 @@
 """
 L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
 
-This file remains dense-only. Basic T2I/T2V/TI2V serving for both dense and MoE
-checkpoints is covered by ``test_lingbot_video.py`` and
-``test_lingbot_video_moe.py``. Distributed feature rows belong in follow-up
+This file remains dense-only. Baseline T2I/T2V/TI2V for the dense checkpoint
+lives in ``test_lingbot_video.py`` (L2/L3). MoE smoke stays in
+``test_lingbot_video_moe.py`` (L4). Distributed feature rows belong in follow-up
 PRs.
 """
 

tests/e2e/online_serving/test_ltx25.py

@@ -7,7 +7,6 @@ import os
 
 import pytest
 
-from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
@@ -24,16 +23,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
 
 
-@pytest.fixture(scope="module", autouse=True)
-def require_ltx25_model_access() -> None:
-    if not os.path.isdir(MODEL):
-        skip_if_gated_repo_inaccessible(
-            MODEL,
-            revision=MODEL_REVISION or None,
-            filename="model_index.json",
-        )
-
-
 def _server(model_class_name: str) -> OmniServerParams:
     return OmniServerParams(
         model=MODEL,

tests/e2e/online_serving/test_minicpmo_4_5_window.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real-model regression for the scheduler-owned Thinker sliding window."""
+
+import asyncio
+
+import pytest
+
+from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+    SERVER_PARAMS,
+    duplex_camera_frames,
+    realtime_url,
+    resolve_ref_audio,
+    validated_input_wav,
+)
+from tests.e2e.online_serving.helpers.minicpmo_window_e2e import run_window_turn
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.omni, pytest.mark.advanced_model]
+
+
+def _assert_complete(result, *, require_audio=False):
+    assert result["created"] == 1
+    assert result["committed"] > 0
+    assert result["done"] > 0
+    assert result["closed"] == 1
+    assert result["errors"] == []
+    if require_audio:
+        assert result["audio_bytes"] > 0
+
+
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("mode,reference", [("basic", False), ("context", False), ("context", True)])
+def test_window_rebuild_and_next_session(omni_server, mode, reference):
+    # Two independent sessions exercise admission after final-unit execution.
+    for _ in range(2):
+        result = asyncio.run(
+            run_window_turn(
+                url=realtime_url(omni_server),
+                model=omni_server.model,
+                input_wav=validated_input_wav(),
+                mode=mode,
+                ref_audio=resolve_ref_audio() if reference else None,
+            )
+        )
+        _assert_complete(result, require_audio=reference)
+
+
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("mode,camera", [("basic", False), ("context", False), ("context", True)])
+def test_window_continuous_input(omni_server, tmp_path, mode, camera):
+    # 24 repeats cover over two minutes and many canonical KV replacements.
+    frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera") if camera else None
+    result = asyncio.run(
+        run_window_turn(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            input_wav=validated_input_wav(),
+            mode=mode,
+            repeats=24,
+            ref_audio=resolve_ref_audio(),
+            video_frames=frames,
+        )
+    )
+    _assert_complete(result, require_audio=True)
+    assert isinstance(result["input_seconds"], (int, float)) and result["input_seconds"] > 120
+    if camera:
+        assert isinstance(result["frames_sent"], int) and result["frames_sent"] >= 120
+
+
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("mode", ["off", "basic", "context"])
+def test_window_buffered_flush(omni_server, mode):
+    # Burst appends cross processor chunk boundaries and leave a final tail.
+    result = asyncio.run(
+        run_window_turn(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            input_wav=validated_input_wav(),
+            mode=mode,
+            ref_audio=resolve_ref_audio(),
+            buffered_flush=True,
+        )
+    )
+    _assert_complete(result, require_audio=True)

tests/e2e/online_serving/test_qwen_image_21.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
+
+MODEL = os.environ.get("QWEN_IMAGE_21_TEST_MODEL", "Qwen/Qwen-Image-2.1")
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(model=MODEL, server_args=["--vae-use-tiling", "--max-num-seqs", "2"] + extra),
+        id=name,
+        marks=hardware_marks(res={"cuda": "H100"}),
+    )
+    for name, extra in [("request", []), ("step", ["--step-execution"])]
+]
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.advanced_model,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(
+        not os.environ.get("QWEN_IMAGE_21_TEST_MODEL"),
+        reason="Set QWEN_IMAGE_21_TEST_MODEL to a provisioned checkpoint until public weights are available.",
+    ),
+]
+
+
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient):
+    online_client.send_diffusion_request(
+        {
+            "model": omni_server.model,
+            "messages": dummy_messages_from_mix_data(content_text="A red ceramic teapot on a wooden table."),
+            "extra_body": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 4,
+                "seed": 42,
+                "negative_prompt": "blurry, low quality",
+                "true_cfg_scale": 4.0,
+            },
+        }
+    )
+
+
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient):
+    image = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    online_client.send_diffusion_request(
+        {
+            "model": omni_server.model,
+            "messages": dummy_messages_from_mix_data(content_text="Make the image blue.", image_data_url=image),
+            "extra_body": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 4,
+                "seed": 42,
+                "negative_prompt": "blurry, low quality",
+                "true_cfg_scale": 4.0,
+            },
+        }
+    )

tests/e2e/online_serving/test_stable_audio_online_expansion.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
 """
 E2E online serving test for Stable Audio Open text-to-audio diffusion.
 
@@ -11,7 +11,6 @@ import os
 
 import pytest
 
-from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
@@ -32,17 +31,10 @@ def _stable_audio_server_cases(model: str):
     ]
 
 
-@pytest.fixture
-def _require_stable_audio_access() -> None:
-    """Skip cleanly (before the server boots) if the gated checkpoint is inaccessible."""
-    skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL)
-
-
 @pytest.mark.slow
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True)
 def test_stable_audio_t2a_online(
-    _require_stable_audio_access: None,
     omni_server: OmniServer,
     openai_client: OpenAIClientHandler,
 ) -> None:

tests/engine/duplex/test_aura_resumable_hooks.py

@@ -0,0 +1,725 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Public duplex hook: resumable defaults True; ephemeral ids when False."""
+
+from __future__ import annotations
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.duplex.config import DuplexCapabilities
+from vllm_omni.engine.duplex.contracts import (
+    DuplexFence,
+    DuplexStageRequestContext,
+    DuplexStageSubmission,
+    duplex_ephemeral_stage_request_id,
+)
+from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_submission_resumable_default_preserves_minicpm() -> None:
+    ctx = DuplexStageRequestContext(
+        request_id="r",
+        session_id="s",
+        fence=DuplexFence("s"),
+        stage_id=0,
+        final_stage_id=2,
+        config_generation=0,
+        sampling_params=(SamplingParams(max_tokens=1),),
+    )
+    submission = DuplexStageSubmission(
+        context=ctx,
+        prompt={"prompt_token_ids": [1, 2]},
+        already_submitted=False,
+    )
+    assert submission.resumable is True
+
+
+def test_stage_request_id_respects_resumable_flag() -> None:
+    fence = DuplexFence("sess", epoch=3, turn_id=9)
+    resumable_id = DuplexSessionManager.stage_request_id(fence, stage_id=0, resumable=True)
+    ephemeral_id = DuplexSessionManager.stage_request_id(fence, stage_id=0, resumable=False)
+    assert ephemeral_id == duplex_ephemeral_stage_request_id(fence, stage_id=0)
+    assert resumable_id != ephemeral_id
+    assert ephemeral_id.endswith("stage0-turn9")
+    from vllm_omni.engine.duplex.contracts import duplex_turn_id_from_request_id
+
+    assert duplex_turn_id_from_request_id(ephemeral_id) == 9
+    assert duplex_turn_id_from_request_id(resumable_id) is None
+
+
+def test_capabilities_expose_concurrent_turn_requests_default_false() -> None:
+    caps = DuplexCapabilities()
+    assert caps.supports_concurrent_turn_requests is False
+    assert caps.allows_video_without_audio() is False
+    assert caps.required_input_modalities == frozenset({"audio"})
+    assert caps.optional_input_modalities == frozenset({"video"})
+    assert caps.supports_core_resumable_request is False
+    payload = caps.as_dict()
+    assert payload["supports_concurrent_turn_requests"] is False
+    assert payload["required_input_modalities"] == ["audio"]
+    assert payload["optional_input_modalities"] == ["video"]
+
+
+def test_next_commit_allowed_soft_opens_on_r4_release() -> None:
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.session import helpers
+
+    caps = DuplexCapabilities(supports_concurrent_turn_requests=True)
+    session = SimpleNamespace(active_response_id="resp", capabilities=caps)
+    tasks = SimpleNamespace(
+        active_response_task=None,
+        has_response_bound_append_tasks=lambda: False,
+    )
+    # Monkeypatch assistant_playback_active via response_in_progress path: active_response_id set.
+    assert helpers.response_in_progress(session, tasks) is True
+    assert helpers.next_commit_allowed(session, tasks, concurrent_turn_requests_released=False) is False
+    assert helpers.next_commit_allowed(session, tasks, concurrent_turn_requests_released=True) is True
+
+    caps_off = DuplexCapabilities(supports_concurrent_turn_requests=False)
+    session_off = SimpleNamespace(active_response_id="resp", capabilities=caps_off)
+    assert helpers.next_commit_allowed(session_off, tasks, concurrent_turn_requests_released=True) is False
+
+
+def test_active_response_accepts_own_turn_when_concurrent_turn_requests() -> None:
+    """Approach A: each response owns its turn; draining TTS uses request→response map."""
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+    session = DuplexEngineSession(
+        session_id="s",
+        config=DuplexSessionConfig(model="m", modalities=["text"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True),
+    )
+    session.begin_response(turn_id=3)
+    assert session.active_response_accepts_model_turn(3)
+    assert not session.active_response_accepts_model_turn(4)
+    assert not session.active_response_accepts_model_turn(2)
+    session.bind_draining_request("req-old", "resp-old")
+    assert session.response_id_for_request("req-old") == "resp-old"
+    assert session.is_draining_request("req-old")
+
+    # Ending the active response must not wipe older draining bindings.
+    session.bind_draining_request("req-newer-tts", session.active_response_id or "resp-active")
+    ended = session.active_response_id
+    session.end_response(commit_text=False)
+    assert session.is_draining_request("req-old")
+    assert session.response_id_for_request("req-old") == "resp-old"
+    assert not session.is_draining_request("req-newer-tts")
+    session.clear_draining_for_response("resp-old")
+    assert not session.is_draining_request("req-old")
+    del ended
+
+    strict = DuplexEngineSession(
+        session_id="s2",
+        config=DuplexSessionConfig(model="m", modalities=["text"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=False),
+    )
+    strict.begin_response(turn_id=3)
+    assert strict.active_response_accepts_model_turn(3)
+    assert not strict.active_response_accepts_model_turn(4)
+
+
+def test_draining_request_exempt_from_completed_turn_filter() -> None:
+    """Draining ownership must win over the model_turn_id < turn_id late-audio guard."""
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+    session = DuplexEngineSession(
+        session_id="s-drain",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    assert r1 is not None
+    session.bind_draining_request("req-r1-tts", r1)
+    # Overlapping R2 becomes active, then finishes — must not wipe R1 draining.
+    session.begin_response(turn_id=2)
+    session.end_response(commit_text=False)
+    session.turn_id = 2
+    assert session.active_response_id is None
+    assert session.is_draining_request("req-r1-tts")
+    assert session.response_id_for_request("req-r1-tts") == r1
+    draining_response_id = session.response_id_for_request("req-r1-tts")
+    model_turn_id = 1
+    drop = (
+        draining_response_id is None
+        and session.active_response_id is None
+        and model_turn_id is not None
+        and model_turn_id < session.turn_id
+    )
+    assert drop is False
+
+
+def test_ephemeral_turn_id_parser_accepts_stage_turn_ids() -> None:
+    from vllm_omni.engine.duplex.contracts import duplex_turn_id_from_request_id
+
+    assert duplex_turn_id_from_request_id("duplex-s.x.e.0.r.stage0-turn12") == 12
+    assert duplex_turn_id_from_request_id("duplex-s.x.e.0.r.stage2-turn3") == 3
+    assert duplex_turn_id_from_request_id("duplex-s.x.e.0.r.stage0_t9") is None
+    assert duplex_turn_id_from_request_id("duplex-s.x.e.0.r.stage0") is None
+
+
+def test_same_turn_request_ids_skip_other_turns_and_epochs() -> None:
+    from vllm_omni.engine.duplex.contracts import duplex_same_turn_request_ids
+
+    stage3 = "duplex-s.x.e.0.r.stage3-turn4"
+    found = duplex_same_turn_request_ids(
+        stage3,
+        [
+            "duplex-s.x.e.0.r.stage0-turn4",
+            "duplex-s.x.e.0.r.stage1-turn4",
+            "duplex-s.x.e.0.r.stage2-turn4",
+            stage3,
+            "duplex-s.x.e.0.r.stage2-turn5",
+            "duplex-s.x.e.1.r.stage0-turn4",
+        ],
+    )
+    assert found == [
+        "duplex-s.x.e.0.r.stage0-turn4",
+        "duplex-s.x.e.0.r.stage1-turn4",
+        "duplex-s.x.e.0.r.stage2-turn4",
+    ]
+    from vllm_omni.engine.duplex.contracts import duplex_session_id_from_request_id
+
+    fence_session = "sess id"
+    import base64
+
+    encoded = base64.urlsafe_b64encode(fence_session.encode()).decode("ascii").rstrip("=")
+    assert duplex_session_id_from_request_id(f"duplex-s.{encoded}.e.0.r.stage1-turn2") == fence_session
+    assert duplex_session_id_from_request_id("not-a-request") is None
+
+
+def test_stale_keys_skip_already_draining_request_ids() -> None:
+    """Rebinding must not move an older draining Stage2/3 id onto the newer response."""
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+    session = DuplexEngineSession(
+        session_id="s-stale",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    assert r1 is not None
+    session.bind_draining_request("req-r1-talker", r1)
+    session.begin_response(turn_id=2)
+    r2 = session.active_response_id
+    assert r2 is not None and r2 != r1
+    # Simulate the concurrent-turn rebind loop: already-draining ids keep R1.
+    for rid in ("req-r1-talker", "req-r2-talker"):
+        if session.is_draining_request(rid):
+            continue
+        session.bind_draining_request(rid, r2)
+    assert session.response_id_for_request("req-r1-talker") == r1
+    assert session.response_id_for_request("req-r2-talker") == r2
+
+
+def test_end_response_clears_only_own_draining_entries() -> None:
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+    session = DuplexEngineSession(
+        session_id="s-end",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    session.bind_draining_request("req-r1", r1)
+    session.begin_response(turn_id=2)
+    r2 = session.active_response_id
+    session.bind_draining_request("req-r2", r2)
+    session.end_response(commit_text=False)
+    assert session.is_draining_request("req-r1")
+    assert session.response_id_for_request("req-r1") == r1
+    assert not session.is_draining_request("req-r2")
+
+
+def test_on_stage_failure_resolves_draining_response_before_active() -> None:
+    """Mirrors runner.on_stage_failure: draining request_id owns the failed done."""
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+    session = DuplexEngineSession(
+        session_id="s-fail",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    assert r1 is not None
+    session.bind_draining_request("req-r1-talker", r1)
+    session.begin_response(turn_id=2)
+    r2 = session.active_response_id
+    assert r2 is not None and r2 != r1
+
+    request_id = "req-r1-talker"
+    draining_response_id = (
+        session.response_id_for_request(request_id) if session.is_draining_request(request_id) else None
+    )
+    response_id = draining_response_id or session.active_response_id
+    assert response_id == r1
+    assert response_id != r2
+    session.clear_draining_for_response(r1)
+    assert not session.is_draining_request(request_id)
+    assert session.active_response_id == r2
+
+
+def test_concurrent_turn_requests_released_resets_when_new_stage0_binds() -> None:
+    """After R4 release, a new ephemeral Stage0 bind must clear the gate flag."""
+    from types import SimpleNamespace
+
+    run = SimpleNamespace(concurrent_turn_requests_released=True)
+    # Simulate the concurrent-turn rebind arm in model_channel._append_via_data_plane.
+    concurrent_turn = True
+    if concurrent_turn:
+        run.concurrent_turn_requests_released = False
+    assert run.concurrent_turn_requests_released is False
+
+
+def test_draining_stage_ids_come_from_the_plugin() -> None:
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    channel = ModelChannel.__new__(ModelChannel)
+    channel._ctx = SimpleNamespace(
+        plugin=SimpleNamespace(draining_stage_ids=lambda *, stage_count: {stage_count - 1}),
+        stage_port=SimpleNamespace(stage_count=5),
+    )
+    assert channel._draining_stage_ids() == frozenset({4})
+
+    channel._ctx.plugin = SimpleNamespace()
+    assert channel._draining_stage_ids() == frozenset()
+
+
+def test_draining_empty_eos_completes_owning_response_before_shortcut() -> None:
+    """R2 finishing first must not swallow R1's empty Stage3 EOS."""
+    import asyncio
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    session = DuplexEngineSession(
+        session_id="s-empty-eos",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"], extra_body={"auto_response": True}),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True, supports_core_resumable_request=False),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    assert r1 is not None
+    r1_request = "duplex-s.x.e.0.r.stage3-turn1"
+    still_running = "duplex-s.x.e.0.r.stage1-turn2"
+    session.bind_draining_request(r1_request, r1)
+    session.begin_response(turn_id=2)
+    session.end_response(commit_text=False)
+    session.turn_id = 2
+    assert session.active_response_id is None
+
+    class _Plane:
+        def is_terminal(self, request_id: str | None) -> bool:
+            del request_id
+            return False
+
+        def close_stream(self, request_id: str) -> None:
+            del request_id
+
+        def mark_terminal(self, request_id: str) -> None:
+            del request_id
+
+    class _Port:
+        def __init__(self) -> None:
+            self.cleanups: list[list[str]] = []
+
+        async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
+            del abort
+            self.cleanups.append(list(request_ids))
+
+    class _Out:
+        def __init__(self) -> None:
+            self.events: list[dict[str, object]] = []
+
+        def auto_responds(self) -> bool:
+            return True
+
+        def emit(self, payload: dict[str, object]) -> None:
+            self.events.append(payload)
+
+    port = _Port()
+    out = _Out()
+
+    async def _noop(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    async def _schedule(*_args: object, **_kwargs: object) -> bool:
+        return False
+
+    channel = ModelChannel(
+        SimpleNamespace(
+            session=session,
+            plugin=SimpleNamespace(data_plane=_Plane()),
+            stage_port=port,
+            model_state=SimpleNamespace(clear_continuation=lambda: None),
+            services=SimpleNamespace(spawn=lambda *_a, **_k: None),
+        ),
+        out,
+        close_from_runtime=_noop,
+        schedule_silence_continuation=_schedule,
+        abort_request=_noop,
+    )
+    asyncio.run(
+        channel._send_one_model_output_event(
+            {
+                "data_plane_request_id": r1_request,
+                "end_of_turn": True,
+                "model_turn_id": 1,
+                "text": "",
+                "stage_role": "tts",
+            }
+        )
+    )
+    done = [event for event in out.events if event.get("type") == "response.done"]
+    assert done and done[-1].get("response_id") == r1
+    assert port.cleanups == [[r1_request]]
+    assert all(still_running not in ids for ids in port.cleanups)
+    assert not session.is_draining_request(r1_request)
+
+
+def test_draining_completion_drops_finished_response_books_only() -> None:
+    """R1 drain finish drops R1 books and leaves the live R2 snapshot."""
+    import asyncio
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    session = DuplexEngineSession(
+        session_id="s-drain-books",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"], extra_body={"auto_response": True}),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True, supports_core_resumable_request=False),
+    )
+    session.begin_response(turn_id=1)
+    r1 = session.active_response_id
+    assert r1 is not None
+    session.append_assistant_text("from-r1")
+    session.snapshot_active_response_for_drain()
+    r1_request = "duplex-s.x.e.0.r.stage3-turn1"
+    session.bind_draining_request(r1_request, r1)
+    session.begin_response(turn_id=2)
+    r2 = session.active_response_id
+    assert r2 is not None and r2 != r1
+    session.append_assistant_text("from-r2")
+    session.snapshot_active_response_for_drain()
+
+    class _Plane:
+        def is_terminal(self, request_id: str | None) -> bool:
+            del request_id
+            return False
+
+        def close_stream(self, request_id: str) -> None:
+            del request_id
+
+        def mark_terminal(self, request_id: str) -> None:
+            del request_id
+
+    class _Port:
+        async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
+            del request_ids, abort
+
+    class _Out:
+        def __init__(self) -> None:
+            self.events: list[dict[str, object]] = []
+
+        def auto_responds(self) -> bool:
+            return True
+
+        def emit(self, payload: dict[str, object]) -> None:
+            self.events.append(payload)
+
+    out = _Out()
+
+    async def _noop(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    async def _schedule(*_args: object, **_kwargs: object) -> bool:
+        return False
+
+    channel = ModelChannel(
+        SimpleNamespace(
+            session=session,
+            plugin=SimpleNamespace(data_plane=_Plane()),
+            stage_port=_Port(),
+            model_state=SimpleNamespace(clear_continuation=lambda: None),
+            services=SimpleNamespace(spawn=lambda *_a, **_k: None),
+        ),
+        out,
+        close_from_runtime=_noop,
+        schedule_silence_continuation=_schedule,
+        abort_request=_noop,
+    )
+    asyncio.run(
+        channel._send_one_model_output_event(
+            {
+                "data_plane_request_id": r1_request,
+                "end_of_turn": True,
+                "model_turn_id": 1,
+                "text": "",
+                "stage_role": "tts",
+            }
+        )
+    )
+    assert r1 not in session._conversation.assistant_response_snapshots
+    assert r1 not in session._playback.by_response
+    assert f"item_{r1}" not in session._conversation.history_item_placeholders
+    assert r2 in session._conversation.assistant_response_snapshots
+    assert session.active_response_id == r2
+    done = [event for event in out.events if event.get("type") == "response.done"]
+    assert done and done[-1].get("response_id") == r1
+
+
+def test_stale_continue_does_not_close_the_new_response() -> None:
+    import asyncio
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    session = DuplexEngineSession(
+        session_id="s-stale-continue",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"], extra_body={"auto_response": True}),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True, supports_core_resumable_request=False),
+    )
+    session.begin_response(turn_id=2)
+    live = session.active_response_id
+    session.epoch = 1
+
+    class _Out:
+        def __init__(self) -> None:
+            self.events: list[dict[str, object]] = []
+
+        def auto_responds(self) -> bool:
+            return True
+
+        def emit(self, payload: dict[str, object]) -> None:
+            self.events.append(payload)
+
+    out = _Out()
+    cleared: list[bool] = []
+
+    async def _noop(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    async def _schedule(*_args: object, **_kwargs: object) -> bool:
+        return False
+
+    channel = ModelChannel(
+        SimpleNamespace(
+            session=session,
+            plugin=SimpleNamespace(data_plane=SimpleNamespace()),
+            stage_port=SimpleNamespace(),
+            model_state=SimpleNamespace(clear_continuation=lambda: cleared.append(True)),
+            services=SimpleNamespace(spawn=lambda *_a, **_k: None),
+            run=SimpleNamespace(closing=False),
+        ),
+        out,
+        close_from_runtime=_noop,
+        schedule_silence_continuation=_schedule,
+        abort_request=_noop,
+    )
+    asyncio.run(channel.maybe_continue_response(expected_epoch=0))
+    assert session.active_response_id == live
+    assert out.events == []
+    assert cleared == []
+
+    session.bind_draining_request("duplex-s.x.e.1.r.stage3-turn2", live or "")
+    asyncio.run(channel.maybe_continue_response(expected_epoch=1))
+    assert session.active_response_id == live
+    assert out.events == []
+
+    session.clear_draining_requests()
+    asyncio.run(channel.maybe_continue_response(expected_epoch=1))
+    assert session.active_response_id is None
+    assert [event.get("type") for event in out.events] == ["response.done"]
+    assert out.events[0].get("response_id") == live
+
+
+def test_silent_listen_with_continuation_releases_ephemeral_request() -> None:
+    import asyncio
+    from collections.abc import Coroutine
+    from types import SimpleNamespace
+    from typing import Any
+
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    session = DuplexEngineSession(
+        session_id="s-silent-continue",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True, supports_core_resumable_request=False),
+    )
+    session.begin_response(turn_id=1)
+    response_id = session.active_response_id
+    request_id = "duplex-s.x.e.0.r.stage1-turn1"
+    session.bind_request(request_id)
+
+    class _Plane:
+        def __init__(self) -> None:
+            self.terminal: list[str] = []
+
+        def is_terminal(self, request_id: str | None) -> bool:
+            return request_id in self.terminal
+
+        def mark_terminal(self, request_id: str) -> None:
+            self.terminal.append(request_id)
+
+    class _Port:
+        def __init__(self) -> None:
+            self.cleanups: list[list[str]] = []
+
+        async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
+            del abort
+            self.cleanups.append(list(request_ids))
+
+    class _Out:
+        def __init__(self) -> None:
+            self.events: list[dict[str, object]] = []
+
+        def auto_responds(self) -> bool:
+            return False
+
+        def emit(self, payload: dict[str, object]) -> None:
+            self.events.append(payload)
+
+    port = _Port()
+    out = _Out()
+    spawned: list[Coroutine[Any, Any, None]] = []
+
+    async def _noop(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    async def _schedule(*_args: object, **_kwargs: object) -> bool:
+        return False
+
+    channel = ModelChannel(
+        SimpleNamespace(
+            session=session,
+            plugin=SimpleNamespace(data_plane=_Plane()),
+            stage_port=port,
+            model_state=SimpleNamespace(
+                continuation_owner_id=None, continuation_units=0, clear_continuation=lambda: None
+            ),
+            services=SimpleNamespace(spawn=lambda coro, **_k: spawned.append(coro)),
+            run=SimpleNamespace(closing=False),
+        ),
+        out,
+        close_from_runtime=_noop,
+        schedule_silence_continuation=_schedule,
+        abort_request=_noop,
+    )
+    asyncio.run(
+        channel._send_one_model_output_event(
+            {
+                "data_plane_request_id": request_id,
+                "is_listen": True,
+                "end_of_turn": True,
+                "model_turn_id": 1,
+                "reason": "model_listen",
+            }
+        )
+    )
+    assert port.cleanups == [[request_id]]
+    assert session.active_request_id is None
+    assert spawned, "silent listen with unused continuation must schedule maybe_continue"
+    asyncio.run(spawned[0])
+    assert port.cleanups == [[request_id]]
+    assert session.active_response_id is None
+    assert [event.get("type") for event in out.events] == ["response.listen", "response.done"]
+    assert out.events[0].get("response_id") == response_id
+    assert out.events[1].get("response_id") == response_id
+
+
+def test_silent_listen_aborts_data_plane_request_with_full_id() -> None:
+    """abort_data_plane_request must pass [request_id], not a bare str."""
+    import asyncio
+    from types import SimpleNamespace
+
+    from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+    from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+    from vllm_omni.engine.duplex.session.model_channel import ModelChannel
+
+    session = DuplexEngineSession(
+        session_id="s-silent-abort",
+        config=DuplexSessionConfig(model="m", modalities=["text", "audio"]),
+        capabilities=DuplexCapabilities(supports_concurrent_turn_requests=True, supports_core_resumable_request=False),
+    )
+    request_id = "duplex-s.x.e.0.r.stage0-turn7"
+    session.bind_request(request_id)
+    aborts: list[list[str]] = []
+
+    class _Plane:
+        def is_terminal(self, request_id: str | None) -> bool:
+            return False
+
+        def mark_terminal(self, request_id: str) -> None:
+            del request_id
+
+    class _Port:
+        async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
+            del request_ids, abort
+
+    class _Out:
+        def auto_responds(self) -> bool:
+            return True
+
+        def emit(self, payload: dict[str, object]) -> None:
+            del payload
+
+    async def _noop(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    async def _abort(request_ids: list[str], *, notify: bool = False) -> None:
+        del notify
+        aborts.append(list(request_ids))
+
+    channel = ModelChannel(
+        SimpleNamespace(
+            session=session,
+            plugin=SimpleNamespace(data_plane=_Plane()),
+            stage_port=_Port(),
+            model_state=SimpleNamespace(
+                continuation_owner_id=None, continuation_units=0, clear_continuation=lambda: None
+            ),
+            services=SimpleNamespace(spawn=lambda coro, **_k: None),
+            run=SimpleNamespace(closing=False),
+        ),
+        _Out(),
+        close_from_runtime=_noop,
+        schedule_silence_continuation=_noop,
+        abort_request=_abort,
+    )
+    asyncio.run(
+        channel._send_one_model_output_event(
+            {
+                "data_plane_request_id": request_id,
+                "is_listen": True,
+                "end_of_turn": True,
+                "model_turn_id": 1,
+                "reason": "model_listen",
+                "abort_data_plane_request": True,
+            }
+        )
+    )
+    assert aborts == [[request_id]]
+    assert all(len(item) > 1 for batch in aborts for item in batch), "must not split request id into chars"

tests/engine/duplex/test_engine_session.py

@@ -242,6 +242,19 @@ def test_history_commit_with_ack_only_policy_defers_unacknowledged_text():
     assert f"item_{response_id}" in session.pending_history_item_ids
 
 
+def test_history_commit_with_ack_only_policy_keeps_only_acknowledged_prefix():
+    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
+    session.begin_response()
+    session.append_assistant_text("hello world")
+    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
+    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
+    session.acknowledge_playback(played_ms=1_000, committed_ms=1_000)
+
+    committed = session.end_response(commit_text=True)
+
+    assert committed == {"role": "assistant", "content": "hello"}
+
+
 @pytest.mark.parametrize("audio_complete", [False, True])
 def test_unaligned_response_keeps_empty_turn_before_later_user_input(audio_complete):
     session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
@@ -469,6 +482,36 @@ def test_cancel_fence_releases_stage_requests_and_advances_identity():
         session.cancel_fence(DuplexFence("sid-other", epoch=1, turn_id=0), DuplexFence("sid-cancel", epoch=2))
 
 
+def test_request_resource_keys_are_stage_id_and_request_id():
+    session = _session("sid-keys")
+    session.bind_stage_request(0, "req-a", fence=session.fence)
+    session.bind_stage_request(1, "req-b", fence=session.fence)
+    session.bind_stage_request(2, "req-tts", fence=session.fence)
+    stale_keys = list(session.request_resources.keys())
+    stale_ids = list(dict.fromkeys(rid for _, rid in stale_keys))
+    assert stale_ids == ["req-a", "req-b", "req-tts"]
+    for sid, rid in stale_keys:
+        if sid < 2:
+            session.request_resources.pop((sid, rid), None)
+    assert session.resource_request_ids() == ["req-tts"]
+    with pytest.raises(TypeError, match="unhashable"):
+        dict.fromkeys(rid for _, rid in session.request_resources.items())
+
+
+def test_release_resources_for_request_ids_keeps_other_ids_on_the_same_fence():
+    session = _session("sid-drain-release")
+    fence = session.fence
+    session.bind_stage_request(2, "drain", fence=fence)
+    session.bind_stage_request(3, "drain", fence=fence)
+    session.bind_stage_request(0, "live", fence=fence)
+
+    released = session.release_resources_for_request_ids(["drain"])
+
+    assert released == ["drain"]
+    assert session.resource_request_ids() == ["live"]
+    assert session.release_resources_for_request_ids([]) == []
+
+
 # ---- public view ----
 
 
@@ -688,3 +731,65 @@ def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment
 
     assert caps["supports_multi_session"] is False
     assert caps["supports_multi_session_same_replica"] is False
+
+
+def test_draining_response_stays_ack_admissible_after_next_begin_response() -> None:
+    from vllm_omni.engine.duplex.events import ErrorEvent
+    from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack
+
+    session = _session()
+    first = session.begin_response(turn_id=1)
+    session.append_assistant_text("hello")
+    session.mark_audio_sent(400, text_chars=5)
+    session.snapshot_active_response_for_drain()
+    second = session.begin_response(turn_id=2)
+    assert second != first
+    assert session.playback.sent_ms == 0
+    events = apply_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": first,
+            "item_id": f"item_{first}",
+            "played_ms": 100,
+            "committed_ms": 100,
+        },
+    )
+    assert not any(isinstance(event, ErrorEvent) and event.code == "playback_item_not_found" for event in events)
+    session.mark_audio_sent(900, text_chars=5, response_id=first)
+    assert session.playback_for_response(first).sent_ms == 900
+    assert session.playback.sent_ms == 0
+
+
+def test_finished_drain_keeps_sent_audio_ackable() -> None:
+    from vllm_omni.engine.duplex.events import ErrorEvent
+    from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack
+
+    session = _session()
+    first = session.begin_response(turn_id=1)
+    session.append_assistant_text("hello")
+    session.mark_audio_sent(400, text_chars=5)
+    session.snapshot_active_response_for_drain()
+    second = session.begin_response(turn_id=2)
+    session.release_finished_drain_response(first)
+    assert first in session._conversation.assistant_response_snapshots
+    assert f"item_{first}" in session._conversation.history_item_placeholders
+    events = apply_playback_ack(
+        session,
+        {
+            "type": "playback.ack",
+            "response_id": first,
+            "item_id": f"item_{first}",
+            "played_ms": 400,
+            "committed_ms": 400,
+        },
+    )
+    assert not any(isinstance(event, ErrorEvent) for event in events)
+    assert session.active_response_id == second
+
+    silent = session.begin_response(turn_id=3)
+    session.snapshot_active_response_for_drain()
+    session.begin_response(turn_id=4)
+    session.release_finished_drain_response(silent)
+    assert silent not in session._conversation.assistant_response_snapshots
+    assert f"item_{silent}" not in session._conversation.history_item_placeholders

tests/engine/duplex/test_input_transcribed.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.engine.duplex.realtime_events import (
+    RealtimeProjectionState,
+    project_internal_event,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_input_transcribed_is_a_user_transcript_event() -> None:
+    events = project_internal_event(
+        RealtimeProjectionState(session_id="s"),
+        {"type": "input.transcribed", "transcript": "  你好  "},
+    )
+    assert len(events) == 1
+    assert events[0].wire_type == "conversation.item.input_audio_transcription.completed"
+    assert events[0].transcript == "你好"
+
+
+def test_input_transcribed_updates_the_committed_item() -> None:
+    state = RealtimeProjectionState(session_id="s")
+    project_internal_event(
+        state,
+        {
+            "type": "input.committed",
+            "realtime_item_id": "item_user_1",
+            "is_speech": True,
+            "message": {"role": "user"},
+        },
+    )
+    events = project_internal_event(
+        state,
+        {"type": "input.transcribed", "transcript": "你好", "realtime_item_id": "item_user_1"},
+    )
+    assert events[0].item_id == "item_user_1"
+    item = state.conversation_items["item_user_1"]
+    assert item["content"][0]["transcript"] == "你好"
+
+
+def test_blank_input_transcribed_emits_nothing() -> None:
+    events = project_internal_event(
+        RealtimeProjectionState(session_id="s"),
+        {"type": "input.transcribed", "transcript": "  "},
+    )
+    assert events == []

tests/engine/duplex/test_minicpmo_window_plugin.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import base64
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.contracts import DuplexFence
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex import plugin as module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _plugin():
+    return module.MiniCPMO45DuplexPlugin(lambda *args: None)
+
+
+@pytest.mark.asyncio
+async def test_window_configuration_is_prepared_by_plugin(monkeypatch):
+    tokenizer = SimpleNamespace(
+        all_special_ids=[90],
+        unk_token_id=-1,
+        eos_token_id=99,
+        encode=lambda text, add_special_tokens=False: [ord(char) for char in text],
+        convert_tokens_to_ids=lambda token: {"<|listen|>": 91}.get(token, -1),
+    )
+    monkeypatch.setattr(module, "_load_tokenizer", lambda config: tokenizer)
+    config = DuplexSessionConfig(
+        modalities=("text",),
+        extra_body={
+            "sliding_window_mode": "basic",
+            "basic_window_high_tokens": 120,
+            "basic_window_low_tokens": 80,
+        },
+    )
+    runtime = await _plugin().prepare_runtime_config(config, model_config=None)
+    assert runtime["duplex_window_config"]["sliding_window_mode"] == "basic"
+    assert "sliding_window_mode" not in config.extra_body
+    assert runtime["duplex_window_prefix_tokens"] > 0
+    assert runtime["duplex_window_suffix_token_ids"]
+    assert runtime["duplex_window_previous_marker_token_ids"]
+
+
+@pytest.mark.asyncio
+async def test_invalid_window_configuration_is_rejected():
+    config = DuplexSessionConfig(
+        modalities=("text",),
+        extra_body={
+            "basic_window_high_tokens": 80,
+            "basic_window_low_tokens": 80,
+        },
+    )
+    with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc:
+        await _plugin().prepare_runtime_config(config, model_config=None)
+    assert exc.value.code == "invalid_sliding_window_config"
+
+
+def test_window_update_cannot_change_created_configuration():
+    config = DuplexSessionConfig(extra_body={"sliding_window_mode": "basic"})
+    with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc:
+        _plugin().runtime_config_for_update(config, {"duplex_window_config": {"sliding_window_mode": "off"}})
+    assert exc.value.code == "sliding_window_update_unsupported"
+
+
+def test_window_internal_configuration_is_server_owned():
+    with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError):
+        _plugin().validate_client_extra_body({"duplex_window_config": {}})
+
+
+@pytest.mark.parametrize("final", [False, True])
+@pytest.mark.parametrize("seq,samples,expected", [(1, 32000, 16), (1, 48000, 28), (2, 16000, 13)])
+def test_first_and_final_append_reserve_exact_window_input(seq, samples, expected, final):
+    prompt = module.build_duplex_data_plane_prompt(
+        request_id="window-request",
+        fence=DuplexFence("sid", turn_id=1),
+        session_config={},
+        runtime_config={"duplex_first_append_context_tokens": 5},
+        seq=seq,
+        turn_seq=seq,
+        payload={
+            "audio": base64.b64encode(bytes(samples * 4)).decode(),
+            "format": "pcm_f32le",
+            "sample_rate_hz": 16000,
+        },
+        final=final,
+    )
+    assert len(prompt["prompt_token_ids"]) == expected
+
+
+@pytest.mark.parametrize("seq", [1, 2, 3])
+def test_final_exact_chunk_append_reserves_one_unit_not_two(seq):
+    """A final append whose audio is already an exact number of chunks must
+    reserve exactly what Stage0 feeds: one unit per chunk (``<unit>`` plus its
+    audio) and the closure pair for every unit after the first. Serving pads
+    the final residual itself, so the old extra 12-slot silent-unit
+    reservation only put 12 pad positions ahead of the audio in the KV."""
+
+    def _budget(*, seq, final):
+        prompt = module.build_duplex_data_plane_prompt(
+            request_id="final-request",
+            fence=DuplexFence("sid", turn_id=1),
+            session_config={},
+            runtime_config={"duplex_first_append_context_tokens": 0},
+            seq=seq,
+            turn_seq=seq,
+            payload={
+                "audio": base64.b64encode(bytes(16000 * 4)).decode(),
+                "format": "pcm_f32le",
+                "sample_rate_hz": 16000,
+            },
+            final=final,
+        )
+        return len(prompt["prompt_token_ids"])
+
+    assert _budget(seq=seq, final=True) == _budget(seq=seq, final=False)

tests/engine/duplex/test_overlap_policy.py

@@ -79,6 +79,39 @@ def test_an_explicit_client_overlap_action_requests_barge_in() -> None:
     assert decision["reason"] == "client_overlap_action"
 
 
+def test_turn_mode_keeps_silent_vision_follow_during_overlap() -> None:
+    """Vision-follow is is_speech=False + a frame; must not drop while TTS plays."""
+    session = _session()
+    session.capabilities = type(session.capabilities)(
+        required_input_modalities=frozenset({"video"}),
+        optional_input_modalities=frozenset({"audio"}),
+    )
+    payload = {"format": "pcm16", "audio": _pcm16(0.0), "video_frames": ["frame"], "is_speech": False}
+    decision = overlap_policy.decide(session, {"is_speech": False}, payload, auto_responds=False)
+    assert decision["action"] == "listen"
+    assert decision["reason"] == "vision_follow"
+    assert decision["buffer_audio"] is True
+    assert decision["defer_runtime_append"] is False
+    assert decision["force_listen"] is False
+
+
+def test_turn_mode_drops_silent_vision_without_capability() -> None:
+    session = _session()
+    assert session.capabilities.allows_video_without_audio() is False
+    payload = {"format": "pcm16", "audio": _pcm16(0.0), "video_frames": ["frame"], "is_speech": False}
+    decision = overlap_policy.decide(session, {"is_speech": False}, payload, auto_responds=False)
+    assert decision["action"] == "drop"
+    assert decision["reason"] == "silence_or_noise"
+
+
+def test_turn_mode_still_drops_silent_overlap_without_vision() -> None:
+    session = _session()
+    payload = {"format": "pcm16", "audio": _pcm16(0.0), "is_speech": False}
+    decision = overlap_policy.decide(session, {"is_speech": False}, payload, auto_responds=False)
+    assert decision["action"] == "drop"
+    assert decision["reason"] == "silence_or_noise"
+
+
 def test_barge_in_is_deferred_when_the_session_cannot_support_it() -> None:
     session = _session()
     session.capabilities = type(session.capabilities)(supports_barge_in=False)

tests/engine/duplex/test_session_manager.py

@@ -51,6 +51,7 @@ from vllm_omni.engine.duplex.plugin import (
     PcmAppendReservation,
 )
 from vllm_omni.engine.duplex.session import manager as session_manager_module
+from vllm_omni.engine.duplex.session.context import DuplexSessionTasks
 from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
 from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity
 from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
@@ -196,7 +197,11 @@ class FakePlugin(DuplexModelPlugin):
 
     def capabilities(self, *, max_sessions: int) -> DuplexCapabilities:
         del max_sessions
-        return DuplexCapabilities(supports_input_append=True, supports_core_resumable_request=True)
+        # Resident Stage0 ids (MiniCPM-shaped); AURA opts out via supports_core_resumable_request=False.
+        return DuplexCapabilities(
+            supports_input_append=True,
+            supports_core_resumable_request=True,
+        )
 
     def validate_client_extra_body(self, extra_body: object) -> None:
         pass
@@ -409,7 +414,8 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non
         assert result.session_id == "sid-open"
         assert result.lease_generation == 0
         assert result.capabilities == DuplexCapabilities(
-            supports_input_append=True, supports_core_resumable_request=True
+            supports_input_append=True,
+            supports_core_resumable_request=True,
         )
         assert result.public_session is not None
         assert result.public_session["id"] == "sid-open"
@@ -1028,6 +1034,40 @@ async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues(
         gate.set()
 
 
+async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
+    """Manager reserves len(audio)+Σlen(frame); video bytes count toward the same limit."""
+    async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=20) as harness:
+        await harness.open("sid-av")
+        harness.events()
+        session = harness.session("sid-av")
+        runner = harness.manager.runners["sid-av"]
+        gate = asyncio.Event()
+        runner._mailbox.put_nowait(_Internal("wait", {"gate": gate}))
+        runner._on_internal = lambda item: gate.wait()  # type: ignore[method-assign]
+        await asyncio.sleep(0)
+
+        # Default caps require audio; attach video to a non-empty audio unit.
+        frame_a = "aaaa"
+        frame_b = "bbbbbb"
+        audio = b"1234"
+        expected = len(audio) + len(frame_a) + len(frame_b)
+        harness.command(
+            "sid-av",
+            AppendAudio(audio=audio, video_frames=(frame_a, frame_b), event_id="evt-av"),
+        )
+        assert session.pending_input_bytes == expected
+        # Second append that would exceed the limit is backpressured.
+        harness.command(
+            "sid-av",
+            AppendAudio(audio=b"x" * 10, video_frames=("yyyyyyyyyy",), event_id="evt-over"),
+        )
+        errors = [event for event in harness.events("sid-av") if isinstance(event, ErrorEvent)]
+        assert [error.code for error in errors] == ["input_backpressure"]
+        assert errors[0].related_event_id == "evt-over"
+        assert session.pending_input_bytes == expected
+        gate.set()
+
+
 async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds() -> None:
     async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness:
         await harness.open("sid-expired")
@@ -1573,3 +1613,34 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo
     finally:
         shutdown.set()
         await asyncio.wait_for(task, timeout=5.0)
+
+
+# --------------------------------------------------------------------------- #
+# Cancellation waits                                                          #
+# --------------------------------------------------------------------------- #
+
+
+async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> None:
+    """The wait after cancelling is allowed to time out; that is not an error.
+
+    On Python 3.10 the timeout arrives as ``asyncio.TimeoutError``, which is a
+    different class from the builtin ``TimeoutError`` until 3.11.
+    """
+
+    async def outlives_the_first_cancel() -> bool:
+        try:
+            await asyncio.sleep(5)
+        except asyncio.CancelledError:
+            await asyncio.sleep(1)
+            raise
+        return True
+
+    tasks = DuplexSessionTasks()
+    task = asyncio.create_task(outlives_the_first_cancel())
+    await asyncio.sleep(0)
+    tasks.track_append_task(task, epoch=0, final=False, response_bound=False)
+    tasks.append_tail = task
+
+    assert await tasks.cancel_append_tasks(timeout_s=0.05) is True
+    assert tasks.append_tail is None
+    await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0)

tests/engine/duplex/test_session_runner.py

@@ -180,7 +180,10 @@ class Harness:
             segment_token_ids=tuple(segment_token_ids),
             segment_output_metadata=dict(segment_output_metadata or {}),
         )
-        return self.runner.on_stage_output(stage_id, output, metrics, request_id=output.request_id, context=context)
+        request_id = getattr(output, "request_id", None)
+        if not isinstance(request_id, str):
+            raise AssertionError("stage output is missing request_id")
+        return self.runner.on_stage_output(stage_id, output, metrics, request_id=request_id, context=context)
 
     async def deliver_and_settle(self, output: object, **kwargs: Any) -> list[DuplexEvent]:
         self.deliver(output, **kwargs)
@@ -520,6 +523,28 @@ async def test_commit_without_audio_is_rejected() -> None:
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_empty_silent_commit_acks_without_opening_response() -> None:
+    """#3: empty is_speech=False commit keeps the silent-ack pair, even auto-respond."""
+    h = await open_harness(auto_response=True)
+    try:
+        # resolve_commit only admits empty silent commits after non-speech was seen.
+        projector = h.runner._require_projector()
+        projector.input_audio_buffer_had_non_speech = True
+        events = await h.run(commands.Commit())
+        assert "error" not in types(events), types(events)
+        assert "input_audio_buffer.committed" in types(events), types(events)
+        assert "response.listen" in types(events), types(events)
+        listen = next(event for event in events if event.type == "response.listen")
+        assert listen.details["reason"] == "silence_or_noise"
+        committed = next(event for event in events if event.type == "input_audio_buffer.committed")
+        assert committed.details.get("empty") is True
+        assert committed.details.get("no_response") is True
+        assert h.port.submissions == []
+    finally:
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_commit_projects_the_user_item_and_does_not_resubmit_consumed_audio() -> None:
     h = await open_harness()
@@ -674,6 +699,45 @@ async def test_stale_epoch_output_is_dropped_after_barge_in() -> None:
         await close_harness(h)
 
 
+async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> None:
+    h = await open_harness()
+    try:
+        await h.run(append_audio())
+        request_id = h.stage0_request_id()
+        await h.deliver_and_settle(tts_output(request_id, samples=24000, text="he"))
+        h.session.bind_draining_request("duplex-drain-tts", "resp-old")
+        # Draining TTS sits on an older turn fence. cancel_fence only drops
+        # the fence being cancelled, so these bindings must be popped by id.
+        from vllm_omni.engine.duplex.contracts import DuplexFence
+        from vllm_omni.engine.duplex.session.engine_session import DuplexRequestResource
+
+        older = DuplexFence(h.session.session_id, epoch=h.session.epoch, turn_id=h.session.turn_id + 5)
+        h.session.request_resources[(2, "duplex-drain-tts")] = DuplexRequestResource(
+            stage_id=2, request_id="duplex-drain-tts", fence=older, submitted=True
+        )
+        h.session.request_resources[(3, "duplex-drain-tts")] = DuplexRequestResource(
+            stage_id=3, request_id="duplex-drain-tts", fence=older, submitted=True
+        )
+        h.session.request_resources[(1, "still-live")] = DuplexRequestResource(
+            stage_id=1, request_id="still-live", fence=older, submitted=True
+        )
+        events = await h.run(commands.BargeIn())
+        assert h.port.aborts == [[request_id, "duplex-drain-tts"]]
+        assert not h.session.is_draining_request("duplex-drain-tts")
+        done_ids = [
+            event.response["id"]
+            for event in events
+            if getattr(event, "wire_type", "") == "response.done" and isinstance(getattr(event, "response", None), dict)
+        ]
+        assert "resp-old" in done_ids
+        assert (2, "duplex-drain-tts") not in h.session.request_resources
+        assert (3, "duplex-drain-tts") not in h.session.request_resources
+        assert (0, request_id) not in h.session.request_resources
+        assert (1, "still-live") in h.session.request_resources
+    finally:
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_cancel_response_without_active_response_is_rejected() -> None:
     h = await open_harness()
@@ -707,6 +771,41 @@ async def test_cancel_response_aborts_the_stage_request_and_reports_playback() -
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_cancel_response_absorbs_a_response_task_that_outlives_the_wait() -> None:
+    """The 0.25 s wait after cancelling the response task may time out.
+
+    On Python 3.10 that timeout is ``asyncio.TimeoutError``, not the builtin
+    ``TimeoutError``, and it must not stop the cancel from completing.
+    """
+
+    async def outlives_the_first_cancel() -> None:
+        try:
+            await asyncio.sleep(5)
+        except asyncio.CancelledError:
+            await asyncio.sleep(1)
+            raise
+
+    h = await open_harness()
+    task = asyncio.create_task(outlives_the_first_cancel())
+    try:
+        await h.run(append_audio())
+        request_id = h.stage0_request_id()
+        await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"))
+        h.runner.tasks.active_response_task = task
+
+        events = await h.run(commands.CancelResponse())
+        # The runner is still inside that wait when the mailbox goes quiet.
+        events += await h.settle(idle_s=0.6, timeout_s=5.0)
+        assert find(events, "response.done").status == "cancelled"
+        assert "error" not in types(events)
+        assert h.session.epoch == 1
+    finally:
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None:
     h = await open_harness()
@@ -778,6 +877,46 @@ async def test_listen_decision_on_a_resumable_request_closes_the_bounded_respons
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_direct_response_listen_still_emits_response_done_after_continuation_clear() -> None:
+    """AURA-like silent DIRECT_RESPONSE: after continuation budget is spent, emit response.done.
+
+    Non-resumable turn-commit still closes the active response on a listen/direct
+    decision so the client is not left without a terminal event.
+    """
+    from dataclasses import replace
+
+    h = await open_harness()
+    try:
+        h.session.capabilities = replace(h.session.capabilities, supports_core_resumable_request=False)
+        await h.run(append_audio())
+        assert h.port.submissions, "ephemeral Stage0 must submit"
+        request_id = h.port.submissions[-1].context.request_id
+        assert "-turn" in request_id
+        await h.deliver_and_settle(tts_output(request_id, samples=24000, text="ok"))
+        response_id = h.session.active_response_id
+        assert response_id is not None
+        model_state = h.runner.model_state
+        model_state.continuation_owner_id = f"response:{response_id}"
+        model_state.continuation_units = h.runner.model._AUTO_RESPONSE_MAX_CONTINUATION_UNITS
+
+        listen = listen_output(request_id)
+        listen.finished = True
+        events = await h.deliver_and_settle(
+            listen,
+            stage_id=0,
+            segment_finished=True,
+            segment_token_ids=[11, 12, LISTEN_TOKEN_ID],
+            segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID},
+        )
+        assert "response.listen" in types(events)
+        done = find(events, "response.done")
+        assert done.response_id == response_id
+        assert h.session.active_response_id is None
+    finally:
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_listen_decision_on_a_resumable_request_keeps_the_turn_going() -> None:
     """Same unfinished listen, budget left: it must schedule the next unit.
@@ -860,6 +999,38 @@ async def test_stage_failure_fails_the_active_response() -> None:
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_stage_failure_on_draining_request_fails_that_response_only() -> None:
+    """Overlapped: a draining older request failure must not wipe the active response."""
+    from dataclasses import replace
+
+    h = await open_harness()
+    try:
+        h.session.capabilities = replace(h.session.capabilities, supports_concurrent_turn_requests=True)
+        await h.run(append_audio())
+        request_id = h.stage0_request_id()
+        await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"))
+        r1 = h.session.active_response_id
+        assert r1 is not None
+        draining_req = "req-r1-talker-drain"
+        h.session.bind_draining_request(draining_req, r1)
+        # Open a newer active response while R1 TTS is still draining.
+        r2 = h.session.begin_response(turn_id=(h.session.turn_id or 0) + 1)
+        assert h.session.active_response_id == r2
+        assert r2 != r1
+
+        h.runner.on_stage_failure(2, RuntimeError("drain-boom"), request_id=draining_req)
+        events = await h.settle()
+        assert types(events)[0] == "error"
+        done = find(events, "response.done")
+        assert done.response_id == r1
+        assert done.status == "failed"
+        assert h.session.active_response_id == r2
+        assert not h.session.is_draining_request(draining_req)
+    finally:
+        await close_harness(h)
+
+
 # --------------------------------------------------------------------------- #
 # Turn mode (no auto response)                                                #
 # --------------------------------------------------------------------------- #
@@ -877,6 +1048,91 @@ async def test_turn_mode_skips_silent_chunks() -> None:
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_turn_mode_keeps_silent_chunks_with_video_frames() -> None:
+    """Engine vision-follow buffers silent+frames; MiniCPM commit stays speech-gated."""
+    from dataclasses import replace
+
+    h = await open_harness(auto_response=False)
+    try:
+        h.session.capabilities = replace(
+            h.session.capabilities,
+            required_input_modalities=frozenset({"video"}),
+            optional_input_modalities=frozenset({"audio"}),
+        )
+        # Minimal 1x1 JPEG (base64) — wire validation only checks non-empty str.
+        frame = (
+            "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkS"
+            "Ew8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJ"
+            "CQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIy"
+            "MjIyMjIyMjIyMjIyMjL/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAA"
+            "AAAAAAAAAAj/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAA"
+            "AAAAAAD/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIQAxAAAAGfAP/EABQQ"
+            "AQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAQUCf//EABQRAQAAAAAAAAAAAAAAAAAA"
+            "AAD/2gAIAQMBAT8Bf//EABQRAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQIBAT8Bf//E"
+            "ABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEABj8Cf//EABQQAQAAAAAAAAAAAAAA"
+            "AAAAAAD/2gAIAQEAAT8hf//Z"
+        )
+        cmd = commands.AppendAudio(
+            audio=pcm_f32(160, value=0.0),
+            format="pcm_f32le",
+            sample_rate_hz=16000,
+            is_speech=False,
+            video_frames=(frame,),
+        )
+        events = await h.run(cmd)
+        assert "response.listen" not in types(events)
+        assert h.runner.model_state.audio_buffer.has_pending()
+        events = await h.run(commands.Commit(create_response=True))
+        # MiniCPM prepare_commit is speech-gated (no frames-only Stage0).
+        assert "input_audio_buffer.committed" in types(events), types(events)
+        assert h.port.submissions == []
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_turn_mode_keeps_silent_vision_while_response_in_progress() -> None:
+    """TTS still playing: vision-follow must buffer, not drop as silence_or_noise."""
+    from dataclasses import replace
+
+    h = await open_harness(auto_response=False)
+    try:
+        h.session.capabilities = replace(
+            h.session.capabilities,
+            required_input_modalities=frozenset({"video"}),
+            optional_input_modalities=frozenset({"audio"}),
+        )
+        h.session._response.active_response_id = "resp-tts"
+        frame = (
+            "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkS"
+            "Ew8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJ"
+            "CQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIy"
+            "MjIyMjIyMjIyMjIyMjL/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAA"
+            "AAAAAAAAAAj/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAA"
+            "AAAAAAD/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIQAxAAAAGfAP/EABQQ"
+            "AQAAAAAAAAAAAAAAAAAAAAD/2gAIAQMBAT8Bf//EABQRAQAAAAAAAAAAAAAAAAAA"
+            "AAD/2gAIAQIBAT8Bf//EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEABj8Cf//E"
+            "ABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAT8hf//Z"
+        )
+        events = await h.run(
+            commands.AppendAudio(
+                audio=pcm_f32(160, value=0.0),
+                format="pcm_f32le",
+                sample_rate_hz=16000,
+                is_speech=False,
+                video_frames=(frame,),
+            )
+        )
+        listen_reasons = [
+            getattr(event, "details", {}) or {} for event in events if getattr(event, "type", None) == "response.listen"
+        ]
+        assert all(details.get("reason") != "silence_or_noise" for details in listen_reasons)
+        assert h.runner.model_state.audio_buffer.has_pending()
+    finally:
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_turn_mode_commit_with_response_create_starts_one_response() -> None:
     h = await open_harness(auto_response=False)
@@ -1033,7 +1289,7 @@ async def test_conversation_items_can_be_injected_and_deleted() -> None:
         await close_harness(h)
 
 
-def _stage_metrics_of(event: object) -> dict[str, dict[str, object]]:
+def _stage_metrics_of(event: DuplexEvent) -> dict[str, dict[str, object]]:
     """Per-stage engine metrics as the client reads them off one wire event."""
     payload = event.to_realtime()
     metadata = payload.get("metadata")
@@ -1081,6 +1337,40 @@ async def test_stage0_metrics_reach_the_response_even_though_its_output_feeds_tt
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_projected_stage1_metrics_reach_the_spoken_audio_event() -> None:
+    """Thinker text is projected, but its TTFT/TPOT still have to ride the audio event."""
+    h = await open_harness(stage_count=3)
+    try:
+        plugin = h.runner.model._ctx.plugin
+
+        def project_stage1(*, stage_id: int, output: object, context: object) -> bool:
+            del output, context
+            return stage_id == 1
+
+        plugin.project_intermediate_output = project_stage1
+        await h.run(append_audio())
+        request_id = h.stage0_request_id()
+        thinker = SimpleNamespace(
+            request_id=request_id,
+            finished=True,
+            outputs=[SimpleNamespace(text="hello", token_ids=[7, 8], multimodal_output={})],
+            multimodal_output={},
+        )
+        forwarded = h.deliver(
+            thinker,
+            stage_id=1,
+            segment_finished=False,
+            metrics=stage_stats(stage_id=1, request_id=request_id, num_tokens_out=2, itls_ms=[6.0]),
+        )
+        assert forwarded is False
+        events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"), stage_id=2)
+        stage_metrics = _stage_metrics_of(find(events, "response.output_audio.delta"))
+        assert stage_metrics["1"]["num_tokens_out"] == 2
+    finally:
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_stage0_metrics_from_several_units_are_summed_into_one_response() -> None:
     """Two pass-through segments before the first audio are one response's tokens."""

tests/engine/duplex/test_session_runner_ephemeral_bind.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Ephemeral Stage0 bind uses r.stage0-turn{N}, not resumable r.stage0."""
+
+from __future__ import annotations
+
+from vllm_omni.engine.duplex.contracts import (
+    DuplexFence,
+    duplex_ephemeral_stage_request_id,
+    duplex_resource_request_id,
+)
+from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
+
+
+def test_stage0_request_id_ephemeral_when_not_resumable() -> None:
+    fence = DuplexFence("duplex-test", epoch=0, turn_id=3)
+    rid = DuplexSessionManager.stage_request_id(fence, stage_id=0, resumable=False)
+    expected = duplex_ephemeral_stage_request_id(fence, stage_id=0)
+    assert rid == expected
+    assert rid.endswith(".r.stage0-turn3")
+    assert not rid.endswith(".r.stage0")
+
+
+def test_stage0_request_id_resumable_keeps_stable_role() -> None:
+    fence = DuplexFence("duplex-test", epoch=1, turn_id=3)
+    rid = DuplexSessionManager.stage_request_id(fence, stage_id=0, resumable=True)
+    expected = duplex_resource_request_id(fence, "stage0")
+    assert rid == expected
+    assert rid.endswith(".r.stage0")

tests/engine/test_arg_utils.py

@@ -22,6 +22,7 @@ from pydantic import ValidationError
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 from vllm.engine.arg_utils import EngineArgs
 
+from tests.helpers.mock import patch_hf_snapshot_download
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.stage_init_utils import build_engine_args_dict
@@ -296,7 +297,7 @@ def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, m
     baseline_config = Mock()
     warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning")
 
-    monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", lambda *args, **kwargs: str(tmp_path))
+    patch_hf_snapshot_download(monkeypatch, lambda *args, **kwargs: str(tmp_path), hf_home=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(

tests/engine/test_async_omni_engine_stage_init.py

@@ -14,6 +14,7 @@ import pytest
 from omegaconf import OmegaConf
 from vllm.v1.engine.utils import EngineZmqAddresses
 
+from tests.helpers.mock import patch_hf_snapshot_download
 from vllm_omni.config.omni_config import OmniStageRuntimeConfig
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine import omni_engine_base as async_omni_engine_module
@@ -1566,22 +1567,20 @@ def test_build_engine_args_pulls_stage_subdirs_missing_from_cached_snapshot(monk
     never materialized. Stage init must fetch exactly those subfolders instead
     of joining onto a path that exists nowhere.
     """
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
     _make_snapshot(snapshot, ["tokenizer"])
     calls = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         calls.append(kwargs)
         if kwargs.get("local_files_only"):
             return str(snapshot)
         _make_snapshot(snapshot, ["language_model"])
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1601,18 +1600,16 @@ def test_build_engine_args_pulls_stage_subdirs_missing_from_cached_snapshot(monk
 
 def test_build_engine_args_skips_hub_call_when_cached_snapshot_is_complete(monkeypatch, tmp_path):
     """A warm cache stays offline-friendly: no outgoing Hub request."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
     _make_snapshot(snapshot, ["language_model", "tokenizer"])
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         assert kwargs.get("local_files_only"), "warm cache must not reach the Hub"
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1633,8 +1630,6 @@ def test_build_engine_args_redownloads_a_partial_subdir(monkeypatch, tmp_path):
     that local path leaves vLLM's loader without any fallback for the missing
     weights.
     """
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
@@ -1644,14 +1639,14 @@ def test_build_engine_args_redownloads_a_partial_subdir(monkeypatch, tmp_path):
     _make_snapshot(snapshot, ["tokenizer"])
     calls = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         calls.append(kwargs)
         if kwargs.get("local_files_only"):
             return str(snapshot)
         _make_snapshot(snapshot, ["language_model"])
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1672,8 +1667,6 @@ def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypat
     The index downloads early; accepting it as proof of weights converts the
     Hub ID into a local path vLLM cannot fetch the remaining shards for.
     """
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
@@ -1693,13 +1686,13 @@ def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypat
     _make_snapshot(snapshot, ["tokenizer"])
     calls = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         calls.append(kwargs)
         if not kwargs.get("local_files_only"):
             (partial / "model-00002-of-00002.safetensors").write_text("x")
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1716,8 +1709,6 @@ def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypat
 
 def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, tmp_path):
     """Shard-named weights always ship an index; one without it is partial."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
@@ -1727,7 +1718,7 @@ def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, t
     _make_snapshot(snapshot, ["tokenizer"])
     calls = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         calls.append(kwargs)
         if not kwargs.get("local_files_only"):
             shards = {f"w{i}.weight": f"model-0000{i}-of-00004.safetensors" for i in range(1, 5)}
@@ -1736,7 +1727,7 @@ def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, t
                 (partial / shard).write_text("x")
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1752,8 +1743,6 @@ def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, t
 
 def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(monkeypatch, tmp_path):
     """Templates and configs download first; alone they are not a tokenizer."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
@@ -1763,13 +1752,13 @@ def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(mon
     (partial / "chat_template.jinja").write_text("{{ messages }}")
     calls = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         calls.append(kwargs)
         if not kwargs.get("local_files_only"):
             _make_snapshot(snapshot, ["tokenizer"])
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1787,19 +1776,17 @@ def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(mon
 def test_build_engine_args_forwards_revision_and_download_dir(monkeypatch, tmp_path):
     """revision/download_dir must shape snapshot selection (they cannot be
     corrected downstream once the repo ID is a local path)."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "pinned"
     _make_snapshot(snapshot, ["language_model", "tokenizer"])
     seen = []
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         seen.append(kwargs)
         return str(snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1824,8 +1811,6 @@ def test_build_engine_args_forwards_revision_and_download_dir(monkeypatch, tmp_p
 
 def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, tmp_path):
     """A tokenizer pinned to a different revision resolves against its own snapshot."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     model_snapshot = tmp_path / "snapshots" / "model-rev"
@@ -1833,12 +1818,12 @@ def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, t
     tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev"
     _make_snapshot(tokenizer_snapshot, ["tokenizer"])
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         if kwargs.get("revision") == "tok-rev":
             return str(tokenizer_snapshot)
         return str(model_snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1861,8 +1846,6 @@ def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, t
 def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypatch, tmp_path):
     """An empty tokenizer_subdir targets the snapshot root and still honors
     its own revision."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     model_snapshot = tmp_path / "snapshots" / "model-rev"
@@ -1875,12 +1858,12 @@ def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypat
     tokenizer_snapshot.mkdir(parents=True, exist_ok=True)
     (tokenizer_snapshot / _SUBDIR_ARTIFACT["tokenizer"]).write_text("x")
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         if kwargs.get("revision") == "tok-rev":
             return str(tokenizer_snapshot)
         return str(model_snapshot)
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1902,19 +1885,17 @@ def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypat
 
 def test_build_engine_args_fails_closed_when_subdir_cannot_be_downloaded(monkeypatch, tmp_path):
     """An undownloadable subfolder raises here instead of reaching HuggingFace."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
     snapshot = tmp_path / "snapshots" / "deadbeef"
     _make_snapshot(snapshot, ["tokenizer"])
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         if kwargs.get("local_files_only"):
             return str(snapshot)
         raise OSError("offline")
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -1927,16 +1908,14 @@ def test_build_engine_args_fails_closed_when_subdir_cannot_be_downloaded(monkeyp
         build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
 
 
-def test_build_engine_args_fails_closed_on_cold_cache_instead_of_joining_repo_id(monkeypatch):
+def test_build_engine_args_fails_closed_on_cold_cache_instead_of_joining_repo_id(monkeypatch, tmp_path):
     """With nothing cached, the stage must not join a subdir onto the repo id."""
-    from huggingface_hub import HfApi
-
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 
-    def fake_snapshot_download(self, repo_id, **kwargs):
+    def fake_snapshot_download(repo_id, **kwargs):
         raise OSError("offline")
 
-    monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+    patch_hf_snapshot_download(monkeypatch, fake_snapshot_download, hf_home=tmp_path)
 
     stage_cfg = types.SimpleNamespace(
         stage_id=0,

tests/engine/test_duplex_orchestrator.py

@@ -55,6 +55,7 @@ def _build(
     stages: int = 1,
     running_counter: FakeRunningCounter | None = None,
     runtime_config: DuplexSessionRuntimeConfig | None = None,
+    plugin=None,
 ) -> tuple[DuplexOrchestrator, list[FakeStageClient], asyncio.Queue, asyncio.Queue]:
     clients = [FakeStageClient(stage_type="llm", final_output=index == stages - 1) for index in range(stages)]
     pools = _build_stage_pools(
@@ -70,7 +71,7 @@ def _build(
         rpc_async_queue=rpc_q,
         stage_pools=pools,
         running_counter=running_counter,
-        plugin=MiniCPMO45DuplexPlugin(_encode_audio),
+        plugin=plugin or MiniCPMO45DuplexPlugin(_encode_audio),
         duplex_session_config=runtime_config or DuplexSessionRuntimeConfig(reaper_interval_s=0.01),
         model_config=None,
     )
@@ -197,7 +198,8 @@ async def test_open_preregisters_the_stage0_request_and_close_releases_it() -> N
     assert request_state.session_owned is True
     assert request_state.session_id == SESSION_ID
     assert request_state.fence == DuplexFence(SESSION_ID)
-    assert request_state.streaming.enabled is True
+    # Set on submit from submission.resumable, not again at preregister.
+    assert request_state.streaming.enabled is False
     # Preregistration reserves the id; nothing is running until an append submits.
     assert counter.value == 0
     assert clients[0].add_request_calls == []
@@ -211,6 +213,18 @@ async def test_open_preregisters_the_stage0_request_and_close_releases_it() -> N
     assert orchestrator.session_manager.active_count() == 0
 
 
+@pytest.mark.asyncio
+async def test_aura_ephemeral_preregister_disables_streaming() -> None:
+    from vllm_omni.model_executor.models.aura_omni.duplex.plugin import AuraDuplexPlugin
+
+    orchestrator, _, rpc_q, _ = _build(plugin=AuraDuplexPlugin(_encode_audio))
+    result = await _open(orchestrator, rpc_q)
+    assert result.ok is True
+    request_state = next(iter(orchestrator.request_states.values()))
+    assert request_state.streaming.enabled is False
+    await orchestrator.session_manager.shutdown()
+
+
 @pytest.mark.asyncio
 async def test_open_failure_rolls_back_session_and_reserved_request() -> None:
     orchestrator, _, rpc_q, _ = _build()
@@ -275,6 +289,7 @@ async def test_append_submits_the_resumable_stage0_request_and_counts_it_running
     assert submitted.resumable is True
     assert submitted.sampling_params.max_tokens == 20
     request_state = orchestrator.request_states[request_id]
+    assert request_state.streaming.enabled is True
     assert request_state.stage_fences[0] == DuplexFence(SESSION_ID)
     assert 0 in request_state.stage_submit_ts
     assert counter.value == 1
@@ -390,6 +405,31 @@ async def test_forward_failure_closes_the_owning_session() -> None:
     assert types[-1] == "session.expired"
 
 
+@pytest.mark.asyncio
+async def test_aura_forward_failure_keeps_the_session() -> None:
+    from vllm_omni.model_executor.models.aura_omni.duplex.plugin import AuraDuplexPlugin
+
+    orchestrator, clients, rpc_q, output_q = _build(stages=2, plugin=AuraDuplexPlugin(_encode_audio))
+    await _open(orchestrator, rpc_q)
+    request_id = next(iter(orchestrator.request_states))
+    await _submit(orchestrator, _append_audio())
+    request_state = orchestrator.request_states[request_id]
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None
+
+    absorbed = await orchestrator._handle_forward_failure(request_id, 1, request_state, ValueError("stale talker"))
+    await _settle(orchestrator)
+
+    assert absorbed is True
+    assert request_id not in orchestrator.request_states
+    assert SESSION_ID in orchestrator.session_manager.runners
+    assert session.state != DuplexSessionState.CLOSED
+    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    assert "error" in types
+    assert "session.expired" not in types
+    await orchestrator.session_manager.shutdown()
+
+
 @pytest.mark.asyncio
 async def test_request_cleanup_failure_is_deferred_and_retried_by_the_reaper(monkeypatch) -> None:
     orchestrator, clients, rpc_q, _ = _build()
@@ -560,3 +600,53 @@ async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(m
         assert not orchestrator.request_states[submitted.request_id].streaming.enabled
     finally:
         await _close(orchestrator, rpc_q)
+
+
+@pytest.mark.asyncio
+async def test_sentence_partial_does_not_legacy_forward_the_full_stage_output() -> None:
+    from vllm_omni.engine.duplex.plugin import PartialStageForward
+
+    orchestrator, *_ = _build(stages=4)
+    forwarded: list[object] = []
+
+    async def record_forward(req_id, stage_id, output, req_state, **kwargs):
+        del req_id, stage_id, req_state, kwargs
+        forwarded.append(output)
+
+    async def no_intercept(*args, **kwargs):
+        del args, kwargs
+        return False
+
+    orchestrator._forward_to_next_stage = record_forward  # type: ignore[method-assign]
+    orchestrator._intercept_stage_output = no_intercept  # type: ignore[method-assign]
+    orchestrator.async_chunk = True
+    orchestrator._stage_receives_async_chunks = lambda stage_id: False  # type: ignore[method-assign]
+
+    full = SimpleNamespace(request_id="r", finished=True, text="FULL")
+    chunk = SimpleNamespace(request_id="r", finished=True, text="CHUNK")
+    req_state = DuplexOrchestratorRequestState(
+        request_id="r",
+        final_stage_id=3,
+        session_owned=True,
+        sampling_params_list=[SimpleNamespace() for _ in range(4)],
+    )
+    orchestrator.plugin.plan_partial_stage_output = lambda *args, **kwargs: PartialStageForward(  # type: ignore[method-assign]
+        output=chunk, is_final_update=True
+    )
+    orchestrator.plugin.partial_stage_followup = lambda *args, **kwargs: None  # type: ignore[method-assign]
+
+    await orchestrator._route_output(1, 0, full, req_state, None)
+
+    assert forwarded == [chunk]
+    assert req_state.skip_legacy_stage_forward is False
+
+    forwarded.clear()
+    orchestrator.plugin.plan_partial_stage_output = lambda *args, **kwargs: None  # type: ignore[method-assign]
+    await orchestrator._route_output(1, 0, full, req_state, None)
+
+    assert forwarded == [full]
+
+
+def test_default_plugin_declares_no_draining_stages() -> None:
+    plugin = MiniCPMO45DuplexPlugin(_encode_audio)
+    assert plugin.draining_stage_ids(stage_count=4) == frozenset()

tests/engine/test_orchestrator_error_handling.py

@@ -13,6 +13,7 @@ import asyncio
 import queue
 import time
 from types import SimpleNamespace
+from unittest.mock import AsyncMock
 
 import janus
 import pytest
@@ -22,6 +23,7 @@ from vllm.v1.serial_utils import MsgpackEncoder
 
 from vllm_omni.engine.messages import (
     AddCompanionRequestMessage,
+    CollectiveRPCRequestMessage,
     EngineQueueMessage,
     ErrorMessage,
     ShutdownRequestMessage,
@@ -1143,3 +1145,64 @@ async def test_diffusion_client_error_output_propagates_non_400_status(
     finally:
         orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
         orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal() -> None:
+    """A control RPC that fails on one replica must come back as that
+    replica's error result instead of raising out of _request_handler.
+    """
+
+    class _FailingClient(FakeStageClient):
+        async def collective_rpc_async(self, *args, **kwargs):
+            raise TimeoutError("outputs did not drain")
+
+    ok = FakeStageClient(stage_type="diffusion")
+    ok.collective_rpc_async = AsyncMock(return_value=None)
+    orchestrator, queues = _build_bare_orchestrator(
+        _build_stage_pools([[ok], [_FailingClient(stage_type="diffusion")]])
+    )
+    try:
+        await orchestrator._handle_collective_rpc(
+            CollectiveRPCRequestMessage(
+                rpc_id="rpc-1",
+                method="pause_scheduler",
+                args=(),
+                kwargs={"mode": "keep"},
+                stage_ids=[0, 1],
+            )
+        )
+
+        result = queues[2].async_q.get_nowait()
+        assert result.rpc_id == "rpc-1"
+        assert result.stage_ids == [0, 1]
+        assert result.results[0] is None
+        assert result.results[1]["supported"] is False
+        assert "outputs did not drain" in result.results[1]["error"]
+    finally:
+        for q in queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+async def test_rpc_failure_capture_is_limited_to_pause_and_resume() -> None:
+    """Only pause/resume failures are reported as a result; every other RPC
+    keeps propagating the way it does today.
+    """
+    orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="diffusion")]]))
+    orchestrator.stage_pools[0].collective_rpc = AsyncMock(side_effect=TimeoutError("worker died"))
+
+    def _msg(rpc_id: str, method: str) -> CollectiveRPCRequestMessage:
+        return CollectiveRPCRequestMessage(rpc_id=rpc_id, method=method, args=(), kwargs={}, stage_ids=[0])
+
+    try:
+        await orchestrator._handle_collective_rpc(_msg("rpc-resume", "resume_scheduler"))
+        result = queues[2].async_q.get_nowait()
+        assert result.results[0]["supported"] is False
+        assert "worker died" in result.results[0]["error"]
+
+        with pytest.raises(TimeoutError):
+            await orchestrator._handle_collective_rpc(_msg("rpc-sleep", "sleep"))
+    finally:
+        for q in queues:
+            q.close()

tests/engine/test_orchestrator_event_driven.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
 """Event-driven orchestration loop (``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``) tests.
 
 Parity suite: re-runs the legacy orchestration scenarios from
@@ -19,7 +19,10 @@ import pytest
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.messages import ShutdownRequestMessage
-from vllm_omni.engine.orchestrator import _event_driven_orch_enabled
+from vllm_omni.engine.orchestrator import (
+    _event_driven_orch_default_for_pipeline,
+    _event_driven_orch_enabled,
+)
 
 from . import test_orchestrator as legacy
 from . import test_orchestrator_error_handling as legacy_errors
@@ -125,6 +128,19 @@ def test_flag_parsing(monkeypatch) -> None:
         assert _event_driven_orch_enabled() is False
 
 
+def test_only_qwen3_tts_has_a_pipeline_default() -> None:
+    assert _event_driven_orch_default_for_pipeline("qwen3_tts") is True
+    for model_type in (None, "qwen3_omni_moe", "minicpmo_4_5", "moss_tts_delay"):
+        assert _event_driven_orch_default_for_pipeline(model_type) is False
+
+
+def test_explicit_event_driven_value_overrides_pipeline_default(monkeypatch) -> None:
+    monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)
+    assert _event_driven_orch_enabled(default=True) is True
+    monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "0")
+    assert _event_driven_orch_enabled(default=True) is False
+
+
 def test_default_is_legacy_loop(monkeypatch) -> None:
     """Without the env flag, the harness runs the legacy poll loop."""
     monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)

tests/engine/test_orchestrator_monitor.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import json
 import time
 
@@ -45,6 +48,8 @@ def test_note_loop_and_flush_json(tmp_path, monkeypatch):
     monitor.register_replica(0, 0)
 
     monitor.note_loop(idle=False)
+    monitor.set_dispatch_queue_size(9)
+    monitor.set_dispatch_queue_size(7)
     monitor.note_loop(idle=True)
 
     now[0] = 1.5
@@ -54,10 +59,37 @@ def test_note_loop_and_flush_json(tmp_path, monkeypatch):
     assert payload["configured_window_s"] == monitor_mod._WINDOW_S
     assert payload["windows"]["loop_idle"] == [1]
     assert payload["windows"]["loop_active"] == [1]
+    assert payload["windows"]["dispatch_queue_size"] == [7]
+    assert payload["windows"]["dispatch_queue_high_water"] == [9]
     assert payload["replicas"]["stage=0,replica=0"]["outputs_queue_size"] == [5]
     assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10]
 
 
+@pytest.mark.parametrize("drain", [False, True])
+def test_dispatch_backlog_carries_across_windows(tmp_path, monkeypatch, drain):
+    out_path = tmp_path / "orch_monitor.json"
+    monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
+    now = [0.0]
+    monkeypatch.setattr(time, "monotonic", lambda: now[0])
+    monitor = OrchestratorMonitor(replica_sampler=lambda: {})
+    monitor.set_dispatch_queue_size(3)
+    monitor.set_dispatch_queue_size(1)
+
+    now[0] = monitor_mod._WINDOW_S
+    monitor.note_loop(idle=False)
+    if drain:
+        monitor.set_dispatch_queue_size(0)
+    now[0] += monitor_mod._WINDOW_S
+    monitor.note_loop(idle=False)
+    now[0] += monitor_mod._WINDOW_S
+    monitor.flush()
+
+    windows = json.loads(out_path.read_text())["windows"]
+    remaining = 0 if drain else 1
+    assert windows["dispatch_queue_size"] == [1, remaining, remaining]
+    assert windows["dispatch_queue_high_water"] == [3, 1, remaining]
+
+
 def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch):
     out_path = tmp_path / "orch_monitor.json"
     monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -14,6 +14,7 @@ import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 
+from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
@@ -249,6 +250,41 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N
     )
 
 
+@pytest.mark.asyncio
+async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> None:
+    """AURA Stage1 has asr2aura: must not be zero-prewarmed under async_chunk.
+
+    The custom-process-input gate lives on DuplexOrchestrator (AURA path), not
+    the turn-based Orchestrator base.
+    """
+    orchestrator = object.__new__(DuplexOrchestrator)
+    stage0 = FakePrewarmPool("sender")
+    stage1 = FakePrewarmPool("receiver")  # would receive chunks if role alone decided
+    stage1.stage_client = SimpleNamespace(custom_process_input_func=lambda *a, **k: None)
+    stage2 = FakePrewarmPool("receiver")
+    orchestrator.stage_pools = [stage0, stage1, stage2]
+    orchestrator._emit_tx_edge = lambda **_kwargs: None
+    orchestrator._on_stage_submitted = MagicMock()
+    req_state = OrchestratorRequestState(
+        request_id="req-prewarm-custom",
+        prompt={"prompt_token_ids": [1, 2]},
+        sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)],
+        final_stage_id=2,
+    )
+
+    prewarmed = await orchestrator._prewarm_async_chunk_stages(
+        "req-prewarm-custom",
+        SimpleNamespace(prompt_token_ids=[1, 2], resumable=True),
+        req_state,
+    )
+
+    assert prewarmed is True
+    assert stage1.submitted == []
+    assert len(stage2.submitted) == 1
+    assert 1 not in req_state.stage_submit_ts
+    assert 2 in req_state.stage_submit_ts
+
+
 @pytest.mark.asyncio
 async def test_async_route_forwards_to_outgoing_only_stage() -> None:
     orchestrator = object.__new__(Orchestrator)

tests/engine/test_single_stage_mode.py

@@ -17,7 +17,12 @@ from vllm.v1.engine.utils import EngineZmqAddresses
 
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig
-from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig
+from vllm_omni.config.stage_config import (
+    DuplexSessionRuntimeConfig,
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_startup import (
@@ -393,6 +398,7 @@ class TestSingleStageModeDetection:
         stage_cfgs: list[Any] | None = None,
         resolved_config_path: str = "/fake/path",
         patch_deploy_config: bool = True,
+        resolved_pipeline: PipelineConfig | None = None,
         **kwargs: Any,
     ) -> AsyncOmniEngine:
         mock_stage_configs = stage_cfgs or [_make_stage_cfg(0)]
@@ -407,10 +413,19 @@ class TestSingleStageModeDetection:
                 "vllm_omni.engine.omni_engine_base.load_deploy_config",
                 return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()),
             )
+
+        def resolve_stage_configs(engine, *args, **kwargs):
+            if resolved_pipeline is not None:
+                engine._config_resolution = SimpleNamespace(
+                    pipeline_config=resolved_pipeline, config_path=resolved_config_path
+                )
+            return resolved_config_path, mock_stage_configs
+
         mocker.patch.object(
             AsyncOmniEngine,
             "_resolve_stage_configs",
-            return_value=(resolved_config_path, mock_stage_configs),
+            autospec=True,
+            side_effect=resolve_stage_configs,
         )
         mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
         mock_thread_cls = mocker.patch("threading.Thread")
@@ -435,6 +450,19 @@ class TestSingleStageModeDetection:
         )
         assert engine.single_stage_mode is True
 
+    @pytest.mark.parametrize("model_type, expected", [("qwen3_tts", True), ("qwen3_omni_moe", False)])
+    def test_event_driven_default_uses_resolved_pipeline(self, mocker, model_type, expected):
+        pipeline = PipelineConfig(
+            model_type=model_type,
+            stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
+        )
+        engine = self._make_engine_no_thread(mocker, resolved_pipeline=pipeline)
+
+        StageConfigFactory.get_pipeline_config.assert_called_once_with(
+            model="fake-model", trust_remote_code=False, deploy_config_path=None
+        )
+        assert engine._event_driven_orch_default is expected
+
     def test_stage_id_kwarg_promotes_to_single_stage_mode(self, mocker: MockerFixture):
         engine = self._make_engine_no_thread(
             mocker,