[Bugfix][Qwen3-TTS] Read full payload from model buffer

Signed-off-by: Sy03 <1370724210@qq.com>

.buildkite/common/ci_source_file_dependencies.yml

@@ -188,6 +188,10 @@ source_file_dependencies:
 
   # --- models ---
   # --- omni / tts jobs ---
+  entrypoints_function:
+    - tests/entrypoints/
+    - vllm_omni/entrypoints/
+
   omni_qwen3_omni_function:
     - *qwen3_omni
     - tests/e2e/offline_inference/test_qwen3_omni.py
@@ -236,6 +240,14 @@ source_file_dependencies:
     - *minicpmo_4_5
     - tests/e2e/accuracy/minicpmo_4_5/
 
+  omni_duplex_eval_ci:
+    - *minicpmo_4_5
+    - *minicpmo_4_5_duplex
+    - tests/e2e/accuracy/omni_duplex_eval/**
+    - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
+    - vllm_omni/benchmarks/duplex/**
+    - vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py
+
   omni_minicpmo_4_5_perf:
     - *minicpmo_4_5
     - tests/dfx/perf/scripts/run_benchmark.py
@@ -390,6 +402,8 @@ source_file_dependencies:
   diffusion_hunyuan_image3_accuracy:
     - *hunyuan_image3
     - tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
+    - tests/e2e/accuracy/test_hunyuan_image3_lora.py
+    - vllm_omni/diffusion/lora/
 
   diffusion_hunyuan_image3_perf:
     - *hunyuan_image3

.buildkite/cuda/test-merge.yml

@@ -60,11 +60,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        # TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per
-        # run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds
-        # 2915 vs 2916), which pushed this suite past 60 min. Revert once the
-        # 0.26 startup regression is understood and fixed.
-        timeout_in_minutes: 90
+        timeout_in_minutes: 60
         commands:
           # Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5
           # (see the P0 thread on #7413); the test stays in the tree and runs again there.
@@ -73,10 +69,25 @@ steps:
 
       - label: "Entrypoint Test with L4"
         timeout_in_minutes: 60
+        # Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
+        # leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
+        # This job does not collect the ready leftover
+        # (test_diffusion_generate_request_id, core_model); isolate the
+        # advanced_model DiT handshake instead. Revert to a single sweep after
+        # the product shutdown fix.
         commands:
-          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
+          - |
+            set -eu
+            pytest -s -v tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model"
+            pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model" --deselect tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake
         mirror_hardwares: l4_1
 
+      - label: "Entrypoint Test with L4 · 2-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_2" --run-level "advanced_model"
+        mirror_hardwares: l4_2
+
   - group: ":card_index_dividers: E2E Test"
     depends_on: upload-merge-pipeline
     steps:

.buildkite/cuda/test-nightly.yml

@@ -70,6 +70,22 @@ steps:
           - export SEED_TTS_EVAL_DEVICE=cuda:1
           - pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and H100 and B200 and cuda and (cards_1 or cards_2)" --run-level full_model
 
+      - label: ":full_moon: Omni · Omni-DuplexEval CI Guard"
+        source_file_dependencies: omni_duplex_eval_ci
+        timeout_in_minutes: 240
+        env:
+          VLLM_DUPLEX_EVAL_JUDGE_URL: "http://127.0.0.1:8001"
+        commands:
+          - |
+            set +e
+            pytest -s -v tests/e2e/accuracy/omni_duplex_eval/test_omni_duplex_eval_ci.py \
+              -m "full_model and H100 and B200 and cuda and cards_2" \
+              --run-level full_model \
+              --basetemp /tmp/omni_duplex_ci
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/e2e/accuracy/omni_duplex_eval/results/*.json"
+            exit $$EXIT
+
       - label: ":full_moon: Omni · Perf Test · No Async Chunk"
         source_file_dependencies: omni_qwen3_omni_perf
         key: nightly-omni-performance-no-async-chunk
@@ -116,6 +132,20 @@ steps:
         commands:
           - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model and H100 and B200 and cards_4" --run-level "core_model"
 
+  - group: ":card_index_dividers: Entrypoints Test"
+    key: nightly-entrypoints-test-group
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    steps:
+
+      - label: ":full_moon: Entrypoint Test with H100"
+        source_file_dependencies: entrypoints_function
+        # Merge collects advanced_model (5 duplex interruptions). Nightly
+        # collects full_model (20 interruptions + BAGEL sleep).
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "full_model and cuda and H100 and B200 and cards_2" --run-level "full_model"
+
   - group: ":card_index_dividers: TTS Model Test"
     key: nightly-tts-test-group
     depends_on: upload-nightly-pipeline
@@ -247,7 +277,7 @@ steps:
         commands:
           - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
           - export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
-          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py tests/e2e/accuracy/test_hunyuan_image3_lora.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU"
         source_file_dependencies: diffusion_qwen_image_perf

.buildkite/cuda/test-ready.yml

@@ -117,11 +117,18 @@ steps:
 
   - label: "Entrypoints Test"
     depends_on: upload-ready-pipeline
+    # Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
+    # leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
+    # Isolate that leftover case in its own process so the driver drops the
+    # mapping; the rest stays a directory sweep with --deselect. Revert to a
+    # single sweep after the product shutdown fix.
     commands:
       - |
         timeout 30m bash -c "
+          set -eu
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model'
+          pytest -sv tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id -m 'core_model and cuda' --run-level 'core_model'
+          pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model' --deselect tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id
         "
     mirror_hardwares: l4_1
 
@@ -204,7 +211,7 @@ steps:
           - |
             timeout 40m bash -c "
               export VLLM_IMAGE_FETCH_TIMEOUT=60
-              pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
+              pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model and cuda' --run-level 'core_model'
             "
         mirror_hardwares: h100_1
 

.buildkite/cuda/test-weekly.yml

@@ -416,6 +416,7 @@ steps:
           - >-
             pytest -sv
             tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/offline_inference/test_bagel_expansion.py
             -m "slow and diffusion and H100 and cards_2"
             --run-level "full_model"
         mirror_hardwares: h100_2

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

@@ -85,6 +85,25 @@ steps:
             EXIT=$$?
             buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json"
             exit $$EXIT
+      - label: ":full_moon: Omni · Omni-DuplexEval CI Guard"
+        source_file_dependencies: omni_duplex_eval_ci
+        timeout_in_minutes: 240
+        mirror_hardwares: a3_npu_2
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+          VLLM_DUPLEX_EVAL_JUDGE_URL: "http://127.0.0.1:8001"
+        commands:
+          - |
+            set +e
+            pytest -s -v tests/e2e/accuracy/omni_duplex_eval/test_omni_duplex_eval_ci.py \
+              -m "full_model and npu and A3 and cards_2" \
+              --run-level full_model \
+              --basetemp /tmp/omni_duplex_ci
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/e2e/accuracy/omni_duplex_eval/results/*.json"
+            exit $$EXIT
       - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test"
         source_file_dependencies: omni_minicpmo_4_5_perf
         key: nightly-omni-performance-minicpmo-4-5-npu
@@ -228,6 +247,42 @@ steps:
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
             exit $$EXIT
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-NPU"
+        source_file_dependencies: diffusion_qwen_image_perf
+        key: nightly-diffusion-x2iat-performance-qwen-image-single-npu
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_2
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3 and cards_1" --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-NPU"
+        source_file_dependencies: diffusion_qwen_image_perf
+        key: nightly-diffusion-x2iat-performance-qwen-image-4npu
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3 and cards_4" --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
     depends_on: upload-nightly-pipeline
@@ -258,13 +313,12 @@ steps:
           HF_HUB_OFFLINE: "1"
           DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
         commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export BENCHMARK_DIR=tests/dfx/perf/results
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
             EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
             exit $$EXIT
       - label: ":full_moon: Diffusion X2V · Wan 2.2 Function Test"
         source_file_dependencies: diffusion_wan22_function

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

@@ -55,8 +55,8 @@ steps:
             -m 'core_model and npu and 310P'
             --run-level 'core_model'
 
-  - group: ":card_index_dividers: Omni Model Test"
-    key: ready-omni-group
+  - group: ":card_index_dividers: E2E Test"
+    key: ready-e2e-group
     depends_on: upload-ready-pipeline
     steps:
       - label: "Omni · MiniCPM-o 4.5 Test"
@@ -66,6 +66,7 @@ steps:
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
+
       - label: "Omni · MiniCPM-o 4.5 Duplex Test"
         mirror_hardwares: a3_npu_2
         env:
@@ -73,6 +74,7 @@ steps:
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
         commands:
           - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
+
       - label: "Omni · MiniCPM-o 4.5 Perf Test"
         timeout_in_minutes: 60
         mirror_hardwares: a3_npu_2
@@ -83,10 +85,15 @@ steps:
           - export BENCHMARK_DIR=tests/dfx/perf/results
           - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni'
 
-  - group: ":card_index_dividers: TTS Model Test"
-    key: ready-tts-group
-    depends_on: upload-ready-pipeline
-    steps:
+      - label: "Omni · Qwen3-Omni Test"
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
+        timeout_in_minutes: 60
+
       - label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
         mirror_hardwares: a3_npu_2
         env:
@@ -95,3 +102,55 @@ steps:
           HF_TOKEN: "${HF_TOKEN}"
         commands:
           - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model and npu' --run-level 'core_model'
+
+      - label: "Diffusion · HunyuanVideo-1.5 Test"
+        mirror_hardwares: a3_npu_2
+        env:
+            VLLM_WORKER_MULTIPROC_METHOD: spawn
+            VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+        commands:
+            - |
+              timeout 40m bash -c "
+                export VLLM_IMAGE_FETCH_TIMEOUT=60
+                pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion and npu' --run-level 'core_model'
+                "
+        timeout_in_minutes: 60
+
+      - label: "Diffusion · Qwen Image Test"
+        mirror_hardwares: a3_npu_2
+        env:
+            VLLM_WORKER_MULTIPROC_METHOD: spawn
+            VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+        commands:
+            - |
+                timeout 40m bash -c "
+                export VLLM_IMAGE_FETCH_TIMEOUT=60
+                pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model and npu' --run-level 'core_model'
+                "
+        timeout_in_minutes: 60
+
+      - label: "Diffusion · Wan22 Test"
+        mirror_hardwares: a3_npu_2
+        env:
+            VLLM_WORKER_MULTIPROC_METHOD: spawn
+            VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+        commands:
+            - |
+                timeout 40m bash -c "
+                export VLLM_IMAGE_FETCH_TIMEOUT=60
+                pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model and npu' --run-level 'core_model'
+                "
+        timeout_in_minutes: 60
+
+      - label: "Diffusion · Hunyuan Image Test"
+        mirror_hardwares: a3_npu_4
+        env:
+            VLLM_WORKER_MULTIPROC_METHOD: spawn
+            VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+        commands:
+            - |
+                timeout 40m bash -c "
+                export VLLM_IMAGE_FETCH_TIMEOUT=60
+                pytest -s -v tests/e2e/online_serving/test_hunyuan_image.py -m 'core_model and diffusion and npu' --run-level 'core_model'
+                "
+        timeout_in_minutes: 60

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,4 +12,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE.
 ## Test Result
 
 **BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
+
+> [!NOTE]
+> To avoid agent spam, vLLM-Omni limits contributors without write access to **5 open PRs**. If this blocks well-intentioned critical work, please contact a committer to request bypass list access.
 (anything written below this line will be removed by GitHub Actions)

.gitignore

@@ -245,7 +245,11 @@ datasets/
 !recipes/cosmos3/negative_prompt.json
 !tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
+# Omni-DuplexEval CI config JSON (e.g. accuracy/omni_duplex_eval)
+!tests/e2e/accuracy/omni_duplex_eval/*.json
 *.jsonl
+# Shared Hunyuan IT2I input for accuracy tests and vllm bench.
+!tests/assets/hunyuan_image3/it2i.jsonl
 *.parquet
 
 # Output files

CONTRIBUTING.md

@@ -11,3 +11,5 @@ how to extend allowlists) are documented in the
 section of that guide.
 
 Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
+
+vLLM-Omni limits contributors without write access to **5 open PRs**; see the [Pull Request Limit](docs/contributing/README.md#pull-request-limit) section of the contributing guide.

benchmarks/accuracy/common.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import base64
@@ -200,7 +203,16 @@ class VllmOmniImageClient:
             )
             last_response = response
             if response.status_code != 404:
-                response.raise_for_status()
+                try:
+                    response.raise_for_status()
+                except requests.HTTPError as error:
+                    # Keep the server's validation detail in the exception;
+                    # otherwise multipart image tests only report a bare
+                    # ``400 Bad Request`` and hide the actionable cause.
+                    raise requests.HTTPError(
+                        f"{error}: {response.text}",
+                        response=response,
+                    ) from error
                 return decode_base64_image(response.json()["data"][0]["b64_json"])
 
         assert last_response is not None

benchmarks/diffusion/README.md

@@ -164,3 +164,22 @@ batch may still pay compile or CUDA-graph capture cost.
 
 For a Qwen-Image continuous-batching replay example, see
 [`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
+
+## HunyuanImage3 reference-prefix reuse
+
+The prefix-cache benchmark uses the unified `vllm bench serve --omni` runner,
+not this legacy diffusion benchmark. It reuses the two-image IT2I input on a
+single DiT stage, comparing dense, paged without caching, and paged with caching.
+See [Shared-reference benchmark](../../docs/design/feature/prefix_caching.md#shared-reference-benchmark)
+for the configuration, commands, and separate partial-hit accuracy coverage.
+
+The accuracy regression compares paged-no-cache, partial-hit and exact-repeat
+outputs against the checked-in official-repository reference-image goldens at
+**50 steps, CFG=2.5, seeds 43/45** (both CFG branches remain on CFGP1).
+It uses the AR-to-DiT image criteria: CLIP ≥ 90, SSIM ≥ 0.26, PSNR ≥ 12.5 dB.
+`HUNYUAN_IMAGE3_INCLUDE_DENSE=1` adds dense as another comparison to the goldens.
+It also requires actual reference-image hits and model-side query slicing,
+and checks bitwise repeatability for repeated requests at the same hit boundary.
+Images, logs, deployment YAMLs, hit traces and quality metrics are saved under
+pytest's temporary output directory. `HUNYUAN_IMAGE3_MODEL` selects a local model;
+DFX uses the repository's normal model/cache resolution.

benchmarks/diffusion/bench_wan_vae_decode.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark and exactness check for the Wan VAE decoder fast path.
+
+Loads a diffusers Wan VAE (by default the Cosmos3 one), decodes seeded latents
+for the requested video size at each ``--fast-path`` level and reports wall
+time, peak memory and, against the ``off`` level, bitwise equality, max abs
+difference and PSNR. ``--profile`` adds a torch.profiler CUDA kernel table with
+the convolution / layout-transpose share.
+
+Single GPU::
+
+    python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano
+    python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+        --fast-path off,lossless --frames 33 --profile
+    python benchmarks/diffusion/bench_wan_vae_decode.py --tiny --frames 9 --size 256x256
+
+Multi-GPU VAE parallel decode (launch one process per GPU with ``torchrun``; the
+decode is timed across all ranks and only rank 0 prints)::
+
+    torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+        --vae-patch-parallel-size 2 --vae-parallel-mode tile
+
+``tile`` distributes the VAE's own spatial tiles over the ranks (it needs the
+video to exceed the tile size, 256x256 pixels by default) and ``spatial_shard_*``
+shards every decoder feature map along height or width with halo exchange. Both
+return the full video on rank 0 only, so equality and PSNR are computed there.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import time
+from collections.abc import Iterable
+
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    DistributedAutoencoderKLWan,
+    OmniAutoencoderKLWan,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    VAE_FAST_PATH_LEVELS,
+    install_wan_vae_fastpath,
+)
+
+TINY_CONFIG = dict(
+    base_dim=32,
+    decoder_base_dim=64,
+    z_dim=48,
+    dim_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    temperal_downsample=[False, True, True],
+    is_residual=True,
+    patch_size=2,
+    in_channels=12,
+    out_channels=12,
+    scale_factor_temporal=4,
+    scale_factor_spatial=16,
+)
+
+DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
+VAE_PARALLEL_MODES = ("tile", "spatial_shard_height", "spatial_shard_width")
+_CONV_TAGS = ("conv", "cudnn", "gemm", "cutlass", "xmma")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model", default="nvidia/Cosmos3-Nano", help="HF id or local path of the pipeline")
+    parser.add_argument("--subfolder", default="vae")
+    parser.add_argument("--tiny", action="store_true", help="Use a small random VAE instead of loading weights")
+    parser.add_argument("--size", default="1280x720", help="Output video WxH")
+    parser.add_argument("--frames", type=int, default=189, help="Output frame count (1 + 4k)")
+    parser.add_argument("--dtype", choices=sorted(DTYPES), default="bf16")
+    parser.add_argument("--fast-path", default=",".join(VAE_FAST_PATH_LEVELS), help="Comma-separated levels")
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--iters", type=int, default=2)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--profile", action="store_true", help="Print a torch.profiler CUDA kernel table")
+    parser.add_argument("--profile-rows", type=int, default=25)
+    parser.add_argument("--save-output", default=None, help="Save the `off` decode output to this .pt path")
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="VAE parallel decode over this many ranks (run under torchrun with as many processes)",
+    )
+    parser.add_argument(
+        "--vae-parallel-mode",
+        choices=VAE_PARALLEL_MODES,
+        default="tile",
+        help="tile: distribute the VAE's spatial tiles; spatial_shard_*: shard feature maps with halo exchange",
+    )
+    return parser.parse_args()
+
+
+def init_distributed() -> tuple[int, int]:
+    """Join the torchrun process group (if any) and return ``(rank, world_size)``."""
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size <= 1:
+        return 0, 1
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
+    return rank, world_size
+
+
+def load_vae(
+    args: argparse.Namespace, dtype: torch.dtype, device: torch.device, parallel: bool
+) -> OmniAutoencoderKLWan:
+    cls = DistributedAutoencoderKLWan if parallel else OmniAutoencoderKLWan
+    if args.tiny:
+        torch.manual_seed(args.seed)
+        vae = cls(**TINY_CONFIG).to(device=device, dtype=dtype).eval()
+        if parallel:
+            vae.init_distributed()
+    else:
+        vae = cls.from_pretrained(args.model, subfolder=args.subfolder, torch_dtype=dtype).to(device).eval()
+    if parallel:
+        # Same setup as the registry: parallel decode needs tiling on.
+        vae.use_tiling = True
+        vae.set_parallel_size(args.vae_patch_parallel_size, mode=args.vae_parallel_mode)
+    return vae
+
+
+def make_latents(vae: OmniAutoencoderKLWan, args: argparse.Namespace, dtype: torch.dtype, device: torch.device):
+    width, height = (int(v) for v in args.size.lower().split("x"))
+    spatial = int(getattr(vae.config, "scale_factor_spatial", 16))
+    temporal = int(getattr(vae.config, "scale_factor_temporal", 4))
+    if width % spatial or height % spatial or (args.frames - 1) % temporal:
+        raise SystemExit(f"size must be a multiple of {spatial} and frames must be 1 + {temporal}k")
+    shape = (1, vae.config.z_dim, (args.frames - 1) // temporal + 1, height // spatial, width // spatial)
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    latents = torch.randn(shape, generator=generator).to(device=device, dtype=dtype)
+    mean = getattr(vae.config, "latents_mean", None)
+    std = getattr(vae.config, "latents_std", None)
+    if mean is not None and std is not None:
+        mean = torch.as_tensor(mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        std = torch.as_tensor(std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        latents = latents * std + mean
+    return latents
+
+
+@torch.inference_mode()
+def decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor:
+    return vae.decode(latents, return_dict=False)[0]
+
+
+def sync() -> None:
+    """Wait for this rank's GPU and, under torchrun, for every other rank."""
+    torch.accelerator.synchronize()
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def psnr(a: torch.Tensor, b: torch.Tensor) -> float:
+    mse = torch.mean((a.float() - b.float()) ** 2).item()
+    return math.inf if mse == 0 else 10 * math.log10(4.0 / mse)  # outputs live in [-1, 1]
+
+
+def bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    if a.shape != b.shape or a.dtype != b.dtype:
+        return False
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return bool((a.contiguous().view(int_dtype) == b.contiguous().view(int_dtype)).all())
+
+
+def profile_decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor, rows: int, verbose: bool) -> None:
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=False) as prof:
+        decode(vae, latents)
+        sync()
+    if not verbose:
+        return
+    events = [e for e in prof.key_averages() if e.device_time_total > 0]
+    events.sort(key=lambda e: e.device_time_total, reverse=True)
+    total = sum(e.device_time_total for e in events if e.device_type.name == "CUDA")
+    if total == 0:
+        total = sum(e.device_time_total for e in events)
+    conv = sum(
+        e.device_time_total
+        for e in events
+        if e.device_type.name == "CUDA" and any(tag in e.key.lower() for tag in _CONV_TAGS)
+    )
+    transpose = sum(e.device_time_total for e in events if e.device_type.name == "CUDA" and "nhwc" in e.key.lower())
+    print(
+        f"  GPU kernel time: {total / 1e3:.1f} ms  conv-like: {100 * conv / total:.1f}%  "
+        f"nchw<->nhwc transposes: {100 * transpose / total:.1f}%"
+    )
+    print(f"  {'kernel':90s} {'ms':>9s} {'%':>6s} {'calls':>7s}")
+    shown = 0
+    for e in events:
+        if e.device_type.name != "CUDA":
+            continue
+        share = 100 * e.device_time_total / total
+        print(f"  {e.key[:90]:90s} {e.device_time_total / 1e3:9.2f} {share:6.1f} {e.count:7d}")
+        shown += 1
+        if shown >= rows:
+            break
+
+
+def run_level(
+    args: argparse.Namespace,
+    level: str,
+    dtype: torch.dtype,
+    device: torch.device,
+    *,
+    rank: int,
+    parallel: bool,
+) -> tuple[torch.Tensor, dict]:
+    vae = load_vae(args, dtype, device, parallel)
+    report = install_wan_vae_fastpath(vae, level=level)
+    latents = make_latents(vae, args, dtype, device)
+    for _ in range(args.warmup):
+        decode(vae, latents)
+    sync()
+    torch.accelerator.reset_peak_memory_stats()
+    timings = []
+    output = None
+    for _ in range(args.iters):
+        sync()
+        start = time.perf_counter()
+        output = decode(vae, latents)
+        sync()
+        timings.append(time.perf_counter() - start)
+    peak = torch.tensor(torch.accelerator.max_memory_allocated() / 2**30, device=device)
+    if dist.is_initialized():
+        dist.all_reduce(peak, op=dist.ReduceOp.MAX)  # the largest rank decides whether the config fits
+    stats = {
+        "installed": report.installed,
+        "reason": report.reason,
+        "fused_silu": report.fused_silu_dtypes,
+        "patched": dict(report.patched),
+        "time_s": min(timings),
+        "peak_gib": peak.item(),
+    }
+    if args.profile:
+        if rank == 0:
+            print(f"[{level}] profile")
+        profile_decode(vae, latents, args.profile_rows, verbose=rank == 0)
+    assert output is not None
+    del vae
+    torch.accelerator.empty_cache()
+    return output.detach(), stats
+
+
+def main() -> None:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA device required")
+    rank, world_size = init_distributed()
+    parallel = args.vae_patch_parallel_size > 1
+    if parallel and world_size < args.vae_patch_parallel_size:
+        raise SystemExit(
+            f"--vae-patch-parallel-size {args.vae_patch_parallel_size} needs torchrun with at least that many "
+            f"processes (WORLD_SIZE={world_size})"
+        )
+    device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda")
+    dtype = DTYPES[args.dtype]
+    levels: Iterable[str] = [level.strip() for level in args.fast_path.split(",") if level.strip()]
+    for level in levels:
+        if level not in VAE_FAST_PATH_LEVELS:
+            raise SystemExit(f"unknown fast-path level {level!r}; choose from {VAE_FAST_PATH_LEVELS}")
+    if "off" not in levels:
+        levels = ["off", *levels]
+
+    if rank == 0:
+        parallel_desc = f" vae_parallel={args.vae_parallel_mode}x{args.vae_patch_parallel_size}" if parallel else ""
+        print(
+            f"model={'tiny' if args.tiny else args.model} size={args.size} frames={args.frames} "
+            f"dtype={args.dtype} world_size={world_size}{parallel_desc}"
+        )
+    golden = None
+    baseline_time = None
+    try:
+        for level in levels:
+            output, stats = run_level(args, level, dtype, device, rank=rank, parallel=parallel)
+            if level == "off":
+                golden = output
+                baseline_time = stats["time_s"]
+                if args.save_output and rank == 0:
+                    torch.save(output.cpu(), args.save_output)
+            assert golden is not None
+            if rank != 0:
+                # Parallel decode assembles the video on rank 0 only; the other
+                # ranks just keep the collectives in lockstep.
+                del output
+                continue
+            speedup = baseline_time / stats["time_s"] if baseline_time else float("nan")
+            line = (
+                f"[{level:13s}] decode {stats['time_s'] * 1e3:9.1f} ms  speedup {speedup:5.2f}x  "
+                f"peak {stats['peak_gib']:5.2f} GiB  installed={stats['installed']}"
+            )
+            if level != "off":
+                equal = torch.equal(output, golden)
+                line += (
+                    f"  torch_equal={equal} bits_equal={bits_equal(output, golden)} "
+                    f"max_abs_diff={(output.float() - golden.float()).abs().max().item():.3e} "
+                    f"psnr={psnr(output, golden):.2f} dB"
+                )
+                if stats["installed"]:
+                    line += f"  fused_silu={stats['fused_silu'] or 'off'}"
+                else:
+                    line += f"  reason={stats['reason']}"
+            print(line)
+            del output
+            torch.accelerator.empty_cache()
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/tts/evidence/breeze_tts_2_20260918.json

@@ -0,0 +1,757 @@
+{
+  "measurement_scope": "Historical measurements for PR7752, not a main-branch comparison or a measurement of subsequent review fixes.",
+  "environment": {
+    "gpu": "NVIDIA H100 80 GB",
+    "gpu_count": 1,
+    "python": "3.12",
+    "pytorch": "2.13.0+cu132",
+    "vllm": "0.29.0",
+    "transformers": "5.14.1",
+    "nsight_systems": "2026.5.1"
+  },
+  "model": {
+    "id": "BreezeBlue/Breeze-TTS-2",
+    "revision": "3e28c5151381a722f1d8661b4118c298caa77aa4"
+  },
+  "workload": {
+    "texts": [
+      "Hello, this is Breeze TTS 2 running on vLLM Omni.",
+      "你好，今天天气很好，我们一起出去走走吧。"
+    ],
+    "alternating_texts": true,
+    "seed": 42,
+    "guidance_scale": 1.0,
+    "repetition_penalty": 1.1,
+    "top_p": 1.0,
+    "sampled": {
+      "temperature": 0.9,
+      "top_k": 50
+    },
+    "greedy": {
+      "temperature": 0.0,
+      "top_k": 0
+    },
+    "stream": false,
+    "response_format": "pcm",
+    "sample_rate": 24000,
+    "channels": 1,
+    "pcm_bits": 16
+  },
+  "cumulative_diagnostic": {
+    "baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99",
+    "combined_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222",
+    "intermediate_variants": "Ordered patch subsets of the combined commit, not independent commits.",
+    "patch_sha256": {
+      "batched_eos": "2f107ea8e890e03cc138ef2f69e693e3474d3aaf0eb4051c00886cbaa9785b3d",
+      "greedy_graph": "4e783313260ddc2e36c8a74ba9d95dc5d78338443612ebf89854e31b8be6933b",
+      "batched_embedding": "90557e4ec73fdedcff51fd717d5e8c6477e71087fd97ac52534542234f5d1d3a"
+    },
+    "sequence_slots": 8,
+    "kv_gib": 2,
+    "codec_chunk_frames": 16,
+    "max_new_tokens": 48,
+    "temperature": 0.0,
+    "cohorts_per_variant": 1,
+    "nsight_injection_present": true,
+    "collection_enabled": false,
+    "warmup": "Two eight-request full-response warmups, then four full-response 64-request cohorts, each preceded by one concurrency-sized warmup.",
+    "variants": [
+      {
+        "variant": "baseline",
+        "makespan_s": 0.9672094583511353,
+        "requests": 8,
+        "concurrency": 8,
+        "audio_s": 30.72,
+        "failures": 0
+      },
+      {
+        "variant": "plus_batched_eos",
+        "makespan_s": 0.9280132218264043,
+        "requests": 8,
+        "concurrency": 8,
+        "audio_s": 30.72,
+        "failures": 0
+      },
+      {
+        "variant": "plus_greedy_graph",
+        "makespan_s": 0.8857334200292826,
+        "requests": 8,
+        "concurrency": 8,
+        "audio_s": 30.72,
+        "failures": 0
+      },
+      {
+        "variant": "plus_batched_embedding",
+        "makespan_s": 0.8435833756811917,
+        "requests": 8,
+        "concurrency": 8,
+        "audio_s": 30.72,
+        "failures": 0
+      }
+    ],
+    "limits": [
+      "Single cumulative diagnostic samples; no repeat spread.",
+      "Nsight injection remains when collection is disabled.",
+      "The final variant used graph tracing configuration; prior variants used node tracing configuration.",
+      "Do not interpret these timings as isolated uninstrumented per-optimization gains."
+    ]
+  },
+  "mechanical_attribution": {
+    "batched_eos": {
+      "stage0_stream_synchronization_calls": [
+        773,
+        438
+      ],
+      "device_to_host_copies": [
+        841,
+        506
+      ],
+      "capture_mode": "node"
+    },
+    "greedy_graph": {
+      "radix_sort_kernels": [
+        735,
+        0
+      ],
+      "capture_mode": "node"
+    },
+    "combined_graph_capture": {
+      "captured_wall_ms": [
+        955.4114495404065,
+        855.6996420957148
+      ],
+      "gpu_envelope_ms": [
+        926.633897,
+        826.396835
+      ],
+      "activity_upper_ms": [
+        651.130806,
+        628.864685
+      ],
+      "idle_lower_ms": [
+        275.50309100000004,
+        197.53215
+      ],
+      "stage0_explicit_kernel_launch_count": [
+        7087,
+        6198
+      ],
+      "stage0_graph_launch_count": [
+        156,
+        156
+      ]
+    },
+    "limits": [
+      "Host API durations overlap GPU activity and are not additional idle time.",
+      "Graph envelopes hide internal gaps: activity is an upper bound and idle is a lower bound.",
+      "Stage 1 has eight versus nine graph launches across the graph-mode traces.",
+      "No separate final-capacity stage breakdown was measured."
+    ]
+  },
+  "combined_uninstrumented_ab": {
+    "baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99",
+    "head_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222",
+    "sequence_slots": 8,
+    "kv_gib": 2,
+    "codec_chunk_frames": 16,
+    "requests_per_group": 64,
+    "max_new_tokens": 384,
+    "baseline_rounds": 2,
+    "head_rounds": 1,
+    "profiler_configured": false,
+    "nsight_injection_present": false,
+    "warmup": "One excluded 64-request full-response wave for each temperature/concurrency family, plus an excluded 2*C-request full-response cohort before every group.",
+    "rows": [
+      {
+        "build": "plain-baseline",
+        "phase": "round-0",
+        "temperature": 0.0,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 10.77213971177116,
+        "req_per_s": 5.941252314994074,
+        "aggregate_rtf": 0.030157166046391827,
+        "audio_s": 357.2,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 305.88,
+        "endpoint_trimmed_aggregate_rtf": 0.03521688149526337,
+        "quiet_tail_at_least_3s": 7,
+        "max_quiet_tail_s": 3.24
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-0",
+        "temperature": 0.0,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 11.076066234149039,
+        "req_per_s": 5.778224745774739,
+        "aggregate_rtf": 0.029092420240988227,
+        "audio_s": 380.72,
+        "caps": 2,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 315.48,
+        "endpoint_trimmed_aggregate_rtf": 0.035108616185333584,
+        "quiet_tail_at_least_3s": 6,
+        "max_quiet_tail_s": 26.04
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-0",
+        "temperature": 0.9,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 8.603743681218475,
+        "req_per_s": 7.438622345259851,
+        "aggregate_rtf": 0.03389435739528236,
+        "audio_s": 253.84,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 239.02,
+        "endpoint_trimmed_aggregate_rtf": 0.035995915325991445,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 1.3
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-0",
+        "temperature": 0.9,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 8.66502623911947,
+        "req_per_s": 7.386013409984042,
+        "aggregate_rtf": 0.033143460216950235,
+        "audio_s": 261.44,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 242.14,
+        "endpoint_trimmed_aggregate_rtf": 0.03578519137325296,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 1.3
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-1",
+        "temperature": 0.0,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 13.904579282738268,
+        "req_per_s": 4.602800178172402,
+        "aggregate_rtf": 0.0392784725501081,
+        "audio_s": 354.0,
+        "caps": 1,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 285.36,
+        "endpoint_trimmed_aggregate_rtf": 0.048726448285457906,
+        "quiet_tail_at_least_3s": 8,
+        "max_quiet_tail_s": 25.38
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-1",
+        "temperature": 0.0,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 10.060411099344492,
+        "req_per_s": 6.361569061941222,
+        "aggregate_rtf": 0.02935460754944121,
+        "audio_s": 342.72,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 289.04,
+        "endpoint_trimmed_aggregate_rtf": 0.03480629359031446,
+        "quiet_tail_at_least_3s": 10,
+        "max_quiet_tail_s": 3.16
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-1",
+        "temperature": 0.9,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 8.535738762002438,
+        "req_per_s": 7.4978864494894575,
+        "aggregate_rtf": 0.033363581777683074,
+        "audio_s": 255.84,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 237.9,
+        "endpoint_trimmed_aggregate_rtf": 0.03587952401009852,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 2.12
+      },
+      {
+        "build": "plain-baseline",
+        "phase": "round-1",
+        "temperature": 0.9,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 8.584469985216856,
+        "req_per_s": 7.4553234049642105,
+        "aggregate_rtf": 0.033355882752629996,
+        "audio_s": 257.36,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 242.57999999999998,
+        "endpoint_trimmed_aggregate_rtf": 0.035388201769382704,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 1.94
+      },
+      {
+        "build": "plain-batched",
+        "phase": "round-0",
+        "temperature": 0.0,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 12.983129982370883,
+        "req_per_s": 4.929473870083891,
+        "aggregate_rtf": 0.03320153943936907,
+        "audio_s": 391.04,
+        "caps": 2,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 312.72,
+        "endpoint_trimmed_aggregate_rtf": 0.041516788124747,
+        "quiet_tail_at_least_3s": 11,
+        "max_quiet_tail_s": 28.4
+      },
+      {
+        "build": "plain-batched",
+        "phase": "round-0",
+        "temperature": 0.0,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 9.772029620129615,
+        "req_per_s": 6.549304749155183,
+        "aggregate_rtf": 0.027698496655696188,
+        "audio_s": 352.8,
+        "caps": 1,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 286.26,
+        "endpoint_trimmed_aggregate_rtf": 0.034136902187275955,
+        "quiet_tail_at_least_3s": 6,
+        "max_quiet_tail_s": 26.04
+      },
+      {
+        "build": "plain-batched",
+        "phase": "round-0",
+        "temperature": 0.9,
+        "concurrency": 8,
+        "requests": 64,
+        "makespan_s": 8.171920725144446,
+        "req_per_s": 7.831696140061215,
+        "aggregate_rtf": 0.03169376638669115,
+        "audio_s": 257.84,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 241.72,
+        "endpoint_trimmed_aggregate_rtf": 0.03380738344011437,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 1.22
+      },
+      {
+        "build": "plain-batched",
+        "phase": "round-0",
+        "temperature": 0.9,
+        "concurrency": 16,
+        "requests": 64,
+        "makespan_s": 8.02551719499752,
+        "req_per_s": 7.974563937124526,
+        "aggregate_rtf": 0.030886380830501534,
+        "audio_s": 259.84000000000003,
+        "caps": 0,
+        "failures": 0,
+        "endpoint_trimmed_audio_s": 243.0,
+        "endpoint_trimmed_aggregate_rtf": 0.03302681973250008,
+        "quiet_tail_at_least_3s": 0,
+        "max_quiet_tail_s": 1.78
+      }
+    ],
+    "limits": [
+      "Ranges compare the head sample with two baseline samples; they are not confidence intervals.",
+      "Full-response greedy output lengths and frame-limit incidence vary, so no general greedy throughput gain is claimed.",
+      "No transcription was run on these historical A/B outputs."
+    ]
+  },
+  "capacity": {
+    "runtime_sha": "7bc0b3e3a2e41ef44e8efb30fd927e93bf6a48a7",
+    "main_base_sha": "6bda549ac23a667efbe44a54cf3f6d57e6e24da2",
+    "sequence_slots": [
+      8,
+      16,
+      32
+    ],
+    "kv_gib": [
+      2,
+      4,
+      8
+    ],
+    "codec_graph_buckets": [
+      [
+        1,
+        2,
+        4,
+        8
+      ],
+      [
+        1,
+        2,
+        4,
+        8,
+        16
+      ],
+      [
+        1,
+        2,
+        4,
+        8,
+        16,
+        32
+      ]
+    ],
+    "backbone_capture_sizes": [
+      1,
+      2,
+      4,
+      8,
+      32,
+      64,
+      128,
+      256
+    ],
+    "codec_chunk_frames": 16,
+    "requests_per_group": 128,
+    "cohorts_per_point": 1,
+    "excluded_warmups_per_configuration": {
+      "cohorts": 4,
+      "requests_per_cohort": 64,
+      "covers": "Both client concurrency levels and both sampling settings"
+    },
+    "profiler_configured": false,
+    "nsight_injection_present": false,
+    "rows": [
+      {
+        "slots": 8,
+        "concurrency": 8,
+        "observed_peak_client_requests": 8,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 6.7812081269297915,
+        "makespan_s": 18.87569259107113,
+        "aggregate_rtf": 0.02841355459879445,
+        "audio_s": 664.32,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 8,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 5.8798058352164615,
+        "makespan_s": 21.76942633604631,
+        "aggregate_rtf": 0.029975526459636356,
+        "audio_s": 726.24,
+        "near_cap_count": 2,
+        "exact_cap_count": 2,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 8,
+        "concurrency": 8,
+        "observed_peak_client_requests": 8,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 7.72991106982811,
+        "makespan_s": 16.559052082709968,
+        "aggregate_rtf": 0.0320365502296664,
+        "audio_s": 516.88,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 8,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 7.945057322591292,
+        "makespan_s": 16.110645248088986,
+        "aggregate_rtf": 0.030958196095482296,
+        "audio_s": 520.4,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 8,
+        "concurrency": 8,
+        "observed_peak_client_requests": 8,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 9.850579526557848,
+        "makespan_s": 12.994159344118088,
+        "aggregate_rtf": 0.026436684863521503,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      },
+      {
+        "slots": 8,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 10.665146583658053,
+        "makespan_s": 12.00171033712104,
+        "aggregate_rtf": 0.02441754218978076,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      },
+      {
+        "slots": 16,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 7.754874623816668,
+        "makespan_s": 16.505747186020017,
+        "aggregate_rtf": 0.023979758231665527,
+        "audio_s": 688.32,
+        "near_cap_count": 1,
+        "exact_cap_count": 1,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 16,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 10.104109318018903,
+        "makespan_s": 12.668113138061017,
+        "aggregate_rtf": 0.018577125085143444,
+        "audio_s": 681.92,
+        "near_cap_count": 1,
+        "exact_cap_count": 1,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 16,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 10.920674745703012,
+        "makespan_s": 11.72088748915121,
+        "aggregate_rtf": 0.021831484669108944,
+        "audio_s": 536.88,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 16,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 11.210975235811437,
+        "makespan_s": 11.417383172083646,
+        "aggregate_rtf": 0.02120613516360261,
+        "audio_s": 538.4,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 16,
+        "concurrency": 16,
+        "observed_peak_client_requests": 16,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 15.14487085057202,
+        "makespan_s": 8.451706274878234,
+        "aggregate_rtf": 0.017195040435543284,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      },
+      {
+        "slots": 16,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 16.563937910502617,
+        "makespan_s": 7.727630995213985,
+        "aggregate_rtf": 0.015721905507840955,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      },
+      {
+        "slots": 32,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 15.246181432777554,
+        "makespan_s": 8.395544849336147,
+        "aggregate_rtf": 0.012390119317202106,
+        "audio_s": 677.6,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 32,
+        "concurrency": 64,
+        "observed_peak_client_requests": 64,
+        "temperature": 0.0,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 11.166074835908939,
+        "makespan_s": 11.463294119108468,
+        "aggregate_rtf": 0.015826284127331107,
+        "audio_s": 724.32,
+        "near_cap_count": 1,
+        "exact_cap_count": 1,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 32,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 17.17219888865503,
+        "makespan_s": 7.453908543102443,
+        "aggregate_rtf": 0.014336645143680651,
+        "audio_s": 519.92,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 32,
+        "concurrency": 64,
+        "observed_peak_client_requests": 64,
+        "temperature": 0.9,
+        "max_new_tokens": 384,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 17.514343665613783,
+        "makespan_s": 7.3082955572754145,
+        "aggregate_rtf": 0.013677750331777613,
+        "audio_s": 534.32,
+        "near_cap_count": 0,
+        "exact_cap_count": 0,
+        "fixed48_verified": null
+      },
+      {
+        "slots": 32,
+        "concurrency": 32,
+        "observed_peak_client_requests": 32,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 25.82908453652114,
+        "makespan_s": 4.955653763841838,
+        "aggregate_rtf": 0.010082303393232905,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      },
+      {
+        "slots": 32,
+        "concurrency": 64,
+        "observed_peak_client_requests": 64,
+        "temperature": 0.0,
+        "max_new_tokens": 48,
+        "requests": 128,
+        "failures": 0,
+        "req_per_s": 28.81314277607874,
+        "makespan_s": 4.442417163401842,
+        "aggregate_rtf": 0.00903812085653044,
+        "audio_s": 491.52,
+        "near_cap_count": 128,
+        "exact_cap_count": 128,
+        "fixed48_verified": true
+      }
+    ],
+    "limits": [
+      "This scales slots, KV reservation and codec graph buckets together.",
+      "One cohort per point; no repeated capacity samples or confidence intervals.",
+      "The 48-frame cohorts deliberately truncate speech and are excluded from full-utterance quality claims.",
+      "Peak device memory was not measured.",
+      "No capacity claim extends to cloning, CFG guidance above one, long prompts or other GPUs."
+    ]
+  },
+  "quality": {
+    "full_response_pcm_checked": 1536,
+    "sampled_full_responses": 768,
+    "sampled_frame_limit_hits": 0,
+    "sampled_quiet_tails_at_least_three_seconds": 0,
+    "greedy_frame_limit_hits": 5,
+    "asr_selected_unique_clips": 47,
+    "asr_selected_sampled_clips": 24,
+    "sampled_asr": "All selected sampled transcripts matched after accounting for TTS 2 versus TTS2 spacing.",
+    "endpoint_method": "Trim leading/trailing quiet using 20 ms RMS windows at or below -40 dBFS; keep internal pauses and original makespans.",
+    "limits": [
+      "ASR selects median-duration and longest English/Chinese clips per group, deduplicated by WAV hash.",
+      "Selected ASR and endpoint measurements do not establish perceptual quality parity.",
+      "Capped greedy clips have extra ASR words with unreliable tail alignment.",
+      "No listening or MOS assessment was performed."
+    ]
+  },
+  "metrics": {
+    "requests_per_second": "Completed requests divided by cohort wall time",
+    "aggregate_rtf": "Cohort wall time divided by total returned audio duration",
+    "timing_exclusions": [
+      "PCM validation",
+      "hashing",
+      "WAV file writes",
+      "post-timing audio analysis"
+    ]
+  }
+}

benchmarks/tts/model_configs.yaml

@@ -7,6 +7,17 @@
 # is tracked here.
 
 models:
+  BreezeBlue/Breeze-TTS-2:
+    supported_tasks: [default_voice, voice_design, voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      default_voice: {}
+      voice_design:
+        task_type: VoiceDesign
+      voice_clone:
+        task_type: Base
+
   IndexTeam/IndexTTS-2.5:
     supported_tasks: [voice_clone]
     backend: openai-audio-speech

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_VERSION=v0.29.0
+ARG VLLM_ASCEND_VERSION=v0.30.0
 ARG VLLM_ASCEND_DEVICE_SUFFIX=
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
 

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_VERSION=v0.29.0
+ARG VLLM_ASCEND_VERSION=v0.30.0
 ARG VLLM_ASCEND_DEVICE_SUFFIX=
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
 

docs/.nav.yml

@@ -98,6 +98,7 @@ nav:
         - Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
         - RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
       - Regional Compilation: user_guide/diffusion/regional_compilation.md
+      - Wan VAE Decoder Fast Path: user_guide/diffusion/vae_fast_path.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Startup and Loading: user_guide/diffusion/startup_and_loading.md
       - LoRA: user_guide/diffusion/lora.md

docs/cli/bench/serve.md

@@ -302,6 +302,21 @@ vllm bench serve --omni \
 selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the
 `/v1/realtime` endpoint.
 
+To replay an existing sample set (for example `sampled_cases.jsonl` from a prior run), pass
+`--omniinteract-video-list` instead of discovering cases. List order is preserved; `--num-prompts` takes the prefix of
+the list (`0` or a value larger than the list length runs every row). Do not combine video-list mode with `--seed`,
+`--dataset-path`, `--omniinteract-scenario-tags`, or `--omniinteract-scenario-focus`.
+
+Use `--omniinteract-scenario-tags` to steer small runs toward paper-table scenarios (`realtime`, `proactive`,
+`nested`, `interrupted`, `1qna`). By default the sampler covers each requested tag with at least one matching case
+(when available under `--omniinteract-subsets`), then fills the remaining `--num-prompts` budget from other cases.
+Add `--omniinteract-scenario-focus` to run only cases that match those tags. Example smoke coverage:
+
+```bash
+  --num-prompts 5 \
+  --omniinteract-scenario-tags realtime proactive nested interrupted 1qna
+```
+
 Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded
 before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use
 explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by
@@ -309,7 +324,9 @@ explicit `--num-prompts 0` only when the client has enough RAM for the full data
 options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result.
 
 Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under
-`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write
+`--omniinteract-output-dir`. The root also contains `batch_summary.json`, `official_eval_manifest.jsonl` (eligible cases only),
+and `sampled_cases.jsonl` (every sampled case as absolute `video_path` / official `output_name` / `subset` rows for Lucky-Lance
+OmniInteract `batch_inference_minicpmo.py --video_list`); failed cases write
 `.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts,
 not answer accuracy. Transcript timestamps are serialized playback-queue times. Playback ACKs report cumulative progress
 incrementally along that serialized clock, like a live listener; the first ack for a response goes out as soon as its audio
@@ -318,6 +335,48 @@ arrives, checkpointing the response's history position so a later committed user
 outputs are ineligible and omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video
 horizon.
 
+Accuracy evaluation is opt-in and requires an already-running text judge with an OpenAI-compatible Chat Completions API. The
+benchmark does not launch or stop the judge server. Add the following options to the command above:
+
+```bash
+vllm serve Qwen/Qwen2.5-7B-Instruct \
+  --served-model-name Qwen2.5-7B-Instruct \
+  --port 8001
+```
+
+Then add to the benchmark command:
+
+```bash
+  --omniinteract-evaluate \
+  --omniinteract-judge-base-url http://127.0.0.1:8001 \
+  --omniinteract-judge-model Qwen2.5-7B-Instruct
+```
+
+`--omniinteract-judge-model` must match that judge process's `--served-model-name` (or its `--model`
+string if `--served-model-name` is omitted). It is not a checkpoint path to load, and it is not the
+Omni DUT `--served-model-name`. Point `--omniinteract-judge-base-url` at the judge, not at the Omni
+realtime endpoint.
+
+The early / core / interrupted-partial judge prompts follow the official English
+templates in [Lucky-Lance/OmniInteract](https://github.com/Lucky-Lance/OmniInteract)
+`eval/evaluation/llm_judge.py` (commit `de304cef35fd9a50a5caadb5090c34cfbf0dd868`).
+They correspond to the OmniInteract paper appendix
+([arXiv:2605.26485](https://arxiv.org/abs/2605.26485)) Listing A.1 (early-stage),
+Listing A.2 (interrupted partial quality), and Listing A.3 (core-stage).
+A local OpenAI-compatible judge is not the paper's GPT-4o judge, so reported
+IA-QTF1 numbers are protocol-compatible rather than official paper-table scores.
+
+After artifact publication, the evaluator builds `[start, t_a, end)` slots, judges early and core response text, writes
+per-case details plus `evaluation/unified_eval_summary.json`, and prints IA-QTF1 columns
+(1Q1A realtime / proactive / nested / Global, 1QnA, All Global), interruption diagnostics (NOR / PAQ / CSM), and nested
+metrics (NCCS, inner / outer IA-QTF1, missed outer). Slices and sections with no slots in the evaluated sample are
+omitted from the terminal report. The header lists how many evaluated cases carry each scenario tag
+(`realtime`, `proactive`, `nested`, `interrupted`, `1qna`). Realtime and proactive
+exclude nested inner/outer slots; 1Q1A Global recomputes F1 from those three TP/FP/FN aggregates; All Global includes every
+scored slot plus unmatched-chunk false positives. Plain `wav_transcript.json` timestamps provide chunk-level approximate
+timing. If transcript chunks include `aligned_words`, the evaluator splits boundary-crossing chunks and derives trigger
+timing from word alignment.
+
 Per-response TTFT and TTFP start when the server begins executing the native model-turn request that owns the response. RTF
 continues to use client receipt of `response.created` through the last audio packet, divided by emitted audio duration. Global
 TTFT, TTFP, and RTF cover the complete input-stream window. TPOT/ITL use engine stage-0 timing; ITL is emitted only when every
@@ -325,9 +384,12 @@ token interval is present within a continuous generation segment. Model-unit pac
 excluded from TPOT/ITL. Raw request metrics retain `response_created_to_first_text_ms` and
 `response_created_to_first_audio_ms` as client-envelope diagnostics.
 
-The checked-in local performance configuration measures four deterministic cases from each OmniInteract subset (12 videos
-total), with no benchmark warmups and a maximum concurrency of two. Each subset also sends one readiness request before its
-measured cases. Run it from the repository root with:
+The checked-in local performance configuration measures four deterministic cases from `1q1a`, `1q1a_math`, and
+`1qna` (12 videos total), with no benchmark warmups and a maximum concurrency of two. Each subset also sends one readiness request before its
+measured cases. The runner starts MiniCPM-o 4.5 on the first visible GPU and a text judge
+(`Qwen/Qwen2.5-7B-Instruct`) on the second (`CUDA_VISIBLE_DEVICES=1` relative to the process). After artifacts land it
+scores All Global IA-QTF1 through that local judge. This is protocol-compatible with the paper metric, not a GPT-4o
+paper-table score. Run it from the repository root with two visible GPUs:
 
 ```bash
 export HF_HOME=/path/to/persistent/huggingface-cache
@@ -338,12 +400,14 @@ bash tools/nightly/run_nightly_jobs.sh \
   --label-substr minicpmo_4_5_omniinteract
 ```
 
-The first run downloads the pinned OmniInteract archive into `HF_HOME`; later runs reuse that cache.
+The first run downloads the pinned OmniInteract archive and the judge weights into `HF_HOME`; later runs reuse that cache.
 
 It requires every case to commit its input, complete any emitted response lifecycles, and publish the expected WAV,
 transcript, event, and result artifacts without errors. A valid LISTEN-only case may have no response audio or transcript
 chunks. Official-manifest eligibility is reported separately because clipped or cancelled output is a benchmark-quality
-signal, not a transport failure. This local performance test does not score answer accuracy.
+signal, not a transport failure. Accuracy must finish with `status=ok` on every subset. After all three subsets finish,
+All Global IA-QTF1 is recomputed from pooled `Global_TP` / `Global_FP` / `Global_FN` and must be at or above
+`omniinteract_aggregate_min_ia_qtf1` (checked in as `0.2`).
 
 ### Video-MME Benchmark
 

docs/contributing/README.md

@@ -295,6 +295,10 @@ If you encounter a bug or have a feature request, please search existing issues
 
 Thank you for your contribution to vLLM-Omni! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM-Omni maintain the code quality and improve the efficiency of the review process.
 
+### Pull Request Limit
+
+To avoid agent spam, vLLM-Omni uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/) for contributors without write access. The current cap is 5 open PRs. If this blocks well-intentioned critical work, contact a committer to request bypass list access.
+
 ### DCO and Signed-off-by
 
 When contributing changes to this project, you must agree to the [DCO](https://developercertificate.org/). Commits must include a `Signed-off-by:` header which certifies agreement with the terms of the DCO.

docs/design/feature/omni_connectors/nixl_connector.md

@@ -1,5 +1,8 @@
 # NixlConnector
 
+For the reusable model-runner API, illustrated control flow, deadline semantics,
+and failure/ownership contracts, see [Cross-Stage Payload Transport](stage_payload.md).
+
 ## When to Use
 
 Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory
@@ -46,6 +49,15 @@ provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime.
 
 ## Configuration
 
+MiniMax-H3's default disaggregated and Turbo deploy configs use shared memory
+for same-host transfer. To opt into NIXL, select
+`vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml` or
+`vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml` with
+`--deploy-config`. These overlays inherit the corresponding model and
+parallel settings and switch both ends of the stage edge. They require NIXL
+and a working UCX backend; initialization failures do not automatically fall
+back to shared memory.
+
 ```yaml
 connectors:
   nixl_connector:
@@ -92,9 +104,44 @@ Parameters:
 - `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete
   (default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's
   buffers and registrations remain owned by the connector until NIXL reports a
-  terminal state; `close()` waits for that state before releasing them.
+  terminal state. `close()` polls remaining transfers once and returns without
+  releasing active DMA resources. A serialized background closer automatically
+  finishes cleanup after completion; no second call is required.
   `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it.
 
+### Diffusion worker integration
+
+Diffusion initializes `OmniConnectorModelRunnerMixin` in synchronous mode,
+borrowing its KV transfer manager's lazy connector. This mode creates no second
+connector, starts no receive/save threads, and leaves KV key builders and
+sharding callbacks unchanged. Connector lifetime remains with the manager.
+
+The shared `recv_stage_payload` entry point and asynchronous `_poll_single_request`
+use the same key and endpoint metadata builder. Diffusion passes its actual
+incoming edge and external request ID; explicit transfer handles take precedence.
+Sender endpoints are passed per request, without mutating KV sender state.
+Diffusion-specific prompt merging, output handles and device placement live in
+`diffusion/worker/stage_payload.py`.
+
+Diffusion supplies ordered TP and SP groups to the shared payload fanout.
+Only the rank leading both groups reads the connector. Each subsequent broadcast
+uses its own group-local rank zero, including ranks that received the payload
+from the preceding group. Tensor payloads use tensor-dictionary broadcasts rather
+than serializing the full payload as a Python object. Receive misses still reach
+the same collectives. Inline fallback is allowed only when all required payload
+keys remain present; otherwise each rank raises an explicit payload error.
+
+The synchronous path passes one monotonic deadline to `get_with_deadline()` for
+discovery and DMA completion waits, rather than calling the ordinary `get()`
+with its potentially much longer transfer timeout. Native calls, serialization,
+and distributed collectives are not preemptible by this deadline.
+
+`GroupCoordinator.broadcast_tensor_dict` preserves the group-local source for
+`broadcast_object` while passing the mapped global source to PyTorch tensor
+collectives. For subgroup `[1, 3]`, local source `0` maps to global source `1`;
+passing `1` back to `broadcast_object` would incorrectly select global rank `3`.
+The nonzero-subgroup regression covers both sender and receiver behavior.
+
 ### Source ownership and failure limits
 
 Claims and expiry use the same lock. Completion identifies both a payload
@@ -105,12 +152,24 @@ transfers retain the complete destination bundle and source claim until terminal
 Zero-element tensors retain their shape/dtype/skeleton positions but never create
 DMA descriptors or registrations.
 
-A lost metadata response, lost completion ACK, or abandoned consumer can retain a
-claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
+After a lost metadata response, retries on the same consumer thread reuse the
+claim ID for that endpoint, key and requested generation until a reply arrives.
+This makes claim acquisition idempotent; subsequent independent reads still get
+distinct claims. Callers retry on the same thread and call `abandon_get()` when
+giving up. Synchronous stage payload receive does this automatically. Background
+recovery reacquires the same claim solely to acknowledge it, without issuing a
+READ. Lost completion ACKs are retried against the exact payload generation.
+A permanently unreachable or dead consumer can still retain a claim indefinitely.
+NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
 `cleanup()` frees those allocations. Producer `close()` rejects new work and
-retains its agent, listener and claimed allocations; a later `close()` finishes
+retains its agent, listener and claimed allocations; the background closer finishes
 teardown after claims drain. Permanently abandoned claims remain until process
-exit. This favors memory safety over bounded shutdown/memory usage. Only trusted
+exit. A consumer with unfinished local transfers likewise retains a strong reference
+to its connector, agent, tensors and registrations after `close()` returns. Its
+background closer polls and releases resources only after a terminal state;
+permanently stuck transfers retain resources until process exit. Closing connectors
+reject new work and report unhealthy. This avoids an unbounded transfer-polling loop
+in shutdown without freeing DMA-owned memory. Only trusted
 peers may access this unauthenticated control plane. Both endpoints must use the
 claim-aware protocol; rolling compatibility with older GET_META clients is not
 provided. Legacy externally-owned metadata without a generation is accepted only

docs/design/feature/omni_connectors/stage_payload.md

@@ -0,0 +1,280 @@
+# Cross-Stage Payload Transport
+
+Stage payload transport moves model outputs such as encoder conditioning between
+workers without routing their tensor contents through the orchestrator. It is
+separate from native diffusion KV transfer: a deployment can use both, either, or
+neither. MiniMax-H3 is one consumer of this interface, not its implementation.
+
+## Architecture
+
+```mermaid
+flowchart TB
+    Config[Topology and deploy configuration] --> Keys[Input and output payload keys]
+    Keys --> Producer[Producer model runner]
+    Keys --> Consumer[Consumer model runner]
+    Producer --> Adapter[DiffusionStagePayloadMixin]
+    Consumer --> Adapter
+    Adapter --> Runtime[Shared OmniConnector runtime]
+    Runtime --> Descriptor[Request key and endpoint resolution]
+    Runtime --> Fanout[TP then SP or HSDP shard then replicate fanout]
+    Runtime --> Manager[OmniKVTransferManager owns lazy connector]
+    Manager --> SHM[SharedMemoryConnector]
+    Manager --> NIXL[NixlConnector]
+    Producer -. small transfer handle .-> Orchestrator[Orchestrator]
+    Orchestrator -. prompt and handle .-> Consumer
+    Consumer --> NativeKV[Separate native diffusion KV path]
+```
+
+The diffusion adapter owns model-facing selection and validation. The shared
+runtime owns request-key construction, endpoint metadata, retry policy, and
+distributed fanout. Backends own storage, transfer progress, and allocation
+lifetime. No pipeline should instantiate a NIXL agent or implement its own
+metadata-query loop.
+
+The synchronous runtime borrows the manager's connector. It does not create
+background AR receive/save loops or replace the manager's KV key callbacks.
+The asynchronous AR path continues using `get()` through its existing poller.
+Both paths share the descriptor and fanout helpers; their scheduling policies
+are deliberately different.
+
+## Configuration and Identity
+
+`stage_output_payload_keys` selects non-`None` entries in a producer's
+`DiffusionOutput.custom_output`. `stage_input_payload_keys` declares the required
+entries in the consumer's `prompt.additional_information`. Both normalize to
+tuples in `OmniDiffusionConfig`, including lists loaded from YAML.
+
+The model topology declares these keys; deploy configuration selects the
+connector for the edge. MiniMax-H3's standard and Turbo disaggregated profiles
+default to SharedMemory. Their `_nixl.yaml` overlays opt into NIXL with
+`--deploy-config`; choosing NIXL does not imply fallback to SharedMemory.
+
+Without an explicit handle, the key is:
+
+```text
+{external_req_id or request_id}_{from_stage}_{chunk_id}
+```
+
+The synchronous complete-payload path uses chunk zero. Source and destination
+IDs come from the configured edge, not from an assumption that stages are
+adjacent or that the producer is stage zero. An explicit handle can override
+the key, source, destination, and backend metadata:
+
+```python
+{
+    "key": "request-42_2_0",
+    "from_stage": "2",
+    "to_stage": "5",
+    "size_bytes": 4096,
+    "metadata": {...},
+    "payload_keys": ["text_encoder_output"],
+}
+```
+
+This descriptor is attached under `_stage_payload_transfer`. Backend metadata
+is opaque to the pipeline. If it is absent, request-scoped sender information
+can supply `source_host` and `source_port`. `payload_sender_info` takes
+precedence over legacy `kv_sender_info`. A flat endpoint is accepted directly;
+a stage-keyed map is resolved using the actual source stage, accepting string
+and integer keys. A single unambiguous candidate remains supported. Multiple
+unmatched candidates are not arbitrarily resolved to stage zero.
+
+## API Contracts
+
+| API | Owner | Contract |
+| --- | --- | --- |
+| `init_omni_connectors(..., synchronous=True)` | Shared runtime | Borrow an existing manager; no AR polling threads or KV callback mutation. |
+| `_stage_payload_recv_spec(...)` | Shared transport | Resolve explicit handle or conventional key plus request-scoped endpoint. |
+| `recv_stage_payload(..., retry_seconds=2.0)` | Shared transport | Leader retries with one monotonic deadline, retires discovery, then fans out success or absence. |
+| `get_with_deadline(..., deadline=...)` | Connector | Bound backend-controlled waits by the supplied absolute monotonic deadline. Return `(payload, bytes)` or `None`. |
+| `abandon_get(key)` | Connector | Retire unresolved discovery on the calling thread after the retry loop. Never free an active READ. Default is a no-op for stateless discovery. |
+| `_maybe_recv_stage_payload(request)` | Diffusion adapter | Skip warmup, receive, merge declared keys on the target device, and reject missing required values. |
+| `_maybe_send_stage_payload(requests, outputs)` | Diffusion adapter | Leader publishes selected complete outputs; all ranks attach the same handle and remove only transferred fields. |
+| `OmniKVTransferManager.close()` | Manager | Drain prefetch, close only the connector already created, and disable later lazy initialization. |
+
+`get_with_deadline()` is intentionally not a wrapper around an arbitrary
+blocking `get()`. The base implementation raises `NotImplementedError`.
+An unsupported backend must implement this capability before being used for
+synchronous remote-only payloads. Existing asynchronous `get()` users are
+unchanged.
+
+The deadline covers retry sleeps, metadata socket waits, and NIXL completion
+polling. It is **not a hard real-time bound** on Python execution, allocation,
+serialization, individual native driver calls, device copies, or payload-group
+collectives. A stuck native call cannot safely be preempted with a Python
+future timeout. The implementation does not start a detached READ and discard
+its ownership when the model-thread budget expires.
+
+## Successful Transfer
+
+The synchronous diffusion producer uses an explicit handle:
+
+```mermaid
+sequenceDiagram
+    participant P as Producer leader
+    participant B as Connector
+    participant O as Orchestrator
+    participant C as Consumer leader
+    participant R as Consumer peer ranks
+    P->>B: put(key, selected custom_output)
+    B-->>P: success, bytes, backend metadata
+    P->>P: Broadcast handle; remove transferred inline keys
+    P->>O: Output carrying small handle
+    O->>C: Prompt carrying handle and endpoint
+    C->>B: get_with_deadline(key, deadline)
+    B-->>C: Complete payload or absence
+    C->>B: abandon_get(key)
+    C->>R: Delivery status, then tensors if present
+    C->>C: Merge and validate required keys
+    R->>R: Merge and validate required keys
+```
+
+MiniMax-H3 instead uses the AR encoder's asynchronous sender. It does not
+forward the diffusion producer's explicit handle:
+
+```mermaid
+sequenceDiagram
+    participant A as AR encoder leader
+    participant S as Background sender
+    participant B as Connector
+    participant O as Orchestrator
+    participant C as Diffusion consumer leader
+    participant R as Consumer peer ranks
+    A->>S: Queue encoder conditioning for conventional key
+    A->>O: Output without inline conditioning
+    O->>C: Request identity and request-scoped sender endpoint
+    par Asynchronous publication
+        S->>B: put(conventional key, conditioning), with retries
+    and Bounded receive
+        C->>B: get_with_deadline(conventional key, endpoint)
+        B-->>C: Payload or absence
+    end
+    C->>B: abandon_get(key)
+    C->>R: Delivery status, then tensors if present
+    C->>C: Merge and validate encoder_output
+    R->>R: Merge and validate encoder_output
+```
+
+Within a TP-by-SP payload grid, exactly one rank, the leader in both groups,
+accesses the connector. The shared broadcaster traverses TP and then SP in
+the same order on every rank: TP first populates the source SP row, then SP
+propagates each TP column.
+
+When initialized HSDP groups span the stage, the adapter instead uses the
+HSDP shard and replicate dimensions. The shard broadcast populates the first
+replica row; the replicate broadcast propagates every shard position to the
+other rows. Both dimensions participate in leader election. This includes
+standalone HSDP, whose TP/SP groups are singletons, as well as pure sharding
+and pure replication. The two HSDP dimensions replace, rather than overlap,
+the TP/SP fanout.
+
+The single-consumer guarantee is scoped to these configured payload groups;
+it is not a global election across arbitrary independent DP/CFG/PP groups or
+stage replicas. Followers do not initialize or query a connector. Absence is
+broadcast too, so peers do not wait for tensor broadcasts that the leader will
+never issue.
+
+Pipelines must gather sharded producer outputs before publication. Selecting
+payload keys does not gather tensor shards automatically.
+
+## Execution and Failure Policy
+
+Full forward receives before pipeline execution and sends after output
+preparation. Stepwise execution receives once when a new cached state is
+created. That state retains the external request ID; final decoded output is
+published exactly once. Intermediate denoising steps and non-final streaming
+chunks are not published under the complete-payload key. Chunk-addressed
+streaming transfer remains the asynchronous path's separate contract.
+
+Dummy request IDs (`dummy_req_id` and its slash-prefixed variants) bypass both
+directions. A pure warmup batch does not lazily initialize the connector,
+publish buffers, wait for remote data, or enter payload fanout.
+
+For the synchronous diffusion producer, failed `put()` calls preserve inline
+fields; only successful publications replace those fields with a handle.
+This guarantee does not apply to H3's AR encoder asynchronous sender. That
+path omits inline conditioning before the background `put()` completes;
+exhausting its send retries does not restore conditioning to the forwarded
+output. It uses the conventional key and request-scoped sender endpoint,
+not the diffusion producer's handle-based handoff.
+
+If a receive fails, inline continuation is valid only when **all required
+keys have non-`None` values**. Missing required values raise an explicit
+stage-payload error on each rank after fanout. H3 with omitted inline
+conditioning therefore fails explicitly if remote publication or delivery
+fails; the original prompt is not equivalent conditioning. With no declared
+input keys, an explicit handle's `payload_keys` defines the required set.
+
+## NIXL Ownership and Recovery
+
+```mermaid
+stateDiagram-v2
+    [*] --> Published: put registers source
+    Published --> Released: unclaimed lease expires
+    Published --> Claimed: atomic metadata query
+    Claimed --> Reading: consumer submits READ
+    Claimed --> Recovering: reply lost and caller abandons
+    Recovering --> Recovering: query or ACK lost
+    Recovering --> Released: same claim recovered and ACKed
+    Reading --> Deferred: deadline expires while DMA may be active
+    Reading --> Released: terminal DMA and completion ACK
+    Deferred --> Deferred: active or unknown DMA state
+    Deferred --> Released: terminal state and completion ACK
+```
+
+Claims are scoped by payload generation and UUID. Query timeouts or lost
+metadata replies retain the cached UUID for retries, since the producer may
+already have granted ownership. An explicit `META_NOT_FOUND` reply clears
+the cached claim; a subsequent query can use a new UUID. At the end of the
+synchronous retry loop, `abandon_get()`
+moves unresolved queries to background recovery. Recovery asks for the same
+claim, obtains the generation, and acknowledges it **without submitting a
+READ**. A lost completion reply is retried against that exact generation,
+not by acquiring a fresh payload with the same key. Existing claims may be
+recovered while a producer is closing; new claims may not.
+
+If a READ was submitted and its state is active or unknown, timeout does not
+release source or destination allocations. The deferred-transfer reaper keeps
+tensors, registrations, descriptor lists, and transfer handles until NIXL
+reports a terminal state. Completion ACK failure also stays in background
+retry. Producer close can therefore remain deferred until ownership drains.
+
+Recovery is eventual, not guaranteed under permanent network partition or
+consumer process death. A producer cannot safely infer that an unreachable
+consumer has stopped DMA. Conservatively retained ownership in that case is
+an explicit safety limit, distinct from a live consumer silently abandoning
+a recoverable discovery claim.
+
+## Reusing the Interface
+
+| Backend | Synchronous support | Backend responsibility |
+| --- | --- | --- |
+| SharedMemory | Implemented | Nonblocking read-lock acquisition; return absence on contention; shared runtime retries. Copy and deserialize after acquiring the lock. |
+| NIXL | Implemented | Deadline-aware socket and completion waits, generation-scoped claims, deferred active-DMA cleanup, discovery recovery. |
+| Other existing connectors | Opt-in required | Implement `get_with_deadline`; implement `abandon_get` only when discovery creates ownership. Do not inherit an unbounded blocking call. |
+
+To add a model, declare input/output keys and use ordinary prompt and output
+structures. To add a backend, implement the connector capability and ownership
+tests, register it through the existing factory, and select it in deploy
+configuration. Neither requires a backend-specific diffusion model runner.
+
+Worker teardown disables offloading, closes the manager after draining
+prefetch, then tears down native KV and distributed resources in `finally`
+blocks. The synchronous mixin does not close the borrowed connector itself.
+Repeated manager close is safe, including for an unused manager.
+
+## Validation Map
+
+| Test module | Main invariants |
+| --- | --- |
+| `tests/diffusion/test_diffusion_stage_payload.py` | Warmup bypass, strict missing-payload errors, valid inline fallback, source-stage endpoints, TP/SP fanout, standalone HSDP election and SHM fanout, handle publication. |
+| `tests/diffusion/test_diffusion_model_runner.py` | New-request receive and final-only stepwise publication with cached external identity. |
+| `tests/distributed/omni_connectors/test_nixl_connector.py` | Lost replies, same-claim recovery, deadlines, generation isolation, active DMA retention, deferred close and ACKs. |
+| `tests/distributed/omni_connectors/test_shm_connector.py` | Key and metadata reads, lock contention, expired deadline, cleanup. |
+| `tests/worker/test_omni_connector_mixin.py` | Borrowed connector ownership and shared synchronous/asynchronous helpers. |
+| `tests/distributed/omni_connectors/test_kv_async_prefetch.py` | Prefetch ordering and idempotent manager close. |
+| `tests/config/test_omni_config.py` | Tuple normalization across topology, deploy, stage CLI, and serialization. |
+
+Mocked ownership tests do not establish native transport performance or
+accelerator correctness. Native NIXL and model E2E runs are separate checks;
+their results must name the tested revision and runtime.

docs/design/feature/prefix_caching.md

@@ -8,6 +8,7 @@
 - [High-Level Approach](#high-level-approach)
 - [Example](#example)
 - [What About Multimodal Inputs?](#what-about-multimodal-inputs)
+- [Diffusion KV Prefix Caching](#diffusion-kv-prefix-caching)
 - [Implementation](#implementation)
 - [Related Files](#related-files)
 
@@ -25,7 +26,7 @@ vLLM implements automatic prefix caching for managing its kv-cache, which is bes
 - Model / stage specific multimodal data
 
 !!! note "Note 1"
-    This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice.
+    The following sections describe caching tensor outputs passed between AR stages. This is distinct from multimodal encoder caching and from [diffusion KV prefix caching](#diffusion-kv-prefix-caching), which reuses GPU KV pages inside a diffusion stage.
 
 ### High-Level Approach
 
@@ -45,7 +46,7 @@ Host footprint: each cached key costs `num_blocks × block_size × D × dtype_by
 ### Example
 
 !!! note "Note 3"
-    Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config.
+    The stage-output tensor cache described here supports AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config. Diffusion KV caching has the separate requirements below.
 
 The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
 
@@ -167,6 +168,65 @@ When we pass our multimodal tensors to the language model component in the same
 
 Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled.
 
+### Diffusion KV Prefix Caching
+
+HunyuanImage3's standalone DiT pipeline can reuse stable text/reference-image KV
+across requests. Enable it on the diffusion stage in the deploy config:
+
+```yaml
+pipeline: hunyuan_image3_dit
+stages:
+  - stage_id: 0
+    diffusion_kv_mode: paged_scheduler
+    enable_prefix_caching: true
+```
+
+The Scheduler uses vLLM's native `KVCacheManager` to look up and retain complete
+prefix blocks. The Worker owns the GPU pages and block tables. On a hit, the first
+denoise forward computes only the uncached query suffix while attending to both
+cached and new KV. Dynamic target-image KV is never published as a reusable prefix.
+
+Cache identity includes token IDs, reference-image content and VAE random state,
+plus model/layout and LoRA context. The same image with different prompts can
+reuse the common leading blocks; it does not imply that every image span or CFG
+branch is interchangeable. Disabling prefix caching skips cache-identity hashing;
+`dense_legacy` remains the default. Enabling prefix caching with a mode other than
+`paged_scheduler`, without a registered model hook, or together with native
+`kv_transfer_config` raises a configuration error.
+Combining it with `enable_sleep_mode: true` is also rejected: sleep discards KV
+pages without invalidating the Scheduler's prefix-cache index. Disable either
+prefix caching or sleep mode.
+
+The current scope is local DiT reuse, not AR-imported KV or missing-page-only
+cross-stage transfer. Prefix-hit accuracy has been exercised with TP4/SP1 and
+TP2/SP2 (Ulysses), both with EP and CFGP1. Other combinations still require validation. See the
+[diffusion compatibility notes](../../user_guide/diffusion_features.md#diffusion-kv-prefix-caching).
+
+#### Shared-reference benchmark
+
+From the repository root, use the unified `vllm bench serve --omni` test runner:
+
+```bash
+pytest tests/dfx/perf/scripts/run_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json
+```
+
+Run GPU tests through your environment's GPU scheduler. The suite reuses the
+two-image IT2I accuracy input on a single DiT stage, comparing dense, paged without
+prefix caching, and paged with prefix caching. Each mode runs two ordinary upstream
+warmups followed by eight identical requests at CFG 2.5, 8 denoise steps and seed 42.
+Warmups are excluded from latency and request throughput. This measures full-prefix
+reuse, not the earlier distinct-prompt workload; its results must be reported separately.
+
+`tests/e2e/accuracy/test_hunyuan_image3_prefix_cache_accuracy.py` uses the same input
+at CFG 2.5, 50 steps and seed 42, comparing all modes against the existing official
+Instruct output image and IT2I thresholds (CLIP ≥90, SSIM ≥0.26, PSNR ≥12.5).
+A different prompt warms the cache; the original prompt then exercises partial
+and repeated hits, with reference-image reuse and query slicing verified.
+Only outputs for the original prompt are scored. The golden comes from AR-to-DiT;
+this single-DiT quality check does not establish identical conditioning between
+the pipelines. The existing AR-to-DiT test remains unchanged.
+
 ### Implementation
 
 The block/slot model is `vllm_omni/core/prefix_cache/`.

docs/design/feature/preprocess_phase_contract.md

@@ -42,7 +42,8 @@ Scheduler / current input batch
        prefill / other spans     single-token decode
           preprocess             preprocess or preprocess_decode_batch
               |                              |
-       prompt embedding              runner-managed talker_mtp
+       prompt embedding              decode embeddings
+              |                      talker_mtp (if present)
                   \                 /
                     model forward
 ```
@@ -96,9 +97,24 @@ computed per-row fields are not promised at that earlier boundary.
 ### Routing guarantee
 
 With `has_preprocess=True`, the batched decode hook is eligible only when it is
-callable, the runner has `talker_mtp`, the scheduled span is one token, and
+callable, the scheduled span is one token, and
 `_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs.
 
+The hook's return format depends on whether the runner manages MTP:
+
+| Runner capability | `preprocess_decode_batch` return tuple |
+| --- | --- |
+| Without `talker_mtp` | `(input_ids, inputs_embeds, updates)` |
+| With `talker_mtp` | `(input_ids, inputs_embeds, hidden_states, text_step, updates)` |
+
+Each tensor's batch dimension and each entry in the `updates` list correspond
+to one physical request row, including separate CFG branches. Non-MTP batches
+contain contiguous single-token decode spans; the runner flushes a batch before
+processing any other span. It copies embeddings into its existing input buffer
+and applies each update dictionary to the corresponding request. Returning the
+original `input_ids` object lets it skip an unchanged token-ID copy. The existing
+five-result MTP format remains supported.
+
 Runner-managed MTP is likewise restricted to single-token decode rows. A prefill
 row, including a one-token tail, is never added to that MTP batch: the runner
 does not consume its `mtp_inputs`, generate codes for it, or overwrite its
@@ -130,6 +146,8 @@ formalizes that existing behavior. Prefer upgrading the matching vLLM and
 vLLM-Omni release pair and consuming the explicit phase keys. Do not silently
 fall back to `input_ids.numel() > 1` when the model requires chunking correctness.
 For a custom or backported runtime, verify both metadata and MTP guards.
+Older runners can restrict `preprocess_decode_batch` to MTP models; a non-MTP
+model must retain its scalar `preprocess` implementation for those runtimes.
 
 ## Correctness and testing plan
 

docs/design/feature/skip_softmax.md

@@ -178,7 +178,7 @@ t = s * u / (1 + (s - 1) * u)
 ```
 
 and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
-`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
+`num_inference_steps = 49` (50 sigma points, `N = 49` denoiser forwards), the sequence
 starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
 
 | `disabled_until_timestep` | Dense steps | Skip-Softmax steps |

docs/design/module/diffusion/diffusion_model_integration.md

@@ -57,3 +57,58 @@ cross-stage routing.
 
 Test registry selection, checkpoint loading, minimal inference, input and output
 contracts, and every declared optional capability.
+
+## Exact conditioning projection reuse
+
+Models with repeated conditioning projections can use `ExactProjectionCache`
+from `vllm_omni.diffusion.cache`. It stores actual projection outputs for identical
+conditioning tensors and unchanged weights, with a byte budget and TP-wide hit
+coordination. It has no model names, sampling schedules, projection layouts, or
+artifact formats built in. MiniMax-H3 is its first model integration.
+
+Keep one cache per model instance and prepare each conditioning tensor once
+before its projections. For a leaf linear with fixed SiLU preprocessing:
+
+```python
+from torch.nn.functional import silu
+
+from vllm_omni.diffusion.cache import ExactProjectionCache
+
+cache = ExactProjectionCache(max_bytes=256 * 1024**2)
+
+# Inside an inference forward, with gradients disabled:
+cache.prepare(conditioning)
+modulation = cache.project(
+    "block.0.modulation",
+    projection,
+    conditioning,
+    lambda: projection(silu(conditioning)),
+)
+```
+
+The original computation must depend only on the supplied conditioning tensor,
+the leaf projection's weights/buffers, and fixed preprocessing. Include prompt,
+guidance, modality, or other request-dependent conditioning in the actual input;
+a timestep-only key is unsafe when those affect the result. Do not mutate the
+input between `prepare` and `project`, or modify outputs returned on cache hits.
+Use separate cache instances for concurrent forwards with different inputs.
+
+Call `clear()` when loading or switching adapters, changing projection behavior,
+or moving the model. Weight-version checks invalidate individual entries.
+Compilation, gradient-enabled execution, projection hooks, and unversioned
+inference weights use the original computation. Validate cold/warm parity,
+invalidation, and rank-local misses when integrating another model.
+Hook checks include global hooks and hooks on wrapped submodules. Numerical
+settings are captured at each projection call, so entering an autocast context
+after `prepare()` cannot reuse an output from a different precision.
+
+Validated offline results can be supplied through `_lookup_precomputed()`;
+this extension is consulted only at TP1 on runtime-cache misses. Artifact
+validation, weight binding, and model-specific layouts remain in the adapter.
+H3's adapter retains its main/Ref2VA sidecars and their eager-only admission.
+
+The cache retains projection results and leaves all original weights available.
+It does not reduce weight residency or cancel layerwise prefetch. A future
+integration that skips weight transfers must coordinate with `OffloadPlan`
+before prefetch and provide a safe miss path; moving weight storage currently
+causes conservative cache invalidation.

docs/features/README.md

@@ -59,6 +59,7 @@ defines the shared configuration and backend extension points.
 | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
 | Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
 | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
+| Speed up Wan VAE video decode | [Wan VAE Decoder Fast Path](../user_guide/diffusion/vae_fast_path.md) | User-facing optimization guide |
 | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
 | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
 | Apply diffusion adapters | [LoRA](../user_guide/diffusion/lora.md) | User-facing extension guide |

docs/getting_started/installation/npu/npu.inc.md

@@ -10,25 +10,25 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 === "A2"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0
     ```
 
 === "A3"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a3
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a3
     ```
 
 === "A5"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a5
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a5
     ```
 
 === "310P (Experimental)"
 
     ```bash
-    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-310p
+    export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-310p
     ```
 
 ```bash
@@ -53,7 +53,7 @@ docker run --rm \
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.29.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.30.0rc1 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -71,10 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 
 # --8<-- [start:installation-main]
 
-You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.29.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
+You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
 
 ```bash
-# Inside the vLLM-Ascend v0.29.0 container, install vLLM-Omni from main
+# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
 cd /vllm-workspace
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -103,7 +103,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 === "A2"
 
     ```bash
-    export IMAGE=quay.io/ascend/vllm-omni:v0.29.0
+    export IMAGE=quay.io/ascend/vllm-omni:v0.30.0
     docker run --rm \
         --name vllm-omni-a2 \
         --shm-size=64g \
@@ -127,7 +127,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 === "A3"
 
     ```bash
-    export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a3
+    export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-a3
     docker run --rm \
         --name vllm-omni-a3 \
         --shm-size=64g \
@@ -151,7 +151,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 === "A5"
 
     ```bash
-    export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a5
+    export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-a5
     docker run --rm \
         --name vllm-omni-a5 \
         --shm-size=64g \
@@ -175,7 +175,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 === "310P (Experimental)"
 
     ```bash
-    export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-310p
+    export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-310p
     docker run --rm \
         --name vllm-omni-310p \
         --shm-size=64g \

docs/models/supported_models.md

@@ -72,7 +72,7 @@ th {
 | `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) |
 | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
 | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
-| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
+| `BreezeForConditionalGeneration` | Breeze-TTS-2 (voice design, voice clone, voice direction, CFG) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/serving/video_stream_api.md

@@ -101,7 +101,7 @@ The server accepts these legacy field names and rewrites them before validation.
 | Variable | Values | Default | Description |
 | -------- | ------ | ------- | ----------- |
 | `VLLM_VIDEO_ASYNC_CHUNK` | `on`, `off` | `on` | Wire-level streaming switch. `off` buffers server-side deltas and emits coalesced outputs at the end of a query. |
-| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Audio delta extraction strategy. `fast` emits only newly produced chunks; `slow` recomputes from accumulated audio and exists for A/B verification. |
+| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Both settings forward every fresh engine audio delta. `slow` is retained as a compatibility alias for `fast`. |
 
 ## EVS Semantics
 

docs/user_guide/diffusion/lora.md

@@ -27,6 +27,21 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
 - `lora_alpha`: LoRA alpha scaling factor
 - `target_modules`: List of module names to apply LoRA to
 
+Every module supplied by a PEFT adapter must bind successfully when the adapter
+is activated. If any module cannot bind, activation raises an error naming the
+unbound modules and resets the LoRA layers to avoid leaving the adapter partly
+active. An adapter may still target only some projections of a fused layer, such
+as Q and V without K, as long as all supplied modules bind.
+
+Previously, some adapters activated even when unmatched modules were silently
+dropped. Those adapters now fail at activation. For example, in-memory tensor
+adapters for Qwen-Image, Boogu-Image, and Wan can contain checkpoint keys such
+as `to_out.0` where the engine expects `to_out`. PEFT checkpoints with extra
+unbound keys also fail at activation if those keys pass the loader's earlier
+validation; unsupported module suffixes may already fail during loading.
+Checkpoint-to-engine name mapping is tracked in [#8001](https://github.com/vllm-project/vllm-omni/issues/8001);
+until it is available, adapter keys must resolve to supported engine modules.
+
 ## Quick Start
 
 ### Offline Inference
@@ -244,14 +259,22 @@ Notes:
 
 ## MiniMax-H3 adapter-declared schedules
 
-MiniMax-H3 supports two few-step mechanisms that must not be conflated:
+MiniMax-H3 counts `num_inference_steps` as denoising updates: requesting `8`
+produces eight updates using nine sigma boundaries, including the terminal zero.
+Earlier uniform-schedule requests counted boundaries instead; to reproduce an
+old N-point schedule, request N-1 steps. This also changes Turbo requests from
+`5`/`9` to `4`/`8`, preserving their sampling grids. Checkpoint-pinned and native
+LoRA schedules already count updates and are unchanged.
+
+MiniMax-H3 supports the following few-step mechanisms:
 
 - **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
   `model_index.json`. Requests must pass `num_inference_steps` as the interval
   count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
 - **Runtime Turbo LoRA**: each LightX2V Turbo artifact carries its own contract,
-  read from its filename. A four-step artifact requests five sigma points and an
-  eight-step one requests nine; the 768p retrains enforce `flow_shift=6` and the
+  read from its filename. A four-step artifact requests `num_inference_steps=4`
+  and an eight-step one requests `8`; the scheduler adds the terminal sigma
+  boundary. The 768p retrains enforce `flow_shift=6` and the
   544p artifacts `flow_shift=12`. `audio_flow_shift=3` across the family. A
   request that does not match the loaded artifact is rejected by name. Alpha
   comes from the artifact's metadata, or 8 when it declares none. `ref2v`

docs/user_guide/diffusion/vae_fast_path.md

@@ -0,0 +1,102 @@
+# Wan VAE Decoder Fast Path
+
+Video pipelines built on the Wan 2.1/2.2 causal VAE (Cosmos3, Wan2.2, LingBot,
+Helios, SANA-Video, LongCat-Video, DreamZero) spend a large share of their
+end-to-end time decoding latents to pixels. The decoder is a chunked 3D
+convolutional network whose reference implementation in diffusers moves every
+activation through memory several times per layer (normalization, activation,
+causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
+installs a fast path on every loaded Wan VAE that fuses this data movement into
+a handful of Triton kernels while leaving the convolutions themselves untouched.
+
+## Levels
+
+The fast path is controlled by `--vae-fast-path` (engine argument
+`vae_fast_path`, deploy-config key `vae_fast_path`):
+
+| Level | Default | Output vs. diffusers | What it does |
+|---|---|---|---|
+| `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
+| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus decoder convolution weights converted to channels-last memory format so cuDNN picks its channels-last kernels, and a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the bias of the preceding `conv1`. |
+| `off` | no | bit-identical | Reference diffusers decoder. |
+
+```bash
+# Default: bit-exact fast path
+vllm serve nvidia/Cosmos3-Nano --omni
+
+# Fastest: channels-last decoder (not bit-exact)
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path channels_last
+
+# Reference implementation, e.g. to bisect a quality issue
+vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path off
+```
+
+The equivalent per-stage deploy configuration is:
+
+```yaml
+stages:
+  - stage_id: 0
+    vae_fast_path: channels_last
+```
+
+## Behavior and limitations
+
+- Only CUDA is supported. On other platforms the reference decoder runs.
+- The fast path is installed once per VAE instance when the pipeline is
+  initialized. It rebinds the forwards of the loaded decoder modules; parameter
+  names, `state_dict` keys and weight loading are unchanged.
+- Installation is skipped, with a logged reason, if a custom forward wrapper
+  or forward hook would be replaced or bypassed by the fused paths. Wrappers
+  and hooks on modules that the fast path still calls normally are preserved.
+- Every fused kernel validates its inputs and falls back to the exact PyTorch
+  expression for anything it does not support (unusual dtypes or layouts, CPU
+  tensors, autograd enabled, `torch.compile` tracing).
+- The SiLU activation is folded into the normalization kernel only after an
+  exhaustive self-test over all bf16/fp16 values proves the fused epilogue
+  bit-identical to `F.silu` on the running toolkit; otherwise SiLU stays a
+  separate operation. The startup log reports `fused_silu=...`.
+- Spatially sharded VAE decode (`--vae-parallel-mode spatial_shard_height` or
+  `spatial_shard_width`) is not combined with the fast path; the installer skips
+  the VAE and logs why. Tiled and tile-parallel decode work unchanged.
+- The `channels_last` level changes the order in which cuDNN accumulates
+  convolutions, so outputs differ from the reference in the last bits. Use
+  `lossless` when bitwise reproducibility against diffusers matters.
+- For VAEs kept in fp32, cuDNN's channels-last convolution algorithms run in
+  TF32 under PyTorch's default `torch.backends.cudnn.allow_tf32 = True`, which
+  dominates the difference to the reference (about 1e-3). Set
+  `torch.backends.cudnn.allow_tf32 = False` if full fp32 convolutions are
+  required.
+
+## Measured
+
+Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_decode.py`
+(best of 2 runs after warmup):
+
+| `--vae-fast-path` | Decode time | Speedup | Output vs. `off` |
+|---|-------------|---------|---|
+| `off` | 6.00 s      | 1.00x   | reference |
+| `lossless` | 3.26 s      | 1.84x   | bit-identical |
+| `channels_last` | 2.51 s      | 2.39x   | PSNR 62.6 dB, max abs diff 4.6e-2 |
+
+Peak decode memory was unchanged (about 13 GiB): it is set by the largest
+activations and the cuDNN workspace, not by the output assembly.
+
+## Benchmarking
+
+`benchmarks/diffusion/bench_wan_vae_decode.py` decodes seeded latents with the
+real VAE at each level and reports decode time, peak memory, bitwise equality
+and PSNR against the `off` level. `--profile` prints a per-kernel table with the
+convolution share and any layout-transpose kernels:
+
+```bash
+python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile
+```
+
+The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
+the decode is timed across all ranks and rank 0 reports:
+
+```bash
+torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
+    --vae-patch-parallel-size 2 --vae-parallel-mode tile
+```

docs/user_guide/diffusion_features.md

@@ -32,6 +32,30 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
 | **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
 | **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
 
+#### Diffusion KV Prefix Caching
+
+[KV prefix caching](../design/feature/prefix_caching.md#diffusion-kv-prefix-caching)
+reuses stable context KV across requests, rather than approximating denoise
+steps. Enable `diffusion_kv_mode: paged_scheduler` and `enable_prefix_caching: true`
+on the HunyuanImage3 standalone DiT stage. This is separate from the AR stage-output
+cache, TeaCache and Cache-DiT.
+
+| Model / combination | Scope |
+|---------------------|-------|
+| HunyuanImage3 standalone DiT | Stable text/reference-image prefix reuse; dynamic image tokens are excluded |
+| TP4/SP1 + EP, CFGP1 | Validated reference-image prefix hits |
+| TP2/SP2 (Ulysses) + EP, CFGP1 | Validated reference-image prefix hits; CFG guidance does not require CFG parallelism |
+| Other SP/TP combinations, CFGP>1 | Prefix-hit E2E accuracy not yet validated |
+| AR-imported KV / cross-stage missing-page transfer | Not covered by local prefix caching |
+| Sleep mode | Rejected with prefix caching; discarded KV pages would leave stale cache hits |
+| TeaCache / Cache-DiT, offload, quantization | Not validated with prefix hits |
+| Step execution | Not supported with HunyuanImage3 paged KV; use request-level execution |
+| Other diffusion models | No prefix-cache model adapter yet |
+
+This table describes prefix-cache scope; the general parallelism tables below do
+not establish prefix-hit compatibility. Floating-point kernel differences still
+require accuracy validation. `dense_legacy` remains the default.
+
 #### Lossless Acceleration
 
 Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
@@ -252,6 +276,7 @@ The Diffusion Acceleration navigation groups the remaining guides as follows:
 | Parallelism | [Parallelism Overview](diffusion/parallelism/overview.md) |
 | Attention | [Attention Backends](diffusion/attention_backends.md) |
 | Compilation | [Regional Compilation](diffusion/regional_compilation.md) |
+| VAE decode | [Wan VAE Decoder Fast Path](diffusion/vae_fast_path.md) |
 | Video extension | [Frame Interpolation](diffusion/frame_interpolation.md) |
 | Startup | [Startup and Loading](diffusion/startup_and_loading.md) |
 | Adapters | [LoRA](diffusion/lora.md) |

docs/user_guide/quantization/fp8.md

@@ -65,7 +65,17 @@ vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
 `QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
 In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
 the image — so the first cold start does not recompile. The engine's startup dummy
-run already exercises the kernels, so with a warm cache the first real request is fast.
+run exercises the kernels, but new shapes, layouts, dtypes, or bias settings may
+still need compilation or tuning. The warmup helper uses inference mode and
+transposed weights without bias. Daemon workers compile candidates in-process
+while retaining autotuning and caching.
+
+Scale validation and Quack/FlashInfer dispatch run inside a PyTorch custom op.
+This keeps layer-specific scale addresses and validation-cache updates out of
+Dynamo tracing without introducing a graph break. Unpopulated scales and Quack
+failures still fall back to FlashInfer at runtime. CUDA graph capture additionally
+requires warming the dispatch with populated scales; Python validation and
+dispatch do not rerun during CUDA graph replay.
 
 To pre-warm specific shapes (e.g. at image build time):
 

examples/offline_inference/text_to_image/README.md

@@ -222,6 +222,8 @@ Anima uses the official model dimensions by default. A checkpoint may provide a
 
 To load Anima, point `--model` to the single-file checkpoint path, pass the native pipeline class name using `--model-class-name`, and supply the converted components directory using `--custom-pipeline-args`:
 
+Download the checkpoint and components using the [Anima recipe](../../../recipes/circlestone-labs/Anima.md#command---prepare-assets), then replace the paths below with their local locations.
+
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
   --model /path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors \

examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py

@@ -14,12 +14,10 @@ from pathlib import Path
 
 import soundfile as sf
 import torch
-from transformers import AutoConfig
+from transformers import AutoTokenizer
 
 from vllm_omni import Omni
-from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import (
-    BreezeTTS2PromptBuilder,
-)
+from vllm_omni.model_executor.models.breeze_tts_2.prompt import DEFAULT_INSTRUCTION, build_breeze_prompt
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 REPO_ROOT = Path(__file__).resolve().parents[4]
@@ -95,8 +93,7 @@ def resolve_template(args) -> str:
 def reference_payload(ref_audio_path: str | None) -> dict:
     """Load the reference clip into the payload keys the prompt builder reads.
 
-    ``BreezeTTS2PromptBuilder`` consumes ``ref_audio`` (a waveform) plus
-    ``ref_audio_sample_rate``; a bare path field is not read.
+    The prompt helper consumes a waveform and its sample rate.
     """
     if ref_audio_path is None:
         return {}
@@ -107,36 +104,17 @@ def reference_payload(ref_audio_path: str | None) -> dict:
 
 
 def build_prompt(args) -> dict:
-    config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
-    reference_encoder = None
-    if args.ref_audio is not None:
-        from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import (
-            BreezeReferenceAudioTokenizer,
-        )
-
-        reference_encoder = BreezeReferenceAudioTokenizer.from_pretrained(
-            args.model,
-            num_codebooks=int(getattr(config, "num_codebooks", 16)),
-            device_map="cpu",
-        )
-    builder = BreezeTTS2PromptBuilder.from_pretrained(
-        args.model,
-        config,
-        reference_audio_encoder=reference_encoder,
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    payload = reference_payload(args.ref_audio)
+    reference = (payload["ref_audio"], payload["ref_audio_sample_rate"]) if payload else None
+    return build_breeze_prompt(
+        tokenizer,
+        args.text,
+        DEFAULT_INSTRUCTION if args.instruction is None else args.instruction,
+        speaker=args.voice,
+        ref_audio=reference,
+        ref_text=args.ref_text,
     )
-    payload = {
-        "text": args.text,
-        "speaker": args.voice,
-        **reference_payload(args.ref_audio),
-    }
-    if args.instruction:
-        payload["instruction"] = args.instruction
-    if args.ref_text:
-        payload["ref_text"] = args.ref_text
-    prompt = builder.build(payload, template=resolve_template(args))
-    if args.max_new_tokens is not None:
-        prompt["additional_information"]["breeze_max_new_frames"] = int(args.max_new_tokens)
-    return prompt
 
 
 def extract_audio(multimodal_output: dict) -> torch.Tensor:
@@ -180,7 +158,12 @@ def main():
     print(f"Prompt len  : {len(prompt['prompt_token_ids'])} tokens")
 
     t_start = time.perf_counter()
-    outputs = engine.generate([prompt])
+    if args.max_new_tokens is not None:
+        sampling = [params.clone() for params in engine.resolve_sampling_params_list(None)]
+        sampling[0].max_tokens = args.max_new_tokens
+        outputs = engine.generate([prompt], sampling_params_list=sampling)
+    else:
+        outputs = engine.generate([prompt])
     elapsed = time.perf_counter() - t_start
 
     request_output = outputs[0]

examples/offline_inference/text_to_speech/moss_tts_nano/end2end.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
 """Offline inference example for MOSS-TTS-Nano via vLLM-Omni.
 
 Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec
@@ -136,21 +136,20 @@ def main(args) -> None:
     )
     params_list = sampling_params
 
-    for stage_outputs in omni.generate(inputs, params_list):
-        for i, req_output in enumerate(stage_outputs):
-            for j, out in enumerate(req_output.outputs):
-                mm = out.multimodal_output
-                if mm is None:
-                    print(f"  [req {i}] No audio output.")
-                    continue
-                audio = mm.get("audio")
-                sr_tensor = mm.get("sr")
-                if audio is None:
-                    print(f"  [req {i}] No waveform in multimodal_output.")
-                    continue
-                sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
-                out_path = str(output_dir / f"output_{i}_{j}.wav")
-                save_audio(audio.cpu(), out_path, sr)
+    for i, req_output in enumerate(omni.generate(inputs, params_list)):
+        for j, out in enumerate(req_output.outputs):
+            mm = out.multimodal_output
+            if mm is None:
+                print(f"  [req {i}] No audio output.")
+                continue
+            audio = mm.get("audio")
+            sr_tensor = mm.get("sr")
+            if audio is None:
+                print(f"  [req {i}] No waveform in multimodal_output.")
+                continue
+            sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
+            out_path = str(output_dir / f"output_{i}_{j}.wav")
+            save_audio(audio.cpu(), out_path, sr)
 
     print("Done.")
 

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import base64
 import concurrent.futures
 import os
@@ -457,7 +460,8 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
                         f.write(audio_data)
                     print(f"Audio saved to {audio_file_path}")
                     count += 1
-                elif choice.message.content:
+                # A choice can contain both audio and text.
+                if choice.message.content:
                     print("Chat completion output from text:", choice.message.content)
     else:
         printed_content = False

examples/online_serving/text_to_speech/README.md

@@ -1061,7 +1061,7 @@ vllm serve BreezeBlue/Breeze-TTS-2 --omni --port 8091
 ./breeze_tts_2/run_server.sh
 ```
 
-The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming, 8-frame inter-stage chunks).
+The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming with 1, 2, 4, then 5-frame inter-stage chunks).
 
 ### Sending requests
 

mkdocs.yml

@@ -139,15 +139,16 @@ plugins:
             show_source: false
           inventories:
           - https://docs.python.org/3/objects.inv
-          - https://typing-extensions.readthedocs.io/en/latest/objects.inv
+          # Read the Docs rate-limits these external inventories in strict builds.
+          # - https://typing-extensions.readthedocs.io/en/latest/objects.inv
           # Disabled because Read the Docs strict builds fail if this
           # external inventory intermittently resets the connection.
           # - https://docs.aiohttp.org/en/stable/objects.inv
-          - https://pillow.readthedocs.io/en/stable/objects.inv
+          # - https://pillow.readthedocs.io/en/stable/objects.inv
           - https://numpy.org/doc/stable/objects.inv
           # Temporarily disabled due to decompression errors
           # - https://pytorch.org/docs/stable/objects.inv
-          - https://psutil.readthedocs.io/stable/objects.inv
+          # - https://psutil.readthedocs.io/stable/objects.inv
 
 markdown_extensions:
   - attr_list

pyproject.toml

@@ -331,6 +331,7 @@ extend-ignore-identifiers-re = [
     ".*arange",
     ".*MoBA",
     ".*temperal_downsample",
+    ".*temperal_upsample",
     ".*_nd",
     ".*nd$",
     ".*EmbedND$",

recipes/BreezeBlue/Breeze-TTS-2.md

@@ -1,210 +1,307 @@
 # Breeze-TTS-2
 
-> Text-to-speech serving (plain / voice design / voice clone / voice direction)
+[Breeze-TTS-2](https://huggingface.co/BreezeBlue/Breeze-TTS-2) generates 24 kHz
+English and Chinese speech. The integration supports Voice Design from text
+instructions, Voice Clone from a reference recording and transcript, and
+Voice Direction combining reference conditioning with instructions.
 
-## Summary
+## Installation and serving
 
-- Vendor: BreezeBlue
-- Model: `BreezeBlue/Breeze-TTS-2` (`BreezeForConditionalGeneration`)
-- Task: Text-to-speech with speaker tags, natural-language style instructions, single-reference voice cloning, and reference+instruction "voice direction"
-- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API (streaming and non-streaming); offline `Omni` example
-- Hardware: 1x NVIDIA L20 48GB (CUDA)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for serving
-Breeze-TTS-2 with vLLM-Omni. Breeze-TTS-2 is a two-stage AR TTS pipeline:
-
-| Stage | Components | Output |
-| ----- | ---------- | ------ |
-| 0 (talker, `LLM_AR`) | T5Gemma2 text encoder + Qwen3 backbone + 2052-way codebook-0 head + 12-layer depth decoder | 16-codebook codec frames |
-| 1 (codec, `LLM_GENERATION`) | Bundled Qwen3-TTS audio tokenizer decoder (Mimi fallback) | 24 kHz mono waveform |
-
-Four prompt modes are selected automatically from the request fields:
-
-| Mode | Trigger | Fields |
-| ---- | ------- | ------ |
-| `tts_plain` | only `input` (default `voice=S0`) | `input`, `voice` |
-| `tts_instruction` (voice design) | `input` + `instructions` | style/emotion/pace directives, not spoken |
-| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript |
-| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery |
-
-Current scope: greedy sampling with `cfg_scale=1.0` and a default
-`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and
-`negative_prompt` are rejected with explicit client errors until companion
-request support lands.
-
-The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py)
-also uses repetition penalty `1.1`, but its
-[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py)
-sample both the backbone and depth decoder with temperature `0.9` and top-k
-`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design)
-use CFG `4`. Those sampling and guidance settings are outside the current
-port's scope, so the examples below do not establish parity with the official
-voice-design outputs.
-
-## Supported model contract
-
-| Item | Value |
-| ---- | ----- |
-| Tasks | Text-to-speech through `/v1/audio/speech` (online) and the offline `Omni` example |
-| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) |
-| Languages | English and Chinese (model card) |
-| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz |
-| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) |
-| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) |
-| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; other guidance values and `negative_prompt` are rejected |
-| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` |
-
-## References
-
-- Issue: [#6656 [New Model]: Breeze TTS 2](https://github.com/vllm-project/vllm-omni/issues/6656)
-- Related examples under `examples/`:
-  [`examples/online_serving/text_to_speech/breeze_tts_2/`](../../examples/online_serving/text_to_speech/breeze_tts_2/),
-  [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../examples/offline_inference/text_to_speech/breeze_tts_2/)
-
-## Hardware
-
-- Accelerator model and per-device memory: 1x NVIDIA L20 48GB
-- Number of devices: 1 (both stages are placed on device 0)
-- Device interconnect: not applicable (single device)
-- Host memory: no special requirement; per-request codec state is small
-- Qualification scope: functional serving of all four prompt modes, streaming PCM, and 4 concurrent mixed-mode requests (greedy, `cfg_scale=1.0`). No other accelerator is qualified by this recipe.
-
-## Software environment
-
-- OS: Linux
-- Python: 3.12.12
-- Driver / runtime: NVIDIA driver 580.82.07, CUDA 13.0 (PyTorch 2.13.0+cu130)
-- vLLM version: 0.28.0
-- vLLM-Omni version or commit: the commit that introduced this recipe (branch based on `main` @ `67ef0b64`)
-- Transformers: 5.14.1
-
-## Command
-
-Start the server from the repository root. The deploy config
-(`vllm_omni/deploy/breeze_tts_2.yaml`, async chunk streaming by default)
-auto-loads from the checkpoint's `model_type`; pass it explicitly if you
-customized it:
+Follow the project's source installation guide for vLLM and vLLM-Omni.
+The integration loads the released checkpoint directly. The original Breeze
+repository and the `qwen-tts` package are not required. The common requirements
+include `soxr` to reproduce reference-audio resampling.
 
 ```bash
-vllm serve BreezeBlue/Breeze-TTS-2 \
-    --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \
-    --omni --port 8091
+vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni --host 127.0.0.1 --port 8091
 ```
 
-## Verification
+The default configuration uses one CUDA GPU, BF16 generation and codec
+computation, and FP32 reference encoding. Decoded waveform tensors are converted
+to FP32 before response formatting. It targets GPUs with at least 16 GiB
+available for the deployment; other workloads consume additional memory.
 
-Plain synthesis (speaker tag `S0`..`S9`, default `S0`):
+On WSL2, set vLLM's supported pinned-memory option before starting the server:
 
 ```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "model": "BreezeBlue/Breeze-TTS-2",
-        "input": "Hello, this is Breeze TTS 2 running on vLLM-Omni.",
-        "voice": "S0",
-        "response_format": "wav",
-        "sample_rate": 24000
-    }' --output breeze_plain.wav
+export VLLM_WSL2_ENABLE_PIN_MEMORY=1
 ```
 
-Voice design (instruction only, no reference audio):
+Native Linux enables pinned memory by default. Startup warmup exercises plain
+generation, CFG, reference conditioning and two-request batches. New text
+lengths or batch sizes can require additional graph
+capture or compilation; warm the workload before measuring steady performance.
+
+## Voice Design
 
 ```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "(laugh) Welcome to our story time.",
-        "instructions": "A warm young woman, clear voice, lively delivery.",
-        "response_format": "wav"
-    }' --output breeze_design.wav
+curl --fail http://127.0.0.1:8091/v1/audio/speech \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "BreezeBlue/Breeze-TTS-2",
+    "input": "你好，欢迎使用语音合成系统。",
+    "instructions": "年轻女性，声音温暖自然，使用标准普通话。",
+    "extra_params": {"guidance_scale": 4.0},
+    "response_format": "wav",
+    "seed": 42,
+    "max_new_tokens": 250
+  }' --output breeze.wav
 ```
 
-Voice cloning (`ref_audio` and `ref_text` must be provided together):
+The `voice` field selects a speaker tag, such as `S0` or `[S0]`; omitting it
+uses `[S0]`. These are checkpoint prompt tags, not a catalog of named voices.
+Omitting `instructions` uses `Speak clearly and naturally.`. An explicitly
+empty instruction remains empty. Vocal-event tokens, including `(laughs)`
+and `[笑]`, are passed to the checkpoint unchanged.
+
+## Voice Clone and Voice Direction
+
+Add `ref_audio` and its exact `ref_text` transcript. Reference audio accepts the
+shared speech API's audio URLs and base64 data URLs, either as a string or
+a one-item list. The shared resolver's
+format, duration and local-file access restrictions apply.
+
+```json
+{
+  "model": "BreezeBlue/Breeze-TTS-2",
+  "input": "Please read this new sentence clearly and naturally.",
+  "ref_audio": "data:audio/wav;base64,<base64-encoded-reference>",
+  "ref_text": "The words actually spoken in the reference recording.",
+  "instructions": "Keep the reference voice. Speak gently and calmly.",
+  "extra_params": {"guidance_scale": 4.0},
+  "response_format": "wav",
+  "seed": 42
+}
+```
+
+Use `guidance_scale: 1.0` with the default instruction for ordinary cloning.
+Reference audio plus a delivery instruction and CFG provides Voice Direction.
+When explicitly specified, `task_type` is `Base` for reference conditioning
+and `VoiceDesign` without reference audio. The legacy `CustomVoice` value
+also remains accepted.
+
+Stereo references are averaged to mono and resampled to 24 kHz with soxr HQ.
+Reference text is encoded independently from the target text. The reference
+codec preserves full causal attention during offline encoding, as used by the
+released Breeze runtime; newer Transformers sliding-window defaults would
+change long references. Reference convolution uses full FP32 precision.
+
+## Sampling and output
+
+Supported `extra_params`:
+
+| Parameter | Default | Meaning |
+| --- | --- | --- |
+| `guidance_scale` | 1.0 | Any finite positive value; values other than 1 use two CFG branches |
+| `temperature` | 0.9 | Zero selects greedy generation |
+| `top_k` | 50 | Zero or -1 disables top-k filtering |
+| `top_p` | 1.0 | Nucleus sampling probability in (0, 1] |
+| `repetition_penalty` | 1.1 | Positive penalty applied to the first codebook's generated history |
+
+`cfg_scale` remains an alias for `guidance_scale`; `guidance_scale` takes
+precedence when both are supplied.
+
+CFG combines conditional and unconditional logits in the Qwen3 output head
+and all 15 remaining depth codebooks. Both branches retain reference audio
+and reference text; the unconditional branch removes only target instructions.
+
+`max_new_tokens` limits generated audio frames, including EOS. Each frame is
+80 ms; the default is 1500 tokens, subject to the remaining 2048-token context.
+Both CFG branches stop at the same generation limit despite unequal prompts.
+Each request owns its seed, generator, audio-code history and codec state.
+Different GPU kernels and batch shapes can still change sampled outputs.
+
+For streamed PCM, use `"stream": true`, `"stream_format": "audio"`, and
+`"response_format": "pcm"`. Output is mono signed 16-bit little-endian PCM at
+24 kHz. Complete WAV responses accumulate the same internal streaming output.
+With the default deployment, the first chunks contain 1, 2, 4 and 5 codec
+frames, then remain at 5 frames.
+
+## Non-streaming throughput profile
+
+For concurrent complete responses, select the throughput deployment from the
+repository root.
 
 ```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "This is new text spoken in the cloned voice.",
-        "ref_audio": "file:///path/to/reference.wav",
-        "ref_text": "The exact transcript of the reference audio.",
-        "response_format": "wav"
-    }' --output breeze_clone.wav
+vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni \
+  --deploy-config vllm_omni/deploy/breeze_tts_2_throughput.yaml \
+  --host 127.0.0.1 --port 8091
 ```
 
-Voice direction (reference + instruction):
+Request a complete WAV with `"stream": false` and omit `stream_format`.
+Setting `"stream_format": "audio"` enables streaming even when `stream` is false.
 
 ```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "We need to discuss what happened last night.",
-        "ref_audio": "file:///path/to/reference.wav",
-        "ref_text": "The exact transcript of the reference audio.",
-        "instructions": "Speak slowly with a restrained, serious tone.",
-        "response_format": "wav"
-    }' --output breeze_direction.wav
+curl --fail http://127.0.0.1:8091/v1/audio/speech \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "BreezeBlue/Breeze-TTS-2",
+    "input": "Hello, welcome to the speech synthesis demonstration.",
+    "stream": false,
+    "response_format": "wav",
+    "seed": 42,
+    "max_new_tokens": 250
+  }' --output breeze-complete.wav
 ```
 
-Streaming raw PCM (`stream_format="audio"`):
+The profile was tested on one H100 80 GB. It sets `max_num_seqs: 32` on both
+stages and reserves 8 GiB of stage-0 KV cache for thirty-two 2048-token contexts.
+These are physical sequence slots: stage 0 can run up to thirty-two plain
+requests (`guidance_scale: 1.0`) or sixteen
+CFG requests, since each CFG request occupies two slots. Mixed workloads share
+the same limit; additional requests queue. The default deployment's memory
+target does not apply to this larger profile.
+
+The internal codec uses fixed 16-frame chunks from the first chunk onward,
+with a shorter final chunk when needed. `cfg_prefill_delay_tokens: 0` removes
+the default eight-frame admission delay, so waiting requests can enter as soon
+as scheduling and KV capacity permit. The stateful codec still requires
+`async_chunk: true` for these complete HTTP responses. Warm the expected text
+lengths and concurrency before measuring the profile.
+
+The shared benchmark defaults to streaming. Override both `stream` and
+`stream_format` to measure complete responses, and set the generation limit
+explicitly in the request body:
 
 ```bash
-curl -N -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Streaming output from Breeze TTS 2.",
-        "stream": true,
-        "stream_format": "audio",
-        "response_format": "pcm",
-        "sample_rate": 24000
-    }' --output breeze_stream.pcm
+python benchmarks/tts/bench_tts.py \
+  --model BreezeBlue/Breeze-TTS-2 --task default_voice --locale en \
+  --dataset-path benchmarks/build_dataset/seed_tts_smoke \
+  --host 127.0.0.1 --port 8091 --concurrency 1 2 4 8 16 32 64 \
+  --num-prompts 128 --num-warmups 64 \
+  --output-dir results/breeze-nonstream -- \
+  --extra-body '{"stream":false,"stream_format":null,"response_format":"pcm","max_new_tokens":384,"seed":42,"extra_params":{"temperature":0.9,"top_k":50,"top_p":1,"repetition_penalty":1.1,"guidance_scale":1}}'
 ```
 
-Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example
-returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta`
-events containing base64-encoded audio. On 1x L20 (greedy, single request after
-warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
+Use the same dataset and sampling settings when comparing deployment profiles.
+Verify completed and failed request counts in each result. For complete
+responses, compare throughput and RTF; first-response timing represents
+completion latency. Run ASR separately: the benchmark's WER mode forces streaming.
 
-## Notes
+## Pipeline and optimization
 
-- **Sample rate**: only `24000` is accepted; other values fail validation
-  before inference.
-- **Speed / seeds / multi-reference**: `speed` must be `1.0`;
-  `task_type=VoiceDesign`, `ref_audio_2`, and speaker embeddings are not
-  supported yet.
-- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and
-  `negative_prompt` return a client error.
-- **Repetition penalty**: the default `1.1` matches upstream's penalty on
-  generated codebook-0 tokens and discourages repeated silence. Override it
-  with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another
-  finite positive value. Text and reference audio are excluded from its history.
-  The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set
-  the same serving override when comparing against those golden frames.
-- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`;
-  synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the
-  12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps.
-- **Checkpoint layout**: the checkpoint must contain the `audio_tokenizer/`
-  subdirectory (bundled Qwen3-TTS codec) for the default streaming path; the
-  Mimi fallback is used only when that directory is absent and does not
-  support async-chunk streaming.
-- **License**: the reference inference code is Apache 2.0; the checkpoints are
-  distributed under the separate BreezeBlue Research and Non-Commercial
-  License. Commercial use requires written authorization from BreezeBlue.
+Stage 0 batches T5Gemma2 text encoding, optional reference-audio encoding,
+vLLM's native Qwen3 backbone with paged KV caching, and Breeze's Llama depth
+decoder. Text segments up to 128 tokens use padded CUDA Graphs with separate
+attention masks. Longer segments use dynamically compiled encoder layers
+without retaining a separate graph workspace for each input length.
+The depth decoder uses fused QKV/MLP projections, compiled layer kernels,
+fixed frame-local KV buffers and a CUDA Graph covering all 15 depth steps.
+Graphs pad requests to power-of-two buckets, keeping CFG branches separate.
+Compiled layers reuse kernels across batch sizes, avoiding a new compilation
+for every shrinking batch. Only real requests advance their random state,
+and returned tensors exclude padded rows.
+Positive temperature, top-k and top-p are mutable GPU inputs to the captured
+sampler. Changing those settings, or changing a positive CFG scale within the
+paired path, reuses the graph. Greedy decoding uses a separate graph that skips
+sorting and random sampling. Each request owns its random-number generator.
 
-## Supported features
+Reference encoding uses dynamic-shape compilation with CUDA Graphs disabled
+for that path, so large reference-convolution workspaces remain reusable by
+the allocator. Every convolution retains each recording's actual boundary;
+the final downsampler repeats its last valid state for partial frames, matching
+independent encoding. Static convolution metadata removes GPU-to-CPU padding
+calculations. Startup compiles reference batches of one and two.
+Waveform storage is aligned to 1920-sample frames, with that frame width
+marked static for compilation. Actual sample lengths still govern every
+convolution boundary. This avoids nested dynamic ceil/mod expressions that
+otherwise make compiler tiling spend minutes simplifying reference shapes.
 
-| Feature | Status | Notes |
-| ------- | ------ | ----- |
-| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` |
-| Voice cloning | ✓ | one reference clip plus transcript per request |
-| Voice design / voice direction | ✓ | `instructions` field, with or without a reference |
-| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) |
-| Non-greedy sampling | ✗ | greedy decoding only |
-| Prefix caching | ✗ | disabled in the deploy config |
-| Tensor / pipeline parallelism | untested | single-GPU profile only |
-| Quantization | ✗ | not supported in this release |
-| CUDA graphs | not tuned | depth decoder and stage-1 codec run eagerly |
-| Platforms | CUDA | qualified on 1x NVIDIA L20 |
+Stage 1 reuses Qwen3-TTS's stateful codec and its CUDA Graph decoder with the
+checkpoint's `audio_tokenizer/` weights. Shared-memory chunks connect stages.
+Only the generated-code decode path is captured: reference conditioning was
+already consumed in stage 0, so the Qwen ICL reference-prefix graphs and their
+large convolution workspaces are unnecessary for Breeze.
+
+The Qwen3 backbone uses compiled piecewise prefill and full decode graphs.
+Stage 0 uses synchronous scheduling to keep CFG branches at the same
+generation step; stage 1 uses asynchronous scheduling. CFG admission reserves
+enough KV pages for both branches' complete generation and preserves their
+actual position IDs. Cancellation finishes both branches.
+
+The default admission policy lets active requests produce eight audio frames
+before adding another prefill. This gives playback a buffer before text and
+reference encoding share the device with an existing stream. It increases a
+queued request's first-audio latency; benchmark it separately from an idle server.
+
+The default `max_num_seqs: 2` counts physical branches: it permits two plain
+requests or one CFG request at a time. Additional CFG requests queue.
+Prefix caching and chunked prefill are disabled because prompt token IDs
+reserve positions for externally encoded embeddings. The deployment requires
+`async_chunk: true`, including for complete HTTP responses.
+
+The accelerated implementation retains the `breeze_tts_2` model package,
+`breeze_tts_2` and `breeze_tts_2_codec` stage names, and default
+`breeze_tts_2.yaml` deployment. The shorter `breeze_tts.yaml` and
+`breeze_tts_throughput.yaml` deployment filenames remain compatibility aliases.
+For a copied custom configuration, use the current `breeze_tts_2.yaml` as
+the starting point: codec chunk controls use `codec_chunk_frames`,
+`initial_codec_chunk_frames` and `codec_chunk_ramp`.
+Defaults now use a 2048-token context and temperature 0.9; set temperature
+to zero explicitly when comparing greedy generation with the earlier runtime.
+
+The supported execution scope is one CUDA GPU, TP=1, PP=1 and unquantized
+released weights. Named voices, speaker embeddings, dual CFG and multiple
+reference recordings are not supported. The original public fast runtime
+also does not expose dual CFG.
+
+## Offline inference and validation
+
+`vllm_omni.model_executor.models.breeze_tts_2.prompt.build_breeze_prompt`
+accepts a checkpoint tokenizer, text, instructions, sampling parameters,
+and optional `ref_audio=(waveform, sample_rate)` plus `ref_text`. Pass its
+result to `Omni.generate()` using the default Breeze deployment. Waveform
+arrays use soundfile's `(samples, channels)` convention.
+
+CPU tests compare cached depth attention and CFG against Hugging Face Llama
+and cover request-owned sampling, payload transport, chunk delivery and CFG
+admission. GPU tests compare padded batched text graphs with independent
+eager encoders and exercise depth graphs across CFG scales and batch shapes.
+
+```bash
+pytest -v tests/model_executor/models/test_breeze_tts_2.py \
+  tests/model_executor/models/test_breeze_tts_2_registration.py \
+  tests/core/sched/test_omni_cfg_ar_scheduler.py -m 'core_model and cpu'
+pytest -v tests/model_executor/models/test_breeze_tts_2_graphs.py -m 'core_model and cuda'
+pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \
+  -m 'core_model and tts' --run-level=core_model
+pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \
+  -m 'advanced_model and tts' --run-level=advanced_model
+```
+
+Serving tests use the shared client for Chinese CFG design, English PCM,
+concurrent CFG requests, reference cloning and direction. Core tests use dummy
+weights and eight-frame requests; advanced tests use real weights and the
+shared content checks. These checks do not establish perceptual quality parity.
+The Chinese content check uses Whisper large-v3 as its primary transcriber:
+small misrecognizes homophones on the fixed test clip. The expected text and
+shared similarity threshold are unchanged; there is no failed-check retry.
+Ensure enough host memory and disk space for the large-v3 test model when
+Whisper runs on CPU alongside the serving process.
+
+The English PCM test uses a 0 dB harmonic-to-noise floor because the original
+eager runtime scores 0.56 dB on that same prompt and seed, below the shared
+1 dB floor. Raw PCM is not transcribed by the shared helper. Separate matched
+evaluations must retain audio and compare intelligibility, speaker similarity,
+instruction effects, first audio, RTF and playback underruns.
+
+The shared serving benchmark supports default speech, Voice Design and
+Voice Clone. For Voice Design:
+
+```bash
+python benchmarks/tts/bench_tts.py \
+  --model BreezeBlue/Breeze-TTS-2 --task voice_design --locale en \
+  --host 127.0.0.1 --port 8091 --concurrency 1 2 \
+  --num-prompts 4 --num-warmups 1 --request-seed 42 --output-len 250 \
+  --output-dir results/breeze
+```
+
+For a local model path, add `--served-model-name /path/to/Breeze-TTS-2`.
+For cloning, select `--task voice_clone --dataset-path /path/to/seed-tts-eval`
+with the desired locale. Reference conditioning with instructions and other
+CFG scales also require identical requests and sampling settings in the
+reference runtime and this deployment.
+
+The reference source and released weights have separate licenses. Consult
+the [reference repository](https://github.com/breezeblue-ai/breeze-tts) and
+checkpoint model card.

recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md

@@ -2,10 +2,40 @@
 
 This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
 audio WVAE encoder together in a vLLM stage. It sends their conditioning
-payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
+payload over shared memory by default, or opt-in NIXL, to an encoder-free
+diffusion stage. The standard MiniMax H3 recipes
 remain single-stage and continue to run all components inside the diffusion
 pipeline.
 
+## Full encoder payload
+
+Both the standard and Turbo builtin deployments use `SharedMemoryConnector`
+for the same-host stage 0 → stage 1 edge. The corresponding `_nixl.yaml`
+overlays opt into `NixlConnector` with the UCX backend. Stage 0 is the sender and stage 1
+is the receiver. The producer hook `encoder2diffusion_full_payload` validates
+the complete conditioning and packages it under `encoder_output`; the diffusion
+stage declares `stage_input_payload_keys: [encoder_output]`.
+
+This is not a text-only handoff. The payload includes:
+
+- `hidden_states.output`: BF16 text embeddings.
+- `meta.token_role_ids`: INT64 token roles aligned with the text embeddings.
+- `embed.embedding`: FP32 visual conditioning latents when present.
+- `embed.speech_feat`: FP32 audio conditioning latents when present.
+- `kv_metadata.minimax_h3_encoder_layout`: packed layout metadata describing
+  the task, dimensions, condition shapes/lengths, keyframes, and reference order.
+
+Text-only requests still carry the full encoder schema, with empty optional
+media tensors. The H3 layout field is model conditioning metadata, **not**
+native paged-KV transfer. Stage 1 does not request native KV-cache reception.
+All participating encoder ranks complete their component collectives before
+the producer sends. The diffusion group leader receives the payload and the
+generic stage-payload path distributes it to the participating diffusion ranks.
+After a successful connector send, the orchestrator carries no duplicate inline
+conditioning; original image/video/audio inputs and temporary encoder media
+buffers are removed from the diffusion prompt. Stage 1 consumes the received
+`encoder_output` without re-encoding text or media.
+
 ## Start the server
 
 Choose the topology explicitly and load its deployment defaults:
@@ -32,18 +62,25 @@ parallelism, and `audio_vae` runs on the group leader. All three role entries
 are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
 the configuration does not create independent role schedulers or world sizes.
 
-For example, this five-GPU topology assigns one GPU to the encoder and four
-to the diffusion stage. `--stage-overrides` keeps placement and parallelism
-scoped to the owning stage rather than broadcasting an override to both:
+For a six-GPU TP2 encoder → TP4 diffusion topology, retain the builtin
+deployment and override only stage 1's TP and Ulysses degrees. VAE patch
+parallelism stays at 4. `--stage-overrides` keeps parallelism scoped to its
+owning stage rather than broadcasting an override to both:
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 \
 vllm-omni serve MiniMaxAI/MiniMax-H3 \
   --omni \
   --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
-  --stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
+  --stage-overrides '{"1":{"tensor_parallel_size":4,"ulysses_degree":1}}'
 ```
 
+The same override is supported with the Turbo builtin deployment. Device IDs
+in the deployment are relative to `CUDA_VISIBLE_DEVICES`; choose visible
+devices and per-stage `devices` overrides for the available hardware. These
+commands describe configuration, not a claim of hardware validation for this
+integrated encoder implementation.
+
 For memory-constrained deployments, start from the CPU-offload or distributed
 layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
 quantization options only to Stage 1 with `--stage-overrides`; retain the
@@ -73,8 +110,35 @@ media conditioning from Stage 0, with no fallback to local encoding.
 
 This H3 topology explicitly keeps its single-replica
 diffusion stage inline, avoiding serialization of decoded video through a
-subprocess. It expects both stages to run in one deployment; cross-node payload
-transport is outside this configuration.
+subprocess. The builtin placement describes a single-host deployment. NIXL's
+transport design can support remote endpoints, but **cross-node H3 execution
+has not been validated** here; these examples are not a tested multi-node
+launch procedure. CPU wiring and payload-contract tests do not establish native
+NIXL, GPU, Turbo-weight, or cross-node execution success.
+
+## Optional NIXL configuration
+
+The base deployments above use `SharedMemoryConnector` and require same-host
+shared-memory access. To use NIXL instead, select the standard overlay:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --omni \
+  --deploy-config vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml
+```
+
+For Turbo, use
+`--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml`
+with the LoRA options shown below. Both overlays require NIXL and a working
+UCX backend; see [NIXL installation](../../docs/design/feature/omni_connectors/nixl_connector.md#installation).
+Initialization failures do not automatically fall back to shared memory.
+
+Each overlay inherits its base deployment and switches **both** edge
+references to NIXL. The encoder roles, `model_loaded` flags, hooks, and full
+`encoder_output` schema stay unchanged. The processor also retains support for
+an inline full-conditioning handoff when supplied; stage 1 still does not fall
+back to local encoders. SharedMemory requires same-host shared-memory access
+and is not a cross-node transport.
 
 The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
 `fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
@@ -99,6 +163,6 @@ that artifact's own contract -- see the
 [Turbo LoRA table](MiniMax-H3.md#turbo-lora) for every published row.
 
 The base deployment intentionally retains 50 inference steps for non-LoRA
-quality. This disaggregated topology is validated on the `fl2v` artifacts,
-which serve T2VA and FL2VA; the `ref2v` artifacts need a `--task-type ref2va`
-server.
+quality. The `fl2v` artifacts serve T2VA and FL2VA; the `ref2v` artifacts need
+a `--task-type ref2va` server. Standard and Turbo configuration coverage does
+not imply that Turbo weights have been exercised on this integrated topology.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -1017,6 +1017,84 @@ balanced switch order.
 > workload. The values above apply to this deployment and are not universal
 > guarantees. `lossless` remains the exact reference path.
 
+## Exact AdaLN reuse
+
+H3 enables exact AdaLN projection reuse by default, independently of the task,
+hardware, sampling schedule, quantization setting, and request `quality`.
+This applies to both original H3 (including its default 50-point schedule) and
+FastH3; no distilled adapter or explicit few-step ladder is required.
+The first occurrence of an input runs the existing projection; later identical
+time embeddings reuse its output only while the projection weights and numerical
+settings remain unchanged. This does not approximate neighboring timesteps or
+change attention, precision, sampling, or the generated audio/video contract.
+
+Each DiT keeps at most 256 MiB of runtime projection outputs. If a long schedule
+exceeds that budget, it retains a reusable subset and computes other entries
+normally, avoiding cyclic eviction on repeated original-H3 requests. All original weights
+remain loaded so new schedules and adapters can compute normally. Adapter changes,
+weight reloads, and model device moves invalidate the cache; parameter versions
+also guard individual entries. TP ranks coordinate hits before skipping a
+projection collective. Gradient-enabled execution, compilation, custom linear
+hooks, and tensors without weight version counters use the original computation.
+Offload paths that replace weight storage can therefore reduce the hit rate.
+
+The runtime cache uses the shared `ExactProjectionCache` implementation; H3 keeps
+only its optional sidecar adaptation. Other models can integrate the same
+[projection cache interface](../../docs/design/module/diffusion/diffusion_model_integration.md#exact-conditioning-projection-reuse).
+This cache retains projection outputs, not offloaded weights, and does not skip
+block weight prefetch.
+
+For an A/B comparison, disable only this reuse at server startup:
+
+```bash
+--cache-config '{"minimax_h3_adaln_cache": false}'
+```
+
+### Optional offline sidecar
+
+An offline sidecar can seed the first projection results; it is not required to
+enable the default cache. Sidecars require `--enforce-eager`, native BF16 TP1
+math, and the same numerical environment as the builder. With the default
+compiled execution, sidecars are rejected before reading their payloads: compiled
+H3 blocks bypass cached projections, so retaining those payloads would waste GPU
+memory. This applies to both the main and Ref2VA sidecars. Other serving
+configurations retain the default runtime-cache behavior described above.
+From the repository root, for a fixed FastH3 adapter and its own four-step schedule:
+
+```bash
+PYTHONPATH=. python tools/minimax_h3/build_adaln_cache.py \
+  --transformer-path "${MODEL_ROOT}/FL2VA/transformer" \
+  --model-variant fl2va --mode t2va \
+  --fasth3-adapter "${FASTH3_ADAPTER}" \
+  --num-inference-steps 4 --flow-shift 12 --audio-flow-shift 3 \
+  --device cuda --output /path/to/h3-adaln.safetensors
+```
+
+Omit `--fasth3-adapter` for base weights and use the serving step count and shifts.
+The builder accepts a native transformer directory with `config.json` and indexed
+or single-file safetensors. It streams the required inputs and refuses to overwrite
+an existing output. It does not instantiate the full DiT.
+
+Pass the resulting local artifact at eager server startup:
+
+```bash
+--enforce-eager \
+--cache-config '{"minimax_h3_adaln_cache_path": "/path/to/h3-adaln.safetensors"}'
+```
+
+Combined servers can also use `minimax_h3_ref_adaln_cache_path` for their Ref2VA
+transformer. Each artifact binds the model architecture, task, schedule, shifts,
+fixed adapter file, effective time-embedding/AdaLN weights, payload checksums, and
+numerical environment. A missing, corrupt, incompatible, or outdated sidecar is
+rejected with a warning; the model still loads every weight and computes through
+the runtime path. A request with different settings similarly falls back.
+Build sidecars from trusted local inputs: checksums verify identity and integrity,
+not the correctness of an untrusted generator.
+
+This implementation saves repeated projection work. It does not remove AdaLN
+weights or claim a GPU memory reduction. End-to-end gains depend on the workload,
+offload behavior, embedding fingerprint cost, and TP coordination overhead.
+
 ## LoRA
 
 ### Turbo LoRA

recipes/cosmos3/Cosmos3-Nano.md

@@ -126,9 +126,16 @@ speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
 `0` removes the streak cap.
 
 To run **without** guardrails (you are responsible for license compliance),
-add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
+add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
+For extra GPUs use
 `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
 `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
+decoder convolutions to channels-last kernels (output no longer bit-identical to
+diffusers; the default `lossless` fast path is bit-exact, see
+[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
 `--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
 generation from ~50 GB to ~36 GB with BF16-level quality (T2V composition can
 shift at the same seed). The pipeline
@@ -724,8 +731,10 @@ vllm serve nvidia/Cosmos3-Nano \
 ```
 
 To run **without** guardrails (you are responsible for license compliance),
-add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
-add `--tensor-parallel-size 8`. `--quantization fp8` and
+add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
+For tensor parallel add `--tensor-parallel-size 8`. `--quantization fp8` and
 `--enable-layerwise-offload` are not supported on NPU.
 The pipeline auto-resolves from `model_index.json`; pass
 `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.

recipes/cosmos3/Cosmos3-Super.md

@@ -57,7 +57,17 @@ vllm serve nvidia/Cosmos3-Super \
 
 Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
 cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
-disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+disable, or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` on either multi-GPU command above (opt-in omni deploy
+overlay, `pipeline: cosmos3_omni_deploy`; without that `pipeline:` key
+`--deploy-config` cannot apply for Cosmos3). Same YAML works for Super / Nano
+(T2I and video); topology `final_output_type` matches the CLI default (`video`).
+`--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
+decoder convolutions to channels-last kernels (output no longer bit-identical to
+diffusers; the default `lossless` fast path is bit-exact, see
+[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
 `--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
 generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
 composition can shift at the same seed).

recipes/inclusionAI/Ming-Image.md

@@ -0,0 +1,141 @@
+# Ming-Image 0.1 Design
+
+> Text-to-image, image editing, and layer decomposition.
+
+## Summary
+
+- Vendor: inclusionAI
+- Models: `inclusionAI/Ming-Image-0.1-Design` and `inclusionAI/Ming-Image-0.1-Design-Layer`
+- Runtime: vLLM-Omni two-stage online serving
+- API: OpenAI-compatible chat completions
+
+## Architecture
+
+Stage 0 uses a Qwen2.5-VL vision tower plus a 20-layer BailingMoeV2 language model.
+It appends 256 learned image-query tokens and exports the final query states plus direct VLM states from layers 5, 12, and 20.
+The query-token checkpoint is in the root sibling `mlp/` directory, so remote stage-0 materialization downloads both `mllm/` and `mlp/`.
+
+Stage 1 projects those conditions through the checkpoint connector and runs a 30-layer Z-Image DiT with a Qwen-Image RGBA VAE.
+
+For Design-Layer, the first returned image is the reconstructed composite and the remaining images are the requested layers.
+
+## CUDA
+
+### Environment
+
+- OS: Linux
+- Python: 3.10+
+- CUDA: 13.0
+- vLLM version: 0.29.0
+- vLLM-Omni version or commit: 59db6a428
+- 2x H100 80GB (1xH100 to be validated)
+
+### Commands
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design
+vllm serve "$MODEL" --omni --deploy-config vllm_omni/deploy/ming_image.yaml --port 8091
+```
+
+For layer decomposition, set `MODEL` to `inclusionAI/Ming-Image-0.1-Design-Layer`.
+
+## Text-to-image
+
+Note that a prompt refiner is expected to describe the prompts with details; we will refine with more example inputs soon.
+
+```bash
+curl -s http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "inclusionAI/Ming-Image-0.1-Design",
+    "messages": [{"role": "user", "content": "A clean editorial botanical poster"}],
+    "modalities": ["image"],
+    "extra_body": {
+      "height": 1024, "width": 1024,
+      "steps": 12, "cfg": 1.0, "seed": 42
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_design_smoke.png
+```
+
+## Image editing
+
+Pass one input image and an editing instruction:
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design
+INPUT_IMAGE=/path/to/input.png
+
+jq -n \
+  --arg model "$MODEL" \
+  --rawfile image <(base64 -w0 "$INPUT_IMAGE") \
+  '{
+    model: $model,
+    messages: [{role: "user", content: [
+      {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
+      {type: "text", text: "Change the background to blue"}
+    ]}],
+    modalities: ["image"],
+    extra_body: {
+      height: 1024, width: 1024,
+      steps: 12, cfg: 1.0, seed: 42
+    }
+  }' |
+curl -sS http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  --data-binary @- \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_edit.png
+```
+
+## Layer decomposition
+
+Use the Design-Layer checkpoint and set `INPUT_IMAGE` to a local flattened design image.
+Note that the prompt should better depict each layer to be decomposed, we will refine with more example inputs soon.
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
+INPUT_IMAGE=/path/to/input.png
+PROMPT="Decompose this design into editable visual layers"
+
+jq -n \
+  --arg model "$MODEL" \
+  --rawfile image <(base64 -w0 "$INPUT_IMAGE") \
+  --arg prompt "$PROMPT" \
+  '{
+    model: $model,
+    messages: [{role: "user", content: [
+      {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
+      {type: "text", text: $prompt}
+    ]}],
+    modalities: ["image"],
+    extra_body: {
+      num_layers: 6,
+      height: 1024, width: 1024,
+      steps: 12, cfg: 2.0, seed: 42
+    }
+  }' |
+curl -sS http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  --data-binary @- > response.json
+
+jq -r '.choices[0].message.content[].image_url.url | split(",")[1]' response.json |
+  nl -v 0 |
+  while read -r index data; do
+    printf "%s" "$data" | base64 -d > "ming_layer_${index}.png"
+  done
+```
+
+## Notes
+
+- The default deployment keeps Stage 0 eager and enables dynamic regional
+  compilation with CUDA Graph Trees for the repeated Stage 1 DiT blocks.
+  The first request for each new image shape or layer count pays compilation
+  and graph-capture cost; warm up every production shape before measuring or
+  serving latency-sensitive traffic.
+- Only one reference image and one request at a time are currently supported.
+- Design-Layer requires a reference image except during warmup.
+- A non-empty `negative_prompt` is rejected; Ming-Image uses zero negative conditioning.
+- Height and width must be divisible by 16.
+- Returned images retain the checkpoint's four RGBA channels.

requirements/common.txt

@@ -15,6 +15,8 @@ kernels==0.16.1
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
+# Breeze reference conditioning must use the released codec's soxr resampling.
+soxr>=0.5.0.post1
 cache-dit==1.5.0
 tqdm>=4.66.0
 torchsde>=0.2.6

tests/assets/hunyuan_image3/it2i.jsonl

@@ -0,0 +1 @@
+{"prompt": "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴", "image_files": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png", "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png"]}

tests/benchmarks/patch/test_patch.py

@@ -2045,3 +2045,79 @@ def test_get_samples_forwards_upstream_multimodal_backends_kwarg(mocker: MockerF
     # No upstream kwargs: unchanged legacy delegate call.
     assert patch.get_samples(args, sentinel) == ["delegated"]
     assert calls[-1] == (args, sentinel, {})
+
+
+@pytest.mark.asyncio
+async def test_benchmark_preserves_stage_metrics_request_order_and_missing_snapshots(monkeypatch):
+    """Persist compact formal-request snapshots in input order, excluding warmups and retaining gaps."""
+    second_finished = asyncio.Event()
+    completion_order = []
+    # Request 1 carries the empty dict that openai-chat-omni initializes before any SSE merge.
+    snapshots = [
+        {"1": {"num_tokens_out": 1536, "finish_reason": "length", "vllm_itls_ms": [8.0, 9.0]}},
+        {},
+        {"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
+    ]
+    expected = [
+        {"1": {"num_tokens_out": 1536, "finish_reason": "length"}},
+        None,
+        {"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
+    ]
+
+    async def request_func(request_func_input, session, pbar=None):
+        request_id = request_func_input.request_id
+        if not request_id:
+            return MixRequestFuncOutput(success=True, stage_metrics={"1": {"num_tokens_out": 999}})
+        index = int(request_id)
+        if index == 0:
+            await second_finished.wait()
+        elif index == 1:
+            second_finished.set()
+        completion_order.append(index)
+        return MixRequestFuncOutput(
+            success=index != 1,
+            stage_metrics=snapshots[index],
+            prompt_len=1,
+            output_tokens=900,
+            ttft=0.01,
+            text_latency=0.1,
+            latency=0.1,
+        )
+
+    monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-stage-metrics", request_func)
+    result = await patch.benchmark(
+        task_type=patch.TaskType.GENERATION,
+        endpoint_type="test-stage-metrics",
+        api_url="http://unused/v1/chat/completions",
+        base_url="http://unused",
+        model_id="test-model",
+        model_name="test-model",
+        tokenizer=None,
+        input_requests=[
+            patch.SampleRequest(prompt="hello", prompt_len=1, expected_output_len=900, request_id=str(i))
+            for i in range(3)
+        ],
+        logprobs=None,
+        request_rate=float("inf"),
+        burstiness=1.0,
+        disable_tqdm=True,
+        num_warmups=2,
+        profile=False,
+        selected_percentile_metrics=[],
+        selected_percentiles=[],
+        ignore_eos=True,
+        goodput_config_dict={},
+        max_concurrency=3,
+        lora_modules=None,
+        extra_headers=None,
+        extra_body={"return_stage_metrics": True},
+        ready_check_timeout_sec=0,
+    )
+    assert completion_order.index(1) < completion_order.index(0)
+    assert result["request_stage_metrics"] == expected
+
+
+@pytest.mark.parametrize("snapshot", [None, {}, "not-a-dict"])
+def test_compact_request_stage_metrics_drops_empty_snapshots(snapshot):
+    """Empty chat-omni snapshots must not make every result persist request_stage_metrics."""
+    assert patch._compact_request_stage_metrics(snapshot) is None

tests/benchmarks/test_accuracy_bench_utils.py

@@ -9,6 +9,7 @@ import sys
 from pathlib import Path
 
 import pytest
+import requests
 from PIL import Image
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -311,7 +312,8 @@ def test_trajectory_judge_payload_collapses_six_frames_into_single_storyboard():
     assert judge_images[0].size == (24, 12)
 
 
-def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
+@pytest.mark.parametrize("num_images", [1, 2])
+def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch, num_images):
     captured = {}
 
     class FakeResponse:
@@ -344,16 +346,46 @@ def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
     output = client.generate_image_edit(
         model="Qwen/Qwen-Image-Edit",
         prompt="edit this image",
-        images=image,
+        images=image if num_images == 1 else [image] * num_images,
         width=512,
         height=512,
+        guidance_scale=2.5,
+        seed=42,
+        bot_task="think_recaption",
+        sys_type="en_unified",
     )
 
     assert output.size == (1, 1)
     assert captured["url"] == "http://127.0.0.1:8093/v1/images/edits"
     assert captured["data"]["prompt"] == "edit this image"
     assert captured["data"]["size"] == "512x512"
-    assert captured["files"][0][0] == "image"
+    assert len(captured["files"]) == num_images
+    assert all(field == ("image" if num_images == 1 else "image[]") for field, _ in captured["files"])
+    assert captured["data"]["guidance_scale"] == "2.5"
+    assert captured["data"]["seed"] == "42"
+    assert captured["data"]["bot_task"] == "think_recaption"
+    assert captured["data"]["sys_type"] == "en_unified"
+
+
+@pytest.mark.parametrize("status_code", [400, 500])
+def test_image_edit_client_preserves_server_error_detail(mocker, status_code):
+    response = requests.Response()
+    response.status_code = status_code
+    response._content = b'{"error": "image input limit exceeded"}'
+    post = mocker.patch("benchmarks.accuracy.common.requests.post", return_value=response)
+    client = VllmOmniImageClient(base_url="http://127.0.0.1:8093")
+
+    with pytest.raises(requests.HTTPError, match="image input limit exceeded") as exc:
+        client.generate_image_edit(
+            model="test-model",
+            prompt="edit this image",
+            images=Image.new("RGB", (2, 2)),
+            width=512,
+            height=512,
+        )
+
+    assert exc.value.response is response
+    post.assert_called_once()
 
 
 def test_text_to_image_client_forwards_output_compression(monkeypatch):

tests/benchmarks/test_omniinteract.py

@@ -249,6 +249,252 @@ def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
     assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
 
 
+def _write_scenario_dataset(root: Path) -> Path:
+    data_root = root / "data"
+    one_q1a = data_root / "1q1a"
+    (one_q1a / "videos").mkdir(parents=True)
+    (one_q1a / "annotations").mkdir()
+    entries = []
+    specs = [
+        ("realtime", "multi_turn", [{"question_time": 0.0, "answer_time": 1.0, "question_type": "realtime"}]),
+        ("proactive", "multi_turn", [{"question_time": 0.0, "answer_time": 1.0, "question_type": "proactive"}]),
+        (
+            "nested",
+            "nested",
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 4.0,
+                    "question_type": "proactive",
+                    "question_text": "outer",
+                    "answer_text": "outer",
+                },
+                {
+                    "question_time": 1.0,
+                    "answer_time": 2.0,
+                    "question_type": "realtime",
+                    "question_text": "inner",
+                    "answer_text": "inner",
+                },
+            ],
+        ),
+        (
+            "interrupted",
+            "multi_turn",
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 1.0,
+                    "question_type": "realtime",
+                    "is_interrupted": True,
+                }
+            ],
+        ),
+    ]
+    for name, scene, annotation in specs:
+        video = one_q1a / "videos" / f"{name}.mp4"
+        ann = one_q1a / "annotations" / f"{name}.json"
+        video.touch()
+        ann.write_text(json.dumps(annotation))
+        entries.append(
+            {
+                "video": f"videos/{name}.mp4",
+                "annotation": f"annotations/{name}.json",
+                "scene_type": scene,
+            }
+        )
+    (one_q1a / "video_json_map.json").write_text(json.dumps({"entries": entries}))
+
+    one_qna = data_root / "1qna"
+    (one_qna / "videos_bench").mkdir(parents=True)
+    (one_qna / "annotations").mkdir(parents=True)
+    (one_qna / "videos_bench" / "guide.mp4").touch()
+    (one_qna / "annotations" / "guide.json").write_text(
+        json.dumps(
+            {
+                "question_time": 0.0,
+                "question_text": "q",
+                "answers": [{"answer_time": 1.0, "answer_text": "a", "label": "step"}],
+            }
+        )
+    )
+    math_root = data_root / "1q1a_math"
+    (math_root / "videos").mkdir(parents=True)
+    (math_root / "annotations").mkdir()
+    (math_root / "videos" / "math.mp4").touch()
+    (math_root / "annotations" / "math.json").write_text(
+        json.dumps([{"question_time": 0.0, "answer_time": 1.0, "question_type": "realtime"}])
+    )
+    (math_root / "video_json_map.json").write_text(
+        json.dumps(
+            {
+                "entries": [
+                    {
+                        "video": "videos/math.mp4",
+                        "annotation": "annotations/math.json",
+                        "scene_type": "multi_turn",
+                    }
+                ]
+            }
+        )
+    )
+    return data_root
+
+
+def test_scenario_cover_and_focus_sampling(tmp_path: Path) -> None:
+    root = _write_scenario_dataset(tmp_path)
+    covered = data.discover_omniinteract_cases(
+        root,
+        ("1q1a", "1q1a_math", "1qna"),
+        num_prompts=5,
+        disable_shuffle=True,
+        scenario_tags=("realtime", "proactive", "nested", "interrupted", "1qna"),
+    )
+    assert len(covered) == 5
+    covered_tags = set().union(*(data.case_scenario_tags(case) for case in covered))
+    assert {"realtime", "proactive", "nested", "interrupted", "1qna"} <= covered_tags
+
+    focused = data.discover_omniinteract_cases(
+        root,
+        ("1q1a", "1q1a_math", "1qna"),
+        num_prompts=0,
+        disable_shuffle=True,
+        scenario_tags=("nested", "1qna"),
+        scenario_focus=True,
+    )
+    assert {case.scene_type for case in focused} == {"nested", "1qna"}
+
+
+def test_load_cases_from_video_list_preserves_order_and_caps(tmp_path: Path) -> None:
+    root = _write_dataset(tmp_path)
+    all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
+    ordered = list(reversed(all_cases))
+    video_list = tmp_path / "sampled_cases.jsonl"
+    video_list.write_text("".join(json.dumps(data.sampled_case_row(case)) + "\n" for case in ordered))
+
+    capped = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=2)
+    assert [case.video_path for case in capped] == [case.video_path for case in ordered[:2]]
+
+    full = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=0)
+    assert [case.video_path for case in full] == [case.video_path for case in ordered]
+    oversized = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=99)
+    assert [case.video_path for case in oversized] == [case.video_path for case in ordered]
+
+
+def test_load_cases_from_video_list_requires_existing_video(tmp_path: Path) -> None:
+    root = _write_dataset(tmp_path)
+    case = data.discover_omniinteract_cases(root, ("1q1a",), num_prompts=1, disable_shuffle=True)[0]
+    row = data.sampled_case_row(case)
+    row["video_path"] = str(tmp_path / "missing.mp4")
+    video_list = tmp_path / "bad.jsonl"
+    video_list.write_text(json.dumps(row) + "\n")
+    with pytest.raises(FileNotFoundError, match="video missing"):
+        data.load_omniinteract_cases_from_video_list(video_list, num_prompts=0)
+
+
+def test_get_samples_from_video_list_without_dataset_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    root = _write_dataset(tmp_path)
+    cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
+    video_list = tmp_path / "sampled_cases.jsonl"
+    video_list.write_text("".join(json.dumps(data.sampled_case_row(case)) + "\n" for case in cases))
+    ref = tmp_path / "ref.wav"
+    ref.touch()
+    parser = TrackingArgumentParser()
+    OmniBenchmarkServingSubcommand.add_cli_args(parser)
+    args = parser.parse_args(
+        [
+            "--backend",
+            "openai-realtime-duplex",
+            "--dataset-name",
+            "omniinteract",
+            "--model",
+            MODEL,
+            "--endpoint",
+            "/v1/realtime",
+            "--num-prompts",
+            "2",
+            "--omniinteract-ref-audio",
+            str(ref),
+            "--omniinteract-video-list",
+            str(video_list),
+        ]
+    )
+    preprocess_serve_args(args)
+    monkeypatch.setattr(
+        "vllm_omni.clients.duplex.reference_audio_data_url",
+        lambda _: "data:audio/wav;base64,ref",
+    )
+    monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
+    samples = benchmark_patch.get_samples(args, None)
+    assert len(samples) == 2
+    assert [sample.omniinteract_case.video_path for sample in samples] == [case.video_path for case in cases[:2]]
+
+
+def _write_overlapping_cover_dataset(root: Path) -> Path:
+    data_root = root / "data"
+    one_q1a = data_root / "1q1a"
+    (one_q1a / "videos").mkdir(parents=True)
+    (one_q1a / "annotations").mkdir()
+    entries = []
+    specs = [
+        (
+            "pair_a",
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 1.0,
+                    "question_type": "realtime",
+                    "is_interrupted": True,
+                }
+            ],
+        ),
+        (
+            "pair_b",
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 1.0,
+                    "question_type": "realtime",
+                    "is_interrupted": True,
+                }
+            ],
+        ),
+        (
+            "proactive_only",
+            [{"question_time": 0.0, "answer_time": 1.0, "question_type": "proactive"}],
+        ),
+    ]
+    for name, annotation in specs:
+        video = one_q1a / "videos" / f"{name}.mp4"
+        ann = one_q1a / "annotations" / f"{name}.json"
+        video.touch()
+        ann.write_text(json.dumps(annotation))
+        entries.append(
+            {
+                "video": f"videos/{name}.mp4",
+                "annotation": f"annotations/{name}.json",
+                "scene_type": "multi_turn",
+            }
+        )
+    (one_q1a / "video_json_map.json").write_text(json.dumps({"entries": entries}))
+    return data_root
+
+
+def test_cover_sampling_skips_tags_already_covered(tmp_path: Path) -> None:
+    root = _write_overlapping_cover_dataset(tmp_path)
+    selected = data.discover_omniinteract_cases(
+        root,
+        ("1q1a",),
+        num_prompts=2,
+        disable_shuffle=True,
+        scenario_tags=("realtime", "interrupted", "proactive"),
+    )
+    names = [case.video_path.stem for case in selected]
+    assert names == ["pair_a", "proactive_only"]
+    selected_tags = set().union(*(data.case_scenario_tags(case) for case in selected))
+    assert {"realtime", "interrupted", "proactive"} <= selected_tags
+
+
 def _archive(path: Path, member: str) -> None:
     with tarfile.open(path, "w") as archive:
         content = b"{}"
@@ -542,6 +788,51 @@ def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_tim
     assert reason in result.official_eval_ineligible_reasons
     oi.write_batch_artifacts(tmp_path / "out", [case], [result])
     assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
+    sampled = [
+        json.loads(line) for line in (tmp_path / "out" / "sampled_cases.jsonl").read_text().splitlines() if line.strip()
+    ]
+    assert sampled == [
+        {
+            "video_path": str(case.video_path.resolve()),
+            "output_name": data.official_output_name(case),
+            "subset": case.subset,
+        }
+    ]
+
+
+def test_official_output_name_matches_minicpmo_batch_layout(tmp_path: Path):
+    one_q1a = _case(tmp_path, subset="1q1a", name="videos/clip_a.mp4")
+    one_qna = _case(tmp_path, subset="1qna", name="videos_bench/nested/guide.mp4")
+    assert data.official_output_name(one_q1a) == "1q1a/videos__clip_a"
+    assert data.official_output_name(one_qna) == "1qna/nested__guide"
+    row = data.sampled_case_row(one_q1a)
+    assert row["video_path"] == str(one_q1a.video_path.resolve())
+    assert row["output_name"] == "1q1a/videos__clip_a"
+    assert row["subset"] == "1q1a"
+
+
+def test_sampled_cases_jsonl_lists_every_case_including_failures(tmp_path: Path):
+    ok_case = _case(tmp_path, name="videos/ok.mp4")
+    bad_case = _case(tmp_path, name="videos/bad.mp4")
+    ok_result = oi.OmniInteractCaseResult(
+        ok_case.subset, str(ok_case.video_path), str(oi._output_dir(tmp_path / "out", ok_case)), success=True
+    )
+    ok_result.eligible_for_official_eval = True
+    bad_result = oi.OmniInteractCaseResult(
+        bad_case.subset, str(bad_case.video_path), str(oi._output_dir(tmp_path / "out", bad_case)), success=False
+    )
+    oi.write_batch_artifacts(tmp_path / "out", [ok_case, bad_case], [ok_result, bad_result])
+    sampled = [
+        json.loads(line) for line in (tmp_path / "out" / "sampled_cases.jsonl").read_text().splitlines() if line.strip()
+    ]
+    assert [row["output_name"] for row in sampled] == ["1q1a/videos__ok", "1q1a/videos__bad"]
+    manifest = [
+        json.loads(line)
+        for line in (tmp_path / "out" / "official_eval_manifest.jsonl").read_text().splitlines()
+        if line.strip()
+    ]
+    assert len(manifest) == 1
+    assert manifest[0]["sample_id"].startswith("1q1a__")
 
 
 def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):

tests/benchmarks/test_omniinteract_eval.py

@@ -0,0 +1,820 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+    OmniInteractCase,
+    OmniInteractEvaluationOptions,
+)
+from vllm_omni.benchmarks.omniinteract import OmniInteractCaseResult
+from vllm_omni.benchmarks.omniinteract_eval import (
+    AlignedWord,
+    TranscriptChunk,
+    _summarize,
+    build_slots,
+    evaluate_batch,
+    evaluate_case,
+    evaluation_inputs_fingerprint,
+    match_slots,
+)
+from vllm_omni.benchmarks.omniinteract_judge import (
+    CoreJudgment,
+    EarlyJudgment,
+    OmniInteractJudge,
+    PartialJudgment,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+
+
+class _FixedJudge:
+    def judge_early(
+        self,
+        slot: dict[str, object],
+        full_context: str,
+        actual_text: str,
+    ) -> EarlyJudgment:
+        del slot, full_context, actual_text
+        return EarlyJudgment("neutral", 1.0, "ok", "{}", "llm_json")
+
+    def judge_core(
+        self,
+        slot: dict[str, object],
+        full_context: str,
+        actual_text: str,
+        future_answers: str,
+    ) -> CoreJudgment:
+        del slot, full_context, actual_text, future_answers
+        return CoreJudgment(1.0, "answer", False, "ok", "{}", "llm_json")
+
+    def judge_interrupted_partial(
+        self,
+        slot: dict[str, object],
+        actual_text: str,
+    ) -> PartialJudgment:
+        del slot, actual_text
+        return PartialJudgment(1.0, False, "ok", "{}", "llm_json")
+
+
+class _GeneratedJudge(OmniInteractJudge):
+    def __init__(self, responses: list[str]) -> None:
+        super().__init__("http://judge", "judge-model")
+        self.responses = iter(responses)
+
+    def _generate(self, system_prompt: str, user_prompt: str) -> str:
+        del system_prompt, user_prompt
+        return next(self.responses)
+
+
+def test_build_slots_and_split_word_aligned_chunk() -> None:
+    slots = build_slots(
+        [
+            {
+                "question_time": 0.0,
+                "answer_time": 2.0,
+                "question_text": "question",
+                "answer_text": "answer",
+                "question_type": "realtime",
+            },
+            {
+                "question_time": 5.0,
+                "answer_time": 6.0,
+                "question_text": "next",
+                "answer_text": "next answer",
+                "question_type": "proactive",
+            },
+        ],
+        "multi_turn",
+    )
+    chunk = TranscriptChunk(
+        source_id=0,
+        text="okay answer",
+        start=0.5,
+        end=2.5,
+        aligned_words=(
+            AlignedWord("okay", 0.5, 0.9),
+            AlignedWord("answer", 2.1, 2.5),
+        ),
+    )
+
+    matched, unmatched = match_slots(slots, [chunk])
+
+    assert not unmatched
+    assert matched[0].early_chunks[0].text == "okay"
+    assert matched[0].core_chunks[0].text == "answer"
+    assert matched[0].core_chunks[0].start == pytest.approx(2.1)
+
+
+def test_judge_parses_json_wrapped_in_model_text() -> None:
+    judge = _GeneratedJudge(
+        [
+            'result: {"flag":"FP_Hallucination","score":1,"rationale":"early"}',
+            '```json\n{"score":0.8,"trigger_phrase":"answer","spoiler":false,"rationale":"good"}\n```',
+            '{"score":0.7,"hallucination":false,"rationale":"partial"}',
+        ]
+    )
+    slot = {"question_text": "q", "gt_answer": "answer"}
+
+    early = judge.judge_early(slot, "context", "guess")
+    core = judge.judge_core(slot, "context", "answer", "(none)")
+    partial = judge.judge_interrupted_partial(slot, "answer")
+
+    assert early.category == "hallucination"
+    assert early.score == 0.0
+    assert core.score == pytest.approx(0.8)
+    assert core.trigger_phrase == "answer"
+    assert partial.score == pytest.approx(0.7)
+
+
+def test_evaluate_case_persists_accuracy_artifact(tmp_path: Path) -> None:
+    annotation = tmp_path / "annotation.json"
+    annotation.write_text(
+        json.dumps(
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 1.0,
+                    "question_text": "question",
+                    "answer_text": "answer",
+                    "question_type": "realtime",
+                }
+            ]
+        )
+    )
+    output_dir = tmp_path / "case-output"
+    output_dir.mkdir()
+    (output_dir / "wav_transcript.json").write_text(
+        json.dumps(
+            {
+                "chunks": [
+                    {
+                        "timestamp": [1.0, 1.5],
+                        "text": "answer",
+                        "aligned_words": [{"text": "answer", "start": 1.0, "end": 1.5}],
+                    }
+                ]
+            }
+        )
+    )
+    video = tmp_path / "video.mp4"
+    video.touch()
+    case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
+    result = OmniInteractCaseResult(
+        subset="1q1a",
+        video=str(video),
+        output_dir=str(output_dir),
+        success=True,
+        eligible_for_official_eval=True,
+    )
+    artifact = tmp_path / "evaluation.json"
+
+    evaluation = evaluate_case(case, result, _FixedJudge(), artifact)
+
+    assert evaluation["timing_precision"] == "word-aligned"
+    assert evaluation["summary"]["IA_QTF1"] == pytest.approx(1.0)
+    paper = evaluation["summary"]["paper_metrics"]["exp_f1"]
+    assert paper["all_global"]["IA_QTF1"] == pytest.approx(1.0)
+    assert paper["realtime"]["IA_QTF1"] == pytest.approx(1.0)
+    assert json.loads(artifact.read_text())["status"] == "ok"
+    fingerprint = evaluation["inputs_fingerprint"]
+    assert fingerprint["annotation_sha256"]
+    assert fingerprint["transcript_sha256"]
+    assert fingerprint["protocol_source"]
+
+
+def _realtime_eval_case(tmp_path: Path, *, transcript_text: str) -> tuple[OmniInteractCase, OmniInteractCaseResult]:
+    annotation = tmp_path / "annotation.json"
+    annotation.write_text(
+        json.dumps(
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 1.0,
+                    "question_text": "question",
+                    "answer_text": "answer",
+                    "question_type": "realtime",
+                }
+            ]
+        )
+    )
+    output_dir = tmp_path / "case-output"
+    output_dir.mkdir(exist_ok=True)
+    chunks: list[dict[str, object]] = []
+    if transcript_text:
+        chunks.append({"timestamp": [1.0, 1.5], "text": transcript_text})
+    (output_dir / "wav_transcript.json").write_text(json.dumps({"chunks": chunks}))
+    video = tmp_path / "video.mp4"
+    video.touch()
+    case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
+    result = OmniInteractCaseResult(
+        subset="1q1a",
+        video=str(video),
+        output_dir=str(output_dir),
+        success=True,
+        eligible_for_official_eval=True,
+    )
+    return case, result
+
+
+def _eval_options(tmp_path: Path, *, skip_existing: bool) -> OmniInteractEvaluationOptions:
+    return OmniInteractEvaluationOptions(
+        judge_base_url="http://127.0.0.1:9",
+        judge_model="cached-judge",
+        judge_api_key="EMPTY",
+        judge_timeout_s=60.0,
+        judge_max_tokens=32,
+        workers=1,
+        output_dir=tmp_path / "evaluation",
+        skip_existing=skip_existing,
+    )
+
+
+class _ConfiguredFixedJudge(_FixedJudge):
+    def __init__(
+        self,
+        base_url: str,
+        model: str,
+        *,
+        api_key: str = "EMPTY",
+        timeout_s: float = 60.0,
+        max_tokens: int = 512,
+    ) -> None:
+        del api_key, timeout_s
+        self.base_url = base_url
+        self.model = model
+        self.max_tokens = max_tokens
+
+
+def test_skip_existing_reuses_matching_fingerprint(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm_omni.benchmarks.omniinteract_eval as eval_mod
+
+    case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
+    options = _eval_options(tmp_path, skip_existing=False)
+    monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
+    first = evaluate_batch([case], [result], options)
+    assert first["summary"]["IA_QTF1"] == pytest.approx(1.0)
+
+    calls: list[str] = []
+
+    def _forbidden(
+        case: OmniInteractCase,
+        result: OmniInteractCaseResult,
+        judge: object,
+        output_path: Path,
+        config: object = None,
+        inputs_fingerprint: object = None,
+    ) -> dict[str, object]:
+        del case, result, judge, output_path, config, inputs_fingerprint
+        calls.append("evaluate_case")
+        raise AssertionError("matching fingerprint must reuse the cached evaluation")
+
+    monkeypatch.setattr(eval_mod, "evaluate_case", _forbidden)
+    reused = evaluate_batch([case], [result], _eval_options(tmp_path, skip_existing=True))
+    assert calls == []
+    assert reused["summary"]["IA_QTF1"] == pytest.approx(1.0)
+
+
+def test_skip_existing_rejects_stale_transcript(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm_omni.benchmarks.omniinteract_eval as eval_mod
+
+    case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
+    options = _eval_options(tmp_path, skip_existing=True)
+    fingerprint = evaluation_inputs_fingerprint(
+        annotation_path=case.annotation_path,
+        transcript_path=Path(result.output_dir) / "wav_transcript.json",
+        judge_model=options.judge_model,
+        judge_base_url=options.judge_base_url,
+        judge_max_tokens=options.judge_max_tokens,
+    )
+    sample_id = f"{case.subset}__{Path(result.output_dir).name}"
+    cache_path = options.output_dir / f"{sample_id}.unified_eval.json"
+    cache_path.parent.mkdir(parents=True)
+    cache_path.write_text(
+        json.dumps(
+            {
+                "status": "ok",
+                "sample_id": sample_id,
+                "inputs_fingerprint": asdict(fingerprint),
+                "summary": {
+                    "IA_QTF1": 1.0,
+                    "num_slots": 1,
+                    "num_unmatched_chunks": 0,
+                    "Global_TP": 1,
+                    "Global_FP": 0,
+                    "Global_FN": 0,
+                },
+                "slots": [],
+            }
+        )
+    )
+    (Path(result.output_dir) / "wav_transcript.json").write_text(json.dumps({"chunks": []}))
+    monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
+    fresh = evaluate_batch([case], [result], options)
+    assert fresh["summary"]["IA_QTF1"] == pytest.approx(0.0)
+
+
+class _MissingCoreScoreJudge(_ConfiguredFixedJudge):
+    def judge_core(
+        self,
+        slot: dict[str, object],
+        full_context: str,
+        actual_text: str,
+        future_answers: str,
+    ) -> CoreJudgment:
+        del slot, full_context, actual_text, future_answers
+        return CoreJudgment(0.0, "", False, "Could not evaluate", "{}", "llm_parse_failed")
+
+
+def test_evaluate_batch_reports_missing_score_as_parse_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm_omni.benchmarks.omniinteract_eval as eval_mod
+
+    case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
+    options = _eval_options(tmp_path, skip_existing=False)
+    monkeypatch.setattr(eval_mod, "OmniInteractJudge", _MissingCoreScoreJudge)
+    evaluation = evaluate_batch([case], [result], options)
+    assert evaluation["status"] == "failed"
+    assert evaluation["evaluated"] == 0
+    assert evaluation["failed"] == 1
+    sample_id = f"{case.subset}__{Path(result.output_dir).name}"
+    cached = json.loads((options.output_dir / f"{sample_id}.unified_eval.json").read_text())
+    assert cached["status"] == "parse_failed"
+
+
+class _CoreRationaleNumberJudge(OmniInteractJudge):
+    def __init__(
+        self,
+        base_url: str,
+        model: str,
+        *,
+        api_key: str = "EMPTY",
+        timeout_s: float = 60.0,
+        max_tokens: int = 512,
+    ) -> None:
+        super().__init__(base_url, model, api_key=api_key, timeout_s=timeout_s, max_tokens=max_tokens)
+
+    def _generate(self, system_prompt: str, user_prompt: str) -> str:
+        del system_prompt, user_prompt
+        return '{"rationale":"Cannot score: 1 required reference is missing."}'
+
+
+def test_evaluate_batch_rejects_core_rationale_number_as_parse_failure(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    import vllm_omni.benchmarks.omniinteract_eval as eval_mod
+
+    case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
+    options = _eval_options(tmp_path, skip_existing=False)
+    monkeypatch.setattr(eval_mod, "OmniInteractJudge", _CoreRationaleNumberJudge)
+    evaluation = evaluate_batch([case], [result], options)
+    assert evaluation["status"] == "failed"
+    assert evaluation["evaluated"] == 0
+    assert evaluation["failed"] == 1
+    sample_id = f"{case.subset}__{Path(result.output_dir).name}"
+    cached = json.loads((options.output_dir / f"{sample_id}.unified_eval.json").read_text())
+    assert cached["status"] == "parse_failed"
+    assert cached["slots"][0]["stage_core"]["parse_source"] == "llm_parse_failed"
+    assert cached["slots"][0]["stage_core"]["S_core"] == 0.0
+
+
+def test_skip_existing_rejects_parse_failed_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm_omni.benchmarks.omniinteract_eval as eval_mod
+
+    case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
+    options = _eval_options(tmp_path, skip_existing=True)
+    fingerprint = evaluation_inputs_fingerprint(
+        annotation_path=case.annotation_path,
+        transcript_path=Path(result.output_dir) / "wav_transcript.json",
+        judge_model=options.judge_model,
+        judge_base_url=options.judge_base_url,
+        judge_max_tokens=options.judge_max_tokens,
+    )
+    sample_id = f"{case.subset}__{Path(result.output_dir).name}"
+    cache_path = options.output_dir / f"{sample_id}.unified_eval.json"
+    cache_path.parent.mkdir(parents=True)
+    cache_path.write_text(
+        json.dumps(
+            {
+                "status": "parse_failed",
+                "sample_id": sample_id,
+                "inputs_fingerprint": asdict(fingerprint),
+                "summary": {
+                    "IA_QTF1": 0.0,
+                    "num_slots": 1,
+                    "num_unmatched_chunks": 0,
+                    "Global_TP": 0,
+                    "Global_FP": 0,
+                    "Global_FN": 1,
+                },
+                "slots": [
+                    {
+                        "stage_core": {"parse_source": "llm_parse_failed", "S_core": 0.0},
+                        "TP_n": 0.0,
+                        "FP_delta": 0,
+                        "FN_delta": 1,
+                    }
+                ],
+            }
+        )
+    )
+    monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
+    fresh = evaluate_batch([case], [result], options)
+    assert fresh["evaluated"] == 1
+    assert fresh["failed"] == 0
+    assert fresh["summary"]["IA_QTF1"] == pytest.approx(1.0)
+    assert json.loads(cache_path.read_text())["status"] == "ok"
+
+
+def _slot_row(
+    *,
+    scene_type: str,
+    question_type: str,
+    tp_n: float,
+    fp: int = 0,
+    fn: int = 0,
+    nested_role: str | None = None,
+    nested_group_id: int | None = None,
+    sample_id: str = "s",
+    tp_core: float | None = None,
+    answer_start: float | None = 1.0,
+) -> dict[str, object]:
+    core_tp = tp_n if tp_core is None else tp_core
+    return {
+        "sample_id": sample_id,
+        "scene_type": scene_type,
+        "question_type": question_type,
+        "nested_role": nested_role,
+        "nested_group_id": nested_group_id,
+        "TP_n": tp_n,
+        "TP_core": core_tp,
+        "FP_delta": fp,
+        "FN_delta": fn,
+        "is_interrupted": False,
+        "stage_core": {"answer_start": answer_start},
+    }
+
+
+def test_paper_metrics_exclude_nested_from_realtime_and_proactive() -> None:
+    summary = _summarize(
+        [
+            _slot_row(scene_type="multi_turn", question_type="realtime", tp_n=1.0),
+            _slot_row(
+                scene_type="nested",
+                question_type="realtime",
+                tp_n=1.0,
+                nested_role="inner",
+                nested_group_id=1,
+                answer_start=0.5,
+            ),
+            _slot_row(
+                scene_type="nested",
+                question_type="proactive",
+                tp_n=1.0,
+                nested_role="outer",
+                nested_group_id=1,
+                answer_start=1.5,
+            ),
+            _slot_row(scene_type="multi_turn", question_type="proactive", tp_n=1.0),
+            _slot_row(scene_type="1QnA", question_type="step", tp_n=1.0, fn=0),
+        ],
+        unmatched=0,
+    )
+    paper = summary["paper_metrics"]["exp_f1"]
+
+    assert paper["realtime"]["num_slots"] == 1
+    assert paper["proactive"]["num_slots"] == 1
+    assert paper["nested"]["num_slots"] == 2
+    assert paper["one_q1a_global"]["num_slots"] == 4
+    assert paper["one_qna"]["num_slots"] == 1
+    assert paper["all_global"]["num_slots"] == 5
+    # All Global includes 1QnA; 1Q1A Global does not.
+    assert paper["all_global"]["Global_TP"] == pytest.approx(5.0)
+    assert paper["one_q1a_global"]["Global_TP"] == pytest.approx(4.0)
+    assert summary["nested"]["NCCS"] == pytest.approx(1.0)
+    assert summary["nested"]["missing_q1_count"] == 0
+    assert summary["paper_metrics"]["exp_nested"]["missed_outer"] == 0
+    assert summary["paper_metrics"]["exp_nested"]["inner_IA_QTF1"] == pytest.approx(1.0)
+    assert summary["paper_metrics"]["exp_nested"]["outer_IA_QTF1"] == pytest.approx(1.0)
+    assert summary["scenario_case_counts"]["realtime"] == 1
+    assert summary["scenario_case_counts"]["proactive"] == 1
+    assert summary["scenario_case_counts"]["nested"] == 1
+    assert summary["scenario_case_counts"]["1qna"] == 1
+
+
+def test_nccs_allows_missing_answer_start() -> None:
+    summary = _summarize(
+        [
+            _slot_row(
+                scene_type="nested",
+                question_type="realtime",
+                tp_n=0.8,
+                tp_core=0.8,
+                nested_role="inner",
+                nested_group_id=1,
+                answer_start=None,
+            ),
+            _slot_row(
+                scene_type="nested",
+                question_type="proactive",
+                tp_n=0.8,
+                tp_core=0.8,
+                nested_role="outer",
+                nested_group_id=1,
+                answer_start=None,
+            ),
+        ],
+        unmatched=0,
+    )
+    assert summary["nested"]["success_pairs"] == 1
+    assert summary["nested"]["NCCS"] == pytest.approx(0.8)
+
+
+def test_missed_outer_counts_zero_outer_core() -> None:
+    summary = _summarize(
+        [
+            _slot_row(
+                scene_type="nested",
+                question_type="realtime",
+                tp_n=1.0,
+                tp_core=1.0,
+                nested_role="inner",
+                nested_group_id=1,
+                answer_start=0.5,
+            ),
+            _slot_row(
+                scene_type="nested",
+                question_type="proactive",
+                tp_n=0.0,
+                tp_core=0.0,
+                fn=1,
+                nested_role="outer",
+                nested_group_id=1,
+                answer_start=1.5,
+            ),
+        ],
+        unmatched=0,
+    )
+    nested = summary["paper_metrics"]["exp_nested"]
+    assert nested["missed_outer"] == 1
+    assert nested["NCCS"] == pytest.approx(0.0)
+    assert nested["inner_IA_QTF1"] == pytest.approx(1.0)
+    assert nested["outer_IA_QTF1"] == pytest.approx(0.0)
+
+
+def test_cross_boundary_core_sets_effective_start_hint() -> None:
+    slots = build_slots(
+        [
+            {
+                "question_time": 0.0,
+                "answer_time": 2.0,
+                "question_text": "question",
+                "answer_text": "answer",
+                "question_type": "realtime",
+            }
+        ],
+        "multi_turn",
+    )
+    chunk = TranscriptChunk(
+        source_id=0,
+        text="okay answer",
+        start=0.5,
+        end=2.5,
+        aligned_words=(
+            AlignedWord("okay", 0.5, 0.9),
+            AlignedWord("answer", 2.1, 2.5),
+        ),
+    )
+    matched, unmatched = match_slots(slots, [chunk])
+    assert not unmatched
+    assert matched[0].core_chunks[0].effective_start_hint == pytest.approx(2.0)
+
+
+class _NonNeutralEarlyJudge(_FixedJudge):
+    def judge_early(
+        self,
+        slot: dict[str, object],
+        full_context: str,
+        actual_text: str,
+    ) -> EarlyJudgment:
+        del slot, full_context, actual_text
+        return EarlyJudgment("other", 1.0, "ignored", "{}", "llm_json")
+
+
+def test_early_tp_ack_only_for_neutral(tmp_path: Path) -> None:
+    annotation = tmp_path / "annotation.json"
+    annotation.write_text(
+        json.dumps(
+            [
+                {
+                    "question_time": 0.0,
+                    "answer_time": 2.0,
+                    "question_text": "question",
+                    "answer_text": "answer",
+                    "question_type": "realtime",
+                }
+            ]
+        )
+    )
+    output_dir = tmp_path / "case-output"
+    output_dir.mkdir()
+    (output_dir / "wav_transcript.json").write_text(
+        json.dumps(
+            {
+                "chunks": [
+                    {
+                        "timestamp": [0.5, 0.9],
+                        "text": "okay",
+                        "aligned_words": [{"text": "okay", "start": 0.5, "end": 0.9}],
+                    },
+                    {
+                        "timestamp": [2.1, 2.5],
+                        "text": "answer",
+                        "aligned_words": [{"text": "answer", "start": 2.1, "end": 2.5}],
+                    },
+                ]
+            }
+        )
+    )
+    video = tmp_path / "video.mp4"
+    video.touch()
+    case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
+    result = OmniInteractCaseResult(
+        subset="1q1a",
+        video=str(video),
+        output_dir=str(output_dir),
+        success=True,
+        eligible_for_official_eval=True,
+    )
+    evaluation = evaluate_case(case, result, _NonNeutralEarlyJudge(), tmp_path / "evaluation.json")
+    assert evaluation["slots"][0]["TP_ack"] == 0.0
+    assert evaluation["slots"][0]["stage_early"]["category"] == "other"
+
+
+def test_paper_metrics_unknown_inner_does_not_underflow_realtime() -> None:
+    summary = _summarize(
+        [
+            _slot_row(scene_type="multi_turn", question_type="realtime", tp_n=0.5),
+            _slot_row(
+                scene_type="nested",
+                question_type="unknown",
+                tp_n=2.0,
+                nested_role="inner",
+                nested_group_id=1,
+                answer_start=0.5,
+            ),
+            _slot_row(
+                scene_type="nested",
+                question_type="proactive",
+                tp_n=1.0,
+                nested_role="outer",
+                nested_group_id=1,
+                answer_start=1.5,
+            ),
+        ],
+        unmatched=0,
+    )
+    paper = summary["paper_metrics"]["exp_f1"]
+    assert paper["realtime"]["Global_TP"] == pytest.approx(0.5)
+    assert paper["realtime"]["num_slots"] == 1
+    assert paper["realtime"]["Precision"] >= 0.0
+    assert paper["one_q1a_global"]["Global_TP"] == pytest.approx(3.5)
+
+
+def test_core_score_float_ignores_leading_counts() -> None:
+    from vllm_omni.benchmarks.omniinteract_judge import _score_float_from_text
+
+    assert _score_float_from_text("The answer covers 3 of the 5 key points, score: 0.6") == pytest.approx(0.6)
+    assert _score_float_from_text("score: 3") is None
+    assert _score_float_from_text("no score here 0.9") is None
+
+    judge = _GeneratedJudge(["The answer covers 3 of the 5 key points, score: 0.6"])
+    core = judge.judge_core({"question_text": "q", "gt_answer": "a"}, "ctx", "answer", "(none)")
+    assert core.score == pytest.approx(0.6)
+    assert core.parse_source == "llm_float_text"
+
+
+def test_core_missing_json_score_does_not_use_rationale_number() -> None:
+    raw = '{"rationale":"Cannot score: 1 required reference is missing."}'
+    core = _GeneratedJudge([raw]).judge_core({"question_text": "q", "gt_answer": "a"}, "ctx", "answer", "(none)")
+    assert core.score == 0.0
+    assert core.parse_source == "llm_parse_failed"
+
+
+def test_early_missing_score_and_unparsed_flag() -> None:
+    missing_score = _GeneratedJudge(['{"flag":"Neutral","rationale":"ok"}'])
+    early = missing_score.judge_early({"question_text": "q", "gt_answer": "a"}, "ctx", "hi")
+    assert early.category == "unparsed"
+    assert early.score == 0.0
+    assert early.parse_source == "llm_parse_failed"
+
+    garbage = _GeneratedJudge(["not json and not a known flag"])
+    unparsed = garbage.judge_early({"question_text": "q", "gt_answer": "a"}, "ctx", "hi")
+    assert unparsed.category == "unparsed"
+    assert unparsed.score == 0.0
+    assert unparsed.parse_source == "llm_parse_failed"
+
+
+@pytest.mark.parametrize(
+    ("path", "response"),
+    [
+        ("early", '{"flag":"Neutral","rationale":"Could not evaluate this answer"}'),
+        ("early", '{"flag":"Neutral","score":"invalid","rationale":"bad"}'),
+        ("core", '{"rationale":"Could not evaluate this answer"}'),
+        ("core", '{"rationale":"Cannot score: 1 required reference is missing."}'),
+        ("core", '{"score":"invalid","trigger_phrase":"","spoiler":false,"rationale":"bad"}'),
+        ("partial", '{"rationale":"Could not evaluate this answer"}'),
+        ("partial", '{"score":"invalid","hallucination":false,"rationale":"bad"}'),
+    ],
+)
+def test_missing_or_invalid_json_scores_are_parse_failures(path: str, response: str) -> None:
+    slot = {"question_text": "q", "gt_answer": "a"}
+    judge = _GeneratedJudge([response])
+    if path == "early":
+        judgment = judge.judge_early(slot, "ctx", "hi")
+    elif path == "core":
+        judgment = judge.judge_core(slot, "ctx", "answer", "(none)")
+    else:
+        judgment = judge.judge_interrupted_partial(slot, "answer")
+    assert judgment.score == 0.0
+    assert judgment.parse_source == "llm_parse_failed"
+
+
+def test_explicit_zero_json_score_is_valid_llm_json() -> None:
+    slot = {"question_text": "q", "gt_answer": "a"}
+    early = _GeneratedJudge(['{"flag":"Neutral","score":0,"rationale":"zero"}']).judge_early(slot, "ctx", "hi")
+    core = _GeneratedJudge(['{"score":0.0,"trigger_phrase":"","spoiler":false,"rationale":"zero"}']).judge_core(
+        slot, "ctx", "answer", "(none)"
+    )
+    partial = _GeneratedJudge(['{"score":0,"hallucination":false,"rationale":"zero"}']).judge_interrupted_partial(
+        slot, "answer"
+    )
+    assert early.category == "neutral"
+    assert early.score == 0.0
+    assert early.parse_source == "llm_json"
+    assert core.score == 0.0
+    assert core.parse_source == "llm_json"
+    assert partial.score == 0.0
+    assert partial.parse_source == "llm_json"
+
+
+@pytest.mark.parametrize("score_literal", ['"NaN"', '"Infinity"', '"-Infinity"', "NaN", "Infinity", "-Infinity"])
+def test_nonfinite_json_scores_are_parse_failures(score_literal: str) -> None:
+    slot = {"question_text": "q", "gt_answer": "a"}
+    early = _GeneratedJudge([f'{{"flag":"Neutral","score":{score_literal},"rationale":"bad"}}']).judge_early(
+        slot, "ctx", "hi"
+    )
+    core = _GeneratedJudge(
+        [f'{{"score":{score_literal},"trigger_phrase":"answer","spoiler":false,"rationale":"bad"}}']
+    ).judge_core(slot, "ctx", "answer", "(none)")
+    partial = _GeneratedJudge(
+        [f'{{"score":{score_literal},"hallucination":false,"rationale":"bad"}}']
+    ).judge_interrupted_partial(slot, "answer")
+
+    assert early.score == 0.0
+    assert early.parse_source == "llm_parse_failed"
+    assert core.score == 0.0
+    assert core.parse_source == "llm_parse_failed"
+    assert core.trigger_phrase == ""
+    assert partial.score == 0.0
+    assert partial.parse_source == "llm_parse_failed"
+
+
+def test_fingerprint_requires_evaluator_schema_version() -> None:
+    from vllm_omni.benchmarks.omniinteract_eval import EvaluationInputsFingerprint
+
+    payload = {
+        "annotation_sha256": "a" * 64,
+        "transcript_sha256": "b" * 64,
+        "judge_model": "m",
+        "judge_base_url": "http://127.0.0.1:9",
+        "judge_max_tokens": 32,
+        "protocol_source": "proto",
+    }
+    assert EvaluationInputsFingerprint.from_mapping(payload) is None
+    payload["evaluator_schema_version"] = 1
+    assert EvaluationInputsFingerprint.from_mapping(payload) is not None
+    payload["evaluator_schema_version"] = 999
+    matched = EvaluationInputsFingerprint.from_mapping(payload)
+    assert matched is not None
+    assert matched.evaluator_schema_version == 999
+
+
+def test_summarize_counts_parse_sources() -> None:
+    rows = [
+        {
+            **_slot_row(scene_type="multi_turn", question_type="realtime", tp_n=1.0),
+            "stage_early": {"parse_source": "llm_json"},
+            "stage_core": {"parse_source": "llm_parse_failed", "trigger_fallback": True},
+        }
+    ]
+    summary = _summarize(rows, unmatched=0)
+    assert summary["judge_parse"]["llm_json"] == 1
+    assert summary["judge_parse"]["llm_parse_failed"] == 1
+    assert summary["judge_parse"]["judge_calls"] == 2
+    assert summary["trigger_fallback_slots"] == 1

tests/benchmarks/test_serve_cli.py

@@ -6,6 +6,7 @@ import json
 import os
 import subprocess
 import textwrap
+import time
 from pathlib import Path
 
 import pytest
@@ -21,6 +22,76 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
 
 
+def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path):
+    """Exercise real dataset/benchmark code; replace only the HTTP backend."""
+    from vllm_omni.benchmarks import serve
+    from vllm_omni.benchmarks.patch import patch
+    from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
+
+    dataset_path = Path(__file__).parents[1] / "assets/hunyuan_image3/it2i.jsonl"
+    sample = json.loads(dataset_path.read_text(encoding="utf-8"))
+    parser = argparse.ArgumentParser()
+    OmniBenchmarkServingSubcommand.add_cli_args(parser)
+    args = parser.parse_args(
+        [
+            "--model",
+            "test-model",
+            "--endpoint",
+            "/v1/images/edits",
+            "--dataset-name",
+            "custom_image",
+            "--dataset-path",
+            str(dataset_path),
+            "--skip-tokenizer-init",
+            "--disable-shuffle",
+            "--disable-tqdm",
+            "--num-prompts",
+            "8",
+            "--num-warmups",
+            "2",
+            "--max-concurrency",
+            "1",
+            "--ready-check-timeout-sec",
+            "0",
+            "--percentile-metrics",
+            "e2el",
+            "--extra-body",
+            '{"seed":42,"guidance_scale":2.5,"bot_task":"think_recaption"}',
+            "--save-result",
+            "--result-dir",
+            str(tmp_path),
+            "--result-filename",
+            "result.json",
+        ]
+    )
+    requests = []
+
+    async def request_func(request_func_input, session, pbar=None):
+        requests.append(request_func_input)
+        if pbar is not None:
+            pbar.update(1)
+        return patch.MixRequestFuncOutput(
+            success=True,
+            prompt_len=1,
+            output_tokens=1,
+            image_count=1,
+            start_time=time.perf_counter(),
+            latency=99.0 if len(requests) <= 2 else 0.1,
+        )
+
+    monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "/v1/images/edits", request_func)
+    result = serve.main(args)
+    assert len(requests) == 10
+    for request in requests:
+        assert request.prompt == sample["prompt"]
+        assert [image["image_url"]["url"] for image in request.multi_modal_content] == sample["image_files"]
+        assert request.extra_body == args.extra_body
+    assert result["completed"] == 8
+    saved = json.loads((tmp_path / "result.json").read_text())
+    assert saved["completed"] == 8
+    assert saved["mean_e2el_ms"] == pytest.approx(100.0)  # Excludes both 99-second warmups.
+
+
 @pytest.mark.parametrize(
     "dataset_name",
     [
@@ -153,6 +224,53 @@ def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
     assert args.num_prompts == expected
 
 
+@pytest.mark.parametrize(
+    ("extra", "match"),
+    [
+        (["--seed", "7"], "cannot be combined with --seed"),
+        (["--dataset-path", "/tmp/omniinteract-data"], "cannot be combined with --dataset-path"),
+        (["--omniinteract-scenario-tags", "realtime"], "cannot be combined with --omniinteract-scenario-tags"),
+        (["--omniinteract-scenario-focus"], "cannot be combined with --omniinteract-scenario-focus"),
+    ],
+)
+def test_preprocess_serve_args_rejects_omniinteract_video_list_conflicts(
+    tmp_path: Path,
+    extra: list[str],
+    match: str,
+) -> None:
+    ref = tmp_path / "ref.wav"
+    ref.touch()
+    video_list = tmp_path / "list.jsonl"
+    video_list.write_text("{}\n")
+    parser = TrackingArgumentParser()
+    parser.add_argument("--dataset-name", default="sharegpt")
+    parser.add_argument("--backend", default="vllm")
+    parser.add_argument("--endpoint", default="/v1/completions")
+    parser.add_argument("--model", default="dummy")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--dataset-path", default=None)
+    parser.add_argument("--num-prompts", type=int, default=1000)
+    parser.add_argument("--max-concurrency", type=int, default=None)
+    add_omni_args(parser)
+    args = parser.parse_args(
+        [
+            "--dataset-name",
+            "omniinteract",
+            "--backend",
+            "openai-realtime-duplex",
+            "--endpoint",
+            "/v1/realtime",
+            "--omniinteract-ref-audio",
+            str(ref),
+            "--omniinteract-video-list",
+            str(video_list),
+            *extra,
+        ]
+    )
+    with pytest.raises(ValueError, match=match):
+        preprocess_serve_args(args)
+
+
 @pytest.mark.parametrize(
     ("argv", "expected_extra_body", "expected_explicit"),
     [
@@ -293,6 +411,19 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
                 async def close(self):
                     return None
 
+            class MockTokenizer:
+                vocab_size = 256
+                all_special_ids = []
+
+                def num_special_tokens_to_add(self, *args, **kwargs):
+                    return 0
+
+                def decode(self, token_ids, *args, **kwargs):
+                    return " ".join(str(token_id) for token_id in token_ids)
+
+                def encode(self, text, *args, **kwargs):
+                    return [int(token_id) for token_id in text.split()]
+
             # Patch globally so modules loaded after this also see the mock
             import aiohttp
             aiohttp.ClientSession = MockClientSession
@@ -303,6 +434,11 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
             patch_mod.aiohttp.ClientSession = MockClientSession
             patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()
 
+            # This test exercises CLI and HTTP routing, not model tokenization.
+            # Keep its subprocess independent of hub access and model caches.
+            import vllm.benchmarks.serve as benchmark_serve
+            benchmark_serve.get_tokenizer = lambda *args, **kwargs: MockTokenizer()
+
             calls_file = os.environ.get("VLLM_OMNI_TEST_POST_CALLS_FILE")
 
             if calls_file:
@@ -319,6 +455,9 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
     env = os.environ.copy()
     env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "")
     env["VLLM_OMNI_TEST_POST_CALLS_FILE"] = str(calls_path)
+    env["HF_HOME"] = str(tmp_path / "hf-cache")
+    env["HF_HUB_OFFLINE"] = "1"
+    env["TRANSFORMERS_OFFLINE"] = "1"
 
     cmd = [
         "vllm",

tests/clients/test_duplex_client.py

@@ -1486,8 +1486,8 @@ def test_pcm16_wav_round_trip(tmp_path):
 
 
 def test_duplex_unit_boundary_and_residual_math():
-    assert duplex_unit_boundary_ms(0) == 1030
-    assert duplex_unit_boundary_ms(2) == 3030
+    assert duplex_unit_boundary_ms(0) == 1000
+    assert duplex_unit_boundary_ms(2) == 3000
     assert has_residual_model_unit(b"\x00" * 32_000, chunk_period_ms=1000) is False
     assert has_residual_model_unit(b"\x00" * 32_002, chunk_period_ms=1000) is True
     created = {"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 500}}}
@@ -1512,9 +1512,13 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
         # A composite belongs to the unit it was captured in, so it rides the
         # same append as that unit's base frame; a unit without one sends the
         # base alone. Frame k rides the append that closes model unit k.
-        assert [event["video_frames"] for event in appends if "video_frames" in event] == [["f0", "s0"], ["f1"]]
-        assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1200, 2200]
-        assert frames_sent == 2
+        assert [event["video_frames"] for event in appends if "video_frames" in event] == [
+            ["f0", "s0"],
+            ["f1"],
+            ["f1"],
+        ]
+        assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1000, 2000, 3000]
+        assert frames_sent == 3
         sock.feed(SESSION_CLOSED)
 
 

tests/config/test_config_factory.py

@@ -694,6 +694,105 @@ class TestCosmos3PolicyPipeline:
         assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position"
 
 
+class TestCosmos3OmniDeployPipeline:
+    """Omni deploy yaml selects an opt-in topology so --deploy-config applies (#6874)."""
+
+    def test_registered_without_auto_capturing_cosmos3_omni(self):
+        assert "cosmos3_omni_deploy" in OMNI_PIPELINES
+        # Same shared HF metadata as policy: must not register cosmos3_omni for
+        # auto-detect, or T2I/video/policy would collide.
+        assert "cosmos3_omni" not in OMNI_PIPELINES
+        pipeline = OMNI_PIPELINES["cosmos3_omni_deploy"]
+        assert pipeline.hf_architectures == ()
+        assert pipeline.diffusers_class_name is None
+        # Align with CLI Cosmos3OmniDiffusersPipeline default (video), not T2I-only.
+        assert pipeline.stages[0].final_output_type == "video"
+
+    def test_deploy_yaml_applies_guardrails_false(self):
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
+        assert deploy.pipeline == "cosmos3_omni_deploy"
+        assert deploy.stages[0].devices is None
+
+        stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_omni_deploy"], deploy)
+        assert len(stages) == 1
+        stage = stages[0].to_omegaconf()
+
+        assert stage.stage_type == "diffusion"
+        assert stage.final_output_type == "video"
+        assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline"
+        assert stage.engine_args.model_config.guardrails is False
+        assert "devices" not in stage.runtime
+
+    def test_deploy_yaml_accepts_documented_multi_gpu_cli(self):
+        from vllm_omni.engine.stage_init_utils import _check_stage_device_layout
+
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            OMNI_PIPELINES["cosmos3_omni_deploy"],
+            {
+                "cfg_parallel_size": 2,
+                "use_hsdp": True,
+                "hsdp_shard_size": 2,
+            },
+            user_deploy_config=deploy,
+        )
+        assert len(stages) == 1
+        omega = stages[0].to_omegaconf()
+        assert omega.engine_args.model_config.guardrails is False
+        assert omega.engine_args.parallel_config.cfg_parallel_size == 2
+        assert omega.engine_args.parallel_config.use_hsdp is True
+        assert omega.engine_args.parallel_config.hsdp_shard_size == 2
+        assert "devices" not in omega.runtime
+        _check_stage_device_layout(omega, dict(omega.engine_args))
+
+    def test_get_pipeline_config_honours_deploy_yaml_pipeline_key(self):
+        deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
+        # No HF download: explicit pipeline: key is highest priority.
+        pipeline = StageConfigFactory.get_pipeline_config(
+            model="nvidia/Cosmos3-Super-Text2Image",
+            trust_remote_code=True,
+            deploy_config_path=deploy_path,
+        )
+        assert pipeline is not None
+        assert pipeline.model_type == "cosmos3_omni_deploy"
+
+    def test_without_deploy_yaml_stays_on_unregistered_fallback(self):
+        # Shared HF metadata must not auto-select a pipeline; only the deploy
+        # yaml ``pipeline:`` key selects cosmos3_omni_deploy and merges guardrails.
+        assert "cosmos3_omni" not in OMNI_PIPELINES
+
+        class FakeCosmos3Config(PretrainedConfig):
+            model_type = "cosmos3_omni"
+
+        fake_config = FakeCosmos3Config()
+        fake_config.architectures = ["Cosmos3OmniDiffusersPipeline"]
+        model = "nvidia/Cosmos3-Super-Text2Image"
+        deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
+
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+        with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config):
+            assert (
+                StageConfigFactory.get_pipeline_config(
+                    model=model,
+                    trust_remote_code=True,
+                    deploy_config_path=None,
+                )
+                is None
+            )
+
+            pipeline = StageConfigFactory.get_pipeline_config(
+                model=model,
+                trust_remote_code=True,
+                deploy_config_path=deploy_path,
+            )
+        assert pipeline is not None
+        assert pipeline.model_type == "cosmos3_omni_deploy"
+
+        stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
+        assert stages[0].to_omegaconf().engine_args.model_config.guardrails is False
+
+
 class TestStagePipelineConfig:
     def test_frozen(self):
         s = StagePipelineConfig(stage_id=0, model_stage="a")
@@ -1588,19 +1687,32 @@ stages:
         assert "text_encoder_tp_size" not in stages[1].runtime_overrides
         assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
 
-    def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    def test_minimax_h3_disaggregated_defaults(self, deploy_name):
         pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
-        deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
+        deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
         stages = merge_pipeline_deploy(pipeline, deploy)
 
+        assert pipeline.model_arch == "MiniMaxH3Encoder"
+        assert stages[0].model_stage == "encoder"
         assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
+        processor = "vllm_omni.model_executor.stage_input_processors.minimax_h3"
+        assert stages[0].yaml_extras["prompt_transform_func"] == f"{processor}.prepare_encoder_prompt"
+        assert (
+            stages[0].yaml_engine_args["custom_process_next_stage_input_func"]
+            == f"{processor}.encoder2diffusion_full_payload"
+        )
+        assert stages[1].custom_process_input_func == f"{processor}.encoder2diffusion"
+        assert tuple(stages[1].yaml_engine_args["stage_input_payload_keys"]) == ("encoder_output",)
+        assert stages[1].yaml_engine_args["omni_kv_config"]["need_recv_cache"] is False
         assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
         assert stages[0].yaml_runtime["num_replicas"] == 1
         assert stages[1].yaml_runtime["num_replicas"] == 1
-        assert stages[1].yaml_engine_args["model_loaded"] == {
-            "text_encoder": False,
-            "vae_encoder": False,
-        }
+        assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False, "vae_encoder": False}
+        assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
         assert stages[0].yaml_engine_args["max_num_seqs"] == 1
         assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
             "text_encoder": {"parallel_mode": "tp"},
@@ -1614,7 +1726,12 @@ stages:
         assert parallel["ulysses_degree"] == 4
         assert parallel["vae_patch_parallel_size"] == 4
         assert stages[1].yaml_engine_args["inline_diffusion"] is True
-        assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
+        sampling = stages[1].yaml_extras["default_sampling_params"]
+        if "turbo" in deploy_name:
+            assert sampling["num_inference_steps"] == 5
+            assert sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
+        else:
+            assert sampling["num_inference_steps"] == 50
 
         turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
         turbo_stages = merge_pipeline_deploy(pipeline, turbo)
@@ -1625,6 +1742,73 @@ stages:
         assert turbo_sampling["num_inference_steps"] == 5
         assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
 
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    @pytest.mark.parametrize("shared_memory", [False, True], ids=["nixl", "shared_memory"])
+    def test_minimax_h3_deploy_connector_wiring_and_roles(self, deploy_name, shared_memory):
+        from vllm_omni.config.stage_config import resolve_deploy_yaml
+        from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
+        from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+        base_config = resolve_deploy_yaml(get_deploy_config_path(deploy_name))
+        if not shared_memory:
+            deploy_name = deploy_name.replace(".yaml", "_nixl.yaml")
+        deploy_path = get_deploy_config_path(deploy_name)
+        config = resolve_deploy_yaml(deploy_path)
+        assert config["async_chunk"] is False
+        connector_name = "shared_memory_connector" if shared_memory else "nixl_connector"
+        assert connector_name in config["connectors"]
+        assert config["stages"][0]["output_connectors"] == {"to_stage_1": connector_name}
+        assert config["stages"][1]["input_connectors"] == {"from_stage_0": connector_name}
+        load_deploy_config(deploy_path)
+        for base_stage, stage in zip(base_config["stages"], config["stages"], strict=True):
+            edge_keys = {"input_connectors", "output_connectors"}
+            assert {key: value for key, value in stage.items() if key not in edge_keys} == {
+                key: value for key, value in base_stage.items() if key not in edge_keys
+            }
+        transfer = load_omni_transfer_config(config_dict=config)
+        assert transfer is not None
+        assert set(transfer.connectors) == {("0", "1")}
+        edge = transfer.connectors[("0", "1")]
+        expected_name = "SharedMemoryConnector" if shared_memory else "NixlConnector"
+        assert edge.name == expected_name
+        assert "role" not in edge.extra
+        for stage_id, role in ((0, "sender"), (1, "receiver")):
+            spec = get_stage_connector_spec(transfer, stage_id, async_chunk=False)
+            assert spec["name"] == expected_name
+            assert spec["extra"]["role"] == role
+            if shared_memory:
+                assert "from_stage" not in spec["extra"]
+                assert "to_stage" not in spec["extra"]
+            else:
+                assert spec["extra"]["from_stage"] == 0
+                assert spec["extra"]["to_stage"] == 1
+                assert spec["extra"]["host"] == "auto"
+                assert spec["extra"]["zmq_port"] == 50071
+                assert spec["extra"]["backends"] == ["UCX"]
+        assert "role" not in edge.extra
+
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    def test_minimax_h3_builtin_stage_one_tp4_ulysses1_override(self, deploy_name):
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            OMNI_PIPELINES["minimax_h3_disaggregated"],
+            {"stage_1_tensor_parallel_size": 4, "stage_1_ulysses_degree": 1},
+            deploy_config_path=get_deploy_config_path(deploy_name),
+        )
+        encoder_args = stages[0].to_omegaconf().engine_args
+        diffusion_args = stages[1].to_omegaconf().engine_args
+        assert encoder_args.tensor_parallel_size == 2
+        assert "ulysses_degree" not in stages[0].runtime_overrides
+        assert diffusion_args.parallel_config.tensor_parallel_size == 4
+        assert diffusion_args.parallel_config.ulysses_degree == 1
+        assert diffusion_args.parallel_config.vae_patch_parallel_size == 4
+        assert tuple(diffusion_args.stage_input_payload_keys) == ("encoder_output",)
+
     def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
         pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
 

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 30,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 59,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
         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: 16,
+        ModelEnvironmentVariableDisposition.INTERNALIZE: 15,
         ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
     }
 

tests/config/test_omni_config.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import warnings
-from dataclasses import fields
+from dataclasses import fields, replace
 from inspect import Parameter, signature
 from multiprocessing.reduction import ForkingPickler
 from pathlib import Path
@@ -246,6 +246,12 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
             assert omni_stage.diffusion_config is not None
             assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
             assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
+            assert omni_stage.diffusion_config.stage_input_payload_keys == engine_args.get(
+                "stage_input_payload_keys", ()
+            )
+            assert omni_stage.diffusion_config.stage_output_payload_keys == engine_args.get(
+                "stage_output_payload_keys", ()
+            )
         elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
             assert isinstance(omni_stage, VllmOmniARStageConfig)
             assert not hasattr(omni_stage, "diffusion_config")
@@ -1572,6 +1578,68 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
     assert cfg.max_cpu_loras == 1
 
 
+@pytest.mark.parametrize("enabled", [False, True])
+def test_diffusion_projection_retains_prefix_caching(enabled):
+    projection = omni_config_module._DiffusionConfigProjection
+    assert projection.from_kwargs().enable_prefix_caching is False
+    config = projection.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=enabled,
+    )
+    assert config.enable_prefix_caching is enabled
+
+
+@pytest.mark.parametrize(
+    "pipeline_value,stage_value,cli_value,expected",
+    [
+        (None, None, None, False),
+        (True, None, None, True),
+        (False, True, None, True),
+        (True, False, None, False),
+        (None, None, True, True),
+        (True, True, False, False),
+        (False, False, True, True),
+    ],
+)
+def test_diffusion_prefix_caching_precedence_and_transport(
+    tmp_path,
+    pipeline_value,
+    stage_value,
+    cli_value,
+    expected,
+):
+    import yaml
+
+    from vllm_omni.engine.stage_init_utils import _project_omni_stage_engine_args
+
+    stage_deploy = {
+        "stage_id": 0,
+        "diffusion_kv_mode": "paged_scheduler",
+        "diffusion_kv_max_rows_per_request": 2,
+    }
+    deploy = {"pipeline": "hunyuan_image3_dit", "async_chunk": False, "stages": [stage_deploy]}
+    if pipeline_value is not None:
+        deploy["enable_prefix_caching"] = pipeline_value
+    if stage_value is not None:
+        stage_deploy["enable_prefix_caching"] = stage_value
+    deploy_path = tmp_path / "prefix.yaml"
+    deploy_path.write_text(yaml.safe_dump(deploy))
+    cli = {} if cli_value is None else {"stage_0_enable_prefix_caching": cli_value}
+    stage = _from_pipeline_key(
+        "hunyuan_image3_dit",
+        deploy_config_path=str(deploy_path),
+        cli_overrides=cli,
+    ).stage_by_id(0)
+
+    assert bool(stage.cache_config.enable_prefix_caching) is expected
+    assert stage.diffusion_config.enable_prefix_caching is expected
+    assert _project_omni_stage_engine_args(stage)["enable_prefix_caching"] is expected
+    serialized = _serialize_stage_config(stage)
+    restored = omni_config_module._DiffusionConfigProjection.from_kwargs(**serialized["diffusion_config"])
+    assert restored.enable_prefix_caching is expected
+
+
 def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     from vllm_omni.platforms import current_omni_platform
 
@@ -1725,6 +1793,58 @@ def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
         )
 
 
+@pytest.mark.parametrize("source", ["default", "topology", "deploy", "stage-cli"])
+@pytest.mark.parametrize("key_container", [list, tuple])
+def test_diffusion_stage_payload_keys_roundtrip(source, key_container):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+    topology_keys = {
+        "stage_input_payload_keys": ("conditioning", "metadata"),
+        "stage_output_payload_keys": ("latents",),
+    }
+    override_keys = {
+        "stage_input_payload_keys": (),
+        "stage_output_payload_keys": ("audio", "video"),
+    }
+    pipeline = _resolve_pipeline_or_skip("dreamzero")
+    topology = pipeline.get_stage(0)
+    if source != "default":
+        topology = replace(topology, **topology_keys)
+    pipeline = replace(pipeline, stages=(topology,))
+    deploy = _load_default_deploy(pipeline)
+    cli_overrides = {}
+    if source == "deploy":
+        deploy = replace(deploy, stages=[StageDeployConfig(stage_id=0, engine_extras=override_keys)])
+    elif source == "stage-cli":
+        cli_overrides = {f"stage_0_{name}": value for name, value in override_keys.items()}
+
+    stage = VllmOmniConfig.from_pipeline_config(
+        pipeline, user_deploy_config=deploy, cli_overrides=cli_overrides
+    ).stage_by_id(0)
+    expected = topology_keys if source == "topology" else override_keys
+    if source == "default":
+        expected = dict.fromkeys(topology_keys, ())
+    legacy_stage = merge_pipeline_deploy(pipeline, deploy)[0]
+    legacy_args = {**legacy_stage.yaml_engine_args, **(override_keys if source == "stage-cli" else {})}
+    restored_stage = ForkingPickler.loads(ForkingPickler.dumps(stage))
+    engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model="test-model")
+    diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs(
+        engine_args, stage_id=restored_stage.stage_id, include_engine_adapter_metadata=True
+    )
+    for name in topology_keys:
+        diffusion_kwargs[name] = key_container(diffusion_kwargs[name])
+    od_config = OmniDiffusionConfig.from_kwargs(**diffusion_kwargs)
+
+    for name, keys in expected.items():
+        assert legacy_args.get(name, ()) == keys
+        assert getattr(stage.diffusion_config, name) == keys
+        assert engine_args[name] == keys
+        assert getattr(od_config, name) == keys
+    for name in topology_keys:
+        assert getattr(topology, name) == (() if source == "default" else topology_keys[name])
+
+
 def test_diffusion_config_field_classification_covers_current_fields():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -1745,6 +1865,8 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "diffusion_kv_cache_dtype",
         "diffusion_kv_mode",
         "diffusion_kv_max_rows_per_request",
+        "stage_input_payload_keys",
+        "stage_output_payload_keys",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
     assert {
         "revision",

tests/config/test_sleep_mode_deploy.py

@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-"""Check the sleep/wake E2E topology before loading large BAGEL weights."""
-
-import pytest
-
-from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config
-from vllm_omni.config.omni_config import (
-    VllmOmniConfig,
-    VllmOmniDiffusionStageConfig,
-    extract_diffusion_stage_config_kwargs,
-)
-from vllm_omni.config.pipeline_registry import resolve_pipeline_config
-from vllm_omni.config.stage_config import load_deploy_config
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)])
-def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size):
-    deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size))
-    pipeline = resolve_pipeline_config(deploy.pipeline or "bagel")
-    config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
-
-    assert len(config.stage_configs) == (1 if diffusion_only else 2)
-    for stage in config.stage_configs:
-        is_diffusion = diffusion_only or stage.stage_id == 1
-        assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion
-        # Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens.
-        assert stage.scheduler_config.max_num_batched_tokens == 32768
-        assert stage.model_config.enable_sleep_mode is True
-        assert stage.parallel_config.tensor_parallel_size == tp_size
-        assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size))
-        if not is_diffusion:
-            assert stage.cache_config.gpu_memory_utilization == 0.8
-        if is_diffusion:
-            engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel")
-            terminal = OmniDiffusionConfig.from_kwargs(
-                **extract_diffusion_stage_config_kwargs(
-                    engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
-                )
-            )
-            assert terminal.enable_sleep_mode is True
-            assert terminal.parallel_config.tensor_parallel_size == tp_size
-
-    if not diffusion_only:
-        assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,)
-        assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector"

tests/core/sched/test_generation_scheduler_restore.py

@@ -133,7 +133,9 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
     monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue)
     scheduler._native_data_plane = True
     scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None)
+    scheduler.input_coordinator = SimpleNamespace(
+        _async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None
+    )
     scheduler.running = []
     scheduler.waiting = create_request_queue(scheduler.policy)
     scheduler.skipped_waiting = create_request_queue(scheduler.policy)

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -1527,7 +1527,9 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.use_v2_model_runner = native
     sched.max_num_running_reqs = 8
     sched.input_coordinator = (
-        SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None
+        SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
+        if native
+        else None
     )
     sched.chunk_transfer_adapter = (
         None

tests/core/sched/test_omni_cfg_ar_scheduler.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+from vllm import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import Request, RequestStatus
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_cfg_ar_scheduler import OmniCFGARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _request(name: str, prompt_length: int) -> Request:
+    return Request(name, [0] * prompt_length, SamplingParams(max_tokens=100), None)
+
+
+def _scheduler(monkeypatch, waiting, running=()):
+    scheduler = OmniCFGARScheduler.__new__(OmniCFGARScheduler)
+    scheduler.policy = SchedulingPolicy.FCFS
+    scheduler.waiting = create_request_queue(scheduler.policy)
+    for request in waiting:
+        scheduler.waiting.add_request(request)
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+    scheduler.running = list(running)
+    scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
+    scheduler.max_model_len = 128
+    scheduler.max_num_running_reqs = 2
+    scheduler.max_num_scheduled_tokens = 256
+    scheduler.block_size = 16
+    scheduler._capacity = 16
+    scheduler._pairs = {"pair": {"cond": "cond", "uncond": "uncond"}}
+    scheduler._request_pairs = {"cond": "pair", "uncond": "pair"}
+    scheduler._prefill_min_output_tokens = 0
+
+    def advance(self, throttle_prefills=False):
+        self.running.extend(self.waiting)
+        self.waiting = create_request_queue(self.policy)
+        scheduled = {}
+        for request in self.running:
+            count = request.num_tokens - request.num_computed_tokens
+            request.num_computed_tokens += count
+            request.append_output_token_ids(0)
+            scheduled[request.request_id] = count
+        output = SchedulerOutput.make_empty()
+        output.num_scheduled_tokens = scheduled
+        output.total_num_scheduled_tokens = sum(scheduled.values())
+        return output
+
+    monkeypatch.setattr(OmniARScheduler, "schedule", advance)
+    return scheduler
+
+
+def test_unequal_prompts_keep_real_positions_and_share_context_limit(monkeypatch):
+    cond, uncond = _request("cond", 80), _request("uncond", 30)
+    scheduler = _scheduler(monkeypatch, [cond, uncond])
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens == {"cond": 80, "uncond": 30}
+    assert (cond.max_tokens, uncond.max_tokens) == (48, 48)
+    for _ in range(4):
+        assert scheduler.schedule().num_scheduled_tokens == {"cond": 1, "uncond": 1}
+    assert cond.num_computed_tokens - uncond.num_computed_tokens == 50
+
+
+def test_waiting_pair_cannot_be_split_or_starved_by_later_single(monkeypatch):
+    cond, uncond = _request("cond", 40), _request("uncond", 20)
+    running, later = _request("running", 10), _request("later", 10)
+    scheduler = _scheduler(monkeypatch, [cond, uncond, later], [running])
+    assert scheduler.schedule().num_scheduled_tokens == {"running": 10}
+    assert list(scheduler.waiting) == [cond, uncond, later]
+    scheduler.running.clear()
+    del scheduler.requests["running"]
+    assert scheduler.schedule().num_scheduled_tokens == {"cond": 40, "uncond": 20}
+    assert list(scheduler.waiting) == [later]
+
+
+def test_cancel_either_branch_finishes_both_and_releases_pair(monkeypatch):
+    cond, uncond = _request("cond", 40), _request("uncond", 20)
+    scheduler = _scheduler(monkeypatch, [cond, uncond])
+
+    def finish(self, ids, status):
+        assert set(ids) == {"cond", "uncond"}
+        assert status == RequestStatus.FINISHED_ABORTED
+        finished = [self.requests.pop(rid) for rid in ids]
+        self.waiting.remove_requests(finished)
+        return finished
+
+    monkeypatch.setattr(OmniARScheduler, "finish_requests", finish)
+    assert len(scheduler.finish_requests("uncond", RequestStatus.FINISHED_ABORTED)) == 2
+    assert not scheduler._pairs and not scheduler._request_pairs and not scheduler.waiting
+
+
+def test_new_prefill_waits_for_initial_streaming_audio(monkeypatch):
+    request = _request("plain", 20)
+    scheduler = _scheduler(monkeypatch, [], [request])
+    scheduler._prefill_min_output_tokens = 8
+    request.append_output_token_ids([0] * 7)
+    assert scheduler._should_defer_waiting_admission()
+    request.append_output_token_ids(0)
+    assert not scheduler._should_defer_waiting_admission()

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -6,8 +6,10 @@ from types import SimpleNamespace
 from typing import Any
 
 import pytest
+from vllm.config import SchedulerConfig, VllmConfig
 from vllm.v1.engine import FinishReason
 
+from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniChunkRecvHandle
@@ -90,6 +92,31 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy
     assert (scheduler.input_coordinator is not None) is enabled
 
 
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize(("stage_id", "required"), [(0, False), (1, False), (1, True)])
+def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_id, required):
+    scheduler = _Scheduler()
+    model_config = object.__new__(OmniModelConfig)
+    model_config.stage_id = stage_id
+    model_config.async_chunk = async_chunk
+    model_config.requires_full_payload_input = required
+    model_config.supports_native_mrv2_data_plane = True
+    model_config.use_v2_model_runner = True
+    scheduler.vllm_config = object.__new__(VllmConfig)
+    scheduler.vllm_config.model_config = model_config
+    scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig)
+    scheduler.vllm_config.scheduler_config.max_num_seqs = 1
+    scheduler._init_omni_io_scheduling_state()
+
+    assert scheduler._native_data_plane
+    if async_chunk or (stage_id > 0 and required):
+        assert scheduler.input_coordinator._async_chunk is async_chunk
+    else:
+        assert scheduler.input_coordinator is None
+    assert scheduler._async_chunk_transport_enabled() is async_chunk
+    assert scheduler.chunk_transfer_adapter is None
+
+
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

@@ -6,12 +6,8 @@
 ``_process_pending_chunk_timeouts`` covers async-chunk. Both share
 ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
 
-Verifies that the mixin correctly *delegates* timed-out requests to the
-base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
-The end-to-end effect (queue removal + status set + per-request cleanup +
-client-facing FINISHED_ERROR emission) is the responsibility of upstream
-vLLM's ``finish_requests`` implementation and is covered by upstream tests;
-this file only asserts the wiring from the mixin to that API.
+Verifies scheduler cleanup and explicit client-facing ERROR outputs. Upstream
+``finish_requests`` frees requests but does not emit their terminal output.
 """
 
 from __future__ import annotations
@@ -23,6 +19,9 @@ import os
 from types import SimpleNamespace
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.request import Request
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
@@ -52,7 +51,7 @@ class _FakeScheduler(OmniSchedulerMixin):
 def test_process_pending_input_timeouts_delegates_to_finish_requests():
     """Timed-out request present in self.requests is forwarded to finish_requests."""
     req_id = "stuck-req"
-    requests = {req_id: SimpleNamespace(request_id=req_id)}
+    requests = {req_id: SimpleNamespace(request_id=req_id, client_index=0)}
     coord = _FakeCoordinator(timed_out_ids={req_id})
     scheduler = _FakeScheduler(requests, coord)
 
@@ -142,7 +141,7 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
     """
     req_id = "stalled-stream"
     adapter = _FakeChunkAdapter(timed_out_ids={req_id})
-    scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
+    scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id, client_index=0)}, adapter)
 
     scheduler._process_pending_chunk_timeouts()
 
@@ -154,6 +153,41 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
     assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
 
 
+@pytest.mark.parametrize("chunked", [False, True])
+@pytest.mark.parametrize("synthesize_abort_outputs", [False, True])
+def test_timeout_emits_error_once_even_without_model_outputs(chunked, synthesize_abort_outputs):
+    from vllm.v1.engine import FinishReason
+
+    request = Request(
+        request_id="stuck",
+        prompt_token_ids=[1],
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        client_index=3,
+    )
+    if chunked:
+        scheduler = _FakeChunkScheduler({"stuck": request}, _FakeChunkAdapter({"stuck"}))
+        scheduler._process_pending_chunk_timeouts()
+    else:
+        scheduler = _FakeScheduler({"stuck": request}, _FakeCoordinator({"stuck"}))
+        scheduler._process_pending_input_timeouts()
+    # Model upstream's finish_requests cleanup. Generation must still emit
+    # ERROR, and AR's synthetic ABORT must not replace or duplicate it.
+    scheduler.requests.clear()
+    scheduler.finished_req_ids_dict = {3: {"stuck"}}
+    outputs: dict[int, EngineCoreOutputs] = {}
+    scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
+
+    [error] = outputs[3].outputs
+    assert error.request_id == "stuck"
+    assert error.finish_reason == FinishReason.ERROR
+    assert "connector input" in error.stop_reason
+    assert outputs[3].finished_requests == {"stuck"}
+    outputs = {}
+    scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
+    assert outputs == {}
+
+
 def test_process_pending_chunk_timeouts_skips_already_freed_request():
     adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
     scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)

tests/core/test_prefix_cache.py

@@ -1363,6 +1363,43 @@ def test_same_step_hit_prefetch_starts_at_save(caplog):
     assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8)))
 
 
+@pytest.mark.parametrize("reuse_blocks", [False, True], ids=["fresh", "reused"])
+@pytest.mark.parametrize("deferred_mm", [False, True], ids=["immediate", "deferred"])
+def test_same_step_hit_refreshes_prefetch_after_write(reuse_blocks, deferred_mm):
+    policy = ModelCachePolicy(deferred_keys=frozenset({"mm"}) if deferred_mm else frozenset())
+    mgr, view = make_manager(policy=policy)
+    try:
+        blocks = [0, 1] if reuse_blocks else [8, 9]
+        sid = run_step(mgr, view, {"old": (blocks, 0, 8)}, mm={"mm": torch.full((8, 2), 100.0)})
+        mgr.materialize(sid, ["old"])
+        mgr.new_step_starts(FakeSchedOut(finished=["old"]))
+
+        mgr.new_step_starts(
+            FakeSchedOut(
+                new_reqs=[FakeNewReq("a", 0, [[0, 1]]), FakeNewReq("b", 8, [[0, 1, 2]])],
+                num_scheduled={"a": 8, "b": 4},
+            )
+        )
+        old_prefetch = dict(mgr._hit_prefetch["b"])
+        # Complete the early reads before A publishes the prefix B actually hits.
+        for future in old_prefetch.values():
+            future.result(timeout=5)
+
+        view.order = ["a", "b"]
+        view.req_blocks.update(a=[0, 1], b=[0, 1, 2])
+        view.computed.update(a=0, b=8)
+        hidden = torch.cat([torch.full((8, HIDDEN), 20.0), torch.full((4, HIDDEN), 30.0)])
+        mm = torch.cat([torch.full((8, 2), 200.0), torch.full((4, 2), 300.0)])
+        sid = mgr.save_outputs(hidden, {"mm": mm}, num_tokens_unpadded=12, num_tokens_padded=12)
+        outs = mgr.materialize(sid, ["a", "b"])
+
+        assert torch.equal(outs.hidden_states["b"][:8], hidden[:8])
+        assert torch.equal(outs.hidden_states["b"][8:], hidden[8:])
+        assert torch.equal(outs.mm_outputs["mm"]["b"], mm)
+    finally:
+        mgr.shutdown()
+
+
 def test_delayed_read_of_reassigned_hit_raises():
     """A version check with no COW copy still raises for live and finished.
     The production path registers the ref first so remap preserves the rows

tests/dfx/perf/scripts/run_benchmark.py

@@ -100,6 +100,71 @@ class _SingleActiveContext:
             stack.close()
 
 
+def _config_for_test(test_name: str) -> dict[str, object] | None:
+    for config in BENCHMARK_CONFIGS:
+        if isinstance(config, dict) and config.get("test_name") == test_name:
+            return config
+    return None
+
+
+def _judge_server_params_for_test(test_name: str) -> dict[str, object] | None:
+    config = _config_for_test(test_name)
+    if config is None:
+        return None
+    raw = config.get("judge_server_params")
+    if raw is None:
+        return None
+    if not isinstance(raw, dict):
+        raise TypeError(f"judge_server_params for {test_name} must be an object")
+    return raw
+
+
+def _resolve_judge_cuda_visible_devices(configured: object) -> str:
+    """Map a judge card index onto the process-visible CUDA device list.
+
+    ``judge_server_params.cuda_visible_devices`` is an index into
+    ``CUDA_VISIBLE_DEVICES`` when that env is set (``1`` with
+    ``CUDA_VISIBLE_DEVICES=2,3`` selects physical GPU 3). Without the env it
+    is used as a raw device list.
+    """
+    if not isinstance(configured, str) or not configured:
+        raise ValueError("judge_server_params.cuda_visible_devices must be a non-empty string")
+    parent = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if parent is None or not parent.strip():
+        return configured
+    visible = [part.strip() for part in parent.split(",") if part.strip()]
+    if configured.isdigit():
+        index = int(configured)
+        if index < 0 or index >= len(visible):
+            raise ValueError(
+                f"judge_server_params.cuda_visible_devices={configured!r} is outside CUDA_VISIBLE_DEVICES={parent!r}"
+            )
+        return visible[index]
+    return configured
+
+
+@contextmanager
+def _start_judge_server(judge_params: dict[str, object]):
+    model = judge_params.get("model")
+    if not isinstance(model, str) or not model:
+        raise ValueError("judge_server_params.model must be a non-empty string")
+    extra = judge_params.get("extra_cli_args") or ()
+    if not isinstance(extra, list | tuple):
+        raise TypeError("judge_server_params.extra_cli_args must be a list")
+    visible = _resolve_judge_cuda_visible_devices(judge_params.get("cuda_visible_devices", "1"))
+    print(f"Starting OmniInteract judge with model: {model} on CUDA_VISIBLE_DEVICES={visible}")
+    with OmniServer(
+        model,
+        [str(item) for item in extra],
+        use_omni=bool(judge_params.get("use_omni", False)),
+        env_dict={"CUDA_VISIBLE_DEVICES": visible},
+    ) as judge:
+        print(f"OmniInteract judge started on {judge.host}:{judge.port}")
+        yield judge
+        print("OmniInteract judge stopping...")
+    print("OmniInteract judge stopped")
+
+
 # OmniServer defaults for flags not already present in JSON ``serve_args`` /
 # ``extra_cli_args``. Add new (flag, value) pairs here rather than special-casing.
 _OMNI_DEFAULT_SERVER_ARGS: tuple[tuple[str, str], ...] = (
@@ -200,10 +265,17 @@ def _start_omni_server(server_param):
         server_args = ["--stage-overrides", stage_overrides] + server_args
     if extra:
         server_args = list(extra) + server_args
+    judge_params = _judge_server_params_for_test(test_name)
     with OmniServer(model, server_args, use_omni=use_omni, env_dict=_omni_server_env()) as server:
         server.test_name = test_name
+        server.judge_base_url = None
         print("OmniServer started successfully")
-        yield server
+        if judge_params is None:
+            yield server
+        else:
+            with _start_judge_server(judge_params) as judge:
+                server.judge_base_url = f"http://{judge.host}:{judge.port}"
+                yield server
         print("OmniServer stopping...")
 
     print("OmniServer stopped")
@@ -260,8 +332,153 @@ def _resolve_num_warmups(params: dict[str, Any], *, default: int) -> int:
     return value
 
 
+def _assert_fixed_stage_workloads(result: dict[str, Any], params: dict[str, Any], num_prompts: int) -> None:
+    """Reject successful responses that did less work than a fixed-length benchmark requested."""
+    extra_body = params.get("extra_body") or params.get("extra-body") or {}
+    fixed_lengths = {
+        str(stage): sampling["max_tokens"]
+        for stage, sampling in enumerate(extra_body.get("sampling_params_list", []))
+        if isinstance(sampling, dict)
+        and isinstance(sampling.get("max_tokens"), int)
+        and sampling["max_tokens"] > 0
+        and sampling.get("min_tokens") == sampling["max_tokens"]
+    }
+    if not fixed_lengths:
+        return
+
+    snapshots = result.get("request_stage_metrics")
+    assert isinstance(snapshots, list) and len(snapshots) == num_prompts, (
+        "Fixed stage workload: missing per-request stage metrics"
+    )
+    for request_index, snapshot in enumerate(snapshots):
+        for stage, expected in fixed_lengths.items():
+            metrics = snapshot.get(stage) if isinstance(snapshot, dict) else None
+            actual = metrics.get("num_tokens_out") if isinstance(metrics, dict) else None
+            assert actual == expected, (
+                f"Fixed stage workload: request {request_index}, stage {stage}: "
+                f"expected {expected} output tokens, got {actual!r}"
+            )
+    stages = ", ".join(f"stage {stage}={expected}" for stage, expected in fixed_lengths.items())
+    print(f"Fixed stage workload OK: {num_prompts} requests, num_tokens_out {stages}")
+
+
+def _is_finite_number(value: object) -> bool:
+    return isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(value)
+
+
+_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1 = "omniinteract_aggregate_min_ia_qtf1"
+_OMNIINTERACT_AGGREGATE_SUBSETS = "omniinteract_aggregate_subsets"
+_OMNIINTERACT_AGGREGATE_GROUP = "omniinteract_aggregate_group"
+_OMNIINTERACT_AGGREGATE_COUNTS: dict[tuple[str, str], tuple[float, float, float]] = {}
+_OMNIINTERACT_MIN_EVALUATED_CASES = 1
+
+
+def _reset_omniinteract_aggregate_counts() -> None:
+    _OMNIINTERACT_AGGREGATE_COUNTS.clear()
+
+
+def _ia_qtf1_from_counts(tp: float, fp: float, fn: float) -> float:
+    precision = tp / (tp + fp) if tp + fp > 0 else 0.0
+    recall = tp / (tp + fn) if tp + fn > 0 else 0.0
+    return (2.0 * precision * recall / (precision + recall)) if precision + recall > 0 else 0.0
+
+
+def _assert_subset_evaluated(accuracy: dict[str, object]) -> None:
+    """Reject a subset whose clipped or cancelled cases left nothing to score.
+
+    Individual ineligible cases may be skipped. A subset with fewer than
+    ``_OMNIINTERACT_MIN_EVALUATED_CASES`` evaluated cases must not pass, or its
+    zero TP/FP/FN would be pooled as a successful accuracy result.
+    """
+
+    evaluated = accuracy.get("evaluated")
+    if not isinstance(evaluated, int) or isinstance(evaluated, bool):
+        raise AssertionError("OmniInteract accuracy evaluated count is missing")
+    if evaluated < _OMNIINTERACT_MIN_EVALUATED_CASES:
+        raise AssertionError(
+            f"OmniInteract subset evaluated {evaluated} cases (skipped={accuracy.get('skipped')}); "
+            f"at least {_OMNIINTERACT_MIN_EVALUATED_CASES} evaluated case is required"
+        )
+
+
+def _finite_accuracy_count(value: object, name: str) -> float:
+    assert _is_finite_number(value), f"OmniInteract accuracy summary {name} is missing"
+    return float(value)
+
+
+def _aggregate_accuracy_params_for_test(test_name: str) -> dict[str, object]:
+    config = _config_for_test(test_name)
+    if config is None:
+        return {}
+    min_ia_qtf1 = config.get(_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1)
+    if min_ia_qtf1 is None:
+        return {}
+    if not _is_finite_number(min_ia_qtf1):
+        raise ValueError("omniinteract_aggregate_min_ia_qtf1 must be a finite number")
+    subsets: list[str] = []
+    raw_params = config.get("benchmark_params") or []
+    if not isinstance(raw_params, list):
+        raise TypeError(f"benchmark_params for {test_name} must be a list")
+    for item in raw_params:
+        if not isinstance(item, dict) or not item.get("omniinteract_evaluate"):
+            continue
+        subset = item.get("omniinteract_subsets")
+        if not isinstance(subset, str) or not subset:
+            raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_subsets on each evaluated case")
+        subsets.append(subset)
+    if not subsets:
+        raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires at least one evaluated OmniInteract subset")
+    if len(set(subsets)) != len(subsets):
+        raise ValueError("omniinteract_aggregate_subsets must not contain duplicates")
+    return {
+        _OMNIINTERACT_AGGREGATE_MIN_IA_QTF1: float(min_ia_qtf1),
+        _OMNIINTERACT_AGGREGATE_SUBSETS: subsets,
+        _OMNIINTERACT_AGGREGATE_GROUP: test_name,
+    }
+
+
+def _maybe_assert_aggregate_ia_qtf1(params: dict[str, object], acc_summary: dict[str, object]) -> None:
+    min_ia_qtf1 = params.get(_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1)
+    if min_ia_qtf1 is None:
+        return
+    assert _is_finite_number(min_ia_qtf1), "omniinteract_aggregate_min_ia_qtf1 must be a finite number"
+    subsets = params.get(_OMNIINTERACT_AGGREGATE_SUBSETS)
+    if not isinstance(subsets, list | tuple) or not subsets:
+        raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_aggregate_subsets")
+    if len(set(subsets)) != len(subsets):
+        raise ValueError("omniinteract_aggregate_subsets must not contain duplicates")
+    if not all(isinstance(item, str) and item for item in subsets):
+        raise ValueError("omniinteract_aggregate_subsets must be non-empty strings")
+    subset = params.get("omniinteract_subsets")
+    if not isinstance(subset, str) or not subset:
+        raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_subsets")
+    if subset not in subsets:
+        raise ValueError(f"omniinteract_subsets {subset!r} is not in omniinteract_aggregate_subsets")
+    group = params.get(_OMNIINTERACT_AGGREGATE_GROUP)
+    if not isinstance(group, str) or not group:
+        raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_aggregate_group")
+    tp = _finite_accuracy_count(acc_summary.get("Global_TP"), "Global_TP")
+    fp = _finite_accuracy_count(acc_summary.get("Global_FP"), "Global_FP")
+    fn = _finite_accuracy_count(acc_summary.get("Global_FN"), "Global_FN")
+    _OMNIINTERACT_AGGREGATE_COUNTS[(group, subset)] = (tp, fp, fn)
+    recorded = [_OMNIINTERACT_AGGREGATE_COUNTS.get((group, name)) for name in subsets]
+    present = [counts for counts in recorded if counts is not None]
+    if len(present) != len(subsets):
+        return
+    total_tp = sum(counts[0] for counts in present)
+    total_fp = sum(counts[1] for counts in present)
+    total_fn = sum(counts[2] for counts in present)
+    ia_qtf1 = _ia_qtf1_from_counts(total_tp, total_fp, total_fn)
+    print(
+        f"OmniInteract aggregate All Global IA-QTF1: {ia_qtf1:.6f} "
+        f"(TP={total_tp:.6f} / FP={total_fp:.6f} / FN={total_fn:.6f})"
+    )
+    assert ia_qtf1 >= float(min_ia_qtf1), f"OmniInteract aggregate All Global IA-QTF1 {ia_qtf1} is below {min_ia_qtf1}"
+
+
 def assert_result(result, params, num_prompt) -> None:
     assert result["completed"] == num_prompt, "Request failures exist"
+    _assert_fixed_stage_workloads(result, params, num_prompt)
     if params.get("dataset_name") == "omniinteract":
         summary = result.get("omniinteract")
         assert isinstance(summary, dict), "OmniInteract summary is missing"
@@ -271,6 +488,23 @@ def assert_result(result, params, num_prompt) -> None:
             0,
         ), "OmniInteract requests did not all succeed"
         assert summary.get("artifacts_complete") is True, "OmniInteract artifacts are incomplete"
+        if params.get("omniinteract_evaluate"):
+            accuracy = summary.get("accuracy")
+            assert isinstance(accuracy, dict), "OmniInteract accuracy is missing"
+            assert accuracy.get("status") == "ok", "OmniInteract accuracy did not complete"
+            assert accuracy.get("failed") == 0, "OmniInteract accuracy reported failed cases"
+            _assert_subset_evaluated(accuracy)
+            acc_summary = accuracy.get("summary")
+            assert isinstance(acc_summary, dict), "OmniInteract accuracy summary is missing"
+            ia_qtf1 = acc_summary.get("IA_QTF1")
+            assert _is_finite_number(ia_qtf1), "OmniInteract All Global IA-QTF1 is missing"
+            min_ia_qtf1 = params.get("omniinteract_min_ia_qtf1")
+            if min_ia_qtf1 is not None:
+                assert _is_finite_number(min_ia_qtf1), "omniinteract_min_ia_qtf1 must be a finite number"
+                assert float(ia_qtf1) >= float(min_ia_qtf1), (
+                    f"OmniInteract All Global IA-QTF1 {ia_qtf1} is below {min_ia_qtf1}"
+                )
+            _maybe_assert_aggregate_ia_qtf1(params, acc_summary)
     baseline = params.get("baseline")
     hardware = result.get("Hardware")
     hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
@@ -304,8 +538,13 @@ def assert_result(result, params, num_prompt) -> None:
 )
 def test_performance_benchmark(omni_server, benchmark_params):
     test_name = benchmark_params["test_name"]
-    params = benchmark_params["params"]
+    params = dict(benchmark_params["params"])
     dataset_name = params.get("dataset_name", "")
+    if params.get("omniinteract_evaluate"):
+        judge_base_url = getattr(omni_server, "judge_base_url", None)
+        if not isinstance(judge_base_url, str) or not judge_base_url:
+            raise ValueError("omniinteract_evaluate requires top-level judge_server_params to start a judge")
+        params["omniinteract_judge_base_url"] = judge_base_url
 
     host = omni_server.host
     port = omni_server.port
@@ -349,6 +588,11 @@ def test_performance_benchmark(omni_server, benchmark_params):
         "eval_phase",
         "trust_remote_code",
         "expected_duplex_audio_turns_per_session",
+        "omniinteract_min_ia_qtf1",
+        _OMNIINTERACT_AGGREGATE_MIN_IA_QTF1,
+        _OMNIINTERACT_AGGREGATE_SUBSETS,
+        _OMNIINTERACT_AGGREGATE_GROUP,
+        "judge_server_params",
         "name",
         "enable_negative_prompt",
         "random_request_config",
@@ -403,7 +647,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
             num_warmups=_resolve_num_warmups(params, default=2),
             benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
         )
-        assert_result(result, params, num_prompt)
+        assert_result(result, {**params, **_aggregate_accuracy_params_for_test(test_name)}, num_prompt)
 
     # concurrency test
     for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
@@ -422,4 +666,4 @@ def test_performance_benchmark(omni_server, benchmark_params):
             num_warmups=_resolve_num_warmups(params, default=max(2, int(concurrency))),
             benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
         )
-        assert_result(result, params, num_prompt)
+        assert_result(result, {**params, **_aggregate_accuracy_params_for_test(test_name)}, num_prompt)

tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json

@@ -0,0 +1,131 @@
+[
+  {
+    "test_name": "test_hunyuan_image3_prefix_dense",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: dense baseline (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "dense_legacy",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": false,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  },
+  {
+    "test_name": "test_hunyuan_image3_prefix_paged_no_cache",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: paged without prefix caching (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "paged_scheduler",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": false,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  },
+  {
+    "test_name": "test_hunyuan_image3_prefix_paged_prefix",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: paged with full prefix reuse (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "paged_scheduler",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": true,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  }
+]

tests/dfx/perf/tests/test_minicpmo_4_5_omniinteract.json

@@ -7,7 +7,7 @@
           "res": {
             "cuda": "H100"
           },
-          "num_cards": 1
+          "num_cards": 2
         }
       },
       "local_model",
@@ -21,6 +21,16 @@
         "--trust-remote-code"
       ]
     },
+    "judge_server_params": {
+      "model": "Qwen/Qwen2.5-7B-Instruct",
+      "use_omni": false,
+      "cuda_visible_devices": "1",
+      "extra_cli_args": [
+        "--served-model-name",
+        "Qwen2.5-7B-Instruct"
+      ]
+    },
+    "omniinteract_aggregate_min_ia_qtf1": 0.2,
     "benchmark_params": [
       {
         "dataset_name": "omniinteract",
@@ -40,6 +50,8 @@
         "omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
         "omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1q1a",
         "omniinteract_max_video_duration_s": 600,
+        "omniinteract_evaluate": true,
+        "omniinteract_judge_model": "Qwen2.5-7B-Instruct",
         "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
       },
       {
@@ -60,6 +72,8 @@
         "omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
         "omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1q1a_math",
         "omniinteract_max_video_duration_s": 600,
+        "omniinteract_evaluate": true,
+        "omniinteract_judge_model": "Qwen2.5-7B-Instruct",
         "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
       },
       {
@@ -80,6 +94,8 @@
         "omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
         "omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1qna",
         "omniinteract_max_video_duration_s": 600,
+        "omniinteract_evaluate": true,
+        "omniinteract_judge_model": "Qwen2.5-7B-Instruct",
         "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
       }
     ]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -1,6 +1,7 @@
 [
   {
     "test_name": "test_qwen3_omni",
+    "description": "random row: fixed workload (Thinker 900 / Talker 1536 tokens, min_tokens == max_tokens) for comparable audio RTF; H100 baseline pending recalibration, see #8100",
     "mark": [
       {
         "hardware_marks": {
@@ -43,17 +44,28 @@
         "random_output_len": 900,
         "ignore_eos": true,
         "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "H100": {
-            "mean_audio_rtf": [
-              0.169,
-              0.2041,
-              0.2382,
-              0.3177,
-              0.4595
-            ]
-          }
-        }
+        "seed": 0,
+        "random_range_ratio": 0.0,
+        "extra_body": {
+          "return_stage_metrics": true,
+          "sampling_params_list": [
+            {
+              "temperature": 0.0,
+              "min_tokens": 900,
+              "max_tokens": 900,
+              "ignore_eos": true
+            },
+            {
+              "temperature": 0.9,
+              "top_k": 50,
+              "repetition_penalty": 1.05,
+              "seed": 0,
+              "min_tokens": 1536,
+              "max_tokens": 1536
+            }
+          ]
+        },
+        "baseline": {}
       },
       {
         "dataset_name": "random-mm",

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -60,6 +60,42 @@
       }
     ]
   },
+  {
+    "test_name": "test_qwen_image_single_device_a3",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "npu": "A3"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (no parallelism)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
   {
     "test_name": "test_qwen_image_single_device_step_execution",
     "mark": [
@@ -144,13 +180,66 @@
       }
     ]
   },
+  {
+    "test_name": "test_qwen_image_single_device_step_execution_a3",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "npu": "A3"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device step execution: sequential 512 plus 512 concurrency sweep. One server (max-num-seqs 8).",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "step-execution": true,
+        "max-num-seqs": 8
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "512x512_steps20_high_concurrency",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": [20, 40, 80, 160],
+        "max-concurrency": [1, 2, 4, 8],
+        "warmup-requests": 8,
+        "warmup-concurrency": 8,
+        "warmup-num-inference-steps": 20,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
   {
     "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
     "mark": [
       {
         "hardware_marks": {
           "res": {
-            "cuda": ["H100", "B200"]
+            "cuda": ["H100", "B200"],
+            "npu": "A3"
           },
           "num_cards": 4
         }
@@ -265,5 +354,55 @@
         }
       }
     ]
+  },
+  {
+    "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit_a3",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "npu": "A3"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "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_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
+        },
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
   }
 ]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -235,6 +235,53 @@ def test_is_diffusion_perf_config():
     assert is_diffusion_perf_config(custom_edits_cfg)
 
 
+def test_buildkite_perf_steps_use_matching_runner_schema():
+    """Every Buildkite step must run a perf JSON with the runner matching its schema.
+
+    Regression for #8074: the NPU nightly HunyuanVideo-1.5 step kept invoking
+    ``run_diffusion_benchmark.py`` after the JSON was migrated to the omni-bench
+    schema (#7737). The runner skipped every case as omni-bench, pytest selected
+    0 tests and the step failed with exit 5.
+    """
+    import json
+    import re
+    from pathlib import Path
+
+    from tests.dfx.conftest import is_diffusion_perf_config
+
+    repo_root = Path(__file__).resolve().parents[4]
+    step_re = re.compile(
+        r"tests/dfx/perf/scripts/(?P<runner>run_diffusion_benchmark|run_benchmark)\.py"
+        r".*?--test-config-file(?:=|\s+)(?P<file>\S+\.json)"
+    )
+    pipelines = sorted(repo_root.glob(".buildkite/**/*.yml")) + sorted(repo_root.glob(".buildkite/**/*.yaml"))
+    assert pipelines, "no Buildkite pipelines found"
+
+    invocations = 0
+    for pipeline in pipelines:
+        for match in step_re.finditer(pipeline.read_text(encoding="utf-8")):
+            invocations += 1
+            config_rel = match.group("file")
+            config_path = repo_root / config_rel
+            assert config_path.exists(), f"{pipeline}: perf config not found: {config_rel}"
+            runner_is_diffusion = match.group("runner") == "run_diffusion_benchmark"
+            cases = json.loads(config_path.read_text(encoding="utf-8"))
+            mismatched = [cfg["test_name"] for cfg in cases if is_diffusion_perf_config(cfg) != runner_is_diffusion]
+            assert not mismatched, (
+                f"{pipeline}: {match.group('runner')} runs {config_rel}, but case(s) {mismatched} "
+                f"are not {'diffusion' if runner_is_diffusion else 'omni-bench'}-schema; the runner "
+                "would skip them and pytest may select 0 tests (exit 5), see issue #8074"
+            )
+
+    assert invocations, "no perf runner invocations found; the scan regex or pipeline layout changed"
+
+    # Pin the issue #8074 scenario itself: the migrated HunyuanVideo-1.5 t2v JSON
+    # is omni-bench schema and must never go back to run_diffusion_benchmark.py.
+    t2v_json = repo_root / "tests" / "dfx" / "perf" / "tests" / "test_hunyuanvideo15_t2v_vllm_omni.json"
+    t2v_cases = json.loads(t2v_json.read_text(encoding="utf-8"))
+    assert t2v_cases and all(not is_diffusion_perf_config(cfg) for cfg in t2v_cases)
+
+
 def test_merge_omni_default_server_args_respects_json():
     from tests.dfx.perf.scripts.run_benchmark import _merge_omni_default_server_args
 
@@ -254,6 +301,38 @@ def test_merge_omni_default_server_args_respects_json():
     ]
 
 
+def test_prefix_benchmark_uses_unified_runner(monkeypatch):
+    from pathlib import Path
+
+    from tests.dfx.conftest import is_diffusion_perf_config, load_benchmark_configs
+    from tests.dfx.perf.scripts import run_benchmark
+
+    configs = load_benchmark_configs(str(Path(__file__).with_name("test_hunyuan_image3_prefix_caching.json")))
+    assert len(configs) == 3
+    assert all(not is_diffusion_perf_config(config) for config in configs)
+    assert all(config["benchmark_params"] == configs[0]["benchmark_params"] for config in configs)
+    monkeypatch.setattr(run_benchmark, "BENCHMARK_CONFIGS", configs)
+    monkeypatch.setattr(run_benchmark, "get_runtime_resource_label", lambda: "H100")
+    calls = []
+
+    def benchmark(**kwargs):
+        calls.append(kwargs)
+        assert "--name" not in kwargs["args"]
+        assert "--warmup-dataset-path" not in kwargs["args"]
+        assert "tests/assets/hunyuan_image3/it2i.jsonl" in kwargs["args"]
+        assert kwargs["num_warmups"] == 2
+        return {"completed": 8}
+
+    monkeypatch.setattr(run_benchmark, "run_benchmark", benchmark)
+    for config in configs:
+        for params in config["benchmark_params"]:
+            run_benchmark.test_performance_benchmark(
+                SimpleNamespace(host="localhost", port=8000, model="test-model"),
+                {"test_name": config["test_name"], "params": params},
+            )
+    assert len(calls) == 3
+
+
 def test_benchmark_param_id_suffix_from_task_eval_phase():
     from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
 
@@ -662,6 +741,204 @@ def test_omniinteract_result_rejects_incomplete_artifacts():
         )
 
 
+def test_omniinteract_result_accepts_accuracy_at_floor():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    assert_result(
+        {
+            "completed": 4,
+            "omniinteract": {
+                "total": 4,
+                "success": 4,
+                "failed": 0,
+                "artifacts_complete": True,
+                "accuracy": {
+                    "status": "ok",
+                    "failed": 0,
+                    "evaluated": 4,
+                    "summary": {"IA_QTF1": 0.25},
+                },
+            },
+        },
+        {"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.2},
+        4,
+    )
+
+
+def test_omniinteract_result_rejects_missing_accuracy():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="accuracy is missing"):
+        assert_result(
+            {
+                "completed": 4,
+                "omniinteract": {"total": 4, "success": 4, "failed": 0, "artifacts_complete": True},
+            },
+            {"dataset_name": "omniinteract", "omniinteract_evaluate": True},
+            4,
+        )
+
+
+def test_omniinteract_result_rejects_failed_accuracy():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="accuracy did not complete"):
+        assert_result(
+            {
+                "completed": 4,
+                "omniinteract": {
+                    "total": 4,
+                    "success": 4,
+                    "failed": 0,
+                    "artifacts_complete": True,
+                    "accuracy": {"status": "failed", "failed": 1, "summary": {"IA_QTF1": 0.9}},
+                },
+            },
+            {"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.0},
+            4,
+        )
+
+
+def test_omniinteract_result_rejects_ia_qtf1_below_floor():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="IA-QTF1"):
+        assert_result(
+            {
+                "completed": 4,
+                "omniinteract": {
+                    "total": 4,
+                    "success": 4,
+                    "failed": 0,
+                    "artifacts_complete": True,
+                    "accuracy": {
+                        "status": "ok",
+                        "failed": 0,
+                        "evaluated": 4,
+                        "summary": {"IA_QTF1": 0.1},
+                    },
+                },
+            },
+            {"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.2},
+            4,
+        )
+
+
+def _omniinteract_accuracy_result(
+    *,
+    tp: float,
+    fp: float,
+    fn: float,
+    ia_qtf1: float = 0.5,
+    evaluated: int = 4,
+    skipped: int = 0,
+) -> dict[str, object]:
+    return {
+        "completed": 4,
+        "omniinteract": {
+            "total": 4,
+            "success": 4,
+            "failed": 0,
+            "artifacts_complete": True,
+            "accuracy": {
+                "status": "ok",
+                "failed": 0,
+                "evaluated": evaluated,
+                "skipped": skipped,
+                "total": 4,
+                "summary": {"IA_QTF1": ia_qtf1, "Global_TP": tp, "Global_FP": fp, "Global_FN": fn},
+            },
+        },
+    }
+
+
+def _omniinteract_aggregate_params(subset: str, *, group: str) -> dict[str, object]:
+    return {
+        "dataset_name": "omniinteract",
+        "omniinteract_evaluate": True,
+        "omniinteract_subsets": subset,
+        "omniinteract_aggregate_min_ia_qtf1": 0.2,
+        "omniinteract_aggregate_subsets": ["1q1a", "1q1a_math", "1qna"],
+        "omniinteract_aggregate_group": group,
+    }
+
+
+def test_omniinteract_aggregate_waits_until_all_subsets():
+    from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
+
+    _reset_omniinteract_aggregate_counts()
+    assert_result(
+        _omniinteract_accuracy_result(tp=6.157502, fp=9, fn=8),
+        _omniinteract_aggregate_params("1q1a", group="wait"),
+        4,
+    )
+    assert_result(
+        _omniinteract_accuracy_result(tp=1.133541, fp=2, fn=2),
+        _omniinteract_aggregate_params("1q1a_math", group="wait"),
+        4,
+    )
+
+
+def test_omniinteract_aggregate_accepts_pooled_ia_qtf1_at_floor():
+    from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
+
+    _reset_omniinteract_aggregate_counts()
+    assert_result(
+        _omniinteract_accuracy_result(tp=6.157502, fp=9, fn=8),
+        _omniinteract_aggregate_params("1q1a", group="pass"),
+        4,
+    )
+    assert_result(
+        _omniinteract_accuracy_result(tp=1.133541, fp=2, fn=2),
+        _omniinteract_aggregate_params("1q1a_math", group="pass"),
+        4,
+    )
+    assert_result(
+        _omniinteract_accuracy_result(tp=0.0, fp=11, fn=18, ia_qtf1=0.0),
+        _omniinteract_aggregate_params("1qna", group="pass"),
+        4,
+    )
+
+
+def test_omniinteract_aggregate_rejects_pooled_ia_qtf1_below_floor():
+    from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
+
+    _reset_omniinteract_aggregate_counts()
+    assert_result(
+        _omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
+        _omniinteract_aggregate_params("1q1a", group="fail"),
+        4,
+    )
+    assert_result(
+        _omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
+        _omniinteract_aggregate_params("1q1a_math", group="fail"),
+        4,
+    )
+    with pytest.raises(AssertionError, match="aggregate All Global IA-QTF1"):
+        assert_result(
+            _omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
+            _omniinteract_aggregate_params("1qna", group="fail"),
+            4,
+        )
+
+
+def test_omniinteract_aggregate_rejects_entirely_skipped_subset():
+    from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
+
+    _reset_omniinteract_aggregate_counts()
+    assert_result(
+        _omniinteract_accuracy_result(tp=4.0, fp=0.0, fn=0.0, ia_qtf1=1.0),
+        _omniinteract_aggregate_params("1q1a", group="skipped-subset"),
+        4,
+    )
+    with pytest.raises(AssertionError, match="evaluated 0 cases"):
+        assert_result(
+            _omniinteract_accuracy_result(tp=0.0, fp=0.0, fn=0.0, ia_qtf1=1.0, evaluated=0, skipped=4),
+            _omniinteract_aggregate_params("1q1a_math", group="skipped-subset"),
+            4,
+        )
+
+
 def test_omni_tpot_baseline_accepts_measured_finite_sample():
     from tests.dfx.perf.scripts.run_benchmark import assert_result
 
@@ -702,3 +979,58 @@ def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples
             {"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
             1,
         )
+
+
+# Mirrors the Thinker/Talker lengths in test_qwen3_omni_no_async_chunk.json.
+_FIXED_SAMPLING = {
+    "sampling_params_list": [{"min_tokens": 900, "max_tokens": 900}, {"min_tokens": 1536, "max_tokens": 1536}]
+}
+_FIXED_OK = {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1536}}
+
+
+@pytest.mark.parametrize("extra_body_key", ["extra_body", "extra-body"])
+def test_fixed_stage_workload_accepts_exact_lengths(extra_body_key, capsys):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    assert_result(
+        {"completed": 2, "failed": 0, "request_stage_metrics": [_FIXED_OK, _FIXED_OK]},
+        {extra_body_key: _FIXED_SAMPLING},
+        2,
+    )
+    assert "Fixed stage workload OK: 2 requests, num_tokens_out stage 0=900, stage 1=1536" in capsys.readouterr().out
+
+
+@pytest.mark.parametrize(
+    "snapshots",
+    [
+        None,
+        [],
+        [_FIXED_OK],
+        [_FIXED_OK, None],
+        [_FIXED_OK, {}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 486}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1537}}],
+        [_FIXED_OK, {"1": {"num_tokens_out": 1536}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 512}, "1": {"num_tokens_out": 1536}}],
+    ],
+)
+def test_fixed_stage_workload_rejects_missing_or_wrong_lengths(snapshots):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="Fixed stage workload"):
+        assert_result({"completed": 2, "request_stage_metrics": snapshots}, {"extra_body": _FIXED_SAMPLING}, 2)
+
+
+def test_fixed_stage_workload_rejects_request_failures():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="Request failures"):
+        assert_result({"completed": 1, "failed": 1}, {}, 2)
+
+
+@pytest.mark.parametrize("extra_body", [{}, {"sampling_params_list": [{"max_tokens": 900}]}])
+def test_variable_stage_workload_does_not_require_stage_metrics(extra_body):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    assert_result({"completed": 2}, {"extra_body": extra_body}, 2)

tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py

@@ -1,79 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Invalid audio format rejection on MiniCPM-o 4.5: ``POST /v1/chat/completions``."""
-
-import os
-
-os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
-
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path
-
-pytestmark = [pytest.mark.slow, pytest.mark.omni, pytest.mark.full_model]
-
-_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
-
-_MINICPMO_SERVER = [
-    pytest.param(
-        OmniServerParams(
-            model="openbmb/MiniCPM-o-4_5",
-            stage_config_path=_MINICPMO_DEPLOY,
-            use_stage_cli=False,
-            server_args=[
-                "--trust-remote-code",
-            ],
-        ),
-        id="minicpmo_4_5",
-    )
-]
-
-
-def _system_prompt() -> dict[str, object]:
-    return {
-        "role": "system",
-        "content": [
-            {
-                "type": "text",
-                "text": (
-                    "You are MiniCPM-o 4.5, a virtual human capable of perceiving "
-                    "auditory and visual inputs, as well as generating text and speech."
-                ),
-            }
-        ],
-    }
-
-
-def _prompt() -> str:
-    return "What is the capital of China? Answer in 20 words."
-
-
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True)
-def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None:
-    """
-    Test that invalid audio format is properly rejected.
-    Deploy Setting: default single GPU
-    Input Modal: text + invalid audio format
-    Output Modal: error
-    """
-    fake_base64 = "AAAA"  # not a real audio file
-
-    responses = openai_client.send_chat_completions_http_request(
-        {
-            "json": {
-                "model": omni_server.model,
-                "messages": dummy_messages_from_mix_data(
-                    system_prompt=_system_prompt(),
-                    audio_data_url=f"data:audio/invalid;base64,{fake_base64}",
-                    content_text=_prompt(),
-                ),
-                "stream": True,
-            },
-        },
-        err_code=400,
-        err_message="Invalid or unsupported audio file",
-    )
-    assert not responses[0].success

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
+"""Invalid inputs on omni chat: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
 
 from __future__ import annotations
 
@@ -10,7 +10,7 @@ from typing import Any
 import pytest
 
 from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
 from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.slow, pytest.mark.omni]
@@ -75,6 +75,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
     elif case_id == "audio_format_unsupported":
         body["modalities"] = ["text", "audio"]
         body["audio"] = {"voice": "alloy", "format": "aac"}
+    elif case_id == "audio_data_url_corrupt":
+        body["messages"] = dummy_messages_from_mix_data(
+            audio_data_url="data:audio/invalid;base64,AAAA",
+            content_text="What is the capital of China? Answer in 20 words.",
+        )
+        body["stream"] = True
     else:
         raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
     return {"json": body, "timeout": 120}
@@ -144,6 +150,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
             ("Invalid audio format", "aac", "Supported formats"),
             id="unsupported_audio_format",
         ),
+        pytest.param(
+            "audio_data_url_corrupt",
+            400,
+            "Invalid or corrupted audio data",
+            id="invalid_audio_data_url",
+        ),
     ],
 )
 @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)

tests/diffusion/ar_diffusion/test_history_staging_gpu.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA regressions for contiguous AR-Diffusion history staging."""
+
+import pytest
+import torch
+
+from tests.diffusion.ar_diffusion.test_paged_attention import (
+    BLOCK,
+    HEAD_DIM,
+    N_HEADS,
+    POS,
+    _commit_video_span,
+    _gpu_flash_attn_usable,
+    make_state,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.experimental.ar_diffusion.kv_cache import paged_write_attn
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.skipif(
+    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
+)
+@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
+@pytest.mark.parametrize("window_chunks", [2, 4])
+@torch.inference_mode()
+def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window_chunks):
+    """Full and reused gathers match fresh attention without touching spare capacity or stale history."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device, dtype = torch.device("cuda"), torch.bfloat16
+    kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
+    stage_key, stage_value = kv.history_staging[0]
+    stage_key.fill_(-7)
+    stage_value.fill_(-7)
+    query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
+    for step in range(2):
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+        assert inputs.reuse_history is (step > 0)
+        key = torch.randn_like(query)
+        value = torch.randn_like(query)
+        actual = paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+        expected = paged_write_attn(
+            inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
+            query,
+            key,
+            value,
+            None,
+            None,
+            HEAD_DIM**-0.5,
+        )
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        ids = inputs.block_table[0].long()
+        for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
+            torch.testing.assert_close(
+                staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
+            )
+            assert (staged[inputs.max_seq_len :] == -7).all()
+
+
+@pytest.mark.skipif(
+    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
+)
+@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
+@pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.parametrize("reuse_history", [False, True])
+@torch.inference_mode()
+def test_staged_attention_replays_inductor_cudagraph(monkeypatch, window_chunks, reuse_history):
+    """Inductor must capture the real mutable op and replay with fresh current K/V."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device, dtype = torch.device("cuda"), torch.bfloat16
+    kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
+    stage_key, stage_value = kv.history_staging[0]
+    stage_key.fill_(-7)
+    stage_value.fill_(-7)
+    query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
+    key, value = torch.randn_like(query), torch.randn_like(query)
+
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.max_video_tokens = 2 * BLOCK
+    ctx.prepare(device, action_len=0, query_len=BLOCK)
+    inputs = ctx.layer_inputs(0)
+    # Populate history before testing the reuse variant, as the first denoising probe does.
+    paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+    if reuse_history:
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+    assert inputs.reuse_history is reuse_history
+
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(paged_write_attn, mode="reduce-overhead", fullgraph=True)
+        for _ in range(3):
+            torch.compiler.cudagraph_mark_step_begin()
+            compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+        current_omni_platform.synchronize()
+
+        previous = None
+        for step in range(2):
+            key.normal_()
+            # Distinct value ranges ensure a stale replay cannot match the reference.
+            value.fill_(2.0 if step == 0 else -2.0)
+            torch.compiler.cudagraph_mark_step_begin()
+            with torch.profiler.profile(
+                activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]
+            ) as prof:
+                actual = compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+                current_omni_platform.synchronize()
+            assert any("cudaGraphLaunch" in event.key for event in prof.key_averages()), (
+                "Inductor did not replay a CUDA graph; mutable staging inputs may have lost their static annotation"
+            )
+            expected = paged_write_attn(
+                inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
+                query,
+                key,
+                value,
+                None,
+                None,
+                HEAD_DIM**-0.5,
+            )
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            if previous is not None:
+                assert not torch.equal(actual, previous)
+            # Copy outside capture: graph-owned output storage is reused on the next step.
+            previous = actual.clone()
+            del actual
+            ids = inputs.block_table[0].long()
+            for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
+                torch.testing.assert_close(
+                    staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
+                )
+                assert (staged[inputs.max_seq_len :] == -7).all()
+    finally:
+        torch._dynamo.reset()

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -244,7 +244,8 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
 
 
 @pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
-def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks):
+@pytest.mark.parametrize("window_chunks", [2, 4])
+def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
     """A second probe of the same AR block restages its current K/V where the table actually holds it.
 
     The padded table always ends in at least one action-capacity block, and while the window is still
@@ -254,11 +255,14 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device = torch.device("cpu")
     dtype = torch.float32
-    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True)
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True)
+    for buffer in kv.history_staging[0]:
+        buffer.fill_(-7)
     if history_chunks:
         _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
 
     ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.max_video_tokens = 2 * BLOCK
     ctx.ensure_video_slots(device)
     key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
 
@@ -275,6 +279,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
         n_blocks = max_seq_len // BLOCK
         block_ids = block_table[0, :n_blocks].to(torch.long)
         stage_k, stage_v = ctx.history_staging(0)
+        # The custom op narrows the manager-owned buffers before calling this helper.
+        stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len]
         paged_attention_module._stage_window(
             stage_k,
             stage_v,
@@ -287,6 +293,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
         )
         full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
         full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
+        for buffer in kv.history_staging[0]:
+            assert (buffer[max_seq_len:] == -7).all()
         return stage_k, stage_v, full_k, full_v, n_blocks
 
     stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
@@ -323,6 +331,24 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl
     assert (stage_k is not None) is gather_enabled
 
 
+@pytest.mark.parametrize("window_chunks", [2, 4])
+def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
+    """Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True)
+    stage_key, stage_value = kv.history_staging[0]
+    for _ in range(2):
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+        assert inputs.stage_key is stage_key
+        assert inputs.stage_value is stage_value
+        assert inputs.max_seq_len == 3 * BLOCK
+        assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
+
+
 @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
@@ -363,8 +389,8 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
         query[0],
         current_k[0],
         current_v[0],
-        action_k[0] if action_len else None,
-        action_v[0] if action_len else None,
+        action_k[0] if action_k is not None else None,
+        action_v[0] if action_v is not None else None,
         HEAD_DIM**-0.5,
     ).unsqueeze(0)
 
@@ -556,8 +582,8 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c
             query[0],
             current_k[0],
             current_v[0],
-            action_k[0] if action_len else None,
-            action_v[0] if action_len else None,
+            action_k[0] if action_k is not None else None,
+            action_v[0] if action_v is not None else None,
             HEAD_DIM**-0.5,
         ).unsqueeze(0)
 

tests/diffusion/cache/test_exact_projection_cache.py

@@ -0,0 +1,270 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact projection reuse without model-specific schedules or artifacts."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def tp1(monkeypatch):
+    import vllm.distributed
+
+    monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 1)
+
+
+@pytest.fixture(autouse=True)
+def deterministic_weights():
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(0)
+        yield
+
+
+@pytest.mark.parametrize("batch,in_features,out_features", [(2, 4, 8), (1, 8, 48), (4, 6, 18)])
+def test_runtime_reuses_exact_outputs_and_tracks_changed_weights(mocker, batch, in_features, out_features):
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(in_features, out_features)
+    x = torch.randn(batch, in_features)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        cache.prepare(x)
+        first = cache.project("block", linear, x, compute)
+        first_copy = first.clone()
+        first.zero_()  # The caller/provider cannot corrupt the owned cache.
+        assert torch.equal(cache.project("block", linear, x, compute), first_copy)
+        assert compute.call_count == 1
+        linear.weight.add_(0.25)
+        assert torch.equal(cache.project("block", linear, x, compute), linear(x))
+        assert compute.call_count == 2
+        other = x.clone()
+        other[0, 0] += 1
+        cache.prepare(other)
+        cache.project("block", linear, other, lambda: linear(other))
+        cache.prepare(x)
+        cache.project("block", linear, x, compute)
+        assert compute.call_count == 2
+    assert cache.hits == 2
+
+
+def test_runtime_budget_retains_reusable_entries_without_schedule_thrashing(mocker):
+    cache = ExactProjectionCache(max_bytes=32)
+    linear = torch.nn.Linear(2, 4)
+    x = torch.ones(2, 2)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        cache.prepare(x)
+        for name in ("first", "second", "first"):
+            cache.project(name, linear, x, compute)
+            assert cache._bytes <= 32
+        assert compute.call_count == 2
+        assert cache.hits == 1
+        cache.clear()
+        assert cache._bytes == 0 and not cache._entries
+
+
+def test_dynamic_lora_buffers_and_suspend_mask_invalidate_entries(mocker):
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    linear.lora_a_stacked = (torch.ones(1, 2),)
+    linear.lora_b_stacked = (torch.ones(2, 1),)
+    linear._diffusion_lora_active_slices = (True,)
+    x = torch.ones(1, 2)
+
+    def compute():
+        result = linear(x)
+        if linear._diffusion_lora_active_slices[0]:
+            result = result + x @ linear.lora_a_stacked[0].T @ linear.lora_b_stacked[0].T
+        return result
+
+    counted = mocker.Mock(side_effect=compute)
+    with torch.no_grad():
+        cache.prepare(x)
+        cache.project("a", linear, x, counted)
+        cache.project("a", linear, x, counted)
+        assert counted.call_count == 1
+        linear.lora_b_stacked[0].add_(1)
+        assert torch.equal(cache.project("a", linear, x, counted), compute())
+        linear._diffusion_lora_active_slices = (False,)
+        assert torch.equal(cache.project("a", linear, x, counted), compute())
+        assert counted.call_count == 3
+
+
+@pytest.mark.parametrize("scope", ["local", "global", "child"])
+@pytest.mark.parametrize("kind", ["pre", "post"])
+def test_linear_hooks_run_even_after_cache_warmup(scope, kind):
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    projection = torch.nn.Sequential(linear) if scope == "child" else linear
+    x = torch.ones(1, 2)
+    with torch.no_grad():
+        cache.prepare(x)
+        expected = projection(x)
+        cache.project("a", projection, x, lambda: projection(x))
+        calls = []
+
+        def pre_hook(module, args):
+            if module is linear:
+                calls.append(module)
+                return (args[0] + 1,)
+
+        def post_hook(module, args, output):
+            if module is linear:
+                calls.append(module)
+                return output + 1
+
+        if kind == "pre":
+            expected = linear(x + 1)
+            register = (
+                torch.nn.modules.module.register_module_forward_pre_hook
+                if scope == "global"
+                else linear.register_forward_pre_hook
+            )
+            hook = register(pre_hook)
+        else:
+            expected = expected + 1
+            register = (
+                torch.nn.modules.module.register_module_forward_hook
+                if scope == "global"
+                else linear.register_forward_hook
+            )
+            hook = register(post_hook)
+        try:
+            for _ in range(2):
+                assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), expected)
+            assert len(calls) == 2
+        finally:
+            hook.remove()
+        assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), projection(x))
+
+
+@pytest.mark.parametrize("autocast_at_prepare", [False, True])
+def test_numerical_settings_follow_projection_context(mocker, autocast_at_prepare):
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(4, 4)
+    x = torch.randn(2, 4)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast_at_prepare):
+            cache.prepare(x)
+        for enabled in (False, True, False, True):
+            with torch.autocast("cpu", dtype=torch.bfloat16, enabled=enabled):
+                actual = cache.project("a", linear, x, compute)
+                expected = linear(x)
+                assert actual.dtype == expected.dtype
+                assert torch.equal(actual, expected)
+    assert compute.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "setting",
+    ["allow_tf32", "allow_bf16_reduced_precision_reduction", "allow_fp16_reduced_precision_reduction"],
+)
+def test_matmul_settings_invalidate_prepared_projection(monkeypatch, mocker, setting):
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    x = torch.ones(1, 2)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        cache.prepare(x)
+        cache.project("a", linear, x, compute)
+        monkeypatch.setattr(torch.backends.cuda.matmul, setting, not getattr(torch.backends.cuda.matmul, setting))
+        assert torch.equal(cache.project("a", linear, x, compute), linear(x))
+    assert compute.call_count == 2
+
+
+def test_disabled_and_grad_paths_keep_original_computation(monkeypatch, mocker):
+    import vllm_omni.diffusion.cache.exact_projection_cache as module
+
+    monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
+    linear = torch.nn.Linear(2, 2)
+    x = torch.ones(1, 2, requires_grad=True)
+    for cache in (ExactProjectionCache(max_bytes=0), ExactProjectionCache()):
+        cache.prepare(x)
+        cache.project("a", linear, x, lambda: linear(x)).sum().backward()
+        assert x.grad is not None
+    with torch.no_grad():
+        cache = ExactProjectionCache(max_bytes=0)
+        cache.prepare(x)
+        assert torch.equal(cache.project("a", linear, x, lambda: linear(x)), linear(x))
+
+
+def test_tp_rank_miss_forces_every_rank_to_recompute(monkeypatch, mocker):
+    import vllm.distributed
+
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    x = torch.ones(1, 2)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        cache.prepare(x)
+        cache.project("a", linear, x, compute)
+        votes = mocker.Mock(side_effect=[torch.tensor([1]), torch.tensor([2])])
+        monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
+        cache.project("a", linear, x, compute)
+        assert compute.call_count == 2
+        cache.project("a", linear, x, compute)
+        assert compute.call_count == 2
+        assert votes.call_count == 2
+
+
+def test_compiled_execution_bypasses_runtime_and_precomputed_results(monkeypatch, mocker):
+    import vllm_omni.diffusion.cache.exact_projection_cache as module
+
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    x = torch.ones(1, 2)
+    compute = mocker.Mock(side_effect=lambda: linear(x))
+    with torch.no_grad():
+        cache.prepare(x)
+        cache.project("modulation", linear, x, compute)
+        monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+        monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
+        seeded = mocker.patch.object(cache, "_lookup_precomputed", side_effect=AssertionError("must not seed"))
+        # An input prepared before compilation must also bypass the warm entry.
+        assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
+        cache.prepare(x)
+        assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
+        assert compute.call_count == 3
+        seeded.assert_not_called()
+
+
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_precomputed_extension_is_used_only_at_tp1(monkeypatch, mocker, world_size):
+    import vllm.distributed
+
+    cache = ExactProjectionCache()
+    linear = torch.nn.Linear(2, 2)
+    x = torch.ones(1, 2)
+    monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: world_size)
+    votes = mocker.Mock(side_effect=[torch.tensor([0]), torch.tensor([world_size])])
+    monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
+    with torch.no_grad():
+        expected = linear(x)
+        seeded = mocker.patch.object(cache, "_lookup_precomputed", return_value=expected)
+        compute = mocker.Mock(side_effect=lambda: linear(x))
+        cache.prepare(x)
+        for _ in range(2):
+            assert torch.equal(cache.project("conditioning.scale_shift", linear, x, compute), expected)
+        assert seeded.call_count == (1 if world_size == 1 else 0)
+        assert compute.call_count == (0 if world_size == 1 else 1)
+        assert cache.hits == 1
+
+
+def test_inference_weights_without_versions_use_safe_runtime_path(mocker):
+    with torch.inference_mode():
+        linear = torch.nn.Linear(2, 2)
+        x = torch.ones(1, 2)
+        cache = ExactProjectionCache()
+        compute = mocker.Mock(side_effect=lambda: linear(x))
+        cache.prepare(x)
+        cache.project("a", linear, x, compute)
+        linear.weight.add_(1)
+        assert torch.equal(cache.project("a", linear, x, compute), linear(x))
+        assert compute.call_count == 2 and not cache._entries

tests/diffusion/cache/test_teacache_extractors.py

@@ -33,6 +33,7 @@ from vllm_omni.diffusion.cache.teacache.extractors import (
     extract_flux2_klein_context,
     extract_flux_context,
     extract_minimax_h3_context,
+    extract_zimage_context,
 )
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
@@ -53,6 +54,62 @@ def setup_tp_group():
             yield
 
 
+@pytest.mark.cpu
+def test_zimage_extractor_accepts_extended_patchify_output():
+    class _Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attention_norm1 = nn.Identity()
+
+        def adaLN_modulation(self, value):
+            return value.new_zeros((value.shape[0], 16))
+
+    class _Model(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = nn.ModuleList([_Block()])
+            self.noise_refiner = nn.ModuleList()
+            self.context_refiner = nn.ModuleList()
+            self.all_x_embedder = nn.ModuleDict({"2-1": nn.Linear(2, 4)})
+            self.cap_embedder = nn.Linear(3, 4)
+            self.x_pad_token = nn.Parameter(torch.zeros(1, 4))
+            self.cap_pad_token = nn.Parameter(torch.zeros(1, 4))
+            self.t_scale = 1.0
+
+        @staticmethod
+        def t_embedder(timestep):
+            return timestep.new_ones((timestep.shape[0], 4))
+
+        @staticmethod
+        def rope_embedder(position_ids):
+            rope = torch.zeros((position_ids.shape[0], 2))
+            return rope, rope
+
+        @staticmethod
+        def patchify_and_embed(*_args):
+            positions = [torch.zeros((32, 3), dtype=torch.int32)]
+            masks = [torch.zeros(32, dtype=torch.bool)]
+            return (
+                [torch.zeros((32, 2))],
+                [torch.zeros((32, 3))],
+                [(1, 1, 1)],
+                positions,
+                positions,
+                masks,
+                masks,
+                [],
+            )
+
+    context = extract_zimage_context(
+        _Model(),
+        x=[torch.zeros((1, 1, 1, 1))],
+        t=torch.zeros(1),
+        cap_feats=[torch.zeros((1, 3))],
+    )
+
+    assert context.modulated_input.shape == (1, 64, 4)
+
+
 class BaseExtractorTest(ABC):
     """Base class for testing TeaCache extractors.
 
@@ -595,6 +652,9 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
 
         model = h3.MiniMaxH3DiTModel(_minimax_h3_small_od_config(), quant_config=None)
+        # Normally loaded from the checkpoint; torch.empty() can contain NaNs.
+        # With one frequency per axis, the RoPE inverse frequency is 1.
+        model.rope.inv_freq.fill_(1.0)
         for submodule in model.modules():
             if isinstance(submodule, h3.MiniMaxH3Attention):
                 submodule.rope._forward_method = submodule.rope.forward_native

tests/diffusion/diffusion_kv/test_config.py

@@ -21,6 +21,86 @@ def test_dense_legacy_is_default() -> None:
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
 
 
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", [None, "dense_legacy", DiffusionKVCacheMode.DENSE_LEGACY])
+def test_prefix_caching_rejects_non_paged_mode(config_cls, mode) -> None:
+    kwargs = {} if mode is None else {"diffusion_kv_mode": mode}
+    with pytest.raises(
+        ValueError,
+        match="set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching",
+    ):
+        config_cls.from_kwargs(enable_prefix_caching=True, **kwargs)
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_prefix_caching_accepts_paged_scheduler(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+
+    assert config.enable_prefix_caching is True
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("build_config", [OmniDiffusionConfig, OmniDiffusionConfig.from_kwargs])
+def test_prefix_caching_rejects_sleep_mode(build_config) -> None:
+    with pytest.raises(ValueError, match="disable enable_prefix_caching or enable_sleep_mode"):
+        build_config(
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            enable_prefix_caching=True,
+            enable_sleep_mode=True,
+        )
+
+
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_sleep_mode_accepts_disabled_prefix_caching(mode) -> None:
+    config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+        enable_sleep_mode=True,
+    )
+
+    assert config.enable_sleep_mode is True
+    assert config.enable_prefix_caching is False
+
+
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+def test_native_kv_transfer_requires_prefix_caching_disabled(enable_prefix_caching) -> None:
+    kwargs = dict(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=enable_prefix_caching,
+        kv_transfer_config={
+            "kv_connector": "MooncakeConnector",
+            "kv_role": "kv_consumer",
+            "engine_id": "test-diffusion",
+        },
+    )
+    if enable_prefix_caching:
+        with pytest.raises(ValueError, match="disable enable_prefix_caching for AR KV import"):
+            OmniDiffusionConfig.from_kwargs(**kwargs)
+    else:
+        config = OmniDiffusionConfig.from_kwargs(**kwargs)
+        assert config.kv_transfer_config.kv_role == "kv_consumer"
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_disabled_prefix_caching_accepts_both_modes(config_cls, mode) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+    )
+
+    assert config.enable_prefix_caching is False
+    assert config.diffusion_kv_mode.value == mode
+
+
 def test_paged_worker_local_is_rejected_until_implemented() -> None:
     with pytest.raises(ValueError, match="reserved but not implemented"):
         OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")

tests/diffusion/diffusion_kv/test_initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -225,6 +225,22 @@ def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -
     assert vllm_config.cache_config.gpu_memory_utilization == 0.42
 
 
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+@pytest.mark.parametrize("requests,rows", [(1, 2), (2, 3)])
+def test_native_config_counts_cfg_rows_once(enable_prefix_caching, requests, rows) -> None:
+    od_config = _od_config(
+        max_num_seqs=requests,
+        diffusion_kv_max_rows_per_request=rows,
+        enable_prefix_caching=enable_prefix_caching,
+    )
+    config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    # Engine and Worker must see the same row capacity, including on reconfiguration.
+    diffusion_vllm_config.configure_diffusion_vllm_config(config, od_config)
+    assert config.scheduler_config.max_num_seqs == requests * rows
+    assert od_config.max_num_seqs == requests
+    assert config.cache_config.enable_prefix_caching is enable_prefix_caching
+
+
 def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
     quantization_config = build_quant_config(
         {

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native multimodal block hashing and diffusion-specific identity helpers."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
+
+import vllm_omni.diffusion.diffusion_engine as engine_module
+import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
+import vllm_omni.diffusion.diffusion_kv.request as request_module
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _request(*, num_rows=1, seed=7, adapter_id=None, scale=1.0):
+    request = OmniDiffusionRequest(
+        request_id="fake", prompt="fake model", sampling_params=OmniDiffusionSamplingParams(seed=seed)
+    )
+    request.sampling_params.lora_request = None if adapter_id is None else SimpleNamespace(lora_int_id=adapter_id)
+    request.sampling_params.lora_scale = scale
+    request.diffusion_kv_requests = tuple(
+        DiffusionKVRequest(f"fake/{row}", sequence_id=row, prefix_len=8, target_len=4, seq_len=12)
+        for row in range(num_rows)
+    )
+    return request
+
+
+def _feature(*, value="image", start=4, end=8):
+    return MultiModalFeatureSpec(
+        data=None,
+        modality="image",
+        identifier=hash_prefix_cache_value((value, start, end)).hex(),
+        mm_position=PlaceholderRange(offset=start, length=end - start),
+    )
+
+
+def _prepare(request, feature=None):
+    namespace = get_cache_namespace("fake-model-v1", request.sampling_params)
+    for row in request.diffusion_kv_requests:
+        row.cache_token_ids = tuple(range(row.prefix_len))
+        row.mm_features = [feature or _feature()]
+        row.cache_namespace = namespace
+
+
+def _hashes(request, feature=None):
+    _prepare(request, feature)
+    row = request.diffusion_kv_requests[0]
+    row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+    return row.block_hashes
+
+
+def test_partial_hits_and_randomness_are_model_opt_in():
+    base = _hashes(_request())
+    assert base == _hashes(_request(seed=999))
+    changed = _hashes(_request(), _feature(value="other image"))
+    assert base[0] == changed[0]
+    assert base[1] != changed[1]
+    changed_geometry = _hashes(_request(), _feature(end=9))
+    assert base[0] == changed_geometry[0]
+    assert base[1] != changed_geometry[1]
+
+
+@pytest.mark.parametrize(("adapter_id", "scale"), [(None, 1.0), (2, 0.5), (1, 1.0)])
+def test_lora_is_common_request_wide_identity(adapter_id, scale):
+    base = _hashes(_request(adapter_id=1, scale=0.5))
+    changed = _hashes(_request(adapter_id=adapter_id, scale=scale))
+    assert all(a != b for a, b in zip(base, changed, strict=True))
+    assert _hashes(_request()) == _hashes(_request(scale=0.5))
+
+
+def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
+    row = DiffusionKVRequest(
+        "native",
+        sequence_id=0,
+        prefix_len=14,
+        target_len=2,
+        seq_len=16,
+        cache_token_ids=range(16),
+        mm_features=[
+            _feature(start=2, end=7),
+            _feature(start=7, end=10),
+        ],
+    )
+    upstream = Mock(wraps=generate_block_hash_extra_keys)
+    monkeypatch.setattr(request_module, "generate_block_hash_extra_keys", upstream)
+    hash_fn = get_hash_fn_by_name("sha256")
+    row.build_block_hashes(4, hash_fn)
+    assert [(call.args[1], call.args[2]) for call in upstream.call_args_list] == [(0, 4), (4, 8), (8, 12)]
+    # Independent expected native extras: a crossing item has a negative
+    # relative offset; a block can include the end of one item and the next.
+    first, second = (feature.identifier for feature in row.mm_features)
+    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
+    parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
+    expected = []
+    for start, keys in zip(range(0, 12, 4), extras, strict=True):
+        parent = hash_block_tokens(hash_fn, parent, list(range(start, start + 4)), keys)
+        expected.append(parent)
+    assert row.block_hashes == expected
+    # Never hash the partial prefix tail or mutable target [14, 16), even
+    # though cache_token_ids deliberately contains the whole sequence.
+    assert len(row.block_hashes) == 3
+
+
+def test_same_identifier_at_different_offsets_changes_block_hash():
+    first = _feature(start=4, end=6)
+    shifted = _feature(start=5, end=7)
+    shifted.identifier = first.identifier
+    a = _hashes(_request(), first)
+    b = _hashes(_request(), shifted)
+    assert a[0] == b[0]
+    assert a[1] != b[1]
+
+
+def test_multimodal_identity_beyond_prefix_does_not_affect_hashes():
+    assert _hashes(_request(), _feature(value="a", start=8, end=12)) == _hashes(
+        _request(), _feature(value="b", start=8, end=12)
+    )
+
+
+@pytest.mark.parametrize(
+    "positions",
+    [
+        [(-1, 2)],
+        [(4, 4)],
+        [(4, 13)],
+        [(5, 7), (4, 5)],
+        [(3, 7), (4, 6)],
+    ],
+)
+def test_invalid_or_overlapping_positions_fail_before_hashing(positions):
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    row.mm_features = [_feature(start=start, end=end) for start, end in positions]
+    fail = Mock(side_effect=AssertionError("invalid positions must fail before hashing"))
+    with pytest.raises(ValueError, match="sorted, non-overlapping"):
+        row.build_block_hashes(4, fail)
+    fail.assert_not_called()
+    assert row.block_hashes == []
+
+
+def test_empty_multimodal_identifier_is_rejected():
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    feature = _feature()
+    feature.identifier = ""
+    row.mm_features = [feature]
+    with pytest.raises(ValueError, match="content identifiers"):
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.int64, torch.bool])
+def test_tensor_identity_includes_dtype_shape_and_values(dtype):
+    tensor = torch.tensor([[0, 1], [1, 0]], dtype=dtype)
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.clone())
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.t().contiguous().t())
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.reshape(4))
+    changed = tensor.clone()
+    changed[0, 0] = 1
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(changed)
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.to(torch.float64))
+
+
+def test_numpy_scalar_shape_and_nested_dict_order():
+    assert hash_prefix_cache_value({"b": {"y": 2, "x": 1}, "a": None}) == hash_prefix_cache_value(
+        {"a": None, "b": {"x": 1, "y": 2}}
+    )
+    assert hash_prefix_cache_value(np.array(1)) != hash_prefix_cache_value(np.array([1]))
+    assert hash_prefix_cache_value(np.int64(1)) == hash_prefix_cache_value(1)
+
+
+@pytest.mark.parametrize(
+    ("first", "second"), [(1, True), (1, 1.0), ("1", b"1"), ([], {}), ([None], []), (["ab", "c"], ["a", "bc"])]
+)
+def test_identity_has_unambiguous_type_and_container_boundaries(first, second):
+    assert hash_prefix_cache_value(first) != hash_prefix_cache_value(second)
+
+
+@pytest.mark.parametrize("value", [object(), {1: "value"}, np.array([object()]), np.longdouble(1)])
+def test_identity_rejects_unsupported_values_without_pickle(value):
+    with pytest.raises(TypeError):
+        hash_prefix_cache_value(value)
+
+
+@pytest.mark.parametrize(
+    ("mode", "enabled", "expect_hook"),
+    [
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, False, False),
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, True, True),
+        (DiffusionKVCacheMode.DENSE_LEGACY, True, False),
+    ],
+)
+def test_engine_gates_hook_loading_and_admission(monkeypatch, mode, enabled, expect_hook):
+    events = []
+    request = _request()
+
+    def preprocess(req):
+        events.append("layout")
+        req.prepared_layout = "fake layout"
+        return req
+
+    def prepare(req):
+        assert req.prepared_layout == "fake layout"
+        assert req.diffusion_kv_requests[0].cache_token_ids == ()
+        events.append("cache_inputs")
+        _prepare(req)
+
+    hook = Mock(side_effect=prepare if expect_hook else AssertionError("disabled hook"))
+    load_hook = Mock(return_value=hook)
+    if not expect_hook:
+        load_hook.side_effect = AssertionError("disabled hook factory")
+        monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", Mock(side_effect=AssertionError("disabled hash")))
+    monkeypatch.setattr(engine_module, "get_diffusion_pre_process_func", lambda _config: preprocess)
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", load_hook)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode, enable_prefix_caching=enabled)
+    engine._init_process_hooks(engine.od_config)
+    engine.prefix_cache_func = hook
+    assert engine._prepare_request_for_admission(request) is request
+    row = request.diffusion_kv_requests[0]
+    if expect_hook:
+        assert events == ["layout", "cache_inputs"]
+        assert row.cache_token_ids == tuple(range(8))
+        load_hook.assert_called_once_with(engine.od_config)
+        hook.assert_called_once_with(request)
+    else:
+        assert events == ["layout"]
+        assert row.cache_token_ids == ()
+        assert row.mm_features == []
+        load_hook.assert_not_called()
+        hook.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "model_class_name,load_format",
+    [(None, "default"), ("UnregisteredPipeline", "default"), ("HunyuanImage3ForCausalMM", "diffusers")],
+)
+def test_engine_rejects_enabled_prefix_caching_without_hook(monkeypatch, model_class_name, load_format):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.from_kwargs(
+        model_class_name=model_class_name,
+        diffusion_load_format=load_format,
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+    with pytest.raises(ValueError, match="requires a registered prefix-cache hook"):
+        engine._init_process_hooks(engine.od_config)
+
+
+def test_engine_rejects_out_of_profile_request_before_hashing():
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=True
+    )
+    engine._diffusion_kv_profile_limits = (1, 8, 4)
+    engine.prefix_cache_func = Mock(side_effect=AssertionError("invalid request must fail before identity work"))
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(_request())
+    engine.prefix_cache_func.assert_not_called()

tests/diffusion/diffusion_kv/test_layout.py

@@ -17,6 +17,7 @@ from vllm_omni.diffusion.diffusion_kv.layout import (
     adopt_kv_cache_layout,
     assert_backend_layout_supported,
     build_kv_cache_tensor,
+    get_connector_required_kv_cache_layout,
     resolve_diffusion_kv_cache_layout,
 )
 
@@ -135,6 +136,51 @@ def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
         assert_backend_layout_supported(worker, _DenseBackend)
 
 
+def test_explicit_connector_layout_is_distinct_from_block_stride_requirement() -> None:
+    config = _config()
+
+    layout = resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
+
+    assert layout is KVCacheLayout.LBHNC
+    assert config.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
+
+
+def test_connector_declared_layout_is_resolved(monkeypatch: pytest.MonkeyPatch) -> None:
+    class _Connector:
+        @classmethod
+        def get_required_kvcache_layout(cls, config):
+            del config
+            return "LBHNC"
+
+    from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+
+    monkeypatch.setattr(KVConnectorFactory, "get_connector_class", lambda config: _Connector)
+    config = _config()
+    config.kv_transfer_config = object()
+
+    assert get_connector_required_kv_cache_layout(config) is KVCacheLayout.LBHNC
+
+
+def test_connector_layout_rejects_same_block_order_but_different_head_order() -> None:
+    config = _config()
+    config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
+    with pytest.raises(ValueError, match="contradicts"):
+        resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
+
+
+def test_connector_layout_cannot_override_backend_requirement() -> None:
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        resolve_diffusion_kv_cache_layout(
+            _config(), indexes_kv_by_block_stride=True, required_layout=KVCacheLayout.LBHNC
+        )
+
+
+def test_connector_layout_without_cache_config() -> None:
+    config = _config()
+    config.cache_config = None
+    assert resolve_diffusion_kv_cache_layout(config, required_layout="LBHNC") is KVCacheLayout.LBHNC
+
+
 def test_tensor_strides_come_from_the_upstream_derivation() -> None:
     spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
 

tests/diffusion/diffusion_kv/test_manager.py

@@ -30,26 +30,159 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
 
 
-def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+def _request(
+    public_id: str,
+    sequence_id: int,
+    *,
+    prefix_len: int = 4,
+    target_len: int = 4,
+    seq_len: int = 8,
+    cache_token_ids=(),
+    mm_features=(),
+    kv_contexts=(),
+) -> DiffusionKVRequest:
     return DiffusionKVRequest(
         f"{public_id}/diffusion-kv/{sequence_id}",
         sequence_id=sequence_id,
-        prefix_len=4,
-        target_len=4,
+        prefix_len=prefix_len,
+        target_len=target_len,
         seq_len=seq_len,
+        cache_token_ids=cache_token_ids,
+        mm_features=mm_features,
         kv_contexts=kv_contexts,
     )
 
 
-def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+def _manager(
+    num_blocks: int,
+    *,
+    max_model_len: int = 64,
+    enable_prefix_caching: bool = False,
+) -> DiffusionKVCacheManager:
     return DiffusionKVCacheManager(
         _config(num_blocks),
         max_model_len=max_model_len,
         scheduler_block_size=BLOCK_SIZE,
         hash_block_size=BLOCK_SIZE,
+        enable_prefix_caching=enable_prefix_caching,
     )
 
 
+def test_successful_request_publishes_prefix_for_a_warm_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+
+    cold_metadata = manager.reserve_request("cold", (cold,))
+    assert cold_metadata is not None
+    assert cold_metadata.sequences[0].cached_prefix_len == 0
+    cold_prefix_block = cold_metadata.sequences[0].block_ids[0][0]
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    warm_metadata = manager.reserve_request("warm", (warm,))
+
+    assert warm_metadata is not None
+    assert warm_metadata.sequences[0].cached_prefix_len == 4
+    assert warm_metadata.sequences[0].block_ids[0][0] == cold_prefix_block
+    # Metadata contains both the shared hit and the newly allocated target.
+    assert len(warm_metadata.sequences[0].block_ids[0]) == 2
+
+    manager.free_request("warm")
+    assert warm.num_computed_tokens == 0
+    assert warm.shared_prefix_boundary == 0
+
+
+def test_prefix_hash_root_is_stable_across_manager_instances() -> None:
+    first_manager = _manager(8, enable_prefix_caching=True)
+    second_manager = _manager(8, enable_prefix_caching=True)
+    first = _request("first", 0, cache_token_ids=range(4))
+    second = _request("second", 0, cache_token_ids=range(4))
+
+    first_manager._prepare_block_hashes((first,))
+    second_manager._prepare_block_hashes((second,))
+
+    assert first.block_hashes == second.block_hashes
+
+
+def test_warm_hit_fits_at_the_same_capacity_as_the_cold_request() -> None:
+    manager = _manager(3, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+
+
+def test_unpublished_request_does_not_become_a_prefix_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    failed = _request("failed", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("failed", (failed,)) is not None
+    manager.free_request("failed")
+
+    retry = _request("retry", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("retry", (retry,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 0
+
+
+def test_partial_prefix_hit_recomputes_from_the_first_changed_block() -> None:
+    manager = _manager(10, enable_prefix_caching=True)
+    cold = _request(
+        "cold",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=range(8),
+    )
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request(
+        "warm",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=(*range(4), 100, 101, 102, 103),
+    )
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+    assert len(metadata.sequences[0].block_ids[0]) == 3
+
+
+def test_cfg_rows_use_the_minimum_warm_hit_as_one_execution_boundary() -> None:
+    manager = _manager(16, enable_prefix_caching=True)
+    for public_id, tokens in (("seed-a", range(8)), ("seed-b", (*range(4), 20, 21, 22, 23))):
+        request = _request(public_id, 0, prefix_len=8, seq_len=12, cache_token_ids=tokens)
+        assert manager.reserve_request(public_id, (request,)) is not None
+        manager.publish_request(public_id)
+        manager.free_request(public_id)
+
+    requests = (
+        _request("warm", 0, prefix_len=8, seq_len=12, cache_token_ids=range(8)),
+        _request(
+            "warm",
+            1,
+            prefix_len=8,
+            seq_len=12,
+            cache_token_ids=(*range(4), 30, 31, 32, 33),
+        ),
+    )
+    metadata = manager.reserve_request("warm", requests)
+
+    assert metadata is not None
+    assert [sequence.cached_prefix_len for sequence in metadata.sequences] == [4, 4]
+
+
 def test_reserve_and_free_multi_cfg_request() -> None:
     manager = _manager(8)
     free_before = manager.native_manager.block_pool.get_num_free_blocks()

tests/diffusion/diffusion_kv/test_native_prefetch.py

@@ -256,7 +256,11 @@ def _scheduler(monkeypatch, num_blocks=24):
         kv_cache_config=config,
         scheduler_block_size=4,
         hash_block_size=4,
-        kv_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64), max_in_flight_tokens=64),
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(enable_prefix_caching=False),
+        ),
     )
     return scheduler, connector
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -113,7 +113,7 @@ class _FakeLayer:
         self.head_size_v = 4
         self.spec = _FakeSpec(non_causal=non_causal)
         self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
-        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = []
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, _FakeNativeMetadata]] = []
         self.native_events: list[str] = []
         self.layer_name = "layer-0"
         self.kv_cache = object()
@@ -122,6 +122,8 @@ class _FakeLayer:
 
 
 class _FakeNativeImpl:
+    vllm_flash_attn_version: int
+
     def __init__(self, layer: _FakeLayer) -> None:
         self.layer = layer
 
@@ -143,7 +145,7 @@ class _FakeNativeImpl:
         key: torch.Tensor,
         value: torch.Tensor,
         _kv_cache,
-        metadata: object,
+        metadata: _FakeNativeMetadata,
         output: torch.Tensor,
     ) -> torch.Tensor:
         self.layer.native_events.append("forward")
@@ -202,7 +204,7 @@ def _make_adapter(
 
     config = SimpleNamespace(
         name="vllm-config",
-        model_config=SimpleNamespace(dtype=torch.float32),
+        model_config=SimpleNamespace(dtype=torch.float32, rswa_window=None),
     )
     adapter = DiffusionPagedAttentionAdapter(
         vllm_config=config,
@@ -1222,9 +1224,15 @@ def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: st
         DiffusionPagedAttentionRow(**values)
 
 
+@pytest.mark.parametrize("pcp_size", [1, 2])
+@pytest.mark.parametrize("selection_fails", [False, True])
 def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
     monkeypatch: pytest.MonkeyPatch,
+    pcp_size: int,
+    selection_fails: bool,
 ) -> None:
+    from vllm.v1.attention import selector
+
     selected_backends = []
     specialized_backends = []
     impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
@@ -1237,28 +1245,39 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
 
     original_backend_per_kind = {"full": object()}
     config = SimpleNamespace(
-        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
-        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
-        model_config=SimpleNamespace(dtype=torch.float16),
-        cache_config=SimpleNamespace(cache_dtype="auto"),
+        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind, use_non_causal=True),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size, decode_context_parallel_size=1),
+        model_config=SimpleNamespace(dtype=torch.float16, rswa_window=None),
+        cache_config=SimpleNamespace(cache_dtype="auto", user_specified_block_size=False),
+        kv_transfer_config=None,
+        speculative_config=None,
     )
 
-    def select_backend(**_kwargs):
-        selected_backends.append(
-            (
-                config.attention_config.backend,
-                config.attention_config.backend_per_kind,
-                config.parallel_config.prefill_context_parallel_size,
-            )
-        )
+    def select_backend(*, backend, attn_selector_config, num_heads):
+        # Exercise the real 0.29 selector, replacing only platform resolution.
+        # Ulysses has already gathered tokens and sharded heads before the
+        # native kernel; the MoE PCP mapping must not request PCP attention.
+        assert not attn_selector_config.use_pcp
+        assert not attn_selector_config.use_dcp
+        assert attn_selector_config.use_non_causal
+        assert num_heads == 4
+        selected_backends.append((backend, config.attention_config.backend_per_kind))
+        if selection_fails:
+            raise ValueError("test backend unavailable")
         return native_backend
 
-    monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
+    monkeypatch.setattr("vllm.config.get_current_vllm_config", lambda: config)
+    monkeypatch.setattr(selector, "_cached_get_attn_backend", select_backend)
     monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+    def specialize_backend(backend, *, ulysses_degree):
+        specialized_backends.append((backend, ulysses_degree))
+        return backend
+
     monkeypatch.setattr(
         adapter_module.current_omni_platform,
         "get_diffusion_paged_kv_attn_backend",
-        lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+        specialize_backend,
     )
     layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
     spec = FullAttentionSpec(
@@ -1269,7 +1288,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         non_causal=True,
     )
 
-    native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(
+    kwargs = dict(
         layer_name="layer-0",
         layer=layer,
         spec=spec,
@@ -1277,12 +1296,21 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         device=torch.device("cpu"),
         ulysses_degree=2,
     )
+    if selection_fails:
+        with pytest.raises(ValueError, match="test backend unavailable"):
+            adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
+    else:
+        native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
-    assert specialized_backends == [(native_backend, 2)]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind
-    assert config.parallel_config.prefill_context_parallel_size == 2
+    assert config.parallel_config.prefill_context_parallel_size == pcp_size
+    if selection_fails:
+        assert not specialized_backends
+        impl_cls.assert_not_called()
+        return
+    assert specialized_backends == [(native_backend, 2)]
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
@@ -1432,7 +1460,12 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
     layer._no_parallel_strategy = object()
     layer._get_active_parallel_strategy = lambda: Strategy()
     layer._with_kv_cache_dtype = lambda metadata: metadata
-    layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query
+
+    def run_local_attention(query, _key, _value, _metadata):
+        events.append("dense")
+        return query
+
+    layer._run_local_attention = run_local_attention
     qkv = torch.zeros(1, 2, 2, 4)
 
     assert not layer.is_paged_kv_active()

tests/diffusion/diffusion_kv/test_request.py

@@ -4,8 +4,10 @@
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
 from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.core.kv_cache_utils import BlockHash, init_none_hash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -49,6 +51,51 @@ def test_request_exposes_native_request_and_diffusion_semantics() -> None:
     assert request.status is RequestStatus.WAITING
 
 
+def test_request_builds_only_full_prefix_block_hashes() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    request = _request(
+        prefix_len=6,
+        target_len=2,
+        cache_token_ids=range(6),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="tail",
+                mm_position=PlaceholderRange(offset=4, length=2),
+            )
+        ],
+    )
+
+    request.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert len(request.block_hashes) == 1
+    assert request.skip_reading_prefix_cache is False
+
+
+def test_changed_multimodal_identity_changes_the_block_hash() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    first = _request(cache_token_ids=range(4))
+    second = _request(
+        cache_token_ids=range(4),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="image-a",
+                mm_position=PlaceholderRange(offset=2, length=2),
+            )
+        ],
+    )
+
+    first.build_block_hashes(BLOCK_SIZE, hash_function)
+    second.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert first.block_hashes != second.block_hashes
+
+
 @pytest.mark.parametrize(
     ("overrides", "message"),
     [

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -69,7 +69,7 @@ def make_executor() -> tuple[MultiprocDiffusionExecutor, list[tuple]]:
     executor._ensure_open = lambda: None
     calls: list[tuple] = []
 
-    def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks):
+    def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks, timeout=None):
         calls.append((method, args, unique_reply_rank, exec_all_ranks))
         return DiffusionOutput(output=None)
 
@@ -111,6 +111,29 @@ def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> Non
     ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
 
 
+def test_runner_builds_suffix_prefill_row_for_a_prefix_hit() -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(
+                sequence_id=0,
+                prefix_len=8,
+                target_len=4,
+                seq_len=12,
+                block_ids=([1, 2, 3],),
+                cached_prefix_len=4,
+            ),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    row = attn_metadata.prefill_rows[0]
+    assert (row.kv_start_pos, row.query_len, row.seq_len) == (4, 8, 12)
+
+
 def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.od_config.parallel_config.cfg_parallel_size = 2

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -146,3 +146,34 @@ def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
             world_size=4,
             sp_group_ranks=[[0, 2]],
         )
+
+
+@pytest.mark.parametrize("rank", [1, 3])
+def test_tensor_dict_broadcast_preserves_subgroup_source(rank, monkeypatch):
+    from vllm_omni.diffusion.distributed.group_coordinator import GroupCoordinator, TensorMetadata
+
+    coordinator = object.__new__(GroupCoordinator)
+    coordinator.rank = rank
+    coordinator.ranks = [1, 3]
+    coordinator.rank_in_group = [1, 3].index(rank)
+    coordinator.world_size = 2
+    coordinator.device_group = object()
+    coordinator.cpu_group = object()
+    metadata = [("payload", TensorMetadata("cpu", torch.float32, (2,)))]
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+    def broadcast_object(value, src):
+        assert src == 0
+        return metadata
+
+    def broadcast(tensor, src, group, async_op):
+        assert src == 1
+        assert group is coordinator.cpu_group
+        assert async_op
+        tensor.fill_(7)
+        return SimpleNamespace(wait=lambda: None)
+
+    monkeypatch.setattr(coordinator, "broadcast_object", broadcast_object)
+    monkeypatch.setattr(torch.distributed, "broadcast", broadcast)
+    result = coordinator.broadcast_tensor_dict({"payload": torch.full((2,), 7.0)} if rank == 1 else None)
+    assert torch.equal(result["payload"], torch.full((2,), 7.0))

tests/diffusion/distributed/test_wan_decoder_fast_path.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The exact Wan decoder fast path must be bit-identical to the plain streaming decode.
+
+Two decoders are driven frame by frame through ``WanStreamingDecoder`` over two chunks of one session and
+a fresh session (the persistent buffers' causal front frames go zeros -> cache -> zeros), on CPU: plain
+diffusers modules in one process, and the spatially sharded wrappers across two gloo ranks with real halo
+exchange, each in fp32 and under bf16 autocast (where the parameter cast applies). Every output must
+``torch.equal`` the untouched decoder's on every rank.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanDecoder3d
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_fast_path import install_wan_decoder_fast_path
+from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
+    WanDistCausalConv3d,
+    WanDistConv2d,
+    install_wan_spatial_shard_decode,
+)
+from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+
+def _make_vae(seed: int = 0) -> SimpleNamespace:
+    torch.manual_seed(seed)
+    # Three levels: the default temporal-upsample list gives one upsample2d and one upsample3d resample.
+    decoder = WanDecoder3d(dim=8, z_dim=4, dim_mult=[1, 2, 4], num_res_blocks=1, attn_scales=[])
+    post_quant_conv = WanCausalConv3d(4, 4, 1)
+    for p in list(decoder.parameters()) + list(post_quant_conv.parameters()):
+        p.data.uniform_(-0.5, 0.5)
+    decoder.eval()
+    count = sum(1 for m in decoder.modules() if m.__class__.__name__ == "WanCausalConv3d")
+    return SimpleNamespace(decoder=decoder, post_quant_conv=post_quant_conv, _cached_conv_counts={"decoder": count})
+
+
+def _run(vae, latents: list[torch.Tensor], autocast: bool) -> list[torch.Tensor]:
+    decoder = WanStreamingDecoder(vae)
+    outs = []
+    ctx = torch.autocast("cpu", dtype=torch.bfloat16) if autocast else torch.no_grad()
+    with torch.no_grad(), ctx:
+        state = decoder.new_decode_state("a")
+        outs.append(decoder.decode_chunk(latents[0], state))
+        outs.append(decoder.decode_chunk(latents[1], state))
+        fresh = decoder.new_decode_state("b")
+        outs.append(decoder.decode_chunk(latents[2], fresh))
+    return [o.float() for o in outs]
+
+
+def _latents() -> list[torch.Tensor]:
+    torch.manual_seed(1)
+    return [torch.randn(1, 4, 2, 6, 10), torch.randn(1, 4, 1, 6, 10), torch.randn(1, 4, 2, 6, 10)]
+
+
+def _check_pair(reference, candidate, autocast: bool, sharded: bool, level: str = "exact") -> None:
+    dtype = torch.bfloat16 if autocast else None
+    counts = install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level)
+    assert counts["level"] == level
+    assert counts["upsamples"] == 2
+    convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, torch.nn.Conv3d | torch.nn.Conv2d)) + 1
+    assert counts["conv_params_cast"] == (convs if autocast else 0)
+    dist_convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d))
+    assert counts["persistent_input_buffers"] == dist_convs
+    assert (dist_convs > 0) == sharded
+    if level == "fused":
+        assert (
+            counts["fused_norm_silu"]
+            == 2 * sum(1 for m in candidate.decoder.modules() if m.__class__.__name__ == "WanResidualBlock") + 1
+        )
+        assert counts["channels_last_convs"] == convs
+        assert counts["resample_forwards"] == 2
+        with pytest.raises(ValueError, match="already installed"):
+            install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level="exact")
+    # Idempotent: a second install is a no-op that reports the same counts.
+    assert install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level) is counts
+    latents = _latents()
+    expected = _run(reference, latents, autocast)
+    actual = _run(candidate, latents, autocast)
+    assert len(expected) == len(actual) == 3
+    for e, a in zip(expected, actual):
+        assert e.shape == a.shape
+        if level == "exact":
+            assert torch.equal(e, a)
+        else:
+            # The fused level changes layouts only on CPU (the kernels need CUDA); the eager norm's reduction
+            # order may differ with the layout, so the outputs are close, not identical.
+            assert torch.allclose(e, a, atol=1e-2 if autocast else 1e-5, rtol=1e-2 if autocast else 1e-4), (
+                (e - a).abs().max()
+            )
+    if sharded:
+        bufs = [m._input_buf for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d)]
+        assert all(b is not None for b in bufs)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("autocast", [False, True])
+def test_fast_path_is_bit_identical(autocast: bool) -> None:
+    reference = _make_vae()
+    candidate = copy.deepcopy(reference)
+    _check_pair(reference, candidate, autocast, sharded=False)
+
+
+def _swap_norms_to_rmsnorm_vae(vae) -> int:
+    """Ensure the decoder uses RMSNormVAE and return the number of matching norms.
+
+    Collecting pipeline tests can already apply patch_wan_rms_norm process-wide before this test runs.
+    """
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
+
+    count = 0
+    for module in list(vae.decoder.modules()):
+        for name, child in list(module.named_children()):
+            if isinstance(child, RMSNormVAE):
+                count += 1
+                continue
+            if child.__class__.__name__ == "WanRMS_norm":
+                images = child.gamma.dim() == 3  # (dim, 1, 1) for the attention block's 4D norm
+                norm = RMSNormVAE(child.gamma.shape[0], channel_first=child.channel_first, images=images, bias=False)
+                norm.gamma.data.copy_(child.gamma.data.reshape(norm.gamma.shape))
+                setattr(module, name, norm)
+                count += 1
+    return count
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("norm", ["default", "RMSNormVAE"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_fused_level_plumbing_on_cpu(autocast: bool, norm: str) -> None:
+    reference = _make_vae()
+    if norm == "RMSNormVAE":
+        assert _swap_norms_to_rmsnorm_vae(reference) > 0
+    candidate = copy.deepcopy(reference)
+    _check_pair(reference, candidate, autocast, sharded=False, level="fused")
+
+
+def _sharded_worker(rank: int, world_size: int, port: str, autocast: bool, return_dict, level: str = "exact") -> None:
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = port
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+    try:
+        torch.manual_seed(0)
+        reference = _make_vae()
+        candidate = copy.deepcopy(reference)
+        for vae in (reference, candidate):
+            install_wan_spatial_shard_decode(vae, dist.group.WORLD, split_dim="width", dst=None)
+        _check_pair(reference, candidate, autocast, sharded=True, level=level)
+        return_dict[rank] = "ok"
+    except BaseException as exc:  # report, the parent asserts
+        return_dict[rank] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("level", ["exact", "fused"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_sharded_fast_path_on_every_rank(autocast: bool, level: str) -> None:
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    port = str(29610 + int(autocast) + 2 * (level == "fused"))
+    mp.spawn(_sharded_worker, args=(2, port, autocast, return_dict, level), nprocs=2, join=True)
+    for rank in range(2):
+        assert return_dict.get(rank) == "ok", f"rank {rank}: {return_dict.get(rank)}"
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_conv_dtype_requires_half() -> None:
+    with pytest.raises(ValueError, match="conv_dtype"):
+        install_wan_decoder_fast_path(_make_vae(), conv_dtype=torch.float32)

tests/diffusion/distributed/test_wan_vae_fastpath_install.py

@@ -0,0 +1,646 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for the Wan VAE decoder fast path installer and its exact PyTorch fallbacks."""
+
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable
+from functools import wraps
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+from torch import nn
+
+from tests.diffusion.distributed.wan_vae_fastpath_helpers import (  # noqa: F401
+    original_wan_rms_norm,
+    unpatched_wan_rms_norm,
+)
+from vllm_omni.diffusion import registry as registry_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    REPORT_ATTR,
+    decode_frames,
+    install_wan_vae_fastpath,
+    is_installed,
+    uninstall_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fastpath_forwards
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cpu,
+    pytest.mark.diffusion,
+    pytest.mark.usefixtures("unpatched_wan_rms_norm"),
+]
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 1],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=True,
+)
+TINY_RESIDUAL_PATCH2 = dict(TINY_RESIDUAL, patch_size=2, in_channels=12, out_channels=12)
+TINY_WAN21 = dict(
+    base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 2],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=False,
+)
+CONFIGS = {
+    "residual": TINY_RESIDUAL,
+    "residual_patch2": TINY_RESIDUAL_PATCH2,
+    "wan21": TINY_WAN21,
+}
+
+
+def _build_pair(config: dict, dtype: torch.dtype) -> tuple[AutoencoderKLWan, AutoencoderKLWan]:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("frames", [1, 3])
+def test_lossless_fallback_paths_are_bitwise_exact(config_name: str, dtype: torch.dtype, frames: int) -> None:
+    """Without CUDA every kernel declines, so this exercises the restructured PyTorch paths."""
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    assert report.fused_silu_dtypes == ()
+
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_decode_frames_matches_reference_decode(config_name: str, dtype: torch.dtype) -> None:
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    install_wan_vae_fastpath(candidate)
+    for frames in (1, 2, 4):
+        torch.manual_seed(2)
+        latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+        with torch.no_grad():
+            expected = reference._decode(latents, return_dict=False)[0]
+            actual = decode_frames(candidate, latents)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_omni_wan_vae_decode_override_is_bitwise_exact(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 3, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        before_install = candidate.decode(latents, return_dict=False)[0]
+        install_wan_vae_fastpath(candidate)
+        actual = candidate.decode(latents, return_dict=False)[0]
+        as_dict = candidate.decode(latents).sample
+    # The Omni wrapper runs bf16 decode under autocast (also on CPU), so compare
+    # the override against the same wrapper before installation, and against the
+    # plain diffusers decode only where no autocast is involved.
+    assert actual.stride() == before_install.stride()
+    assert torch.equal(actual, before_install)
+    assert torch.equal(as_dict, before_install)
+    if dtype is torch.float32:
+        assert torch.equal(before_install, expected)
+
+    # Tiling still dispatches to the diffusers implementation.
+    candidate.use_tiling = True
+    candidate.tile_sample_min_height = candidate.tile_sample_min_width = 16
+    candidate.tile_sample_stride_height = candidate.tile_sample_stride_width = 16
+    reference.use_tiling = True
+    reference.tile_sample_min_height = reference.tile_sample_min_width = 16
+    reference.tile_sample_stride_height = reference.tile_sample_stride_width = 16
+    with torch.no_grad():
+        expected_tiled = reference.decode(latents, return_dict=False)[0]
+        actual_tiled = candidate.decode(latents, return_dict=False)[0]
+    if dtype is torch.float32:
+        assert torch.equal(actual_tiled, expected_tiled)
+    assert actual_tiled.shape == expected_tiled.shape
+
+
+def test_omni_wan_vae_decode_preserves_autograd() -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+
+    reference_latents = torch.randn(1, 4, 2, 6, 8, requires_grad=True)
+    candidate_latents = reference_latents.detach().clone().requires_grad_()
+    expected = reference.decode(reference_latents, return_dict=False)[0]
+    actual = candidate.decode(candidate_latents, return_dict=False)[0]
+
+    torch.testing.assert_close(actual, expected)
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_latents.grad, reference_latents.grad)
+
+
+def test_install_is_idempotent_and_reversible() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    keys_before = set(vae.state_dict())
+    report = install_wan_vae_fastpath(vae)
+    assert report.installed and is_installed(vae)
+    assert report.patched["WanDecoder3d"] == 1
+    assert report.patched["WanCausalConv3d"] > 1
+    assert report.patched["WanRMS_norm"] > 1
+    assert set(vae.state_dict()) == keys_before
+    assert "forward" in vae.decoder.__dict__
+
+    assert install_wan_vae_fastpath(vae) is report
+    assert install_wan_vae_fastpath(vae, level="channels_last") is report  # warns, keeps the first level
+
+    uninstall_wan_vae_fastpath(vae)
+    assert not is_installed(vae)
+    assert "forward" not in vae.decoder.__dict__
+    assert all("forward" not in module.__dict__ for module in vae.decoder.modules())
+    assert all(fastpath_forwards.CFG_ATTR not in module.__dict__ for module in vae.decoder.modules())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("failure_stage", ["binding", "conversion", "report"])
+def test_failed_install_restores_original_state(failure_stage: str, monkeypatch) -> None:
+    installer = importlib.import_module(install_wan_vae_fastpath.__module__)
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    # Preserve mixed original layouts, an existing forward, gradients and a
+    # buffer: Module.to can replace tensors as well as change their storage.
+    vae.decoder.conv_in.to(memory_format=torch.channels_last_3d)
+    vae.decoder.forward = vae.decoder.forward
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, object())
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    vae.post_quant_conv.register_buffer("rollback_probe", torch.randn(2, 3, 4, 5, 6))
+    attributes = {
+        module: {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        for module in vae.modules()
+    }
+    parameters = dict(vae.named_parameters())
+    buffers = dict(vae.named_buffers())
+    gradients = {name: parameter.grad for name, parameter in parameters.items()}
+    tensors = [*parameters.values(), *buffers.values(), *(grad for grad in gradients.values() if grad is not None)]
+    originals = [(tensor, tensor.data_ptr(), tensor.stride(), tensor.clone()) for tensor in tensors]
+    latents = torch.randn(1, 4, 2, 6, 8)
+    expected = vae.decode(latents, return_dict=False)[0]
+
+    with monkeypatch.context() as patch:
+        if failure_stage == "binding":
+            conv_type = type(vae.post_quant_conv)
+            original_setattr: Callable[..., None] = conv_type.__setattr__
+
+            def fail_binding(module, name, value):
+                original_setattr(module, name, value)
+                if module is vae.post_quant_conv and name == fastpath_forwards.CFG_ATTR:
+                    raise RuntimeError("injected installation failure")
+
+            patch.setattr(conv_type, "__setattr__", fail_binding)
+        elif failure_stage == "conversion":
+            original_to = vae.post_quant_conv.to
+
+            def fail_conversion(*args, **kwargs):
+                original_to(*args, **kwargs)
+                assert vae.decoder.conv_out.weight.is_contiguous(memory_format=torch.channels_last_3d)
+                raise torch.OutOfMemoryError("injected installation failure")
+
+            patch.setattr(vae.post_quant_conv, "to", fail_conversion)
+        else:
+            original_info = installer.logger.info
+
+            def fail_report(message, *args, **kwargs):
+                if message.startswith("Wan VAE fast path (%s) installed"):
+                    assert is_installed(vae)
+                    raise RuntimeError("injected installation failure")
+                original_info(message, *args, **kwargs)
+
+            patch.setattr(installer.logger, "info", fail_report)
+
+        with pytest.raises(RuntimeError, match="injected installation failure"):
+            install_wan_vae_fastpath(vae, level="channels_last")
+
+    assert not is_installed(vae)
+    assert not hasattr(vae, REPORT_ATTR)
+    assert not hasattr(vae, installer._UNDO_ATTR)
+    for module, original in attributes.items():
+        actual = {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        assert actual == original
+    assert all(dict(vae.named_parameters())[name] is parameter for name, parameter in parameters.items())
+    assert all(dict(vae.named_buffers())[name] is buffer for name, buffer in buffers.items())
+    assert all(parameters[name].grad is grad for name, grad in gradients.items())
+    for tensor, pointer, strides, values in originals:
+        assert tensor.data_ptr() == pointer
+        assert tensor.stride() == strides
+        assert torch.equal(tensor, values)
+    uninstall_wan_vae_fastpath(vae)  # Safe and unnecessary after automatic rollback.
+    assert torch.equal(vae.decode(latents, return_dict=False)[0], expected)
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    uninstall_wan_vae_fastpath(vae)
+
+
+def test_load_state_dict_after_install_updates_patched_modules() -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, torch.float32)
+    install_wan_vae_fastpath(candidate)
+    with torch.no_grad():
+        for parameter in reference.parameters():
+            parameter.mul_(0.5)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert torch.equal(actual, expected)
+
+
+def test_channels_last_level_converts_conv_weights_and_restores_them() -> None:
+    reference, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    report = install_wan_vae_fastpath(vae, level="channels_last")
+    assert report.installed and report.channels_last
+    conv3d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv3d)]
+    conv2d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv2d)]
+    assert conv3d and conv2d
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last_3d) for m in conv3d)
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last) for m in conv2d)
+    assert vae.post_quant_conv.weight.is_contiguous(memory_format=torch.channels_last_3d)
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        out = vae.decode(latents, return_dict=False)[0]
+    # One temporal upsampler: 1 + 2 output frames for 2 latent frames.
+    assert expected.shape == (1, 3, 3, 12, 16)
+    assert out.shape == expected.shape
+    torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
+
+    uninstall_wan_vae_fastpath(vae)
+    assert all(m.weight.is_contiguous() for m in conv3d + conv2d)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("assign", [False, True])
+def test_uninstall_restores_mixed_layouts_and_keeps_weight_updates(assign: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    convs = [module for module in vae.decoder.modules() if isinstance(module, (nn.Conv2d, nn.Conv3d))]
+    convs.append(vae.post_quant_conv)
+    for index, conv in enumerate(convs):
+        if index % 3 == 1:
+            conv.to(memory_format=torch.channels_last_3d if isinstance(conv, nn.Conv3d) else torch.channels_last)
+        elif index % 3 == 2:
+            conv.weight.data = conv.weight.data.transpose(-1, -2)
+    vae.post_quant_conv.register_buffer("layout_probe", torch.randn(2, 3, 4, 5, 6).transpose(-1, -2))
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    grad_strides = vae.decoder.conv_out.weight.grad.stride()
+    strides = {name: tensor.stride() for name, tensor in vae.state_dict().items()}
+    # An explicit alias is safe to replace, but must remain an instance
+    # attribute after uninstall, with its exact original identity.
+    original_forward = vae.decoder.forward
+    vae.decoder.forward = original_forward
+    original_cfg = object()
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, original_cfg)
+
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    updated = {name: tensor.contiguous() + 0.25 for name, tensor in vae.state_dict().items()}
+    vae.load_state_dict(updated, assign=assign)
+    uninstall_wan_vae_fastpath(vae)
+
+    assert vae.decoder.__dict__["forward"] is original_forward
+    assert getattr(vae.decoder, fastpath_forwards.CFG_ATTR) is original_cfg
+    for name, tensor in vae.state_dict().items():
+        assert tensor.stride() == strides[name], name
+        assert torch.equal(tensor, updated[name]), name
+    if not assign:
+        assert vae.decoder.conv_out.weight.grad.stride() == grad_strides
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("level", ["lossless", "channels_last"])
+@pytest.mark.parametrize(
+    "target",
+    [
+        "decoder",
+        "decoder.conv_in",
+        "decoder.up_blocks.0.upsampler",
+        "decoder.up_blocks.0.upsampler.resample",
+        "decoder.up_blocks.0.upsampler.resample.1",
+        "decoder.up_blocks.0.avg_shortcut",
+        "decoder.norm_out",
+        "decoder.nonlinearity",
+        "post_quant_conv",
+    ],
+)
+def test_installer_declines_forward_wrappers_it_would_bypass(target: str, level: str) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    module = vae.get_submodule(target)
+    original_forward = module.forward
+    calls = []
+
+    @wraps(original_forward)
+    def wrapped(*args, **kwargs):
+        calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    module.forward = wrapped
+    originals = {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()}
+    report = install_wan_vae_fastpath(vae, level=level)
+    assert not report.installed and target in report.reason and "custom forward" in report.reason
+    assert not is_installed(vae) and not hasattr(vae, REPORT_ATTR)
+    assert all(fastpath_forwards.CFG_ATTR not in child.__dict__ for child in vae.modules())
+    assert {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()} == originals
+    uninstall_wan_vae_fastpath(vae)
+    assert module.forward is wrapped
+    vae.decode(torch.randn(1, 4, 2, 6, 8))
+    assert calls
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("pre_hook", [False, True])
+def test_installer_declines_convolution_hooks_it_would_bypass(pre_hook: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    conv = vae.decoder.conv_in
+    calls = []
+
+    def hook(*args):
+        calls.append(True)
+
+    handle = conv.register_forward_pre_hook(hook) if pre_hook else conv.register_forward_hook(hook)
+    try:
+        report = install_wan_vae_fastpath(vae, level="channels_last")
+        assert not report.installed and "decoder.conv_in" in report.reason and "forward hooks" in report.reason
+        vae.decode(torch.randn(1, 4, 2, 6, 8))
+        assert calls
+    finally:
+        handle.remove()
+
+
+@torch.no_grad()
+def test_installer_preserves_wrappers_and_hooks_on_modules_it_still_calls() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    original_forward = vae.decoder.mid_block.forward
+    wrapper_calls = []
+    hook_calls = []
+
+    def wrapped(*args, **kwargs):
+        wrapper_calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    def hook(*args):
+        hook_calls.append(True)
+
+    vae.decoder.mid_block.forward = wrapped
+    handle = vae.decoder.register_forward_hook(hook)
+    try:
+        latents = torch.randn(1, 4, 2, 6, 8)
+        expected = vae.decode(latents).sample
+        wrapper_calls.clear()
+        hook_calls.clear()
+        assert install_wan_vae_fastpath(vae).installed
+        actual = vae.decode(latents).sample
+        assert torch.equal(actual, expected)
+        assert len(wrapper_calls) == len(hook_calls) == 2
+        uninstall_wan_vae_fastpath(vae)
+        assert vae.decoder.mid_block.forward is wrapped
+    finally:
+        handle.remove()
+
+
+def test_installer_refuses_unsupported_targets() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+
+    assert install_wan_vae_fastpath(vae, level="off").installed is False
+    assert not is_installed(vae)
+
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial-shard" in report.reason
+    del vae._vllm_omni_wan_spatial_shard_installed
+
+    vae.distributed_executor = SimpleNamespace(parallel_mode="spatial_shard_height")
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial_shard_height" in report.reason
+    del vae.distributed_executor
+
+    report = install_wan_vae_fastpath(nn.Linear(2, 2))
+    assert not report.installed and "AutoencoderKLWan" in report.reason
+
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        install_wan_vae_fastpath(vae, level="bogus")
+
+
+@pytest.mark.parametrize(("batch", "frames"), [(1, 1), (1, 2), (2, 1)])
+def test_resample_views_keep_channels_last_recognizable(batch: int, frames: int) -> None:
+    """Regression: ``reshape`` gives a size-1 batch dim a stride the layout heuristic rejects."""
+    import torch.nn.functional as F
+
+    x = torch.randn(batch, 8, frames, 6, 10).contiguous(memory_format=torch.channels_last_3d)
+    merged = fastpath_forwards._merge_batch_and_frames(x)
+    reference = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, 8, 6, 10)
+    assert merged.data_ptr() == x.data_ptr()
+    assert torch.equal(merged, reference)
+    upsampled = F.interpolate(merged, scale_factor=(2.0, 2.0), mode="nearest-exact")
+    assert upsampled.stride(1) == 1, "nearest upsample must keep channels_last for the following Conv2d"
+    assert torch.equal(upsampled, F.interpolate(reference.contiguous(), scale_factor=(2.0, 2.0), mode="nearest-exact"))
+
+    split = fastpath_forwards._split_batch_and_frames(upsampled, batch, frames)
+    assert split.shape == (batch, 8, frames, 12, 20)
+    assert split.data_ptr() == upsampled.data_ptr()
+    assert split.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(split, upsampled.view(batch, frames, 8, 12, 20).permute(0, 2, 1, 3, 4))
+
+    plain = torch.randn(batch, 8, frames, 6, 10)
+    assert torch.equal(
+        fastpath_forwards._merge_batch_and_frames(plain), plain.permute(0, 2, 1, 3, 4).reshape(-1, 8, 6, 10)
+    )
+
+
+@torch.no_grad()
+def test_pending_conv_bias_is_added_exactly_when_no_kernel_consumes_it() -> None:
+    """Without the channels-last kernel the un-added conv bias is applied with ATen's rounding."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    for dtype in (torch.float32, torch.bfloat16):
+        norm = WanRMS_norm(8, images=False).to(dtype)
+        act = nn.SiLU()
+        x = torch.randn(1, 8, 2, 4, 4).to(dtype)
+        bias = torch.randn(8).to(dtype)
+        biased = x.clone()
+        biased.add_(bias.view(1, -1, 1, 1, 1))
+        expected = act(norm(biased))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        setattr(norm, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig(channels_last=True))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        out = fastpath_forwards.rms_norm_fastpath(norm, x, bias=bias)
+        assert out is not None and torch.equal(out, norm(biased))
+
+
+@torch.no_grad()
+def test_resample_return_bias_is_none_without_kernels() -> None:
+    """On CPU the Conv2d keeps its bias, so the up block adds nothing twice."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample, WanResidualUpBlock
+
+    torch.manual_seed(0)
+    resample = WanResample(8, mode="upsample2d").eval()
+    assert fastpath_forwards._is_upsample_conv_pair(resample.resample)
+    assert not fastpath_forwards._is_upsample_conv_pair(WanResample(8, mode="downsample2d").resample)
+    x = torch.randn(1, 8, 2, 6, 10)
+    out, bias = fastpath_forwards.resample_forward(resample, x, [None], [0], return_bias=True)
+    assert bias is None
+    assert torch.equal(out, WanResample.forward(resample, x, feat_cache=[None], feat_idx=[0]))
+
+    # First chunk of a temporal up block: one input frame, the time conv is skipped ("Rep").
+    block = WanResidualUpBlock(8, 8, num_res_blocks=1, temperal_upsample=True, up_flag=True).eval()
+    for module in block.modules():
+        setattr(module, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig())
+    x = torch.randn(1, 8, 1, 6, 10)
+    cache_len = 6
+    expected = WanResidualUpBlock.forward(block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True)
+    actual = fastpath_forwards.residual_up_block_forward(
+        block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True
+    )
+    assert torch.equal(actual, expected)
+
+
+def test_upsample_forward_only_fuses_nearest_2x() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanUpsample
+
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"))
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=2, mode="nearest"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 3.0), mode="nearest-exact"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="bilinear"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(size=(12, 20), mode="nearest-exact"))
+
+    module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+    for dtype in (torch.float32, torch.bfloat16):
+        x = torch.randn(3, 8, 6, 10).to(dtype)
+        for layout in (torch.contiguous_format, torch.channels_last):
+            x = x.contiguous(memory_format=layout)
+            expected = WanUpsample.forward(module, x)
+            actual = fastpath_forwards.upsample_forward(module, x)
+            assert actual.stride() == expected.stride()
+            assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_rms_norm_vae_substitute_is_not_matched(dtype: torch.dtype, monkeypatch) -> None:
+    from diffusers.models.autoencoders import autoencoder_kl_wan
+
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
+
+    assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False))
+    _, vae = _build_pair(TINY_RESIDUAL, dtype)
+    assert fastpath_forwards.is_diffusers_rms_norm(vae.decoder.norm_out)
+
+    # Production Wan pipelines use this substitute. The installer must leave
+    # its different epsilon/upcast behavior intact while optimizing the rest.
+    monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", RMSNormVAE)
+    reference, candidate = _build_pair(TINY_RESIDUAL, dtype)
+    report = install_wan_vae_fastpath(candidate)
+    assert report.installed
+    assert "WanRMS_norm" not in report.patched
+    assert "RMSNormVAE" not in report.patched
+    assert "forward" not in candidate.decoder.norm_out.__dict__
+    latents = torch.randn(1, 4, 3, 6, 8).to(dtype)
+    expected = reference.decode(latents, return_dict=False)[0]
+    actual = candidate.decode(latents, return_dict=False)[0]
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+def test_rms_norm_fastpath_declines_tensor_bias_and_dtype_mismatch() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    norm = WanRMS_norm(8, images=False, bias=True)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4)) is None
+
+    norm = WanRMS_norm(8, images=False)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4).to(torch.bfloat16)) is None
+
+    x = torch.randn(1, 8, 2, 4, 4)
+    out = fastpath_forwards.rms_norm_fastpath(norm, x)
+    assert out is not None and torch.equal(out, norm(x))
+
+
+def test_rms_norm_forward_preserves_autograd() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    reference = WanRMS_norm(8, images=False)
+    candidate = WanRMS_norm(8, images=False)
+    candidate.load_state_dict(reference.state_dict())
+    reference_input = torch.randn(1, 8, 1, 4, 4, requires_grad=True)
+    candidate_input = reference_input.detach().clone().requires_grad_()
+
+    expected = reference(reference_input)
+    assert fastpath_forwards.rms_norm_fastpath(candidate, candidate_input) is None
+    actual = fastpath_forwards.rms_norm_forward(candidate, candidate_input)
+    torch.testing.assert_close(actual, expected)
+
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_input.grad, reference_input.grad)
+    torch.testing.assert_close(candidate.gamma.grad, reference.gamma.grad)
+
+
+def test_omni_diffusion_config_validates_vae_fast_path() -> None:
+    assert OmniDiffusionConfig(model="x").vae_fast_path == "lossless"
+    assert OmniDiffusionConfig(model="x", vae_fast_path="channels_last").vae_fast_path == "channels_last"
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        OmniDiffusionConfig(model="x", vae_fast_path="fast")
+
+
+class _StubPipeline(nn.Module):
+    def __init__(self, vae: nn.Module) -> None:
+        super().__init__()
+        self.vae = vae
+
+
+@pytest.mark.parametrize(("level", "expected"), [("lossless", True), ("channels_last", True), ("off", False)])
+def test_registry_hook_installs_on_cuda_platform(mocker, level: str, expected: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = True
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path=level))
+    assert is_installed(vae) is expected
+    if expected:
+        assert getattr(vae, REPORT_ATTR).level == level
+
+
+def test_registry_hook_skips_non_cuda_platform_and_non_wan_vaes(mocker) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = False
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path="lossless"))
+    assert not is_installed(vae)
+
+    platform.is_cuda.return_value = True
+    other = nn.Linear(2, 2)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(other), SimpleNamespace(vae_fast_path="lossless"))
+    assert not hasattr(other, REPORT_ATTR)

tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py

@@ -0,0 +1,908 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""GPU bit-exactness tests for the Wan VAE decoder fast path kernels and forwards."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    WanCausalConv3d,
+    WanResample,
+    WanResidualBlock,
+    WanRMS_norm,
+    WanUpsample,
+)
+from torch import nn
+
+from tests.diffusion.distributed.wan_vae_fastpath_helpers import (  # noqa: F401
+    original_wan_rms_norm,
+    unpatched_wan_rms_norm,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    install_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_data_movement as dm
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm as rn
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm_cl as cl
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_upsample as up
+
+pytestmark = [
+    pytest.mark.usefixtures("unpatched_wan_rms_norm"),
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+LOW_PRECISION = [torch.bfloat16, torch.float16]
+ALL_DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+# Cosmos3 720p decoder stages, spatially scaled down so the suite stays fast.
+NORM_SHAPES = [
+    (1, 1024, 1, 45, 80),
+    (1, 1024, 2, 45, 80),
+    (1, 512, 4, 45, 80),
+    (1, 256, 4, 90, 160),
+]
+
+
+def _bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return a.shape == b.shape and bool((a.view(int_dtype) == b.view(int_dtype)).all())
+
+
+def _same_dense_strides(a: torch.Tensor, b: torch.Tensor) -> bool:
+    """Stride equality ignoring size-1 dims, whose stride carries no information."""
+    return a.shape == b.shape and all(
+        sa == sb for sa, sb, n in zip(a.stride(), b.stride(), a.shape, strict=True) if n > 1
+    )
+
+
+def _with_negative_zeros(x: torch.Tensor) -> torch.Tensor:
+    flat = x.view(-1)
+    flat[:: max(1, flat.numel() // 97)] = -0.0
+    return x
+
+
+def _reference_rms_norm(norm: WanRMS_norm, x: torch.Tensor) -> torch.Tensor:
+    return WanRMS_norm.forward(norm, x)
+
+
+def _make_norm(channels: int, dtype: torch.dtype, images: bool = False) -> WanRMS_norm:
+    norm = WanRMS_norm(channels, images=images).to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        norm.gamma.copy_(torch.randn_like(norm.gamma) * 0.5 + 1.0)
+    return norm
+
+
+# --------------------------------------------------------------------------- #
+# RMSNorm
+# --------------------------------------------------------------------------- #
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+def test_vector_norm_from_low_precision_matches_fp32_normalize_denominator(dtype: torch.dtype, shape) -> None:
+    """The fast path relies on ATen running the same reduction for both inputs."""
+    torch.manual_seed(0)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    expected = x.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(x, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+    frame_major = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    expected = frame_major.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(frame_major, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_rms_norm_fastpath_is_bitwise_exact(dtype: torch.dtype, shape, layout: str, autocast: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype)
+    x = _with_negative_zeros(torch.randn(shape, device="cuda").to(dtype))
+    if layout == "frame_major":
+        x = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    with torch.autocast("cuda", dtype=dtype if dtype is not torch.float32 else torch.bfloat16, enabled=autocast):
+        expected = _reference_rms_norm(norm, x)
+        actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None
+    assert actual.dtype == expected.dtype
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected), "RMSNorm fast path is not bit-exact (signed zeros included)"
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_rms_norm_fastpath_handles_attention_block_4d_input(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(1024, dtype, images=True)
+    x = torch.randn(2, 1024, 45, 80, device="cuda").to(dtype)
+    expected = _reference_rms_norm(norm, x)
+    actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+def test_fused_silu_epilogue_matches_silu_of_reference(dtype: torch.dtype) -> None:
+    device = torch.device("cuda")
+    if not rn.silu_epilogue_is_exact(device, dtype):
+        pytest.skip(f"fused SiLU epilogue is not bit-exact for {dtype} on this toolkit; fusion stays disabled")
+    torch.manual_seed(0)
+    norm = _make_norm(256, dtype)
+    x = _with_negative_zeros(torch.randn(1, 256, 4, 90, 160, device="cuda").to(dtype))
+    expected = F.silu(_reference_rms_norm(norm, x))
+    actual = fp.rms_norm_fastpath(norm, x, silu=True)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", [(1, 1024, 2, 45, 80), (1, 256, 4, 90, 160), (2, 512, 45, 80)])
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_is_close_and_keeps_layout(dtype: torch.dtype, shape, silu: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype, images=len(shape) == 4)
+    memory_format = torch.channels_last_3d if len(shape) == 5 else torch.channels_last
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=memory_format)
+    expected = _reference_rms_norm(norm, x)
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=memory_format)
+    assert actual.dtype == expected.dtype
+    # Fast-math kernel: approximate divide/exp and a single final rounding.
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    # Declines channels-first input and too many channels.
+    assert cl.rms_norm_channels_last(x.contiguous(), norm.gamma, norm.scale) is None
+    wide = torch.randn(1, 2048, 1, 4, 4, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last_3d)
+    assert cl.rms_norm_channels_last(wide, torch.ones(2048, device="cuda", dtype=dtype), 1.0) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_absorbs_conv_bias(dtype: torch.dtype, silu: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 90, 160)
+    norm = _make_norm(shape[1], dtype)
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=torch.channels_last_3d)
+    bias = (torch.randn(shape[1], device="cuda") * 0.1).to(dtype)
+    # Reference: ATen adds the conv bias in place (one rounding) before the norm.
+    expected = _reference_rms_norm(norm, x + bias.view(1, -1, 1, 1, 1))
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu, bias=bias)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias[:-1]) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.view(1, -1)) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.cpu()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_channels_last_residual_block_folds_conv1_bias(dtype: torch.dtype, monkeypatch) -> None:
+    """The conv1 bias is left un-added and consumed by norm2's kernel (rounded like ATen); output stays close."""
+    torch.manual_seed(0)
+    block = WanResidualBlock(64, 64, dropout=0.0).eval().to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        for conv in (block.conv1, block.conv2):
+            conv.bias.copy_(torch.randn_like(conv.bias) * 0.1)
+    x = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=dtype)
+    expected = WanResidualBlock.forward(block, x, feat_cache=[None, None], feat_idx=[0])
+
+    for conv in (block.conv1, block.conv2):  # as the installer does: conv weights only
+        conv.to(memory_format=torch.channels_last_3d)
+    cfg = fp.FastPathConfig(channels_last=True)
+    for module in (block, block.norm1, block.norm2):
+        setattr(module, fp.CFG_ATTR, cfg)
+    seen: list[torch.Tensor | None] = []
+    original = fp.rms_norm_fastpath
+
+    def spy(norm, x, *, silu=False, bias=None):
+        seen.append(bias)
+        return original(norm, x, silu=silu, bias=bias)
+
+    monkeypatch.setattr(fp, "rms_norm_fastpath", spy)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    actual = fp.residual_block_forward(block, x_cl, feat_cache=[None, None], feat_idx=[0])
+    assert seen[0] is None and seen[1] is block.conv1.bias, "norm2 should receive conv1's un-added bias"
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)  # channels-last convs may run TF32
+    else:
+        # Fast-math norm plus channels-last cuDNN accumulation order: a few outputs differ by one bf16 ulp.
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+def test_rms_norm_scale_declines_bad_operands() -> None:
+    x = torch.randn(1, 8, 16, device="cuda", dtype=torch.bfloat16)
+    denom = torch.ones(1, 1, 16, device="cuda")
+    gamma = torch.ones(8, device="cuda", dtype=torch.bfloat16)
+    assert rn.rms_norm_scale(x, denom, gamma, 1.0) is not None
+    assert rn.rms_norm_scale(x.cpu(), denom.cpu(), gamma.cpu(), 1.0) is None
+    assert rn.rms_norm_scale(x, denom.to(torch.bfloat16), gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, denom, gamma.float(), 1.0) is None
+    assert rn.rms_norm_scale(x.transpose(1, 2), denom, gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, torch.ones(1, 1, 8, device="cuda"), gamma, 1.0) is None
+
+
+# --------------------------------------------------------------------------- #
+# Data movement
+# --------------------------------------------------------------------------- #
+
+
+def _dense_5d(shape, dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
+    tensor = torch.randn(shape, device="cuda", dtype=dtype)
+    if channels_last:
+        return tensor.contiguous(memory_format=torch.channels_last_3d)
+    return tensor.contiguous()
+
+
+def _reference_cat_pad(x: torch.Tensor, cache: torch.Tensor | None, padding) -> torch.Tensor:
+    reference_padding = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        reference_padding[4] -= cache.shape[2]
+    return F.pad(x, reference_padding)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,padding",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        # Non-power-of-two channels and a ragged number of rows per program.
+        (12, 3, 7, 9, 1, (1, 1, 1, 1, 2, 0)),
+        # Width wider than one block so the in-kernel width loop runs.
+        (8, 2, 3, 1300, 2, (1, 1, 1, 1, 2, 0)),
+        # Cosmos3 stage-3 shape, scaled down in height.
+        (256, 4, 8, 640, 2, (1, 1, 1, 1, 2, 0)),
+    ],
+)
+def test_cat_pad_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, padding
+) -> None:
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = None
+    if cache_frames:
+        pad_height, pad_width = padding[2], padding[0]
+        buffer = _dense_5d(
+            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width), dtype, channels_last
+        )
+        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
+
+    reference = _reference_cat_pad(x, cache, padding)
+    output = dm.cat_pad_5d(x, cache, padding)
+    assert output is not None
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert torch.equal(output_with_cache, reference)
+    pad_height, pad_width = padding[2], padding[0]
+    keep_frames = min(CACHE_T, reference.shape[2])
+    expected_cache = reference[:, :, -keep_frames:, pad_height : pad_height + height, pad_width : pad_width + width]
+    assert torch.equal(next_cache, expected_cache)
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("cache_layout", ["contiguous", "channels_last"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames",
+    [
+        (96, 1, 10, 14, 0),
+        (96, 1, 10, 14, 2),
+        (12, 3, 7, 9, 1),
+        (8, 2, 3, 1300, 2),
+        (256, 4, 8, 640, 2),
+        (64, 4, 70, 70, 2),
+    ],
+)
+def test_cat_pad_5d_assembles_channels_last_output_from_channels_first_input(
+    dtype, layout, cache_layout, channels, frames, height, width, cache_frames
+) -> None:
+    """The conv_out path: channels-first activations, channels-last padded input, cache in the input layout."""
+    torch.manual_seed(0)
+    padding = (1, 1, 1, 1, 2, 0)
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype)
+    cache = None
+    if cache_frames:
+        cache = _dense_5d((1, channels, cache_frames, height, width), dtype, cache_layout == "channels_last")
+    reference = _reference_cat_pad(x, cache, padding)
+
+    output = dm.cat_pad_5d(x, cache, padding, channels_last_output=True)
+    assert output is not None
+    assert output.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T, channels_last_output=True)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert output_with_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output_with_cache, reference)
+    keep_frames = min(CACHE_T, reference.shape[2])
+    assert next_cache.is_contiguous()  # follows x's layout, not the output's
+    assert torch.equal(next_cache, reference[:, :, -keep_frames:, 1 : 1 + height, 1 : 1 + width])
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+def test_conv_out_channels_last_verdict_is_recorded_and_exact(layout: str) -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 12, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    if layout == "frame_major":
+        chunk = torch.randn(1, 4, 64, 24, 40, device="cuda", dtype=torch.bfloat16).permute(0, 2, 1, 3, 4)
+    else:
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16)
+    cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16)
+    reference = WanCausalConv3d.forward(conv, chunk, cache)
+    fast_cache = [cache.clone()]
+    out = fp._run_conv_out_channels_last(conv, chunk, fast_cache, 0)
+    assert out is not None  # the verdict decides which formulation's result is returned, both are exact
+    assert out.is_contiguous()
+    assert _bits_equal(out, reference)
+    verdicts = fp._CONV_OUT_LAYOUT_VERDICTS[conv]
+    assert len(verdicts) == 1
+    assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+    assert fast_cache[0].is_contiguous()
+    again = fp._run_conv_out_channels_last(conv, chunk, [cache.clone()], 0)
+    if next(iter(verdicts.values())):
+        assert again is not None and _bits_equal(again, reference)
+    else:
+        assert again is None  # routed back to the standard path
+    # Channels-last activations are not this path's business.
+    assert (
+        fp._run_conv_out_channels_last(conv, chunk.contiguous(memory_format=torch.channels_last_3d), [None], 0) is None
+    )
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,pad_front",
+    [(96, 1, 10, 14, 0, 2), (96, 1, 10, 14, 1, 2), (96, 4, 10, 14, 2, 2), (12, 3, 7, 9, 2, 2), (256, 4, 8, 640, 2, 2)],
+)
+def test_cat_time_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, pad_front
+):
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = _dense_5d((1, channels, cache_frames, height, width), dtype, channels_last) if cache_frames else None
+
+    parts = [torch.zeros_like(x[:, :, :1]).expand(-1, -1, pad_front - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2)
+    if channels_last:
+        reference = reference.contiguous(memory_format=torch.channels_last_3d)
+    else:
+        reference = reference.contiguous()
+
+    pair = dm.cat_time_5d(x, cache, pad_front, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert assembled.stride() == reference.stride()
+    assert torch.equal(assembled, reference)
+    assert torch.equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+    assert torch.equal(dm.cat_time_5d(x, cache, pad_front), reference)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("cache_frames", [0, 1, 2])
+def test_cat_time_5d_exceeds_grid_y_limit(monkeypatch: pytest.MonkeyPatch, cache_frames: int) -> None:
+    # Smaller blocks reproduce the large-plane launch failure without allocating
+    # 4K feature maps. The final block is partial, exercising the plane boundary.
+    monkeypatch.setattr(dm, "_PLANE_BLOCK", 32)
+    batch, channels, frames, height, width = 2, 8, 1, 257, 1021
+    plane_size = channels * height * width
+    assert dm.triton.cdiv(plane_size, dm._PLANE_BLOCK) > 65_535
+    assert plane_size % dm._PLANE_BLOCK != 0
+
+    x = _dense_5d((batch, channels, frames, height, width), torch.bfloat16, True)
+    cache = _dense_5d((batch, channels, cache_frames, height, width), x.dtype, True) if cache_frames else None
+    parts = [torch.zeros_like(x).expand(-1, -1, 2 - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2).contiguous(memory_format=torch.channels_last_3d)
+
+    pair = dm.cat_time_5d(x, cache, pad_front=2, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert _bits_equal(assembled, reference)
+    assert assembled.stride() == reference.stride()
+    assert _bits_equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(dm.cat_time_5d(x, cache, pad_front=2), reference)
+
+
+@torch.no_grad()
+def test_cached_conv_spatial_padding_verdict_is_recorded_and_exact() -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 64, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    for layout in (torch.contiguous_format, torch.channels_last_3d):
+        if layout is torch.channels_last_3d:
+            conv.to(memory_format=torch.channels_last_3d)
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        reference_cache = [cache.clone()]
+        reference = WanCausalConv3d.forward(conv, chunk, cache)  # upstream pre-padded path
+        fast_cache = [cache.clone()]
+        out = fp._run_cached_causal_conv(conv, chunk, fast_cache, 0)
+        verdicts = fp._SPATIAL_PAD_VERDICTS[conv]
+        assert len(verdicts) >= 1
+        assert _bits_equal(out, reference)  # exact whichever path the verdict selected
+        again = fp._run_cached_causal_conv(conv, chunk, [cache.clone()], 0)
+        assert _bits_equal(again, reference)
+        assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+        del reference_cache
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize("with_bias", [False, True])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
+    [
+        # Cosmos3 / Wan2.2 block shapes: repeats 8 (same channels), 4 (halved), 2 (spatial only, halved).
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),
+        (64, 32, 2, 10, 14, 1, 2, False),
+        (64, 64, 4, 10, 14, 1, 2, False),
+        # More output than input channels (repeats 8 with a 2x channel increase) and a 4x spatial factor.
+        (32, 64, 2, 10, 14, 1, 2, False),
+        (16, 16, 1, 5, 6, 1, 4, False),
+        # Odd sizes and a row wider than one column block.
+        (12, 12, 3, 7, 9, 2, 2, True),
+        (8, 8, 1, 3, 700, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_is_bitwise_and_layout_exact(
+    dtype,
+    main_layout,
+    with_bias,
+    in_channels,
+    out_channels,
+    frames,
+    height,
+    width,
+    factor_temporal,
+    factor_spatial,
+    drop_first,
+) -> None:
+    torch.manual_seed(0)
+    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
+    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
+    if main_layout == "frame_major":
+        # The previous block's output is frame-major too in the lossless level.
+        source = torch.randn((1, frames, in_channels, height, width), device="cuda", dtype=dtype)
+        source = source.permute(0, 2, 1, 3, 4)
+    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
+    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
+    if main_layout == "channels_last":
+        main = _dense_5d(shape, dtype, True)
+    elif main_layout == "frame_major":
+        main = torch.randn((shape[0], shape[2], shape[1], shape[3], shape[4]), device="cuda", dtype=dtype)
+        main = main.permute(0, 2, 1, 3, 4)
+    else:
+        main = _dense_5d(shape, dtype, False)
+
+    shortcut = source.repeat_interleave(repeats, dim=1)
+    shortcut = shortcut.view(1, out_channels, factor_temporal, factor_spatial, factor_spatial, frames, height, width)
+    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
+    if drop_first:
+        shortcut = shortcut[:, :, factor_temporal - 1 :]
+    bias = None
+    biased_main = main
+    if with_bias:
+        bias = (torch.randn(out_channels, device="cuda") * 0.1).to(dtype)
+        biased_main = main.clone()
+        biased_main.add_(bias.view(1, -1, 1, 1, 1))  # ATen's conv bias add
+    reference = biased_main + shortcut
+
+    output = dm.dup_up3d_add(main, source, factor_temporal, factor_spatial, repeats, drop_first, main_bias=bias)
+    assert output is not None
+    assert _same_dense_strides(output, reference)
+    assert torch.equal(output, reference)
+
+
+def test_dup_up3d_add_declines_unsupported_cases() -> None:
+    main = torch.randn(1, 8, 2, 10, 14, device="cuda", dtype=torch.bfloat16)
+    source = torch.randn(1, 32, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    # repeats (1) below the spatial factor: the source channel would change along a row.
+    assert dm.dup_up3d_add(main, source, 1, 2, 1, False) is None
+    source = torch.randn(1, 8, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False) is not None
+    # Mixed layouts, wrong bias shape, wrong main shape.
+    assert dm.dup_up3d_add(main.contiguous(memory_format=torch.channels_last_3d), source, 1, 2, 4, False) is None
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False, main_bias=torch.ones(7, device="cuda")) is None
+    assert dm.dup_up3d_add(main[:, :, :1], source, 1, 2, 4, False) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "frames,channels,height,width",
+    [
+        (1, 96, 10, 14),
+        (4, 64, 10, 14),
+        # Non-power-of-two channels and ragged rows/columns for every block size.
+        (3, 12, 7, 9),
+        # Width wider than one column block so the in-kernel loop runs.
+        (2, 8, 3, 1300),
+        # Cosmos3 720p stage-3 input, scaled down in height.
+        (4, 256, 8, 640),
+        # Channels equal to and above the channels-last block: in-kernel channel loop.
+        (2, 1024, 5, 160),
+        (1, 1100, 3, 20),
+    ],
+)
+def test_upsample_nearest_2x_is_bitwise_and_layout_exact(dtype, layout, frames, channels, height, width) -> None:
+    torch.manual_seed(0)
+    if layout == "frame_major":
+        # The strided ``(t, c, h, w)`` view ``WanResample`` hands the upsampler in the lossless level.
+        x = fp._merge_batch_and_frames(torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype))
+        assert frames == 1 or not x.is_contiguous()
+    elif layout == "channels_last":
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+        x = x.contiguous(memory_format=torch.channels_last)
+    else:
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+    for mode in ("nearest-exact", "nearest"):
+        reference = F.interpolate(x, scale_factor=(2.0, 2.0), mode=mode)
+        output = up.upsample_nearest_2x(x)
+        assert output is not None
+        assert output.stride() == reference.stride()
+        assert torch.equal(output, reference)
+
+
+def test_upsample_nearest_2x_declines_unsupported_inputs() -> None:
+    x = torch.randn(2, 8, 6, 10, device="cuda", dtype=torch.bfloat16)
+    assert up.upsample_nearest_2x(x) is not None
+    assert up.upsample_nearest_2x(x.cpu()) is None
+    assert up.upsample_nearest_2x(x.unsqueeze(2)) is None
+    assert up.upsample_nearest_2x(x.transpose(2, 3)) is None
+    assert up.upsample_nearest_2x(x[..., ::2]) is None
+    assert up.upsample_nearest_2x(x.double()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_upsample_forward_matches_wan_upsample(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    x = torch.randn(3, 64, 10, 14, device="cuda", dtype=dtype)
+    for layout in (torch.contiguous_format, torch.channels_last):
+        x = x.contiguous(memory_format=layout)
+        module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+        expected = WanUpsample.forward(module, x)
+        actual = fp.upsample_forward(module, x)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+        # Anything but a nearest 2x upsample takes the reference path.
+        other = WanUpsample(scale_factor=(3.0, 3.0), mode="nearest-exact")
+        assert torch.equal(fp.upsample_forward(other, x), WanUpsample.forward(other, x))
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_forward_can_leave_conv_bias_pending(dtype: torch.dtype, mode: str) -> None:
+    torch.manual_seed(0)
+    resample = WanResample(32, mode=mode).eval().to(device="cuda", dtype=dtype)
+    setattr(resample, fp.CFG_ATTR, fp.FastPathConfig())
+    x = torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)
+    cache = [torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)] if mode == "upsample3d" else [None]
+    expected = WanResample.forward(
+        resample, x, feat_cache=[c.clone() if c is not None else None for c in cache], feat_idx=[0]
+    )
+    out, bias = fp.resample_forward(
+        resample, x, [c.clone() if c is not None else None for c in cache], [0], return_bias=True
+    )
+    assert bias is resample.resample[1].bias
+    assert out.stride() == expected.stride()
+    assert _bits_equal(out + bias.view(1, -1, 1, 1, 1), expected)
+    plain = fp.resample_forward(resample, x, [c.clone() if c is not None else None for c in cache], [0])
+    assert _bits_equal(plain, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("force_prepad", [False, True])
+def test_cached_conv_deferred_bias_matches_autocast(dtype: torch.dtype, force_prepad: bool, monkeypatch) -> None:
+    """FP32 parameters and inputs must defer the bias in the convolution's output dtype."""
+    torch.manual_seed(8)
+    conv = WanCausalConv3d(32, 32, 3, padding=1).cuda().eval()
+    conv.bias.copy_(torch.randn_like(conv.bias))
+    if force_prepad:
+        monkeypatch.setattr(dm, "cat_time_5d", lambda *args, **kwargs: None)
+    cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+            expected = conv(x, cache[0])
+            out, bias = fp._run_cached_causal_conv(conv, x, cache, 0, return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = dm.add_bias_residual(out, bias, torch.zeros_like(out), None)
+            assert actual is not None
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("channels_last", [False, True])
+def test_residual_block_deferred_bias_matches_autocast(dtype: torch.dtype, channels_last: bool) -> None:
+    """Exercise conv1/conv2 and a non-identity shortcut with FP32 weights under autocast."""
+    torch.manual_seed(8)
+    block = WanResidualBlock(32, 64, dropout=0.0).cuda().eval()
+    for conv in (block.conv1, block.conv2, block.conv_shortcut):
+        conv.bias.copy_(torch.randn_like(conv.bias))
+        if channels_last:
+            conv.to(memory_format=torch.channels_last_3d)
+    x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+    with torch.autocast("cuda", dtype=dtype):
+        expected = block(x, feat_cache=[None, None], feat_idx=[0])
+        setattr(block, fp.CFG_ATTR, fp.FastPathConfig(channels_last=channels_last))
+        actual = fp.residual_block_forward(block, x, feat_cache=[None, None], feat_idx=[0])
+    assert actual.dtype == expected.dtype == dtype
+    assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_deferred_bias_matches_autocast(dtype: torch.dtype, mode: str) -> None:
+    """The upsampler's FP32 bias must be rounded before its deferred addition."""
+    torch.manual_seed(8)
+    resample = WanResample(32, mode=mode).cuda().eval()
+    resample.resample[1].bias.copy_(torch.randn_like(resample.resample[1].bias))
+    reference_cache = [None]
+    candidate_cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 1, 6, 10, device="cuda")
+            expected = resample(x, feat_cache=reference_cache, feat_idx=[0])
+            out, bias = fp.resample_forward(resample, x, candidate_cache, [0], return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = out + bias.view(1, -1, 1, 1, 1)
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("with_bias_x", [True, False])
+@pytest.mark.parametrize("with_bias_h", [True, False])
+def test_add_bias_residual_matches_aten_rounding(dtype: torch.dtype, with_bias_x: bool, with_bias_h: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 45, 80)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    h = torch.randn(shape, device="cuda").to(dtype)
+    bias_x = torch.randn(256, device="cuda").to(dtype) if with_bias_x else None
+    bias_h = torch.randn(256, device="cuda").to(dtype) if with_bias_h else None
+    expected_x = x.clone()
+    expected_h = h.clone()
+    if bias_x is not None:
+        expected_x.add_(bias_x.view(1, -1, 1, 1, 1))
+    if bias_h is not None:
+        expected_h.add_(bias_h.view(1, -1, 1, 1, 1))
+    expected = expected_x + expected_h
+    actual = dm.add_bias_residual(x, bias_x, h, bias_h)
+    assert actual is not None
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    h_cl = h.contiguous(memory_format=torch.channels_last_3d)
+    actual_cl = dm.add_bias_residual(x_cl, bias_x, h_cl, bias_h)
+    assert actual_cl is not None and actual_cl.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(actual_cl.contiguous(), expected)
+    assert dm.add_bias_residual(x_cl, bias_x, h, bias_h) is None
+    assert dm.add_bias_residual(x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4), bias_x, h, bias_h) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize("temporal_only", [False, True])
+@pytest.mark.parametrize("start", [None, "Rep"])
+def test_cached_conv_chunk_loop_matches_reference_including_mixed_paths(
+    monkeypatch, channels_last: bool, temporal_only: bool, start
+) -> None:
+    torch.manual_seed(0)
+    channels = 64
+    if temporal_only:
+        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(channels, channels, 3, padding=1)
+    conv = conv.to(device="cuda", dtype=torch.float32)
+    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
+
+    decline = lambda *a, **k: None  # noqa: E731
+    original_time, original_pad = fp.dm.cat_time_5d, fp.dm.cat_pad_5d
+
+    def run(mode: list[str]):
+        cache = [start]
+        outputs = []
+        for chunk, m in zip(chunks, mode, strict=True):
+            monkeypatch.setattr(fp.dm, "cat_time_5d", decline if m in ("reference", "padded") else original_time)
+            monkeypatch.setattr(fp.dm, "cat_pad_5d", decline if m == "reference" else original_pad)
+            outputs.append(fp._run_cached_causal_conv(conv, chunk, cache, 0))
+        monkeypatch.setattr(fp.dm, "cat_time_5d", original_time)
+        monkeypatch.setattr(fp.dm, "cat_pad_5d", original_pad)
+        return outputs
+
+    reference = run(["reference"] * 4)
+    for pattern in (
+        ["time"] * 4,
+        ["padded"] * 4,
+        ["time", "reference", "padded", "time"],
+        ["reference", "padded", "time", "reference"],
+    ):
+        fused = run(pattern)
+        for out, ref in zip(fused, reference, strict=True):
+            assert torch.equal(out, ref), pattern
+
+
+# --------------------------------------------------------------------------- #
+# Whole decoder
+# --------------------------------------------------------------------------- #
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=16,
+    z_dim=8,
+    dim_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    temperal_downsample=[False, True, True],
+    is_residual=True,
+    patch_size=2,
+    in_channels=12,
+    out_channels=12,
+)
+TINY_WAN21 = dict(base_dim=8, z_dim=4, dim_mult=[1, 2], num_res_blocks=1, temperal_downsample=[False, True])
+
+
+def _build_pair(config: dict, dtype: torch.dtype):
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("config", [TINY_RESIDUAL, TINY_WAN21], ids=["residual_patch2", "wan21"])
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("frames", [1, 2, 5])
+def test_lossless_decoder_is_bitwise_exact_on_cuda(config: dict, dtype: torch.dtype, frames: int) -> None:
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 10, device="cuda").to(dtype)
+    autocast = dtype is not torch.float32
+    with torch.autocast("cuda", dtype=dtype, enabled=autocast):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_omni_wan_vae_decode_is_bitwise_exact_on_cuda(dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 5, 6, 10, device="cuda").to(dtype)
+    with torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32):
+        expected = reference.decode(latents, return_dict=False)[0]
+    actual = candidate.decode(latents, return_dict=False)[0]  # applies its own autocast context
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    ("dtype", "allow_tf32", "tolerance"),
+    [
+        pytest.param(torch.bfloat16, True, 3e-2, id="bf16"),
+        # With TF32 disabled the only difference is cuDNN's channels-last algorithm
+        # choice and the norm kernel's reduction order.
+        pytest.param(torch.float32, False, 2e-4, id="fp32"),
+        # PyTorch's default allows cuDNN to run fp32 convolutions in TF32; the
+        # channels-last algorithms use it while the NCDHW reference may not, so
+        # the deviation is dominated by TF32's 10-bit mantissa.
+        pytest.param(torch.float32, True, 1e-2, id="fp32-tf32"),
+    ],
+)
+def test_channels_last_decoder_keeps_layout_and_stays_close(
+    dtype: torch.dtype, allow_tf32: bool, tolerance: float
+) -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, dtype)
+    report = install_wan_vae_fastpath(candidate, level="channels_last")
+    assert report.installed and report.channels_last
+
+    violations: list[str] = []
+
+    def audit(name: str):
+        def hook(module: nn.Module, inputs) -> None:
+            x = inputs[0]
+            if x.dim() == 5 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last_3d):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+            if x.dim() == 4 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+
+        return hook
+
+    handles = [
+        module.register_forward_pre_hook(audit(name))
+        for name, module in candidate.decoder.named_modules()
+        if isinstance(module, (nn.Conv2d, nn.Conv3d))
+    ]
+    torch.manual_seed(1)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 3, 6, 10, device="cuda").to(dtype)
+    with (
+        torch.backends.cudnn.flags(enabled=True, allow_tf32=allow_tf32),
+        torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32),
+    ):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    for handle in handles:
+        handle.remove()
+    assert not violations, "\n".join(violations)
+    assert actual.shape == expected.shape
+    difference = (actual.float() - expected.float()).abs()
+    error = difference.max().item()
+    mse = (difference**2).mean().item()
+    psnr = float("inf") if mse == 0 else 10 * torch.log10(torch.tensor(4.0 / mse)).item()
+    assert error <= tolerance, f"max_abs_diff={error:.3e} psnr={psnr:.2f} dB tolerance={tolerance}"

tests/diffusion/distributed/wan_vae_fastpath_helpers.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Isolate the upstream numerical reference from Wan's global norm patch."""
+
+from __future__ import annotations
+
+import importlib.util
+
+import pytest
+from diffusers.models.autoencoders import autoencoder_kl_wan
+from torch import nn
+
+
+@pytest.fixture(scope="session")
+def original_wan_rms_norm() -> type[nn.Module]:
+    # Wan's package import replaces the public class, including aliases held
+    # by already collected tests. Load the installed upstream source under a
+    # private name instead of reloading (and changing) the live module/classes.
+    spec = importlib.util.spec_from_file_location(
+        "diffusers.models.autoencoders._wan_fastpath_reference", autoencoder_kl_wan.__file__
+    )
+    assert spec is not None and spec.loader is not None
+    reference = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(reference)
+    return reference.WanRMS_norm
+
+
+@pytest.fixture
+def unpatched_wan_rms_norm(monkeypatch, request, original_wan_rms_norm):
+    # Import the patching package before temporarily restoring the reference,
+    # so importing RMSNormVAE in a test cannot change it halfway through.
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
+
+    assert original_wan_rms_norm is not RMSNormVAE
+    monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", original_wan_rms_norm)
+    if hasattr(request.module, "WanRMS_norm"):
+        monkeypatch.setattr(request.module, "WanRMS_norm", original_wan_rms_norm)

tests/diffusion/lora/test_lora_manager.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.lora.lora_model import LoRAModel
 from vllm.lora.lora_weights import LoRALayerWeights
 from vllm.lora.utils import get_supported_lora_modules
 
@@ -22,6 +23,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyLoRALayer:
+    base_layer: torch.nn.Module
+
     def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
         self.n_slices = n_slices
         self.output_slices = output_slices
@@ -292,7 +295,9 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
     assert torch.allclose(torch.cat([b0, b1, b2], dim=0), B * 0.5)
 
 
-def test_lora_manager_activates_packed_lora_from_sublayers():
+@pytest.mark.parametrize("name_prefix", ["transformer.blocks.0.attn.", "blocks.0.attn.", ""])
+@pytest.mark.parametrize("target_modules", [("to_q",), ("to_q", "to_v"), ("to_q", "to_k", "to_v")])
+def test_lora_manager_activates_packed_lora_from_sublayers(name_prefix, target_modules):
     pipeline = torch.nn.Module()
     bound_names = None
 
@@ -320,8 +325,10 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     rank = 2
     loras: dict[str, LoRALayerWeights] = {}
     for name, out_dim in zip(["to_q", "to_k", "to_v"], [2, 1, 1]):
-        loras[f"transformer.blocks.0.attn.{name}"] = LoRALayerWeights(
-            module_name=f"transformer.blocks.0.attn.{name}",
+        if name not in target_modules:
+            continue
+        loras[f"{name_prefix}{name}"] = LoRALayerWeights(
+            module_name=f"{name_prefix}{name}",
             rank=rank,
             lora_alpha=rank,
             lora_a=torch.ones((rank, 4)) * (1 if name == "to_q" else 2),
@@ -342,9 +349,15 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     assert len(lora_a_list) == 3
     assert len(lora_b_list) == 3
     # Scale should apply to B only.
-    assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
-    assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
-    assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+    for index, name in enumerate(("to_q", "to_k", "to_v")):
+        if name in target_modules:
+            weights = loras[f"{name_prefix}{name}"]
+            assert torch.equal(lora_a_list[index], weights.lora_a)
+            assert torch.equal(lora_b_list[index], weights.lora_b * 2.0)
+        else:
+            assert lora_a_list[index] is None
+            assert lora_b_list[index] is None
+    assert manager._active_adapter_id == 1
     assert bound_names == frozenset(loras)
 
 
@@ -421,6 +434,7 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer():
     )
     layer = _DummyLoRALayer(n_slices=1, output_slices=(2,))
     manager._lora_modules = {"transformer.blocks.0.attn.to_q": layer}
+    manager._expected_lora_modules = {"to_q"}
 
     # Adapter-side name the engine does not expose, e.g. a diffusers-style
     # checkpoint against a differently named engine layout.
@@ -443,11 +457,13 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer():
         )()
     }
 
-    with pytest.raises(ValueError, match="applies to no layer") as excinfo:
+    with pytest.raises(ValueError, match="binding is incomplete") as excinfo:
         manager._activate_adapter(3, scale=1.0)
 
     # The message must name what was received so the mismatch is diagnosable.
+    assert "bound=0/1" in str(excinfo.value)
     assert "unet.down_blocks.0.attn.to_q" in str(excinfo.value)
+    assert "expected target modules in ['to_q']" in str(excinfo.value)
 
     # Nothing was bound and the adapter must not be left marked active.
     assert manager._active_adapter_id is None
@@ -455,6 +471,97 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer():
     assert layer.reset_calls >= 1
 
 
+def test_lora_manager_rejects_empty_adapter():
+    """An empty adapter must not be considered successfully bound."""
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+    layer = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+    manager._lora_modules = {"transformer.attn.to_q": layer}
+    manager._registered_adapters = {1: LoRAModel(1, rank=2, loras={})}
+
+    with pytest.raises(ValueError, match="bound=0/0"):
+        manager._activate_adapter(1, scale=1.0)
+
+    assert manager._active_adapter_id is None
+    assert manager._suspended_adapter_id is None
+    assert 1 not in manager._adapter_scales
+    assert layer.set_calls == []
+    assert layer.active_slices == ()
+    assert layer.reset_calls >= 1
+
+
+@pytest.mark.parametrize("suspend_previous", [False, True])
+@pytest.mark.parametrize("unbound_name", ["transformer.attn.to_out.0", "transformer.attn.to_qkv"])
+def test_lora_manager_rejects_partial_binding_and_rolls_back(unbound_name, suspend_previous):
+    """A partial bind must clear uploaded weights and invalidate fast paths."""
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+    layer = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
+    manager._lora_modules = {
+        "transformer.attn.to_out": layer,
+        "transformer.attn.to_qkv": packed_layer,
+    }
+    manager._expected_lora_modules = {"to_out", "to_qkv"}
+
+    def weights(name, value):
+        return LoRALayerWeights(
+            module_name=name,
+            rank=2,
+            lora_alpha=2,
+            lora_a=torch.full((2, 2), value),
+            lora_b=torch.full((2, 2), value),
+        )
+
+    valid_name = "transformer.attn.to_out"
+    previous = LoRAModel(1, rank=2, loras={valid_name: weights(valid_name, 1.0)})
+    # The extra weights either use an unmatched checkpoint name or have a B
+    # shape that cannot be split across the fused layer's four output rows.
+    partial = LoRAModel(
+        2,
+        rank=2,
+        loras={valid_name: weights(valid_name, 2.0), unbound_name: weights(unbound_name, 2.0)},
+    )
+    manager._registered_adapters = {1: previous, 2: partial}
+    manager._activate_adapter(1, scale=0.5)
+    if suspend_previous:
+        manager._deactivate_all_adapters()
+
+    with pytest.raises(ValueError, match="binding is incomplete") as excinfo:
+        manager._activate_adapter(2, scale=1.0)
+
+    assert "LoRA adapter 2" in str(excinfo.value)
+    assert unbound_name in str(excinfo.value)
+    assert valid_name not in str(excinfo.value).replace(unbound_name, "")
+    if unbound_name.endswith("to_qkv"):
+        assert "lora_b.shape[0]=2" in str(excinfo.value)
+        assert "sum(output_slices)=4" in str(excinfo.value)
+        assert "output_slices=(2, 1, 1)" in str(excinfo.value)
+    else:
+        assert "expected target modules in ['to_out', 'to_qkv']" in str(excinfo.value)
+    assert len(layer.set_calls) == 2
+    assert layer.reset_calls == 1
+    assert layer.active_slices == ()
+    assert packed_layer.active_slices == ()
+    assert manager._active_adapter_id is None
+    assert manager._suspended_adapter_id is None
+    assert 2 not in manager._adapter_scales
+
+    manager._activate_adapter(1, scale=0.5)
+    assert len(layer.set_calls) == 3
+    assert layer.resume_calls == 0
+    assert manager._active_adapter_id == 1
+    assert torch.equal(layer.set_calls[-1][1], previous.loras[valid_name].lora_b * 0.5)
+
+
 def _dummy_lora_request(adapter_id: int) -> LoRARequest:
     return LoRARequest(
         lora_name=f"adapter_{adapter_id}",
@@ -550,6 +657,146 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
     assert set(manager.list_adapters()) == {1, 2}
 
 
+def _make_hunyuan_image3_pipeline(num_heads=4, num_kv_heads=2, head_dim=2, hidden_size=8):
+    """Use production hooks and the production base-weight splitter without
+    allocating the full checkpoint or initializing distributed workers.
+    """
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import HunyuanImage3Model
+    from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import HunyuanImage3Pipeline
+
+    model = HunyuanImage3Model.__new__(HunyuanImage3Model)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(
+        num_attention_heads=num_heads,
+        num_key_value_heads=num_kv_heads,
+        attention_head_dim=head_dim,
+        hidden_size=hidden_size,
+    )
+    pipeline = HunyuanImage3Pipeline.__new__(HunyuanImage3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.model = model
+    pipeline.transformer = model
+    return pipeline
+
+
+def _make_hunyuan_image3_qkv_base():
+    """An un-initialized QKVParallelLinear instance for isinstance() checks.
+
+    Supply the type and checkpoint head metadata for binding-only tests.
+    The PEFT forward test below constructs the real distributed projection.
+    """
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    layer = object.__new__(QKVParallelLinear)
+    layer.total_num_heads = 4
+    layer.total_num_kv_heads = 2
+    layer.head_size = layer.v_head_size = 2
+    return layer
+
+
+def test_lora_manager_activates_hunyuan_image3_fused_qkv_lora():
+    """A fused HI3 ``qkv_proj`` LoRA-B (GQA-interleaved in the adapter) is
+    de-interleaved to the block [Q; K; V] layout the QKV output slices expect,
+    with the namespace alias resolving the PEFT adapter tensors."""
+    manager = DiffusionLoRAManager(
+        pipeline=_make_hunyuan_image3_pipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    wrapper = "transformer.layers.0.self_attn.qkv_proj"
+    packed_layer = _DummyLoRALayer(
+        n_slices=3,
+        output_slices=(8, 4, 4),
+        tp_size=1,
+    )
+    # Use QKV type identity and checkpoint head metadata for the manager's
+    # de-interleave branch; the forward test below constructs real layers.
+    packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
+    manager._lora_modules = {wrapper: packed_layer}
+
+    rank = 3
+    # GQA-interleaved fused QKV rows: [Q-group0, K0, V0, Q-group1, K1, V1].
+    interleaved = torch.arange(16, dtype=torch.bfloat16).unsqueeze(1).repeat(1, rank)
+    lora = LoRALayerWeights(
+        module_name="model.layers.0.self_attn.qkv_proj",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=torch.ones((rank, 8)),
+        lora_b=interleaved,
+    )
+    manager._registered_adapters = {
+        11: type(
+            "LM",
+            (),
+            {
+                "id": 11,
+                "loras": {"model.layers.0.self_attn.qkv_proj": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+
+    manager._activate_adapter(11, 0.5)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert len(lora_a_list) == 3 and len(lora_b_list) == 3
+    # De-interleaved full [Q; K; V] slices, scaled by the lora scale.
+    expected_q = torch.cat([interleaved[0:4], interleaved[8:12]], dim=0) * 0.5
+    expected_k = torch.cat([interleaved[4:6], interleaved[12:14]], dim=0) * 0.5
+    expected_v = torch.cat([interleaved[6:8], interleaved[14:16]], dim=0) * 0.5
+    assert torch.equal(lora_b_list[0], expected_q)
+    assert torch.equal(lora_b_list[1], expected_k)
+    assert torch.equal(lora_b_list[2], expected_v)
+
+
+def test_lora_manager_rejects_hunyuan_image3_qkv_lora_with_unexpected_rows():
+    """If the fused HI3 QKV layout cannot be established (wrong row count), the
+    manager must fail closed instead of applying the LoRA to wrong slices."""
+    manager = DiffusionLoRAManager(
+        pipeline=_make_hunyuan_image3_pipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    wrapper = "transformer.layers.0.self_attn.qkv_proj"
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(8, 4, 4))
+    packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
+    manager._lora_modules = {wrapper: packed_layer}
+
+    rank = 3
+    # Row count (12) does not match the HI3 QKV layout (16): fail closed.
+    lora = LoRALayerWeights(
+        module_name="model.layers.0.self_attn.qkv_proj",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=torch.ones((rank, 8)),
+        lora_b=torch.arange(12 * rank, dtype=torch.bfloat16).view(-1, rank),
+    )
+    manager._registered_adapters = {
+        12: type(
+            "LM",
+            (),
+            {
+                "id": 12,
+                "loras": {"model.layers.0.self_attn.qkv_proj": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+
+    with pytest.raises(ValueError, match="cannot establish HunyuanImage-3 fused-QKV layout"):
+        manager._activate_adapter(12, 0.5)
+
+    assert packed_layer.set_calls == []
+    assert packed_layer.reset_calls >= 1
+    assert manager._active_adapter_id is None
+
+
 def test_lora_manager_applies_multiple_scales_correctly(monkeypatch):
     """Ensure that the LoRA manager applies scales correctly when the
     active adapter receives a different scale, i.e., the rank is unchanged.
@@ -894,3 +1141,96 @@ def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
     assert manager._suspended_adapter_id is None
     assert layer.reset_calls == 1, "the upload must be torn down"
     assert layer.suspended_slices is None
+
+
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+def test_hunyuan_image3_peft_qkv_forward(tmp_path, monkeypatch, tp_size):
+    """Disk PEFT -> real manager/hooks -> vLLM sharding -> numerical forward.
+
+    TP4 also covers replicated KV heads. Only distributed rank discovery is
+    stubbed: projection construction, adapter loading and all LoRA operations
+    execute production code. The oracle indexes checkpoint rows explicitly.
+    """
+    import json
+
+    from safetensors.torch import save_file
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+
+    adapter_dir = tmp_path / "adapter"
+    adapter_dir.mkdir()
+    rank = 2
+    a = torch.arange(16, dtype=torch.float32).reshape(rank, 8) / 16
+    b = torch.arange(32, dtype=torch.float32).reshape(16, rank) / 32
+    prefix = "base_model.model.model.layers.0.self_attn.qkv_proj"
+    save_file(
+        {f"{prefix}.lora_A.weight": a, f"{prefix}.lora_B.weight": b},
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": rank, "lora_alpha": rank, "target_modules": ["qkv_proj"]})
+    )
+    request = LoRARequest(lora_name="hi3", lora_int_id=6411, lora_path=str(adapter_dir))
+    x = torch.arange(24, dtype=torch.float32).reshape(3, 8) / 24
+    q_rows = [0, 1, 2, 3, 8, 9, 10, 11]
+    k_rows = [4, 5, 12, 13]
+    v_rows = [6, 7, 14, 15]
+    for tp_rank in range(tp_size):
+        monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: tp_size)
+        monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: tp_rank)
+        monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: tp_rank)
+        monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: tp_size)
+        with set_current_vllm_config(VllmConfig()):
+            pipeline = _make_hunyuan_image3_pipeline()
+            base = linear.QKVParallelLinear(8, 2, 4, 2, bias=False, params_dtype=torch.float32)
+            base.weight.data.zero_()
+            block = torch.nn.Module()
+            block.self_attn = torch.nn.Module()
+            block.self_attn.qkv_proj = base
+            pipeline.model.layers = torch.nn.ModuleList([block])
+            manager = DiffusionLoRAManager(pipeline, torch.device("cpu"), torch.float32)
+            manager.set_active_adapter(request, lora_scale=0.5)
+            wrapped = pipeline.model.layers[0].self_attn.qkv_proj
+            assert wrapped is not base
+            q_size = 8 // tp_size
+            kv_size = 4 // min(tp_size, 2)
+            kv_rank = tp_rank // max(tp_size // 2, 1)
+            rows = (
+                q_rows[tp_rank * q_size : (tp_rank + 1) * q_size]
+                + k_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
+                + v_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
+            )
+            expected = (x @ a.T) @ b[rows].T * 0.5
+            with torch.inference_mode():
+                torch.testing.assert_close(wrapped.apply(x), expected)
+                manager.set_active_adapter(None)
+                torch.testing.assert_close(wrapped.apply(x), torch.zeros_like(expected))
+                manager.set_active_adapter(request, lora_scale=0.5)
+                torch.testing.assert_close(wrapped.apply(x), expected)
+
+
+@pytest.mark.parametrize("projection", ["qkv_proj", "o_proj"])
+def test_hunyuan_image3_lora_namespace_is_model_scoped(projection):
+    from vllm.lora.lora_model import LoRAModel
+
+    name = f"model.layers.0.self_attn.{projection}"
+    weights = LoRALayerWeights(name, 2, 2, torch.ones(2, 8), torch.ones(8, 2))
+    adapter = LoRAModel(6411, 2, {name: weights})
+    wrapper = f"transformer.layers.0.self_attn.{projection}"
+    hi3 = DiffusionLoRAManager(_make_hunyuan_image3_pipeline(), torch.device("cpu"), torch.float32)
+    other = DiffusionLoRAManager(_DummyPipeline(), torch.device("cpu"), torch.float32)
+    assert hi3._get_lora_weights(adapter, wrapper) is weights
+    assert other._get_lora_weights(adapter, wrapper) is None
+
+
+@pytest.mark.parametrize("head_dim_key", ["head_dim", "attention_head_dim", None])
+def test_hunyuan_image3_real_qkv_hook_head_dimension(head_dim_key):
+    pipeline = _make_hunyuan_image3_pipeline()
+    del pipeline.model.config.attention_head_dim
+    if head_dim_key is not None:
+        setattr(pipeline.model.config, head_dim_key, 2)
+    b = torch.arange(32, dtype=torch.float32).reshape(16, 2)
+    expected_rows = [0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15]
+    torch.testing.assert_close(pipeline._deinterleave_fused_qkv_lora_b(b), b[expected_rows])
+    assert pipeline._deinterleave_fused_qkv_lora_b(b[:-1]) is None

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -75,6 +75,248 @@ def _tiny_cosmos3_edge_config(**overrides):
     return config
 
 
+@pytest.mark.parametrize(
+    ("config_kind", "expected_language", "expected_gen"),
+    [
+        ("flat", "transformer", "transformer"),
+        ("transformer", "transformer", "transformer"),
+        ("other_component", None, None),
+        ("language_model", "language", None),
+        ("gen_layers", None, "gen"),
+        ("transformer_language_override", "language", "transformer"),
+        ("transformer_gen_disabled", "transformer", None),
+        ("independent_subcomponents", "language", "gen"),
+    ],
+)
+def test_transformer_resolves_global_and_subcomponent_quant_configs(
+    monkeypatch: pytest.MonkeyPatch,
+    config_kind: str,
+    expected_language: str | None,
+    expected_gen: str | None,
+) -> None:
+    """Resolve pipeline-level defaults and Cosmos3 subcomponent overrides."""
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    received: list[tuple[str, object]] = []
+    configs = {
+        "transformer": object(),
+        "language": object(),
+        "gen": object(),
+    }
+
+    class _StubLanguageModel(nn.Module):
+        def __init__(self, *, quant_config, **kwargs) -> None:
+            del kwargs
+            super().__init__()
+            received.append(("language_model", quant_config))
+            self.layers = nn.ModuleList()
+
+    class _StubGenDecoderLayer(nn.Module):
+        def __init__(self, *, quant_config, **kwargs) -> None:
+            del kwargs
+            super().__init__()
+            received.append(("gen_layers", quant_config))
+
+    monkeypatch.setattr(Cosmos3VFMTransformer, "_language_model_cls", _StubLanguageModel)
+    monkeypatch.setattr(transformer_cosmos3, "Cosmos3GenDecoderLayer", _StubGenDecoderLayer)
+
+    if config_kind == "flat":
+        top_level_config = configs["transformer"]
+    elif config_kind == "transformer":
+        top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"]})
+    elif config_kind == "other_component":
+        top_level_config = ComponentQuantizationConfig({"vae": configs["transformer"]})
+    elif config_kind == "language_model":
+        top_level_config = ComponentQuantizationConfig({"language_model": configs["language"]})
+    elif config_kind == "gen_layers":
+        top_level_config = ComponentQuantizationConfig({"gen_layers": configs["gen"]})
+    elif config_kind == "transformer_language_override":
+        top_level_config = ComponentQuantizationConfig(
+            {"transformer": configs["transformer"], "language_model": configs["language"]}
+        )
+    elif config_kind == "transformer_gen_disabled":
+        top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"], "gen_layers": None})
+    else:
+        top_level_config = ComponentQuantizationConfig(
+            {"language_model": configs["language"], "gen_layers": configs["gen"]}
+        )
+
+    Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
+            dtype=torch.float32,
+            quantization_config=top_level_config,
+            custom_pipeline_args={},
+            model_config={},
+        )
+    )
+
+    assert received == [
+        ("language_model", None if expected_language is None else configs[expected_language]),
+        ("gen_layers", None if expected_gen is None else configs[expected_gen]),
+    ]
+
+
+def test_transformer_component_quant_reaches_real_linear(monkeypatch: pytest.MonkeyPatch) -> None:
+    """#7323 regression: unwrap ComponentQuantizationConfig before ColumnParallelLinear.
+
+    Passing the router through unchanged makes ``resolve(language_model.*)`` miss
+    a ``transformer``-only map and raise ``ValueError: All linear layers should
+    support quant method``. A leaf config must reach the real linear path after
+    resolution. Attention backends are stubbed so this stays CPU-safe; MLP and
+    QKV linears remain real ``ColumnParallelLinear`` instances.
+    """
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3GatedMLP,
+        Cosmos3VFMTransformer,
+        _resolve_cosmos3_quant_configs,
+    )
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    class _LeafQuantConfig(QuantizationConfig):
+        def get_name(self) -> str:
+            return "test_leaf"
+
+        def get_quant_method(self, layer, prefix):  # noqa: ANN001
+            del layer, prefix
+            return UnquantizedLinearMethod()
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        def get_min_capability(self) -> int:
+            return 0
+
+        @classmethod
+        def from_config(cls, config):  # noqa: ANN001
+            raise NotImplementedError
+
+        def get_config_filenames(self) -> list[str]:
+            return []
+
+    class _StubFrameworkAttention(nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            del args, kwargs
+            super().__init__()
+
+    monkeypatch.setattr(transformer_cosmos3, "FrameworkAttention", _StubFrameworkAttention)
+
+    leaf = _LeafQuantConfig()
+    router = ComponentQuantizationConfig({"transformer": leaf})
+
+    with pytest.raises(ValueError, match="All linear layers should support quant method"):
+        Cosmos3GatedMLP(
+            hidden_size=8,
+            intermediate_size=16,
+            quant_config=router,
+            prefix="language_model.layers.0.mlp",
+        )
+
+    language_qc, gen_qc = _resolve_cosmos3_quant_configs(router)
+    assert language_qc is leaf
+    assert gen_qc is leaf
+
+    model = Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
+            dtype=torch.float32,
+            quantization_config=router,
+            custom_pipeline_args={},
+            model_config={},
+        )
+    )
+    mlp = model.language_model.layers[0].mlp
+    assert isinstance(mlp.gate_proj.quant_method, UnquantizedLinearMethod)
+    assert isinstance(model.gen_layers[0].mlp.gate_proj.quant_method, UnquantizedLinearMethod)
+    assert isinstance(
+        model.language_model.layers[0].self_attn.to_q.quant_method,
+        UnquantizedLinearMethod,
+    )
+
+
+def test_transformer_pathway_default_keeps_nested_layer_override() -> None:
+    """Pathway leaf plus a more-specific key must still longest-prefix match."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3GatedMLP,
+        _resolve_cosmos3_quant_configs,
+    )
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    class _TaggedLinearMethod(UnquantizedLinearMethod):
+        def __init__(self, tag: str) -> None:
+            super().__init__()
+            self.tag = tag
+
+    class _TaggedQuantConfig(QuantizationConfig):
+        def __init__(self, tag: str) -> None:
+            super().__init__()
+            self.tag = tag
+
+        def get_name(self) -> str:
+            return self.tag
+
+        def get_quant_method(self, layer, prefix):  # noqa: ANN001
+            del layer, prefix
+            return _TaggedLinearMethod(self.tag)
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        def get_min_capability(self) -> int:
+            return 0
+
+        @classmethod
+        def from_config(cls, config):  # noqa: ANN001
+            raise NotImplementedError
+
+        def get_config_filenames(self) -> list[str]:
+            return []
+
+    pathway = _TaggedQuantConfig("pathway")
+    mlp_override = _TaggedQuantConfig("mlp")
+    language_qc, gen_qc = _resolve_cosmos3_quant_configs(
+        ComponentQuantizationConfig(
+            {
+                "language_model": pathway,
+                "gen_layers": pathway,
+                "language_model.layers.0.mlp": mlp_override,
+            }
+        )
+    )
+
+    assert isinstance(language_qc, ComponentQuantizationConfig)
+    assert language_qc.resolve("language_model.layers.0.mlp.gate_proj") is mlp_override
+    assert language_qc.resolve("language_model.layers.0.self_attn.to_q") is pathway
+    assert language_qc.resolve("language_model.layers.1.mlp.gate_proj") is pathway
+    assert gen_qc is pathway
+
+    mlp0 = Cosmos3GatedMLP(
+        hidden_size=8,
+        intermediate_size=16,
+        quant_config=language_qc,
+        prefix="language_model.layers.0.mlp",
+    )
+    mlp1 = Cosmos3GatedMLP(
+        hidden_size=8,
+        intermediate_size=16,
+        quant_config=language_qc,
+        prefix="language_model.layers.1.mlp",
+    )
+    assert mlp0.gate_proj.quant_method.tag == "mlp"
+    assert mlp1.gate_proj.quant_method.tag == "pathway"
+
+
 def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
         compute_mrope_position_ids_action,

tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py

@@ -89,6 +89,23 @@ def test_launch_failure_disables_runtime_key(monkeypatch):
     assert len(fused_rope._FAILED_KEYS) == 1
 
 
+def test_non_cuda_platform_does_not_launch_fused_kernel(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    monkeypatch.setattr(fused_rope, "HAS_TRITON", True)
+    monkeypatch.setattr(fused_rope.current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+
+    def unexpected_launch(*args):
+        raise AssertionError("non-CUDA platforms must remain on the native path")
+
+    monkeypatch.setattr(fused_rope, "_launch_fused_qk_rotary_emb", unexpected_launch)
+    query = torch.zeros((1, 1, 1, 128), dtype=torch.bfloat16)
+    freqs = torch.zeros((1, 1, 64), dtype=torch.float32)
+    assert fused_rope.try_fused_qk_rotary_emb(query, query, freqs, freqs) is None
+    assert not fused_rope._FAILED_KEYS
+
+
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_failed_runtime_key_cache_is_bounded():
     from vllm_omni.diffusion.models.ernie_image import fused_rope

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -2,13 +2,24 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from transformers.utils.generic import ModelOutput
+from vllm.lora.request import LoRARequest
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+import vllm_omni.diffusion.diffusion_engine as engine_module
+import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
+import vllm_omni.diffusion.models.hunyuan_image3.request_layout as layout_module
+from tests.helpers.kv_layout import build_kv_cache_tensor
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
@@ -27,6 +38,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
     hunyuan_num_special_tokens,
     normalize_hunyuan_cot_text,
     prepare_hunyuan_layout,
+    prepare_hunyuan_prefix_cache,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -68,18 +80,26 @@ class _FakeTokenizerWrapper:
         seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
         padded_seq_len = max(seq_lens)
         positions = torch.arange(padded_seq_len)
-        joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
+        # Like the real tokenizer, joint spans enclose VAE + separator + ViT,
+        # while all_image_slices contains the individual RoPE spans. Leave one
+        # full text block before the image so partial reuse can be tested.
+        joint_image_slices = [[slice(5, 10)] if has_cond_image else [] for _ in range(rows)]
+        cond_vae_image_slices = [[slice(5, 7)] if has_cond_image else [] for _ in range(rows)]
+        cond_vit_image_slices = [[slice(8, 10)] if has_cond_image else [] for _ in range(rows)]
         gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
-        all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
+        all_image_slices = [
+            vae + vit + generated
+            for vae, vit, generated in zip(cond_vae_image_slices, cond_vit_image_slices, gen_image_slices)
+        ]
         cond_vae_image_mask = (
-            torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 5) & (positions < 7) for _ in range(rows)]) if has_cond_image else None
         )
         cond_vit_image_mask = (
-            torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 8) & (positions < 10) for _ in range(rows)]) if has_cond_image else None
         )
         sections = []
         for _ in range(rows):
-            row_sections = []
+            row_sections: list[dict[str, object]] = []
             if has_cond_image:
                 row_sections.append(
                     {
@@ -101,6 +121,8 @@ class _FakeTokenizerWrapper:
                 ),
                 cond_vae_image_mask=cond_vae_image_mask,
                 cond_vit_image_mask=cond_vit_image_mask,
+                cond_vae_image_slices=cond_vae_image_slices,
+                cond_vit_image_slices=cond_vit_image_slices,
                 joint_image_slices=joint_image_slices,
                 gen_image_slices=gen_image_slices,
                 real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
@@ -155,6 +177,7 @@ def _request(
     guidance_scale: float,
     prompt="draw a cat",
     request_id: str = "req",
+    seed: int | None = None,
 ) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompt=prompt,
@@ -163,6 +186,7 @@ def _request(
             width=1024,
             guidance_scale=guidance_scale,
             num_inference_steps=4,
+            seed=seed,
         ),
         request_id=request_id,
     )
@@ -205,6 +229,8 @@ def test_builds_kv_request_lengths_without_model_execution() -> None:
     assert tokenizer.calls[0]["sequence_template"] == "instruct"
     assert tokenizer.calls[0]["cfg_factor"] == 1
     assert kv_requests[0].block_hashes == []
+    assert kv_requests[0].cache_token_ids == ()
+    assert kv_requests[0].mm_features == []
     assert kv_requests[0].kv_contexts == ()
     assert kv_requests[0].skip_reading_prefix_cache is True
 
@@ -226,12 +252,14 @@ def _reference_image() -> JointImageInfo:
     return JointImageInfo(
         vae_image_info=ImageInfo(
             image_type="vae",
+            image_tensor=torch.arange(12, dtype=torch.float32).reshape(1, 3, 2, 2),
             token_width=8,
             token_height=8,
             image_token_length=64,
         ),
         vision_image_info=ImageInfo(
             image_type="siglip2",
+            image_tensor=torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2),
             token_width=4,
             token_height=4,
             image_token_length=16,
@@ -243,6 +271,342 @@ def _reference_image() -> JointImageInfo:
     )
 
 
+def _prepare_cached_requests(request, tokenizer, image_processor):
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    prepare_hunyuan_prefix_cache(request)
+    return request.diffusion_kv_requests
+
+
+def test_reference_prefix_identity_includes_image_content_and_vae_seed() -> None:
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id: str, image: JointImageInfo, seed: int):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={
+                "prompt": "edit this image",
+                "additional_information": {"batch_cond_image_info": [image]},
+            },
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)[0]
+
+    first = build("first", _reference_image(), 7)
+    same = build("same", _reference_image(), 7)
+    changed_seed = build("changed-seed", _reference_image(), 8)
+    changed_image_value = _reference_image()
+    changed_image_value.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+    changed_image = build("changed-image", changed_image_value, 7)
+
+    assert first.mm_features == same.mm_features
+    assert [(feature.mm_position.offset, feature.mm_position.length) for feature in first.mm_features] == [
+        (5, 2),
+        (8, 2),
+    ]
+    for changed in (changed_seed, changed_image):
+        # Both subspans must carry the identity, not just the first one.
+        assert all(a.identifier != b.identifier for a, b in zip(first.mm_features, changed.mm_features, strict=True))
+
+
+def test_reference_prefix_rejects_missing_inputs_with_nested_spans() -> None:
+    tokenizer, image_processor = _components([20])
+    request = _request(
+        guidance_scale=1.0,
+        prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+
+
+def test_reference_content_is_hashed_once_across_cfg_rows(monkeypatch):
+    tokenizer, image_processor = _components([20, 20])
+    image = _reference_image()
+    hasher = Mock(wraps=layout_module.hash_prefix_cache_value)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", hasher)
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=5.0,
+            request_id=request_id,
+            seed=7,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)
+
+    first = build("first")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 1
+    assert len(first[0].mm_features) == len(first[1].mm_features) == 2
+    assert first[0].mm_features == first[1].mm_features
+    image.vae_image_info.image_tensor[0, 0, 0, 0] += 1
+    changed = build("changed")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 2
+    assert all(a.identifier != b.identifier for a, b in zip(first[0].mm_features, changed[0].mm_features, strict=True))
+
+
+def test_failed_cfg_cache_preparation_does_not_install_partial_inputs():
+    tokenizer, image_processor = _components([20, 20])
+    request = _request(
+        guidance_scale=5.0,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    # First row prepares successfully; the second exposes inconsistent inputs.
+    layout.rope_image_info[0] = []
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+
+
+@pytest.mark.parametrize("generator_list", [False, True])
+def test_reference_identity_uses_explicit_generator_state_before_seed(generator_list):
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id, seed, advance):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        generator = torch.Generator(device="cpu").manual_seed(123)
+        if advance:
+            torch.rand(1, generator=generator)
+        request.sampling_params.generator = [generator] if generator_list else generator
+        before = generator.get_state().clone()
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        torch.testing.assert_close(before, generator.get_state())
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first", 7, False)
+    assert first == build("same-generator-different-unused-seed", 8, False)
+    changed = build("advanced-generator-same-seed", 7, True)
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+def test_unseeded_reference_identity_is_request_local():
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        # Real requests normally receive an automatic seed in __post_init__.
+        request.sampling_params.seed = None
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first")
+    assert first == build("first")
+    changed = build("second")
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_text_prefix_hashes_include_lora_identity_and_scale(guidance_scale: float) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, adapter_id: int | None, scale: float):
+        request = _request(guidance_scale=guidance_scale, request_id=request_id)
+        if adapter_id is not None:
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = scale
+        kv_requests = _prepare_cached_requests(request, tokenizer, image_processor)
+        for kv_request in kv_requests:
+            kv_request.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return kv_requests
+
+    base = build("base", None, 1.0)
+    same_base = build("same-base", None, 0.5)  # Scale has no effect without an adapter.
+    first = build("first", 1, 0.5)
+    same = build("same", 1, 0.5)
+    changed_id = build("changed-id", 2, 0.5)
+    changed_scale = build("changed-scale", 1, 1.0)
+
+    for row in range(len(first)):
+        assert base[row].block_hashes == same_base[row].block_hashes
+        assert first[row].block_hashes == same[row].block_hashes
+        assert len(first[row].block_hashes) == 5
+        for changed in (base, changed_id, changed_scale):
+            assert first[row].cache_token_ids == changed[row].cache_token_ids
+            # LoRA changes even the text prefix, so every block must differ.
+            assert all(a != b for a, b in zip(first[row].block_hashes, changed[row].block_hashes, strict=True))
+
+
+@pytest.fixture
+def prefix_cache_manager(request):
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+    num_blocks = 64
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    manager = DiffusionKVCacheManager(
+        config,
+        max_model_len=64,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        enable_prefix_caching=getattr(request, "param", True),
+    )
+    yield manager
+    manager.close()
+
+
+@pytest.mark.parametrize("prefix_cache_manager", [False], indirect=True)
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_paged_cache_off_prepares_and_allocates_without_identity_work(
+    monkeypatch, prefix_cache_manager, guidance_scale
+):
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+    config = SimpleNamespace(
+        model="fake",
+        model_class_name="HunyuanImage3ForCausalMM",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        enable_prefix_caching=False,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_config",
+        lambda *_args, **_kwargs: SimpleNamespace(
+            vae_downsample_factor=(8, 8),
+            patch_size=2,
+            image_base_size=1024,
+        ),
+    )
+    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
+    monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
+    monkeypatch.setattr(pipeline_module.GenerationConfig, "from_pretrained", lambda _model: SimpleNamespace())
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    fail = Mock(side_effect=AssertionError("disabled prefix cache must not prepare identity"))
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", fail)
+    monkeypatch.setattr(layout_module, "prepare_hunyuan_prefix_cache", fail)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils.MultiModalHasher, "hash_kwargs", fail)
+    monkeypatch.setattr(torch.Tensor, "tolist", fail)
+    monkeypatch.setattr(prefix_cache_manager, "_hash_function", fail)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = config
+    engine._init_process_hooks(config)
+    # Guard admission too, even if someone injects a hook after initialization.
+    engine.prefix_cache_func = fail
+    request = _request(
+        guidance_scale=guidance_scale,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    request.sampling_params.seed = None
+    request.sampling_params.generator = SimpleNamespace(get_state=fail)
+    engine._prepare_request_for_admission(request)
+    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+    metadata = prefix_cache_manager.reserve_request(request.request_id, request.diffusion_kv_requests)
+    assert metadata is not None
+    assert all(row.cached_prefix_len == 0 for row in metadata.sequences)
+    prefix_cache_manager.publish_request(request.request_id)
+    prefix_cache_manager.free_request(request.request_id)
+    fail.assert_not_called()
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+@pytest.mark.parametrize(
+    ("change", "expected_hit"),
+    [
+        ("same", 20),
+        ("prompt_suffix", 12),
+        ("target", 20),
+        ("vae", 4),
+        ("vit", 4),
+        ("vision_kwargs", 4),
+        ("rope", 4),
+        ("seed", 4),
+        ("lora_id", 0),
+        ("lora_scale", 0),
+        ("base", 0),
+    ],
+)
+def test_prefix_lookup_respects_reference_and_lora_identity(
+    prefix_cache_manager: DiffusionKVCacheManager, guidance_scale: float, change: str, expected_hit: int
+) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, change: str):
+        image = _reference_image()
+        if change == "vae":
+            image.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vit":
+            image.vision_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vision_kwargs":
+            image.vision_encoder_kwargs["pixel_attention_mask"][0] = False
+        request = _request(
+            guidance_scale=guidance_scale,
+            request_id=request_id,
+            seed=8 if change == "seed" else 7,
+            prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        if change != "base":
+            adapter_id = 2 if change == "lora_id" else 1
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = 1.0 if change == "lora_scale" else 0.5
+        layout = _prepare(request, tokenizer, image_processor)
+        if change == "rope":
+            for spans in layout.rope_image_info:
+                spans[0] = (slice(5, 7), (2, 1))
+        elif change == "prompt_suffix":
+            layout.tokenizer_output.tokens[:, 12] += 100
+        elif change == "target":
+            layout.tokenizer_output.tokens[:, 20:] += 100
+        request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+        prepare_hunyuan_prefix_cache(request)
+        return request.diffusion_kv_requests
+
+    manager = prefix_cache_manager
+    cold = build("cold", "same")
+    cold_metadata = manager.reserve_request("cold", cold)
+    assert cold_metadata is not None
+    assert all(sequence.cached_prefix_len == 0 for sequence in cold_metadata.sequences)
+    # CPU-only: simulate successful materialization to test native cache lookup.
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = build("warm", change)
+    warm_metadata = manager.reserve_request("warm", warm)
+    assert warm_metadata is not None
+    for cold_row, warm_row, cold_sequence, warm_sequence in zip(
+        cold, warm, cold_metadata.sequences, warm_metadata.sequences, strict=True
+    ):
+        if change == "prompt_suffix":
+            assert cold_row.cache_token_ids[:12] == warm_row.cache_token_ids[:12]
+            assert cold_row.cache_token_ids[12] != warm_row.cache_token_ids[12]
+        else:
+            assert cold_row.cache_token_ids == warm_row.cache_token_ids
+        assert warm_sequence.cached_prefix_len == expected_hit
+        hit_blocks = expected_hit // 4
+        assert cold_sequence.block_ids[0][:hit_blocks] == warm_sequence.block_ids[0][:hit_blocks]
+        assert cold_row.block_hashes[:hit_blocks] == warm_row.block_hashes[:hit_blocks]
+        assert all(
+            a != b for a, b in zip(cold_row.block_hashes[hit_blocks:], warm_row.block_hashes[hit_blocks:], strict=True)
+        )
+
+
 def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
     tokenizer, image_processor = _components([20])
     joint_image = _reference_image()
@@ -366,13 +730,17 @@ def test_prepared_model_inputs_match_local_tokenization(
         # Preparation runs without a forward context. Coverage must come from
         # the scheduler metadata copied onto this request, including every CFG row.
         reused = pipeline.prepare_model_inputs(
-            **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(5,) * len(prefix_lens)
+            # The image span is [5, 10); native KV must cover its end, not
+            # merely reach the first image token.
+            **common_kwargs,
+            prepared_layout=prepared_layout,
+            kv_computed_tokens=(10,) * len(prefix_lens),
         )
         assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None
         partial = pipeline.prepare_model_inputs(
             **common_kwargs,
             prepared_layout=prepared_layout,
-            kv_computed_tokens=(5,) * (len(prefix_lens) - 1) + (4,),
+            kv_computed_tokens=(10,) * (len(prefix_lens) - 1) + (9,),
         )
         assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None
 
@@ -496,6 +864,9 @@ def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(trans
 
     tokenizer, image_processor = _components([12, 14])
     request = _request(guidance_scale=5.0)
+    # Prompt token IDs are only materialized for the native AR -> DiT
+    # transfer path; local prefix caching does not need them.
+    request.kv_transfer_params = {"num_transfer_tokens": 0}
     layout = _prepare(request, tokenizer, image_processor)
     layout.tokenizer_output.think_recaption_end_pos = [[7], [5]]
     layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]]

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -147,6 +147,80 @@ def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
     torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
 
 
+def test_slice_cached_prefix_inputs_keeps_full_kv_axis() -> None:
+    inputs_embeds = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2)
+    attention_mask = torch.arange(36).reshape(1, 1, 6, 6)
+    position_ids = torch.arange(6).reshape(1, 6)
+    custom_pos_emb = (torch.arange(6).reshape(1, 6), torch.arange(10, 16).reshape(1, 6))
+    image_mask = torch.tensor([[False, False, True, True, True, True]])
+    gen_timestep_scatter_index = torch.tensor([[2]])
+
+    sliced = HunyuanImage3Pipeline._slice_cached_prefix_inputs(
+        inputs_embeds,
+        attention_mask,
+        position_ids,
+        custom_pos_emb,
+        image_mask,
+        gen_timestep_scatter_index,
+        [6],
+        [6],
+        2,
+    )
+
+    (
+        sliced_embeds,
+        sliced_mask,
+        sliced_positions,
+        sliced_rope,
+        sliced_image_mask,
+        sliced_scatter,
+        sliced_query_lens,
+        sliced_seq_len,
+    ) = sliced
+    torch.testing.assert_close(sliced_embeds, inputs_embeds[:, 2:])
+    torch.testing.assert_close(sliced_mask, attention_mask[:, :, 2:, :])
+    torch.testing.assert_close(sliced_positions, position_ids[:, 2:])
+    torch.testing.assert_close(sliced_rope[0], custom_pos_emb[0][:, 2:])
+    torch.testing.assert_close(sliced_rope[1], custom_pos_emb[1][:, 2:])
+    torch.testing.assert_close(sliced_image_mask, image_mask[:, 2:])
+    torch.testing.assert_close(sliced_scatter, torch.tensor([[0]]))
+    assert sliced_query_lens == [4]
+    assert sliced_seq_len == 4
+
+
+@pytest.mark.parametrize("local_prefix_hit", [True, False])
+def test_ar_reuse_does_not_consume_local_prefix_hits(local_prefix_hit):
+    from vllm_omni.diffusion.forward_context import set_forward_context
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+        HunyuanImage3Text2ImagePipeline,
+    )
+
+    pipe = object.__new__(HunyuanImage3Text2ImagePipeline)
+    input_ids = torch.arange(12).reshape(1, 12)
+    cond_image = torch.ones(1, 1)
+    kwargs = dict(
+        query_lens=[12],
+        attention_mask=torch.ones(1, 1, 12, 12),
+        position_ids=input_ids.clone(),
+        image_mask=torch.zeros_like(input_ids, dtype=torch.bool),
+        gen_timestep_scatter_index=torch.tensor([[8]]),
+        cond_vae_images=cond_image,
+    )
+    runtime = SimpleNamespace(metadata=SimpleNamespace(prefill_rows=[SimpleNamespace(kv_start_pos=4)]))
+    with set_forward_context(paged_kv_runtime=runtime, paged_kv_cached_prefix_len=4 if local_prefix_hit else 0):
+        output, ar_reuse_len = pipe._maybe_handle_ar_kv_reuse(input_ids, kwargs, 1, False, None, torch.device("cpu"))
+    if local_prefix_hit:
+        assert ar_reuse_len == 0
+        assert output is input_ids
+        assert kwargs["query_lens"] == [12]
+        assert kwargs["cond_vae_images"] is cond_image
+    else:
+        assert ar_reuse_len == 4
+        torch.testing.assert_close(output, input_ids[:, 4:])
+        assert kwargs["query_lens"] == [8]
+        assert "cond_vae_images" not in kwargs
+
+
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
     pipeline = _pipeline()
     states = [_state("req-0", 1), _state("req-1", 3)]
@@ -416,7 +490,7 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
             }
         ]
 
-    captured = {}
+    captured: dict[str, object] = {}
 
     def fake_restore_prompt_kv_cache(states_arg, row_state_indexes, row_branches):
         del states_arg
@@ -443,7 +517,9 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
     assert captured["row_state_indexes"] == [0, 1, 0, 1]
     assert captured["row_branches"] == [0, 0, 1, 1]
     assert captured["input_ids"] is None
+    assert isinstance(captured["images"], torch.Tensor)
     assert tuple(captured["images"].shape) == (4, 1)
+    assert isinstance(captured["timestep"], torch.Tensor)
     assert captured["timestep"].tolist() == [0.5, 0.0, 0.5, 0.0]
     assert captured["merged_attention_mask_shape"] == (4, 1, 2, 6)
     assert captured["merged_full_attn_spans"] == [[(4, 6)], [(4, 6)], [(4, 6)], [(4, 6)]]

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -8,7 +8,7 @@ Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP,
 from __future__ import annotations
 
 import math
-from contextlib import ExitStack, contextmanager
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -59,7 +59,7 @@ class MockAttention(nn.Module):
 @contextmanager
 def patched_mgr_env(sp_size=1):
     target = _TRANSFORMER_MODULE
-    patches = [
+    patches: list[AbstractContextManager] = [
         patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_allgather_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_ulysses_parallel_world_size", return_value=1, create=True),
@@ -242,6 +242,39 @@ def test_legacy_tensor_prefix_still_uses_dense_cache_and_reuse() -> None:
     assert mgr.attn.calls[-1][1].shape == (1, prefix_len + first_len, NUM_HEADS, HEAD_DIM)
 
 
+@pytest.mark.parametrize("prefix_len", [0, 19])
+@pytest.mark.parametrize("prompt_query_len", [0, 3])
+def test_paged_sp_splits_uncached_queries(prefix_len, prompt_query_len) -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr(sp_size=2)
+    mgr.attn.paged_kv_active = True
+    shard_image_size = IMAGE_TOKEN_LEN // 2
+    query_len = prompt_query_len + shard_image_size
+    query = torch.randn(2 * query_len, NUM_HEADS, HEAD_DIM)
+    key, value = _make_known_kv(2 * query_len)
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        output = mgr(
+            query,
+            key,
+            value,
+            first_step=True,
+            query_lens=[query_len] * 2,
+            seq_lens=[prefix_len + query_len] * 2,
+            shard_image_size=shard_image_size,
+            full_attn_spans=[[(prefix_len + prompt_query_len, prefix_len + query_len)]] * 2,
+        )
+
+    torch.testing.assert_close(output, query)
+    metadata = mgr.attn.paged_metadata[0]
+    torch.testing.assert_close(
+        metadata.joint_key, key.reshape(2, query_len, NUM_KV_HEADS, HEAD_DIM)[:, :prompt_query_len]
+    )
+    assert metadata.joint_query.shape[1] == prompt_query_len
+    assert mgr.attn.paged_calls[0][0].shape[1] == shard_image_size
+    assert mgr.image_kv_cache_map is None
+
+
 def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
     from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
 
@@ -391,6 +424,7 @@ def test_no_ar_kv(bs):
         gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
     assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
@@ -478,12 +512,14 @@ def test_ar_kv_no_cfg(sp_size):
         gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
         # cached = ar(5) + prompt(3) = 8
         expected_cached_len = 8
     else:
         # cached = seq_len - shard_image_size = 17 - 4 = 13
+        assert shard_image_size is not None
         expected_cached_len = seq_len - shard_image_size
 
     assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -26,6 +26,7 @@ def _make_pipeline():
     pipeline.default_image_negative_prompt = "default image negative"
     pipeline.img_prompt_template = "<image>"
     pipeline.od_config = SimpleNamespace(flow_shift=None)
+    pipeline.enable_diffusion_pipeline_profiler = False
     return pipeline
 
 
@@ -572,3 +573,72 @@ def test_load_weights_rejects_external_weight_stream():
     assert pipeline.load_weights([]) == set()
     with pytest.raises(RuntimeError, match="components are loaded directly"):
         pipeline.load_weights([("transformer.weight", torch.zeros(1))])
+
+
+@pytest.mark.parametrize("enabled", [False, True], ids=["disabled", "enabled"])
+def test_pipeline_profiler_records_executed_stages_per_request(monkeypatch, enabled):
+    from vllm_omni.diffusion.models.lingbot_video import pipeline_lingbot_video as module
+    from vllm_omni.diffusion.profiler import diffusion_pipeline_profiler as profiler_module
+    from vllm_omni.diffusion.worker.utils import consume_pipeline_stage_durations
+
+    class TextEncoder(nn.Module):
+        def forward(self):
+            return torch.ones(1, 2, 4), torch.ones(1, 2, dtype=torch.long)
+
+    class VAE(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(()))
+            self.config = SimpleNamespace(latents_mean=[0.0], latents_std=[1.0])
+
+        def encode(self, value):
+            return value
+
+        def decode(self, latents):
+            return (latents.repeat(1, 3, 1, 1, 1),)
+
+    components = [
+        (module.LingBotVideoTransformer3DModel, _RecordingTransformer()),
+        (module.Qwen3VLForConditionalGeneration, TextEncoder()),
+        (module.Qwen3VLProcessor, SimpleNamespace()),
+        (module.AutoencoderKLWan, VAE()),
+        (module.FlowUniPCMultistepScheduler, _CorruptingScheduler()),
+    ]
+    for component_class, component in components:
+        monkeypatch.setattr(component_class, "from_pretrained", lambda *args, value=component, **kwargs: value)
+    monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(module.LingBotVideoPipeline, "encode_prompt", lambda self, *args, **kwargs: self.text_encoder())
+    monkeypatch.setattr(profiler_module.current_omni_platform, "is_available", lambda: False)
+    ticks = iter(range(100))
+    monkeypatch.setattr(profiler_module.time, "perf_counter", lambda: float(next(ticks)))
+    pipeline = module.LingBotVideoPipeline(
+        od_config=SimpleNamespace(
+            model="unused", dtype=torch.float32, enable_diffusion_pipeline_profiler=enabled, flow_shift=None
+        )
+    )
+    request = _make_request_batch(
+        {"prompt": "a robot", "modalities": ["video"]},
+        height=16,
+        width=16,
+        num_frames=9,
+        num_inference_steps=2,
+        guidance_scale=1.0,
+        seed=42,
+    )
+    prefix = "LingBotVideoPipeline."
+    expected = {prefix + "text_encoder.forward": 1.0, prefix + "transformer.forward": 2.0, prefix + "vae.decode": 1.0}
+    first = pipeline(request)
+    assert first.stage_durations == (expected if enabled else {})
+    assert first.output["video"].shape == (3, 2, 2, 3)
+    # A latent-only request must not inherit the previous request's VAE timing.
+    request.requests[0].sampling_params.output_type = "latent"
+    second = pipeline(request)
+    latent_expected = {key: value for key, value in expected.items() if not key.endswith("vae.decode")}
+    assert second.stage_durations == (latent_expected if enabled else {})
+    assert first.stage_durations == (expected if enabled else {})
+    consumed = consume_pipeline_stage_durations(pipeline)
+    if enabled:
+        assert consumed == {**latent_expected, prefix + "forward": 7.0}
+    else:
+        assert consumed == {}
+    assert consume_pipeline_stage_durations(pipeline) == {}