[Core] Validate HSDP shard size before initialization

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

.buildkite/amd/scripts/bootstrap-amd-omni.sh

@@ -1,4 +1,7 @@
 #!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # vllm-omni customized version
 # Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
 # Last synced: 2025-12-15
@@ -6,6 +9,9 @@
 
 set -euo pipefail
 
+# The bootstrap runs from the repository root; ShellCheck does not follow
+# sourced files unless invoked with -x.
+# shellcheck disable=SC1091
 source .buildkite/common/scripts/resolve_skip_ci.sh
 
 if [[ -z "${RUN_ALL:-}" ]]; then
@@ -63,6 +69,8 @@ upload_pipeline() {
     # Install minijinja
     ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
     curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
+    # Installed by the minijinja bootstrap above and only present on the CI agent.
+    # shellcheck disable=SC1091
     source /var/lib/buildkite-agent/.cargo/env
 
     if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
@@ -90,22 +98,82 @@ upload_pipeline() {
     cd .buildkite/amd
 
     # Select test definition file: merge suite for main, ready suite for PRs.
-    # For debugging, DEBUG_TEST_YAML can override the selection — accepts
-    # "merge" or "ready" (case-insensitive).
+    # For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing
+    # "merge" and/or "ready" (case-insensitive). Multiple suites are combined
+    # before rendering so they share one amd-build step.
     if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
-        case "${DEBUG_TEST_YAML,,}" in
-            merge)
-                TEST_YAML="test-amd-merge.yml"
-                ;;
-            ready)
-                TEST_YAML="test-amd-ready.yml"
-                ;;
-            *)
-                echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
+        declare -a DEBUG_TEST_SPECS=()
+        declare -A SEEN_DEBUG_TESTS=()
+        IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}"
+
+        for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do
+            # Permit readable values such as "ready, merge," and ignore the
+            # empty item produced by a trailing comma.
+            requested_test="${requested_test//[[:space:]]/}"
+            [[ -z "$requested_test" ]] && continue
+
+            if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then
+                echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2
                 exit 1
-                ;;
-        esac
-        echo "DEBUG_TEST_YAML override: using $TEST_YAML"
+            fi
+            SEEN_DEBUG_TESTS[$requested_test]=1
+
+            case "$requested_test" in
+                ready)
+                    DEBUG_TEST_SPECS+=("READY_TESTS:test-amd-ready.yml")
+                    ;;
+                merge)
+                    DEBUG_TEST_SPECS+=("MERGE_TESTS:test-amd-merge.yml")
+                    ;;
+                *)
+                    echo "ERROR: DEBUG_TEST_YAML entries must be 'merge' or 'ready', got '$requested_test'" >&2
+                    exit 1
+                    ;;
+            esac
+        done
+
+        if [[ ${#DEBUG_TEST_SPECS[@]} -eq 0 ]]; then
+            echo "ERROR: DEBUG_TEST_YAML did not contain a test suite" >&2
+            exit 1
+        elif [[ ${#DEBUG_TEST_SPECS[@]} -eq 1 ]]; then
+            TEST_YAML="${DEBUG_TEST_SPECS[0]#*:}"
+        else
+            TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-debug-tests.XXXXXX.yml")
+            python - "$TEST_YAML" "${DEBUG_TEST_SPECS[@]}" <<'PY'
+import sys
+
+import yaml
+
+
+output_path, *suite_specs = sys.argv[1:]
+combined = {"env": {}, "steps": []}
+
+for suite_spec in suite_specs:
+    group_name, input_path = suite_spec.split(":", 1)
+    with open(input_path, encoding="utf-8") as test_file:
+        suite = yaml.safe_load(test_file)
+
+    for name, value in (suite.get("env") or {}).items():
+        previous = combined["env"].get(name, value)
+        if previous != value:
+            raise ValueError(
+                f"Conflicting environment value for {name}: {previous!r} != {value!r}"
+            )
+        combined["env"][name] = value
+
+    suite_steps = []
+    for entry in suite.get("steps") or []:
+        if "group" in entry:
+            suite_steps.extend(entry.get("steps") or [])
+        else:
+            suite_steps.append(entry)
+    combined["steps"].append({"group": group_name, "steps": suite_steps})
+
+with open(output_path, "w", encoding="utf-8") as output_file:
+    yaml.safe_dump(combined, output_file, sort_keys=False)
+PY
+        fi
+        echo "DEBUG_TEST_YAML override: using ${DEBUG_TEST_SPECS[*]}"
     elif [[ $BUILDKITE_BRANCH == "main" ]]; then
         TEST_YAML="test-amd-merge.yml"
     else
@@ -130,19 +198,20 @@ upload_pipeline() {
             > pipeline.yaml
     )
     cat pipeline.yaml
+    if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-debug-tests."*.yml ]]; then
+        rm -f -- "$TEST_YAML"
+    fi
     buildkite-agent artifact upload pipeline.yaml
     buildkite-agent pipeline upload pipeline.yaml
     exit 0
 }
 
 get_diff() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
+    git diff --name-only --diff-filter=ACMDR "$(git merge-base origin/main HEAD)"
 }
 
 get_diff_main() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
+    git diff --name-only --diff-filter=ACMDR HEAD~1
 }
 
 file_diff=$(get_diff)
@@ -181,7 +250,7 @@ for file in $file_diff; do
     # First check if file matches any pattern
     matches_pattern=0
     for pattern in "${patterns[@]}"; do
-        if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
+        if [[ $file == "$pattern"* ]] || [[ $file == "$pattern" ]]; then
             matches_pattern=1
             break
         fi
@@ -191,7 +260,7 @@ for file in $file_diff; do
     if [[ $matches_pattern -eq 1 ]]; then
         matches_ignore=0
         for ignore in "${ignore_patterns[@]}"; do
-            if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
+            if [[ $file == "$ignore"* ]] || [[ $file == "$ignore" ]]; then
                 matches_ignore=1
                 break
             fi

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

@@ -21,7 +21,7 @@ steps:
 
     - label: "Simple · Diffusion Test · Shard %N/%t"
       agent_pool: mi300_1
-      parallelism: 2
+      parallelism: 4
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
       grade: Blocking
@@ -71,26 +71,30 @@ steps:
       commands:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        - |
-          bash -c '
-            set +e
-            pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-            EXIT1=\$?
-            pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
-            EXIT2=\$?
-            exit \$((EXIT1 | EXIT2))
-          '
+        - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+        - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
 
-    - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
+    - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2"
       agent_pool: mi300_4
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
+      timeout_in_minutes: 40
       grade: Blocking
       commands:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
-        - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
+        - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
+
+    - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 2/2"
+      agent_pool: mi300_4
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      timeout_in_minutes: 40
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
 
     # merge-only tests
     - label: "Diffusion Tensor Parallelism Test"
@@ -175,17 +179,41 @@ steps:
 
     - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
       agent_pool: mi300_1
-      parallelism: 2
+      parallelism: 3
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
       grade: Blocking
+      timeout_in_minutes: 80
       commands:
         - |
-          bash -c '
+          timeout --signal=TERM --kill-after=2m 80m bash -c '
             export VLLM_LOGGING_LEVEL=DEBUG
             export VLLM_WORKER_MULTIPROC_METHOD=spawn
             export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+            TESTS=(
+              "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_default_cuda_graph_startup[async_chunk]"
+              "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_002[async_chunk]"
+              "tests/e2e/offline_inference/test_qwen3_tts_customvoice.py::test_text_to_audio_001[no_cuda_graph]"
+            )
+            TEST="$${TESTS[$$BUILDKITE_PARALLEL_JOB]}"
+            pytest -s -v "$$TEST" -m "advanced_model" --run-level "advanced_model"
+          '
+
+    # TODO: Restore this case to the blocking sharded step once its long-running
+    # Whisper validation on AMD CI is resolved.
+    - label: "Qwen3-TTS CustomVoice async_chunk Test (quarantined)"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: NonBlocking
+      timeout_in_minutes: 80
+      commands:
+        - |
+          timeout --signal=TERM --kill-after=2m 75m bash -c '
+            export VLLM_LOGGING_LEVEL=DEBUG
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+            pytest -s -v "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_001[async_chunk]" -m "advanced_model" --run-level "advanced_model"
           '
 
     - label: "Qwen3-TTS Base E2E Test"

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

@@ -191,10 +191,7 @@ steps:
       mirror_hardwares: [amdproduction]
       grade: Blocking
       commands:
-        - |
-          timeout 20m bash -c '
-            pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
-          '
+        - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs

.buildkite/cuda/test-merge.yml

@@ -166,6 +166,7 @@ steps:
           - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
           - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
           - vllm_omni/deploy/minicpmo_4_5.yaml
           - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -216,3 +217,15 @@ steps:
               pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
             "
         mirror_hardwares: h100_1
+
+      - label: "Diffusion · HunyuanVideo-1.5 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
+          - vllm_omni/diffusion/models/hunyuan_video/
+        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 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            "
+        mirror_hardwares: h100_1

.buildkite/cuda/test-nightly.yml

@@ -272,6 +272,12 @@ steps:
           - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test"
+        timeout_in_minutes: 90
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion X2V · Doc Test"
         timeout_in_minutes: 60
         commands:
@@ -322,6 +328,23 @@ steps:
             exit $$EXIT2
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5"
+        key: nightly-diffusion-x2v-performance-hunyuanvideo15
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
+            EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json
+            EXIT2=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
+            exit $$EXIT2
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3"
         key: nightly-diffusion-x2v-performance-minimax-h3
         timeout_in_minutes: 180

.buildkite/cuda/test-ready.yml

@@ -205,6 +205,7 @@ steps:
           - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
           - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
           - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
           - vllm_omni/deploy/minicpmo_4_5.yaml
           - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -306,6 +307,24 @@ steps:
             "
         mirror_hardwares: h100_1
 
+      - label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test"
+        timeout_in_minutes: 90
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
+          - tests/helpers/assertions.py
+          - tests/helpers/runtime.py
+          - vllm_omni/diffusion/models/minimax_h3/
+          - vllm_omni/diffusion/offloader/
+          - vllm_omni/diffusion/distributed/
+          - vllm_omni/entrypoints/openai/
+        commands:
+          - |
+            timeout 80m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py -m 'core_model and cuda and diffusion' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_2
+
       - label: "Diffusion · Qwen Image Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen_image.py
@@ -319,3 +338,15 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
             "
         mirror_hardwares: h100_1
+
+      - label: "Diffusion · HunyuanVideo-1.5 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
+          - vllm_omni/diffusion/models/hunyuan_video/
+        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' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_1

docker/Dockerfile.ci

@@ -1,4 +1,4 @@
-# Target vllm version: v0.27.0 (releases/v0.27.0 branch).
+# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
 # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
 #
 # Release-boundary form: there is no precompiled-wheel pin here, because the
@@ -14,11 +14,14 @@
 # core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
 # under numba's ceiling. All of it was compensation for a base image and a vLLM
 # that disagreed. At a real release boundary they agree, so none of it is
-# needed. If you pin an untagged commit again, restore the reinstall block —
-# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the
-# Dockerfile silently ship the wrong vLLM.
+# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
+# rode on the v0.27.0 base image with the full reinstall block until the
+# official v0.28.0 image shipped, at which point this file dropped back to the
+# form you are reading. If you pin an untagged commit again, restore the
+# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather
+# than letting the Dockerfile silently ship the wrong vLLM.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.27.0
+ARG VLLM_BASE_TAG=v0.28.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docs/.nav.yml

@@ -53,6 +53,7 @@ nav:
       - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
       - MXFP8 W8A8: user_guide/quantization/mxfp8.md
       - MXFP4 W4A4: user_guide/quantization/mxfp4.md
+      - SVDQuant W4A4: user_guide/quantization/svdquant.md
       - ModelOpt: user_guide/quantization/modelopt.md
       - AutoRound: user_guide/quantization/autoround.md
       - msModelSlim: user_guide/quantization/msmodelslim.md

docs/api/README.md

@@ -114,6 +114,10 @@ Configuration classes.
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
 - [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
 - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
 - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
 

docs/configuration/environment_variables.md

@@ -83,6 +83,7 @@ depends on the installed kernels and model path.
 | `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
 | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
 | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
+| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
 | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
 | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
 | `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |

docs/configuration/stage_configs.md

@@ -59,7 +59,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
 | `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
 | `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
 | `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
-| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
+| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
 | `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
 | `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
 | `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
@@ -94,12 +94,14 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
 | `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
 | `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
 | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
-| `devices` | str \| null | optional | `null` | Device list assigned to this stage. |
+| `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. |
 | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
 | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
 | `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
 | `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. |
 
+**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage.
+
 ### Connector schema
 
 Each entry under top-level `connectors:` follows this shape:

docs/contributing/profiling.md

@@ -322,6 +322,15 @@ Each 1-second window records:
 
 On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
 
+`loop_idle` and `loop_active` count loop iterations, so their absolute scale is
+tied to which orchestration loop is running. Under the default poll loop an idle
+orchestrator records roughly one iteration per millisecond. Under the
+event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see
+[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle
+orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
+records one iteration per routed output. Window counts and the `loop_active_pct`
+summary are therefore not comparable across the two modes.
+
 ### Relationship to other diagnostics
 
 This monitor is intentionally separate from the existing profiling tools:

docs/design/feature/host_weight_runtime.md

@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
 
 ## Status
 
-The first implementation provides contracts and a CPU local-filesystem store.
-The initial diffusion consumer contract additionally defines typed,
+The implementation provides contracts and a CPU local-filesystem store. The
+diffusion consumer additionally defines typed,
 representation-independent final-layout identity/restoration mechanics plus a
-concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
-library-only: no loader or DLO path selects, publishes, restores, or transports
-that artifact yet.
+concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
+selects, publishes, restores, and transfers these artifacts.
 
 V1 includes:
 
@@ -33,13 +33,12 @@ V1 includes:
 
 V1 does not include:
 
-- a public CLI, loader activation, or default-on model integration;
-- lease handoff to DLO or another transport consumer;
-- FP8, quantized, merged-adaptation, or additional model producers;
-- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
+- default-on activation or consumers outside no-AllGather DLO;
+- online FP8, quantized, merged-adaptation, or additional model producers;
+- HWR interaction with DLO AllGather;
 - a remote artifact provider or cross-node coordination;
 - automatic eviction; or
-- a change to DLO AllGather or no-AllGather behavior.
+- a change to DLO collective or execution behavior.
 
 ## Motivation and use cases
 
@@ -255,7 +254,8 @@ The contract is intentionally separate from loader activation:
 - `FinalLayoutBF16Producer` accepts only the matching identity context and a
   finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
   coordinate. Its BF16 policy preserves model-declared FP32 parameters and
-  buffers and revalidates MiniMax H3 mixed-precision invariants.
+  buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
+  FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
 
 Other representations reuse source identity, typed parallel identity, tensor
 ownership, and exact restoration only when their policy proves those semantics.

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -220,9 +220,15 @@ model revision and parallel layout with `required`. TP coordinates have distinct
 identities, while equivalent DP replicas share them.
 
 The current producer/consumer boundary is the model-declared final-layout BF16
-contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
-and TP2 rank identities plus SP layout identities; online quantization, HSDP,
-LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
+contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
+contract covers both transformer block stacks, constructor-stable packed QKV
+mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
+buffers. The shared HWR machinery supports ordinary-loader final layouts for
+TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
+in this change is TP1-only; its TP2/SP layouts rely on the existing
+per-coordinate identity mechanics and remain unmeasured. Online quantization,
+HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
+HWR mode
 `disabled`, DLO-disabled, and DLO AllGather configurations stop before source
 identity or store construction and retain the existing checkpoint-mmap or
 ordinary-loader path.

docs/design/module/engine_orchestration.md

@@ -66,6 +66,12 @@ the boundary affected by the in-flight stage client/process refactor in
 [#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
 responsibilities proposed only by that PR are not current contracts.
 
+The orchestration loop also has an opt-in event-driven mode
+(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
+[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
+cadence only: the routing, ordering, and terminal-state contracts below hold
+identically on both loops.
+
 ## Ownership boundary
 
 This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,

docs/design/module/host_weight_runtime.md

@@ -257,6 +257,53 @@ Lookup checks the marker both before validation and after constructing the
 lease so a weaker overlapping lookup cannot return a lease after a stronger
 validation has denied the artifact.
 
+If publishing the external deny marker fails, lookup returns a retryable
+storage failure and marks the artifact directory with a persistent invalidation
+mode before releasing the shared lock. Cooperating readers check that mode both
+while opening and immediately before returning a lease. Lookup then makes a
+bounded attempt to take the exclusive artifact lock and quarantine the entry;
+the quarantine transition removes the internal mode marker. If a cooperating
+builder replaces the marked inode while the locks are exchanged, containment
+preserves that new publication instead of quarantining it. If neither the
+external marker nor the directory-mode marker can be persisted, cross-process
+exclusion cannot be guaranteed on an unavailable filesystem and the domain
+must remain failed rather than treating the artifact as an ordinary invalid
+cache entry.
+
+New publications carry the internal invalidation mode through the staging
+rename. Hardening clears it and is the logical authority transition for a
+process crash; the temporary- and artifact-parent syncs then make the rename
+durable, with an artifact-inode sync first making the hardened mode durable. A
+process that dies before hardening leaves a persistently denied
+entry, while one that dies after successful hardening leaves an authoritative
+ready entry. Any synchronous failure after rename is marked invalid and
+contained in quarantine before `PUBLICATION_FAILED` is returned, so that
+nonretryable result cannot later resolve as a ready hit. If the filesystem
+refuses both the inode marker and the quarantine move, the store instead returns
+a retryable storage failure with `outcome_uncertain` details; as with deny-marker
+double failure, cross-process exclusion cannot be guaranteed until storage
+recovers. Once the artifact-parent sync succeeds, a later artifact- or
+build-lock release error is logged and lookup resolves the committed entry
+instead of reporting a contradictory publication failure. The same authority
+rule applies when a competing publisher already installed byte-identical
+semantic content; staging-cleanup errors are retried without downgrading that
+validated ready entry to publication failure.
+
+Before widening a hardened artifact for a lifecycle move, the store durably
+sets the same invalidation mode. A process crash before rename therefore leaves
+the writable source denied. After rename, the store re-hardens the destination
+but retains the marker until both parent-directory syncs make the move durable;
+only then does it durably clear the marker. An interrupted move therefore leaves
+a marked transition that the next exact-key lifecycle operation re-hardens.
+Explicit cleanup uses the locked move and retries exact-key
+`.cleanup.*` tombstones left by an interrupted removal; a later build performs
+the same tombstone reconciliation before producing replacement content.
+
+An operational failure while inspecting a noncooperative competing publication
+is a retryable storage failure. The competing entry remains authoritative and
+is not mislabeled as corrupt or quarantined merely because its manifest could
+not be read transiently.
+
 `inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots.
 Inspection validation does not create a deny marker. Reported lock activity is
 point-in-time evidence only; kernel locks are not a persistent owner registry.

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
 
 vLLM-Omni supports the following hardware platforms:
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
 
 ```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
 ```
 
-The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.26.0
+git checkout v0.28.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.26.0 \
+    vllm/vllm-omni:v0.28.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
   -t vllm-omni-cuda .
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
 
 ```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
 ```
 
-The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.26.0
+git checkout v0.28.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.26.0 \
+  vllm/vllm-omni-rocm:v0.28.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.26.0
+  vllm/vllm-omni-rocm:v0.28.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/models/supported_models.md

@@ -69,6 +69,8 @@ th {
 | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
 | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
 | `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
+| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
 | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |

docs/serving/speech_api.md

@@ -811,6 +811,51 @@ If you encounter OOM errors:
 
 Use `/v1/audio/voices` to list available voices for the loaded model.
 
+## Orchestration Loop (experimental)
+
+Multi-stage omni deployments route stage outputs through a single orchestrator
+loop. By default that loop polls every stage replica on a 1 ms cadence. An
+opt-in event-driven mode replaces the poll with one reader task per live stage
+replica awaiting its client directly, and switches the serving-side
+final-output drain to a condition-variable wakeup at the same time.
+
+**Configuration (environment variables):**
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
+
+Set it on the process that runs the orchestrator (stage 0 of an omni
+deployment) before starting the server:
+
+```bash
+export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --omni \
+    --port 8091
+```
+
+The server logs the selected loop mode and its reader/poller counts once at
+startup, so you can confirm which loop is live.
+
+Routing, output ordering, and terminal-state behavior are identical on both
+loops; only the poll cadence changes. Leaving the variable unset keeps the
+legacy poll loop, which is the supported default.
+
+**Known limitations:**
+
+- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at
+  concurrency 8) predates the rebuild on the per-replica fault-isolation work
+  in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work
+  changed dead-replica handling and reader/poller lifecycle rather than the
+  steady-state output path, and the parity suite covers it, but the serving
+  A/B has not been re-run on the current head.
+- The diffusion-poller branch is covered by unit tests only. Deployments whose
+  stages all run as standard engine cores never exercise it, including GLM-TTS,
+  which deploys its DiT without `stage_type: diffusion`.
+- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the
+  latency is admission-bound, which this mode does not address.
+
 ## Development
 
 Enable debug logging:

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -112,6 +112,12 @@ perform a collective and uses the same rank-local transfer behavior.
 HWR is an opt-in startup optimization for models that declare the final-layout
 BF16 restore contract. Use it only with no-AllGather DLO:
 
+The validated BF16 model contracts currently cover MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
+class but has not been validated against this contract. FLUX.2-dev, online
+FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
+the validated scope.
+
 ```bash
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \

docs/user_guide/diffusion_features.md

@@ -129,6 +129,7 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --deploy-config /path/to/deploy_config_file
 ```
 
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
 ### Launch individual stages (stage-based CLI)
 
 Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.

docs/user_guide/quantization/svdquant.md

@@ -0,0 +1,50 @@
+# SVDQuant W4A4
+
+## Overview
+
+[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
+and activations with a small low-rank branch that corrects part of the
+quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
+it does not calibrate the model while loading.
+
+## Checkpoint contract
+
+Place the following entry in the diffusion transformer's `config.json`:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "svdquant",
+    "rank": 32,
+    "precision": "nvfp4",
+    "act_unsigned": false,
+    "modules_to_not_convert": []
+  }
+}
+```
+
+For a quantized linear with input size `K`, output size `N`, and correction
+rank `R`, the checkpoint stores:
+
+| Suffix | Shape | dtype |
+| --- | --- | --- |
+| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
+| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
+| `proj_down` | `(K, R)` | `bfloat16` |
+| `proj_up` | `(N, R)` | `bfloat16` |
+| `smooth_factor` | `(K,)` | `bfloat16` |
+| `wcscales` | `(N,)` | `bfloat16` |
+| `wtscale` | `(1,)` | `bfloat16` |
+
+`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
+ones when no per-output correction is needed. Modules listed in
+`modules_to_not_convert` keep their checkpoint precision.
+
+## Runtime support
+
+The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
+by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
+NVFP4 tensor layouts; incompatible forced backends fail during model loading.
+SM103 is the currently validated and enabled hardware target. Native fusion of
+the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
+contract.

examples/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""vLLM-Omni example scripts, importable as a package.
+
+Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
+places an ``examples`` package on ``sys.path``), a regular package at this repo
+root wins over the upstream ``examples`` namespace during ``import examples``.
+"""

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,117 @@
+"""
+Offline inference example: MammothModa2 image summarization (single AR stage).
+
+Example:
+  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+    --model path/to/MammothModa2-Preview \
+    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
+    --image /path/to/input.jpg \
+    --question "Please summarize the content of this image."
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from PIL import Image
+from vllm import SamplingParams
+from vllm.multimodal.image import convert_image_mode
+
+from vllm_omni import Omni
+
+DEFAULT_SYSTEM = "You are a helpful assistant."
+DEFAULT_QUESTION = "Please summarize the content of this image."
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
+    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
+    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
+    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
+    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
+    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Max new tokens to generate.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    return parser.parse_args()
+
+
+def build_prompt(system: str, question: str) -> str:
+    return (
+        f"<|im_start|>system\n{system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image file not found: {args.image}")
+
+    pil_image = Image.open(args.image)
+    image_data = convert_image_mode(pil_image, "RGB")
+    prompt = build_prompt(args.system, args.question)
+
+    omni = Omni(
+        model=args.model,
+        deploy_config=args.deploy_config,
+        trust_remote_code=args.trust_remote_code,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
+    try:
+        sp = SamplingParams(
+            temperature=float(args.temperature),
+            top_p=float(args.top_p),
+            top_k=-1,
+            max_tokens=int(args.max_tokens),
+            seed=int(args.seed),
+            detokenize=True,
+        )
+        # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
+        # Consume it inside the try block so the worker isn't closed early.
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": image_data},
+                        "additional_information": {"omni_task": ["chat"]},
+                    }
+                ],
+                [sp],
+            )
+        )
+    finally:
+        omni.close()
+
+    lines: list[str] = []
+    for stage_outputs in outputs:
+        ro = stage_outputs
+        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+        lines.append("answer:\n")
+        lines.append(text.strip() + "\n")
+        lines.append("\n")
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_image/text_to_image.py

@@ -599,6 +599,17 @@ def main():
             if args.seed is not None and hasattr(params, "seed"):
                 params.seed = args.seed
 
+            # MammothModa2's AR stage emits one visual token per grid cell,
+            # one EOL token per row, and one final look-ahead token whose hidden
+            # state is unavailable. Size the first stage from the prompt metadata
+            # instead of SamplingParams' default of 16 tokens.
+            prompt_info = prompt_dict.get("additional_information", {})
+            if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
+                ar_width = int(prompt_info.get("ar_width", [0])[0])
+                ar_height = int(prompt_info.get("ar_height", [0])[0])
+                if ar_width > 0 and ar_height > 0:
+                    params.max_tokens = ar_height * (ar_width + 1) + 1
+
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]
 

examples/offline_inference/x_to_text/README.md

@@ -10,10 +10,12 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
 | --- | --- | --- | --- |
 | `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
 | `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
+| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
+| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
 
 The script recognizes these model families from `config.json`, applies the
-model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
-BAGEL uses its registered default deploy.
+model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
+and MammothModa2. BAGEL uses its registered default deploy.
 
 ## Text-To-Text
 
@@ -33,6 +35,22 @@ python examples/offline_inference/x_to_text/x_to_text.py \
   --prompt "Explain multimodal inference in three concise sentences."
 ```
 
+### MammothModa2
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Preview \
+  --prompt "Explain multimodal inference in three concise sentences."
+```
+
+The same command supports the Dev checkpoint:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Dev \
+  --prompt "Explain multimodal inference in three concise sentences."
+```
+
 ## Image-To-Text
 
 Pass one image with `--image`. The default question can be replaced with any
@@ -54,6 +72,16 @@ python examples/offline_inference/x_to_text/x_to_text.py \
   --prompt "Describe this image in detail."
 ```
 
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Preview \
+  --image image.png \
+  --prompt "Describe this image in detail."
+```
+
+Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
+the Dev checkpoint.
+
 ## Key Arguments
 
 | Argument | Default | Description |
@@ -74,4 +102,4 @@ python examples/offline_inference/x_to_text/x_to_text.py \
 
 The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
 different topology, pass a custom AR-only YAML through `--deploy-config`.
-BAGEL uses its repository deploy default unless overridden.
+BAGEL and MammothModa2 use their repository deploy defaults unless overridden.

examples/offline_inference/x_to_text/x_to_text.py

@@ -15,6 +15,7 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
 
 _REPO_ROOT = Path(__file__).resolve().parents[3]
 _HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
+_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
 
 
 def parse_args() -> argparse.Namespace:
@@ -61,7 +62,7 @@ def main() -> None:
     omni_kwargs: dict[str, Any] = {
         "model": args.model,
         "mode": "image-to-text" if image is not None else "text-to-text",
-        "trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
+        "trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
         "log_stats": args.log_stats,
     }
     if args.enforce_eager is not None:
@@ -70,6 +71,8 @@ def main() -> None:
         omni_kwargs["deploy_config"] = args.deploy_config
     elif family == "hunyuan_image3":
         omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
+    elif family == "mammoth_moda2":
+        omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
 
     omni = Omni(**omni_kwargs)
     try:

examples/online_serving/minicpmo/README.md

@@ -159,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
     --output-dir /tmp/minicpmo_realtime_duplex_demo
 ```
 
+Video input uses the same session. PyAV demuxes JPEG frames at `--video-fps`
+(default 1.0) and vLLM `load_audio` extracts a 16 kHz mono WAV unless
+`--input-wav` overrides the soundtrack. Frames keep their capture size
+(`--frame-max-side 0`); the server `process_image` normalizes at 448.
+
+`--stack-frames N` raises the visual refresh rate the way official duplex does:
+each 1 s unit also samples the `N-1` sub-frames captured inside it and tiles
+them into a single composite image sent next to that unit's base frame. The
+audio cadence never changes — a unit is always one second — and the wire always
+carries 2 images per unit however large `N` is, because the sub-frames share one
+composite. Official uses 5 for high refresh rate mode:
+
+```bash
+python examples/online_serving/minicpmo/realtime_duplex_demo.py \
+    --url ws://localhost:8099/v1/realtime?duplex=1 \
+    --model openbmb/MiniCPM-o-4_5 \
+    --input-video /path/to/clip.mp4 \
+    --stack-frames 5 \
+    --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
+    --output-dir /tmp/minicpmo_realtime_duplex_video_demo
+```
+
+Detail inside a composite is capped by `scale_resolution=448` at
+`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for
+stacked frames, which the duplex adapter does not implement yet. Reading small
+text or digits out of a wide scene is limited by that, not by frame timing.
+
 ## Open the experimental browser client
 
 The browser UI serves the page and proxies the same-origin Realtime WebSocket to
@@ -192,6 +219,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
 two-response WAV, its `--input-sha256`, and an
 `--expect-second-response-substring` value.
 
+## Run Omni-DuplexEval
+
+Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
+endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
+served from the same vLLM-Omni environment. This validates the vLLM-Omni
+duplex client and local judge integration; it does not reproduce the paper's
+original MiniCPM-o inference implementation.
+
+Start the duplex generation server:
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
+    --served-model-name openbmb/MiniCPM-o-4_5 \
+    --host 0.0.0.0 --port 8099
+```
+
+Start a separate multimodal judge. The allowed path must contain any local
+videos passed with `--judge-video-mode video_url`:
+
+```bash
+vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
+    --served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
+    --allowed-local-media-path /data/omni-duplex-eval \
+    --host 0.0.0.0 --port 8000
+```
+
+Generate, evaluate, and summarize:
+
+```bash
+vllm bench omni-duplex-eval --omni generate \
+    --url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
+    --model openbmb/MiniCPM-o-4_5 \
+    --ref-audio /data/ref.wav \
+    --dataset Hothan/Omni-DuplexEval \
+    --concurrency 2 \
+    --response-root /data/omni-duplex-eval/responses
+
+vllm bench omni-duplex-eval --omni evaluate \
+    --dataset Hothan/Omni-DuplexEval \
+    --response-root /data/omni-duplex-eval/responses \
+    --score-root /data/omni-duplex-eval/scores \
+    --judge-base-url http://127.0.0.1:8000 \
+    --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+    --judge-video-mode video_url \
+    --eval-workers 4
+
+vllm bench omni-duplex-eval --omni summarize \
+    --score-root /data/omni-duplex-eval/scores
+```
+
+The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
+splits, not dataset configurations. Pass one with `--split`, or omit it to
+run all nine splits. `--limit 1` is useful for a smoke test. Realtime
+generation records `clock=media`; artifacts generated with
+`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
+them unless `--allow-invalid-clock` is explicit.
+
 ## Related examples
 
 - [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)

examples/online_serving/minicpmo/realtime_duplex_demo.py

@@ -12,7 +12,9 @@ from __future__ import annotations
 import argparse
 import asyncio
 import base64
+import io
 import json
+import math
 import sys
 import time
 from pathlib import Path
@@ -22,8 +24,8 @@ if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
 from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
-    PCM16_BYTES_PER_SAMPLE,  # noqa: F401 - compatibility export
-    PCM16_SAMPLE_RATE,  # noqa: F401 - compatibility export
+    PCM16_BYTES_PER_SAMPLE,
+    PCM16_SAMPLE_RATE,
     RealtimeDuplexClient,
     RealtimeEventCollector,
     build_realtime_url,
@@ -38,6 +40,10 @@ from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
 from vllm_omni.experimental.fullduplex.client import (  # noqa: E402
     reference_audio_data_url as _ref_audio_data_url,
 )
+from vllm_omni.experimental.fullduplex.video_stacking import (  # noqa: E402
+    concat_frames_b64,
+    unit_subframe_offsets,
+)
 
 
 class _StreamingOutputWriter:
@@ -159,14 +165,258 @@ def _event_count_after(
     return sum(event.get("type") == event_type for event in events[index + 1 :])
 
 
-async def run_demo(args: argparse.Namespace) -> dict[str, object]:
-    input_pcm16 = read_pcm16_wav(Path(args.input_wav))
-    if not input_pcm16:
-        raise ValueError("input WAV has no audio")
+def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]:
+    """Frame count and average FPS via PyAV (same approach as Daily-Omni)."""
+    import av
 
+    with av.open(str(video_path)) as container:
+        stream = container.streams.video[0]
+        avg_fps = float(stream.average_rate) if stream.average_rate else 0.0
+        num_frames = stream.frames or 0
+        if num_frames <= 0:
+            duration: float | None = None
+            if stream.duration is not None and stream.time_base is not None:
+                duration = float(stream.duration * stream.time_base)
+            elif container.duration is not None:
+                duration = container.duration / av.time_base
+            if duration and avg_fps:
+                num_frames = int(math.floor(duration * avg_fps))
+        if num_frames <= 0:
+            stream.thread_type = "AUTO"
+            num_frames = sum(1 for _ in container.decode(stream))
+    if num_frames <= 0:
+        raise ValueError(f"No decodable video frames in {video_path}")
+    return num_frames, avg_fps or 1.0
+
+
+def _upright(image, rotation: int):
+    """Apply a container display-matrix rotation, in degrees, to ``image``.
+
+    PyAV hands back the stored raster and leaves the display matrix for the
+    caller, unlike the ffmpeg CLI which autorotates. Phone clips routinely
+    carry +/-90, so skipping this feeds the model a sideways frame: arrows and
+    digits stop being readable and the model mirrors the motion it reports.
+    """
+    from PIL import Image
+
+    quarter_turns = {
+        90: Image.Transpose.ROTATE_90,
+        180: Image.Transpose.ROTATE_180,
+        270: Image.Transpose.ROTATE_270,
+    }
+    turn = quarter_turns.get(int(rotation) % 360)
+    return image.transpose(turn) if turn is not None else image
+
+
+def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list:
+    """Decode presentation-order ``frame_idx`` into upright RGB ``PIL.Image``s."""
+    import av
+
+    wanted = sorted({int(i) for i in frame_idx})
+    if not wanted:
+        return []
+
+    decoded: dict[int, object] = {}
+    with av.open(str(video_path)) as container:
+        stream = container.streams.video[0]
+        stream.thread_type = "AUTO"
+        cursor = 0
+        for position, frame in enumerate(container.decode(stream)):
+            if wanted[cursor] > position:
+                continue
+            image = _upright(frame.to_image(), getattr(frame, "rotation", 0) or 0)
+            while cursor < len(wanted) and wanted[cursor] <= position:
+                decoded[wanted[cursor]] = image
+                cursor += 1
+            if cursor >= len(wanted):
+                break
+
+    if not decoded:
+        raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}")
+    last_image = decoded[max(decoded)]
+    return [decoded.get(int(i), last_image) for i in frame_idx]
+
+
+def _resize_max_side(image, max_side: int):
+    width, height = image.size
+    longest = max(width, height)
+    if longest <= max_side or max_side <= 0:
+        return image
+    scale = max_side / float(longest)
+    new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale))))
+    return image.resize(new_size)
+
+
+def _jpeg_bytes(image, *, quality: int = 95) -> bytes:
+    buf = io.BytesIO()
+    image.convert("RGB").save(buf, format="JPEG", quality=quality)
+    return buf.getvalue()
+
+
+def _to_pcm16_bytes(audio_np) -> bytes:
+    import numpy as np
+
+    arr = np.asarray(audio_np).reshape(-1)
+    if arr.dtype != np.int16:
+        arr = (np.clip(arr, -1.0, 1.0) * 32767.0).astype(np.int16)
+    return arr.tobytes()
+
+
+def _pcm16_wav_duration_s(path: Path) -> float:
+    pcm16 = read_pcm16_wav(path)
+    return len(pcm16) / float(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+
+
+def _extract_video_frames(
+    video_path: Path,
+    work_dir: Path,
+    *,
+    fps: float,
+    max_side: int,
+    duration_s: float,
+    stack_frames: int = 1,
+) -> tuple[list[str], list[str | None]]:
+    """Demux JPEG frames only. Frame count follows ``duration_s``, not soundtrack.
+
+    Returns the base frames (one per model unit) and the parallel stacked track.
+    With ``stack_frames=N`` each unit also samples its ``N-1`` interior
+    sub-frames and tiles them into one composite, which is how official duplex
+    raises the visual refresh rate while keeping one second of audio per unit.
+    A unit whose interior falls past the end of the clip stacks nothing.
+    """
+    num_frames, avg_fps = _probe_video_frames_and_fps(video_path)
+    sample_fps = max(float(fps), 1e-6)
+    unit_s = 1.0 / sample_fps
+    timeline_s = max(float(duration_s), unit_s)
+    num_samples = max(1, int(math.ceil(timeline_s * sample_fps - 1e-9)))
+    offsets = unit_subframe_offsets(stack_frames, unit_s=unit_s)
+
+    def frame_at(seconds: float) -> int:
+        # Official get_video_frame_audio_segments uses truncating int(), not
+        # round(), so 29.97 fps lands on the same presentation index.
+        return min(int(seconds * avg_fps), num_frames - 1)
+
+    base_idx = [frame_at(unit * unit_s) for unit in range(num_samples)]
+    sub_times = [
+        [unit * unit_s + offset for offset in offsets if unit * unit_s + offset < timeline_s]
+        for unit in range(num_samples)
+    ]
+    sub_idx = [frame_at(seconds) for times in sub_times for seconds in times]
+    images = _decode_video_frames_rgb(video_path, base_idx + sub_idx)
+    sub_images = iter(images[len(base_idx) :])
+
+    frame_dir = work_dir / "frames"
+    frame_dir.mkdir(parents=True, exist_ok=True)
+    for stale in frame_dir.glob("*.jpg"):
+        stale.unlink()
+    frames_b64: list[str] = []
+    stacked_b64: list[str | None] = []
+    for index, (image, times) in enumerate(zip(images[: len(base_idx)], sub_times), start=1):
+        jpeg = _jpeg_bytes(_resize_max_side(image, max_side))
+        (frame_dir / f"frame_{index:05d}.jpg").write_bytes(jpeg)
+        frames_b64.append(base64.b64encode(jpeg).decode("ascii"))
+        interior = [_resize_max_side(next(sub_images), max_side) for _ in times]
+        if not interior:
+            stacked_b64.append(None)
+            continue
+        composite = concat_frames_b64(interior)
+        (frame_dir / f"stack_{index:05d}.jpg").write_bytes(base64.b64decode(composite))
+        stacked_b64.append(composite)
+    return frames_b64, stacked_b64
+
+
+def _extract_video_soundtrack(video_path: Path, work_dir: Path) -> Path:
+    """Write a 16 kHz mono WAV from the video soundtrack."""
+    try:
+        from vllm.multimodal.media.audio import load_audio
+    except ImportError as exc:  # pragma: no cover
+        raise SystemExit("--input-video needs vLLM audio helpers (load_audio)") from exc
+
+    audio_np, sr = load_audio(str(video_path), sr=PCM16_SAMPLE_RATE, mono=True)
+    work_dir.mkdir(parents=True, exist_ok=True)
+    audio_path = work_dir / "input_16k.wav"
+    write_pcm16_wav(audio_path, _to_pcm16_bytes(audio_np), sample_rate_hz=int(sr))
+    return audio_path
+
+
+def _extract_video_input(
+    video_path: Path,
+    work_dir: Path,
+    *,
+    fps: float,
+    max_side: int,
+    stack_frames: int,
+) -> tuple[Path, list[str], list[str | None]]:
+    """Demux a video into duplex inputs: 16 kHz mono WAV plus JPEG frames.
+
+    Use this only when the video soundtrack is the audio source. When the
+    caller already has ``--input-wav``, extract frames against that WAV
+    duration instead so a silent clip cannot fail ``load_audio``.
+    """
+    audio_path = _extract_video_soundtrack(video_path, work_dir)
+    frames_b64, stacked_b64 = _extract_video_frames(
+        video_path,
+        work_dir,
+        fps=fps,
+        max_side=max_side,
+        duration_s=_pcm16_wav_duration_s(audio_path),
+        stack_frames=stack_frames,
+    )
+    return audio_path, frames_b64, stacked_b64
+
+
+def _resolve_duplex_av_inputs(
+    *,
+    input_wav: str | None,
+    input_video: str | None,
+    work_dir: Path,
+    fps: float,
+    max_side: int,
+    stack_frames: int = 1,
+) -> tuple[str, list[str], list[str | None]]:
+    """Return ``(wav_path, frames_b64, stacked_b64)``. External WAV wins over video audio."""
+    if input_video:
+        video_path = Path(input_video)
+        if input_wav:
+            frames, stacked = _extract_video_frames(
+                video_path,
+                work_dir,
+                fps=fps,
+                max_side=max_side,
+                duration_s=_pcm16_wav_duration_s(Path(input_wav)),
+                stack_frames=stack_frames,
+            )
+            return input_wav, frames, stacked
+        extracted_wav, frames, stacked = _extract_video_input(
+            video_path,
+            work_dir,
+            fps=fps,
+            max_side=max_side,
+            stack_frames=stack_frames,
+        )
+        return str(extracted_wav), frames, stacked
+    if not input_wav:
+        raise SystemExit("provide --input-wav, --input-video, or both")
+    return input_wav, [], []
+
+
+async def run_demo(args: argparse.Namespace) -> dict[str, object]:
     output_dir = Path(args.output_dir)
     stream_writer = _StreamingOutputWriter(output_dir)
 
+    input_wav, video_frames, stacked_video_frames = _resolve_duplex_av_inputs(
+        input_wav=args.input_wav,
+        input_video=args.input_video,
+        work_dir=output_dir / "video_input",
+        fps=args.video_fps,
+        max_side=args.frame_max_side,
+        stack_frames=args.stack_frames,
+    )
+
+    input_pcm16 = read_pcm16_wav(Path(input_wav))
+    if not input_pcm16:
+        raise ValueError("input WAV has no audio")
+
     url = build_realtime_url(
         args.url,
         args.model,
@@ -184,10 +434,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
             timeout_s=args.timeout_s,
         )
         stream_event_cursor = len(client.events.events)
-        await client.stream_pcm16(
+        frames_sent = await client.stream_pcm16(
             input_pcm16,
             chunk_ms=args.chunk_ms,
             realtime=not args.no_realtime_pacing,
+            video_frames=video_frames,
+            stacked_video_frames=stacked_video_frames,
         )
         commit_event_cursor = len(client.events.events)
         stream_decision = _latest_model_decision(client.events.events, stream_event_cursor)
@@ -310,6 +562,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
                     committed_index,
                 ),
             },
+            "input_video": args.input_video,
+            "video_frames_extracted": len(video_frames),
+            "video_frames_sent": frames_sent,
+            "video_stacked_frames_extracted": sum(frame is not None for frame in stacked_video_frames),
             "audio_bytes": len(audio),
             "audio_chunk_count": len(stream_writer.audio_chunk_paths),
             "audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths],
@@ -356,7 +612,48 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
     parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
     parser.add_argument("--session-id")
-    parser.add_argument("--input-wav", required=True)
+    parser.add_argument(
+        "--input-wav",
+        help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
+    )
+    parser.add_argument(
+        "--input-video",
+        help=(
+            "Video to stream as omni duplex input. PyAV demuxes frames and "
+            "vLLM load_audio extracts a 16 kHz mono WAV (unless --input-wav overrides the audio)."
+        ),
+    )
+    parser.add_argument(
+        "--video-fps",
+        type=float,
+        default=1.0,
+        help="Frame extraction rate. The wire contract carries ~1 frame per second of audio.",
+    )
+    parser.add_argument(
+        "--frame-max-side",
+        type=int,
+        default=0,
+        help=(
+            "Optional client-side downscale of the longest side, in pixels. "
+            "0 (default) keeps the capture resolution so the server "
+            "process_image can normalize at scale_resolution=448 from a "
+            "higher-quality source. The web demo never resizes on the client."
+        ),
+    )
+    parser.add_argument(
+        "--stack-frames",
+        type=int,
+        default=1,
+        help=(
+            "Frames sampled per 1 s audio unit. 1 (default) sends the unit's "
+            "frame alone. N > 1 also samples the N-1 sub-frames inside the unit "
+            "and tiles them into one composite sent next to the base frame, so "
+            "motion (arrows, doors, digits changing) is visible at N fps while "
+            "the audio stays 1 s per unit. Official uses 5 for high refresh "
+            "rate mode; the wire carries 2 images per unit at any N because the "
+            "sub-frames share one composite."
+        ),
+    )
     parser.add_argument(
         "--ref-audio",
         required=True,
@@ -372,7 +669,7 @@ def parse_args() -> argparse.Namespace:
         "--temperature",
         type=float,
         default=None,
-        help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.",
+        help="Stage 0 sampling temperature; omit to preserve the model default.",
     )
     parser.add_argument("--no-realtime-pacing", action="store_true")
     parser.add_argument("--require-audio", action="store_true")

examples/online_serving/qwen3_omni/README.md

@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --deploy-config /path/to/your_deploy_config.yaml
 ```
 
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
 For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
 use `--stage-overrides` on top of the default config:
 

pyproject.toml

@@ -71,11 +71,11 @@ dev = [
     "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
     "torchdiffeq==0.2.5", # for covo_audio Test
     "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
-    "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
     "s3tokenizer==0.3.0", # for tts test
     # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
     # whose top-level stepaudio2 package conflicts with this distribution.
     "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
+    "silero-vad==6.1", # Realtime server-VAD duplex E2E
 ]
 
 demo = [
@@ -85,6 +85,10 @@ demo = [
     "cn2an>=0.5.22",  # Chinese number normalization for TTS (e.g. "1次" -> "一次")
 ]
 
+server-vad = [
+    "silero-vad==6.1",
+]
+
 forced-aligner = [
     "qwen-asr>=0.0.6",  # faithful Japanese/Korean word segmentation for the TTS forced aligner
 ]

recipes/LTX/LTX-2.5.md

@@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized
 48 kHz stereo audio.
 
 The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
-repository is not directly loadable with `--model`. It supplies the official
-upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
-both model licenses and authenticate before first use.
+repository is not directly loadable with `--model`. It supplies the canonical
+Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
+the Full/SFT two-stage pipeline, so accept both model licenses and authenticate
+before first use.
+
+The vLLM-Omni integration code is provided under the project's
+[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate
+artifacts covered by the
+[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md),
+not the vLLM-Omni license. Users are responsible for reviewing and complying
+with the model license before downloading or using the weights.
 
 ## Choose a pipeline
 
 | Pipeline | Mode | Default size | Schedule |
-|---|---|---:|---:|
+| --- | --- | ---: | ---: |
 | `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
 | `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
 | `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
@@ -28,6 +36,40 @@ latent upsampler, and run a three-step refinement stage. Select the class with
 `--model-class-name`; no `--task-type` flag is required. Supplying one initial
 image selects I2V, while omitting it selects T2V.
 
+## Video decoder
+
+LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in
+to the legacy convolutional VAE, set the startup-only stage-0 model extra
+`ltx2_use_conv_vae: true`. The choice applies to all requests and all four
+pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub
+during startup. `decode_timestep` and `decode_noise_scale` condition only the
+legacy ConvVAE and have no effect when DiffVAE is selected.
+
+For offline Python usage:
+
+```python
+omni = Omni(
+    model="Lightricks/LTX-2.5-Diffusers",
+    model_class_name="LTX2Pipeline",
+    stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
+)
+```
+
+For online serving:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+  --omni \
+  --model-class-name LTX2Pipeline \
+  --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'
+```
+
+Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving
+serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial
+dimension. DiffVAE also supports distributed VAE decode; see the
+[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md).
+DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder.
+
 ## Prerequisites
 
 ```bash
@@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
 ## Hardware
 
 | GPU | Status | Recommended scope |
-|---|---|---|
+| --- | --- | --- |
 | NVIDIA B300 | Verified | All four canonical pipelines |
 | NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
 | NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |

recipes/MammothModa2/MammothModa2.md

@@ -0,0 +1,212 @@
+# MammothModa2
+
+> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
+
+## Summary
+
+- Vendor: ByteDance Research
+- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
+- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run MammothModa2-Preview through shared task-oriented
+offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
+The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
+and forwards MammothModa2-specific generation parameters through the
+pipeline-declared `extra_body` contract.
+
+MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
+interface (not `OmniDiffusionRequest`), so its generation knobs
+(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
+`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
+flags. Image size uses the standard `--height` / `--width` flags.
+
+## References
+
+- Upstream model:
+  [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
+- Dev model:
+  [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
+- Related offline example:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- Related T2T/I2T example:
+  [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
+- Declared parameters:
+  [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
+- Deploy config:
+  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
+
+## Hardware Support
+
+The default deploy config runs both the AR and DiT stages on a single GPU
+(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
+stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
+after rebalancing the split so the AR weights (~23 GB) leave room for the KV
+cache — see the note under *1x L40S 48GB*.
+
+## GPU
+
+### 1x L40S 48GB
+
+> **48 GB config adjustment:** the committed
+> `vllm_omni/deploy/mammoth_moda2.yaml` uses
+> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
+> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
+> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
+  ~80 GB GPU for the default config
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Offline Commands
+
+Download the model:
+
+```bash
+hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+Run text-to-image with the shared offline example from the repository root. The
+deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
+MammothModa2 generation parameters as a JSON object through `--extra-body`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model ./MammothModa2-Preview \
+  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+  --output mammoth_t2i.png
+```
+
+The `--extra-body` JSON forwards MammothModa2-specific parameters into
+`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
+declared `extra_body_params` (see
+[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
+so unknown keys for MammothModa2 are silently dropped:
+
+- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
+  (default `9.0`; CFG is active only when `> 1.0`).
+- `cfg_range` — relative step range `[start, end]` over which CFG is applied
+  (default `[0.0, 1.0]`).
+- `num_inference_steps` — number of DiT denoising steps (default `50`).
+
+`--height` and `--width` must be multiples of 16.
+
+Run text-to-text through the shared understanding example. It recognizes the
+MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Preview \
+  --prompt "Explain multimodal generation in three sentences."
+```
+
+Add an image for image-to-text or image summarization. The shared example
+uses MammothModa2's chat and vision-token template:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Preview \
+  --image /path/to/input.jpg \
+  --prompt "Please summarize the content of this image."
+```
+
+#### Verification
+
+The example writes the generated image to the `--output` path. Confirm the file
+exists and is a valid image:
+
+```bash
+ls -lh mammoth_t2i.png
+python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
+```
+
+### 1x AMD MI300X, MammothModa2 Preview
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Offline Commands
+
+The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+    --model bytedance-research/MammothModa2-Preview \
+    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+    --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+    --height 1024 \
+    --width 1024 \
+    --seed 42 \
+    --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+    --enable-diffusion-pipeline-profiler \
+    --log-stats \
+    --output mammoth_t2i.png
+```
+
+#### Verification
+
+The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
+
+The output was a valid 1024 by 1024 RGB PNG.
+
+## MammothModa2-Dev unified inference
+
+MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
+Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
+`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
+required.
+
+Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
+Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
+then sends the generated visual tokens and hidden states to the DiT stage.
+
+Download the checkpoint:
+
+```bash
+hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
+```
+
+Run text-to-text through the shared understanding example. It recognizes the
+Dev checkpoint as MammothModa2 and automatically selects
+`mammoth_moda2_ar.yaml`:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Dev \
+  --prompt "Explain multimodal generation in three sentences."
+```
+
+Add an image for image-to-text or image summarization:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Dev \
+  --image ./image.png \
+  --prompt "Please summarize the content of this image."
+```
+
+The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
+run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
+and encoder caches. Allow additional GPU memory for those caches and the input
+image.

recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md

@@ -0,0 +1,82 @@
+# MiniMax H3 disaggregated text encoder
+
+This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
+its hidden states and token-role metadata to an encoder-free diffusion stage.
+The standard MiniMax H3 recipes remain single-stage and continue to load the
+text encoder inside the diffusion pipeline.
+
+## Start the server
+
+Choose the topology explicitly and load its deployment defaults:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --omni \
+  --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml
+```
+
+The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size
+2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel
+size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the
+`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a
+deployment override for the available hardware. Diffusion quantization,
+layerwise offload, distributed layerwise offload, VAE parallelism, and USP
+settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
+
+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:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
+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}}'
+```
+
+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
+encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
+one offload strategy per deployment:
+
+```bash
+# Stage 1 model-level CPU offload.
+--stage-overrides '{"1":{"enable_cpu_offload":true}}'
+
+# Stage 1 distributed layerwise offload. Tune resident layers for available RAM.
+--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}'
+
+# Stage 1 online FP8 quantization of the DiT only.
+--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}'
+```
+
+The Stage 1 VAE patch-parallel options remain independent of offload and
+quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
+hardware-qualified profiles before combining these options.
+
+Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
+text-encoder weights. 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.
+
+The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
+`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
+
+## Turbo LoRA
+
+MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
+`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
+requests that omit sampling controls do not inherit the 50-step base schedule:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --omni \
+  --lora-path /path/to/MiniMax-H3-Turbo \
+  --deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
+```
+
+The base deployment intentionally retains 50 inference steps for non-LoRA
+quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
 export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
 ```
 
-Start from a non-offloaded FL2VA server command and add
+Start from a non-offloaded or DLO FL2VA server command and add
 `--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
 `--lora-path` preloads the adapter; each request still activates it and uses
 the published sampling settings:
@@ -853,9 +853,10 @@ the published sampling settings:
 
 For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
 ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
-dynamic-only and does not support prefusion, DLO, or LoRA composition.
-It also rejects model-level, layerwise, and distributed layerwise offload;
-the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
+dynamic-only and does not support prefusion or LoRA composition. DLO is
+supported by keeping the request-switchable LoRA A/B buffers resident on the
+accelerator while DLO streams only the base blocks; budget for this additional
+fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
 The five requested sigma points produce the four denoiser evaluations expected
 by the Turbo artifact.
 

recipes/NVIDIA/NemotronLabs-VoiceChat.md

@@ -1,7 +1,7 @@
 # NemotronLabs VoiceChat 11B
 
-> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
-> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
+> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
+> inference and experimental model-native, frame-locked Realtime serving.
 
 ## Summary
 
@@ -11,8 +11,8 @@
   timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
   per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
   31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
-- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
-  function calling, and batch>1 are documented follow-ups.
+- Mode: Offline single-turn inference (batch=1), plus experimental native
+  duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
 - Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
 
 ## When to use this recipe
@@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
   [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
 - Staged pipeline config:
   [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
+- Native duplex config and E2E probe:
+  [`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
+  [`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
+- Nightly model-level gate:
+  [`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
 - Upstream model card:
   [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
 - Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
@@ -60,7 +65,7 @@ output stayed within one word of the reference in testing.
 
 - OS: Linux
 - Python: 3.12
-- vLLM version: 0.26.0
+- vLLM version: 0.27.0
 - vLLM-Omni version or commit: this PR / current `main`
 
 #### Command
@@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in
 the input WAV, and the WAV should transcribe to (approximately) the same text
 with any ASR model.
 
+#### Native duplex serving
+
+The checkpoint does not contain the underlying Nemotron text tokenizer. Set
+`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
+`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
+
+```bash
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
+
+vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
+  --omni \
+  --served-model-name nemotron-voicechat \
+  --deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
+```
+
+The single-GPU profile runs three eager fp32 engine processes on one device.
+
+In another shell, stream only the user channel of one of NVIDIA's bundled
+stereo conversation fixtures:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+  --model nemotron-voicechat \
+  --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
+  --input-channel 0 \
+  --output-dir results/nemotron_voicechat_duplex \
+  --timeout-s 600
+```
+
+The probe requires a completed response, non-silent 22.05 kHz output, and the
+advertised native 80 ms append capabilities. Add `--expect-function-call` and
+use `tool_call.wav` to validate the function-call channel. To validate the
+Realtime tool round trip, also provide the expected arguments and return a
+tool result:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+  --model nemotron-voicechat \
+  --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
+  --output-dir results/nemotron_voicechat_tool_call \
+  --expect-function-call \
+  --expected-function-name generate_random_number \
+  --expected-function-arguments '{"min":1,"max":50}' \
+  --function-output 20 \
+  --expected-post-tool-text "random number"
+```
+
 #### Notes
 
 - Memory usage: the shipped yaml runs all three stages on one GPU
@@ -109,5 +163,14 @@ with any ASR model.
   instead of pausing.
 - The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
   token, and the stage context is 16384).
-- Known limitations: batch=1 only; offline single turn (no online/duplex
-  serving yet); the function-call channel is decoded but not yet acted on.
+- The 80 ms frame period is a model/protocol contract, not a throughput
+  guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
+  roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
+  stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
+  The current native duplex path is functionally streaming but not wall-clock
+  realtime.
+- Known limitations: batch=1 only. The native duplex deployment allows one
+  active session and does not support barge-in. Tool execution remains
+  client-owned: the server emits function-call events, accepts a validated
+  `function_call_output`, and resumes the live model with the returned result;
+  it does not execute arbitrary tools itself.

recipes/README.md

@@ -85,8 +85,9 @@ recipes/
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
+| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
-| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
+| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 | [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
 

recipes/XiaomiMiMo/MiMo-Audio.md

@@ -0,0 +1,161 @@
+# MiMo-Audio for omni audio generation and understanding
+
+> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
+
+## Summary
+
+- Vendor: Xiaomi MiMo
+- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
+- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for running
+MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
+The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
+FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
+Stage-1 GPU pressure.
+
+## References
+
+- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
+- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
+- Offline example:
+  [`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
+- Online example:
+  [`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
+- Default deploy config (generic single GPU):
+  [`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
+- 5090D-tuned deploy config:
+  [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+- Pipeline / deploy docs:
+  [Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
+  [PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
+
+Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
+offline with `tts_sft`.
+
+#### Environment
+
+- OS: Ubuntu (Linux x86_64)
+- Python: 3.12
+- PyTorch: 2.11.0+cu129
+- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
+- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
+- vLLM: 0.26.0+cu129
+- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
+
+Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
+On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
+PyTorch 2.11+cu129.
+
+Download or cache:
+
+- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
+
+Export the runtime knobs below before offline or online runs:
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
+export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
+export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
+export VLLM_USE_FLASHINFER_SAMPLER=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+```
+
+Optional (mirror / download stability):
+
+```bash
+export HF_HUB_DISABLE_XET=1
+export HF_ENDPOINT=https://hf-mirror.com   # if you use a Hugging Face mirror
+```
+
+Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
+`attention_backend: TRITON_ATTN` in the deploy YAML instead.
+
+#### Command
+
+Offline TTS validation from the repository root:
+
+```bash
+cd examples/offline_inference/mimo_audio
+
+python3 -u end2end.py \
+  --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --text "The weather is so nice today."
+```
+
+Online serving (requires MiMo chat template):
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+
+vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
+  --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+  --served-model-name "MiMo-Audio-7B-Instruct" \
+  --port 18091 \
+  --chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
+```
+
+Pass the same environment exports to the server process when using a CPU
+tokenizer or disabling FlashInfer sampling.
+
+#### Verification
+
+Offline:
+
+```bash
+ls -lh output_audio/tts_sft/*.wav
+```
+
+Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
+with the decoded text. A successful run completes both stages without a
+Stage-1 `decode_vq` device mismatch.
+
+Online:
+
+```bash
+curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
+```
+
+Then use the client under
+[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
+
+#### Notes
+
+- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
+  with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
+- Key flags:
+  - `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
+    and `TRITON_ATTN`.
+  - `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
+  - `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
+    [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
+    fix) so code2wav decode uses the tokenizer device.
+  - `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
+    issues on consumer GPUs.
+- Known limitations:
+  - Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
+    tooling; the tested path uses TRITON attention and cu129 wheels.
+  - Without a working `flash-attn` build, audio quality may be slightly
+    metallic; install FlashAttention when a compatible wheel is available.
+  - First request triggers Triton JIT for rotary / attention kernels (latency
+    spike); subsequent requests are faster.
+  - For longer outputs or batching, you may need to raise `max_model_len` in
+    the deploy YAML and match `max_position_embeddings` in the model config
+    (see the offline example README).

recipes/zai-org/GLM-TTS.md

@@ -1,4 +1,4 @@
-# GLM-TTS for Chinese/English TTS on 1x GPU
+# GLM-TTS for Chinese/English TTS
 
 ## Summary
 
@@ -103,3 +103,137 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
 - Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
 - Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
 - Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
+
+### 2x non-standard RTX 4090 48GB
+
+This configuration was personally validated on two non-standard RTX 4090 cards
+reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not
+validated.
+
+#### Environment
+
+- OS: Ubuntu 22.04.1 LTS
+- NVIDIA driver: 580.76.05
+- Python: 3.12.13
+- PyTorch: 2.11.0+cu130
+- CUDA runtime and JIT toolkit: 13.0
+- vLLM: 0.26.0
+- vLLM-Omni: commit `66a9b0c8`, version
+  `0.26.0rc2.dev5+g66a9b0c84`
+- FlashInfer: 0.6.14
+
+The host's default `nvcc` was 11.8. The working CUDA toolkit came from the
+Python environment, so `CUDA_HOME` was set to
+`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`.
+
+#### Command
+
+From the repository root, start the async-chunk server with:
+
+```bash
+export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13"
+export PATH="$CUDA_HOME/bin:$PATH"
+export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"
+
+CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \
+  --omni \
+  --trust-remote-code \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/glm_tts.yaml \
+  --stage-overrides '{"1":{"devices":"1"}}' \
+  --allowed-local-media-path "${PWD}/tests/assets"
+```
+
+Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync
+variant, add `--no-async-chunk` to the same command.
+
+#### Verification
+
+Run the Chinese voice-cloning check:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --model zai-org/GLM-TTS \
+  --text "今天天气很好，我们一起学习语音模型推理。" \
+  --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
+  --ref-text "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" \
+  --response-format wav \
+  --output glm_tts_zh.wav
+```
+
+Run the English voice-cloning check with the same Chinese reference:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+  --model zai-org/GLM-TTS \
+  --text "Hello, this is a voice cloning test with GLM-TTS." \
+  --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
+  --ref-text "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" \
+  --response-format wav \
+  --output glm_tts_en.wav
+```
+
+Inspect either output:
+
+```bash
+ffprobe -v error -select_streams a:0 \
+  -show_entries stream=codec_name,sample_rate,channels \
+  -of default=noprint_wrappers=1 glm_tts_zh.wav
+```
+
+Expected output:
+
+```text
+codec_name=pcm_s16le
+sample_rate=24000
+channels=1
+```
+
+The English smoke test produced a valid WAV with an actual duration of 5.020
+seconds.
+
+#### Observed warm-cache results
+
+These observations are from one controlled run, not a hardware or performance
+guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard
+deviation.
+
+##### Measurement procedure
+
+- For the single-GPU server, use the documented command with
+  `CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages
+  remain on GPU 0. For the dual-GPU server, use the command as shown.
+- Async chunking is the default. Add `--no-async-chunk` for sync mode.
+- For non-streaming measurements, use the exact Chinese client command above.
+  For streaming, add `--stream` to that command and change the output filename
+  to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV.
+- The fixed Chinese prompt and reference were exactly those shown in
+  Verification and yielded 3.900 seconds of audio. Each case used one warm-up
+  followed by three measured requests at concurrency 1.
+- Wall time starts immediately before the HTTP POST and stops after the complete
+  response body. TTFA runs from immediately before the POST to the first
+  non-empty audio bytes. RTF is wall time divided by audio duration.
+- Peak total device memory used was sampled every 0.2 seconds and reported by
+  NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and
+  baseline use and are not baseline-subtracted.
+
+| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF |
+|---:|---|---|---:|---:|---:|
+| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 |
+| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 |
+| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 |
+| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 |
+| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 |
+| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 |
+
+Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It
+reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync
+wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only
+within matching execution and response modes.
+
+| Configuration | GPU 0 | GPU 1 |
+|---|---:|---:|
+| Single-GPU async | 31.52 GiB | -- |
+| Dual-GPU async | 29.23 GiB | 2.91 GiB |
+| Single-GPU sync | 31.47 GiB | -- |
+| Dual-GPU sync | 29.23 GiB | 2.91 GiB |

requirements/common.txt

@@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.40.0
+# Hub kernels, including LTX-2.5 NATTEN
+kernels==0.14.1
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1

tests/benchmarks/duplex/test_omni_duplex_eval.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
+    extract_timed_sentences,
+    normalize_response_items,
+    split_text,
+    validate_clock,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
+    DuplexSample,
+    canonical_task_type,
+    family_for_split,
+    load_samples,
+    task_type_for_split,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
+    PROTOCOL_PIN,
+    build_content_prompt,
+    build_reminder_prompt,
+    build_temporal_prompt,
+    parse_judge_json,
+    reminder_window,
+    summarize_pr_results,
+    temporal_window,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dataset_mapping_and_manifest(tmp_path):
+    manifest = tmp_path / "samples.json"
+    manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
+    sample = load_samples(manifest, media_root=tmp_path)[0]
+    assert sample.family == "rtd"
+    assert sample.video == str(tmp_path / "clip.mp4")
+    assert family_for_split("PR_event_reminder") == "pr"
+    assert canonical_task_type("post-event-reminder") == "post_event_reminder"
+    assert task_type_for_split("PR_correction") == "correction"
+    assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
+
+
+def test_hugging_face_subset_names_are_splits(monkeypatch):
+    calls = []
+
+    def fake_load_dataset(name, **kwargs):
+        calls.append((name, kwargs))
+        return [{"id": "pr", "question_text": "Correct this."}]
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
+    assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
+    assert sample.split == "PR_correction"
+    assert sample.task_type == "correction"
+
+
+def test_response_aliases_and_clock_guard():
+    assert split_text("One. Two!") == ["One.", "Two!"]
+    assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
+        {"sentence": "x", "start": 800.0, "end": 800.0}
+    ]
+    with pytest.raises(ValueError, match="clock=invalid"):
+        validate_clock({"clock": "invalid"})
+    validate_clock({"clock": "invalid"}, allow_invalid=True)
+    timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
+    assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
+    with pytest.raises(ValueError, match="only clock=media"):
+        extract_timed_sentences([], clock="wall")
+
+
+def test_protocol_windows_and_parsing():
+    assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
+    assert temporal_window(4, 5, 10) == (2.0, 3.0)
+    assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
+    assert reminder_window(5) == (5.0, 15.0)
+    assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
+        "content_score": 2.35,
+        "is_relevant": 1,
+    }
+    assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
+    assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
+
+
+def test_protocol_prompt_wording_is_pinned():
+    prompts = [
+        build_temporal_prompt(1, 2, "response", "question"),
+        build_content_prompt("response", "question", ["one", "two"]),
+        build_reminder_prompt("instruction", "response", "proactive_reminder"),
+        build_reminder_prompt("instruction", "response", "post_event_reminder"),
+        build_reminder_prompt("instruction", "response", "correction", "answer"),
+    ]
+    assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
+        "5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
+        "1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
+        "d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
+        "08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
+        "3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
+    ]
+
+
+def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
+    response = tmp_path / "response.json"
+    response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
+    response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
+    video = tmp_path / "video.mp4"
+    video.write_bytes(b"video")
+    monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
+
+    class Judge:
+        def temporal(self, *args, **kwargs):
+            return '{"temporal_score": 3, "is_relevant": 1}'
+
+        def content(self, prompt, video, frames=None, *, mode="video_url"):
+            assert mode == "frame-sample"
+            assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
+            return '{"content_score": 3}'
+
+    sample = DuplexSample(
+        id="rtd",
+        split="RTD_OCR",
+        family="rtd",
+        task_type=None,
+        video=video,
+        video_duration=4,
+    )
+    score = eval_module.evaluate_sample(
+        sample,
+        response,
+        tmp_path / "score.json",
+        Judge(),
+        judge_video_mode="frame-sample",
+    )
+    assert score["content"]["frame_count"] == 2

tests/benchmarks/test_daily_omni_pack_modes.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
 
 The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
@@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
 
 if _MODULE_NAME not in sys.modules:
     _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    if _spec is None or _spec.loader is None:
+        raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}")
     _mod = importlib.util.module_from_spec(_spec)
     sys.modules[_MODULE_NAME] = _mod
     _spec.loader.exec_module(_mod)
@@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
     def fake_extract(video_path, *, audio_path, include_audio, **_kw):
         counter.append(1)
         frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
-        segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+        segments: list[np.ndarray] = (
+            [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+        )
         return frames, segments
 
     monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)

tests/config/test_config_factory.py

@@ -15,11 +15,7 @@ from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 
 from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
 from vllm_omni.config import config_factory as config_factory_module
-from vllm_omni.config.config_factory import (
-    StageConfigFactory,
-    _materialize_object_storage_configs,
-    _name_match_candidate,
-)
+from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
 from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -38,6 +34,7 @@ from vllm_omni.config.stage_config import (
     build_stage_runtime_overrides,
     load_deploy_config,
     merge_pipeline_deploy,
+    normalize_pipeline_cli_overrides,
     pipeline_cfg_resolver,
 )
 from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
@@ -433,6 +430,20 @@ class TestStageConfigFactory:
         assert cfg["final_output"] is True
         assert cfg["final_output_type"] == "image"
 
+    @pytest.mark.parametrize(
+        "model_class_name",
+        [
+            "MiniMaxH3Pipeline",
+            "WanPipeline",
+            "LTX2Pipeline",
+            "Cosmos3OmniDiffusersPipeline",
+        ],
+    )
+    def test_default_video_diffusion_output_type(self, model_class_name):
+        configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
+
+        assert configs[0]["final_output_type"] == "video"
+
     def test_default_diffusion_with_parallel_config(self):
         """Test diffusion config calculates devices from parallel_config."""
 
@@ -538,6 +549,7 @@ class TestPipelineDiscovery:
         """Check that specific models are in OMNI_PIPELINES."""
         assert "qwen2_5_omni" in OMNI_PIPELINES
         assert "qwen3_omni_moe" in OMNI_PIPELINES
+        assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
     def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -562,6 +574,18 @@ class TestPipelineDiscovery:
         assert pipeline.default_deploy_config_name is None
         assert len(pipeline.stages) == 1  # thinker only
 
+    def test_registry_qwen3_omni_thinker_only_key_is_static(self):
+        """Explicit registry key must not re-enter the HF-config resolver."""
+        pipeline = resolve_pipeline_config(
+            "qwen3_omni_moe_thinker_only",
+            Q3_OMNI_ALL_STAGES_HF_CONFIG,
+        )
+        assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
+        assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
+        assert len(pipeline.stages) == 1
+        assert pipeline.stages[0].engine_output_type == "text"
+        assert pipeline.stages[0].final_output_type == "text"
+
     @pytest.mark.parametrize(
         "pipeline",
         [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -1229,14 +1253,22 @@ stages:
         assert deploy.connectors is not None
         assert deploy.platforms is not None
 
-    def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
+    @pytest.mark.parametrize(
+        ("hf_config", "model"),
+        [
+            (Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
+            (Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
+        ],
+        ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
+    )
+    def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
         with patch(
             "vllm_omni.config.config_factory.get_config",
-            return_value=Q3_OMNI_THINKER_HF_CONFIG,
+            return_value=hf_config,
         ):
             pipeline = StageConfigFactory.get_pipeline_config(
-                "Qwen/Qwen3-Omni-30B-A3B-Thinking",
+                model,
                 trust_remote_code=False,
                 deploy_config_path=str(deploy_path),
             )
@@ -1245,6 +1277,7 @@ stages:
         assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
 
         deploy = load_deploy_config(deploy_path)
+        assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
         stages = merge_pipeline_deploy(pipeline, deploy)
 
         assert len(stages) == 1
@@ -1296,6 +1329,8 @@ stages:
     @pytest.mark.parametrize(
         ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
         [
+            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
+            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
             ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
             ("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
             ("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
@@ -1304,6 +1339,13 @@ stages:
             ("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
             ("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
             ("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
+            (
+                "minimax_h3_disaggregated.yaml",
+                "minimax_h3_disaggregated",
+                2,
+                "video",
+                "minimax_h3_disaggregated",
+            ),
         ],
     )
     def test_load_new_registry_backed_deploy_configs(
@@ -1324,9 +1366,91 @@ stages:
         assert len(stages) == stage_count
         assert stages[-1].final_output is True
         assert stages[-1].final_output_type == final_output_type
+        if pipeline_name == "minimax_h3_disaggregated":
+            assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
+            assert (
+                stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
+            )
+            assert "model_path_resolver" not in stages[0].yaml_extras
+            assert "model_path_resolver" not in stages[1].yaml_extras
         if deploy.trust_remote_code is not None:
             assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
 
+    def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
+        from vllm_omni.engine import stage_init_utils
+
+        model_root = tmp_path / "MiniMax-H3"
+        (model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
+        (model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
+        monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
+
+        deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
+        assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
+        stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
+        resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
+
+        assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
+        assert resolved[1]["model"] == str(model_root / "FL2VA")
+        assert all("model_path_resolver" not in engine_args for engine_args in resolved)
+
+    def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
+        pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            pipeline,
+            {"text_encoder_tp_size": 4},
+        )
+
+        assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
+        assert "text_encoder_tp_size" not in stages[0].runtime_overrides
+        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):
+        pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+        deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert stages[0].yaml_engine_args["max_num_seqs"] == 1
+        assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
+        assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
+        parallel = stages[1].yaml_engine_args["parallel_config"]
+        assert parallel["tensor_parallel_size"] == 1
+        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
+
+        turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
+        turbo_stages = merge_pipeline_deploy(pipeline, turbo)
+        turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
+        assert turbo_sampling["num_inference_steps"] == 5
+        assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
+
+    def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
+        pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+        assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
+        assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
+            "stage_0_tensor_parallel_size": 4
+        }
+
+    def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
+        pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+        with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
+            normalized = normalize_pipeline_cli_overrides(
+                pipeline,
+                {"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
+            )
+
+        assert normalized == {"stage_0_tensor_parallel_size": 2}
+
+    def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
+        pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+        with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
+            normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
+
     @pytest.mark.parametrize(
         ("config_json", "model_index", "expected_pipeline"),
         [
@@ -1362,6 +1486,29 @@ stages:
         assert pipeline is not None
         assert pipeline.model_type == expected_pipeline
 
+    def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
+        def get_model_file(filename, _model, revision=None):
+            del revision
+            if filename == "config.json":
+                return None
+            if filename == "model_index.json":
+                return {"_class_name": "MiniMaxH3Pipeline"}
+            return None
+
+        with (
+            patch.object(StageConfigFactory, "get_hf_config", return_value=None),
+            patch(
+                "vllm_omni.config.config_factory.get_hf_file_to_dict",
+                side_effect=get_model_file,
+            ),
+        ):
+            pipeline = StageConfigFactory.get_pipeline_config(
+                model="/models/unrelated-checkpoint-name",
+                trust_remote_code=False,
+            )
+
+        assert pipeline is None
+
     @pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
     def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
         deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
@@ -2936,44 +3083,3 @@ class TestObjectStorageConfigResolution:
 
         matching = "s3://any-bucket/my-cosyvoice3-model"
         assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
-
-
-class TestNameMatchCandidateSnapshotPaths:
-    """Name-based matching must survive resolved HF cache snapshot paths."""
-
-    _SNAPSHOT_REV = "29e01c4e8d000f4bcd70751be16fa94bf3d85a18"
-
-    def test_candidate_recovers_repo_name_from_hf_cache_layout(self):
-        path = f"/data/hub/models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512/snapshots/{self._SNAPSHOT_REV}"
-        assert _name_match_candidate(path) == "Fun-CosyVoice3-0.5B-2512"
-        assert _name_match_candidate(path + "/") == "Fun-CosyVoice3-0.5B-2512"
-
-    def test_candidate_keeps_basename_for_non_cache_paths(self):
-        assert _name_match_candidate("org/My-Model") == "My-Model"
-        assert _name_match_candidate("/data/snapshots-of/My-Model") == "My-Model"
-        # A ``snapshots`` parent without the ``models--<org>--<name>`` encoding
-        # falls back to the basename, same as before.
-        assert _name_match_candidate(f"/data/notcache/snapshots/{self._SNAPSHOT_REV}") == self._SNAPSHOT_REV
-
-    def test_candidate_recovers_bare_repo_name_from_hf_cache_layout(self):
-        # Legacy un-namespaced repos (e.g. ``gpt2``) cache as ``models--<name>``.
-        path = f"/data/hub/models--my-cosyvoice3-model/snapshots/{self._SNAPSHOT_REV}"
-        assert _name_match_candidate(path) == "my-cosyvoice3-model"
-
-    def test_candidate_ignores_malformed_cache_repo_dirs(self):
-        # Hub validation rejects ``--`` inside repo ids, so more than three
-        # ``--``-separated segments cannot be a real cache dir; keep basename.
-        assert _name_match_candidate("/hub/models--a--b--c/snapshots/abc123") == "abc123"
-        assert _name_match_candidate("/hub/models--/snapshots/abc123") == "abc123"
-
-    def test_try_infer_model_type_matches_cosyvoice3_snapshot_dir(self, tmp_path):
-        """Regression for the HF-cache path shape: CosyVoice3 ships an empty
-        config.json and relies on name matching, but a resolved snapshot dir
-        basename is a bare revision hash, so basename-only matching (#5036)
-        misclassified it as a diffusion pipeline and stage init crashed with
-        "Diffusers pipeline index not found"."""
-        snapshot = tmp_path / "hub" / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / self._SNAPSHOT_REV
-        snapshot.mkdir(parents=True)
-        (snapshot / "config.json").write_text("{}")
-
-        assert StageConfigFactory.try_infer_model_type(model=str(snapshot), trust_remote_code=False) == "cosyvoice3"

tests/config/test_environment_variables.py

@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     """Make an inventory expansion an explicit review decision."""
     category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
     assert category_counts == {
-        EnvironmentVariableCategory.PUBLIC_OMNI: 23,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 24,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 56,

tests/config/test_omni_config.py

@@ -92,10 +92,23 @@ def _from_pipeline_key(
     )
 
 
+def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
+    config = _from_pipeline_key(
+        "minimax_h3_disaggregated",
+        cli_overrides={"text_encoder_tp_size": 4},
+    )
+
+    assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
+    diffusion = config.stage_by_id(1)
+    assert diffusion.parallel_config.tensor_parallel_size == 1
+    assert diffusion.parallel_config.text_encoder_tp_size == 1
+
+
 def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
     omni_config = _from_pipeline_key("personaplex")
 
     assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
+    assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
 
 
 def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
@@ -105,6 +118,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
     omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
 
     assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
+    assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
 
 
 @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
@@ -132,6 +146,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
 
         engine_args = legacy_stage.yaml_engine_args
         assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
+        assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
         assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
         assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
         assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -579,6 +594,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "active_stream_window",
+        "session_mode",
         "duplex_max_sessions",
         "enable_sleep_mode",
         "default_sampling_params",
@@ -639,8 +655,8 @@ def test_sub_config_fields_match_structured_scopes():
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
-        "text_encoder_tp_size",
         "vae_parallel_mode",
+        "text_encoder_tp_size",
         "use_hsdp",
         "mask_sp_padding",
         "hsdp_shard_size",
@@ -936,6 +952,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
+def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
+    deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    extras:",
+                "      ltx2_use_conv_vae: true",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
+def test_stage_override_routes_ltx2_conv_vae_extra():
+    stage = _from_pipeline_key(
+        "dreamzero",
+        cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
+    ).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():
     with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")

tests/config/test_pipeline_registry.py

@@ -66,3 +66,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
     """Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
     with pytest.raises(ValueError):
         register_pipeline(custom_resolver)
+
+
+def test_minimax_h3_disaggregation_is_explicit_opt_in():
+    assert "minimax_h3" not in OMNI_PIPELINES
+    pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+    assert isinstance(pipeline, PipelineConfig)
+    assert pipeline.model_type == "minimax_h3_disaggregated"

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
     return sched

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -25,10 +25,10 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
+def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode))
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
@@ -93,6 +93,7 @@ def _run_resumable_segment_stop(
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 
@@ -198,6 +199,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
         inter_stage_output,
         session,
         True,
+        new_token_ids=[42],
         confirmed_num_computed_tokens=26,
     )
 
@@ -226,6 +228,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 
@@ -427,3 +430,71 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
     assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
     assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
     assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+
+
+def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
+    """A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
+    session = _make_request()
+    session.resumable = True
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+    def queue(*requests):
+        result = MagicMock()
+        result.requests = set(requests)
+        result.add_request.side_effect = result.requests.add
+        result.remove_requests.side_effect = result.requests.difference_update
+        return result
+
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
+    sched.running = []
+    sched.waiting = queue()
+    sched.skipped_waiting = queue(session)
+    sched.num_waiting_for_streaming_input = 1
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=True,
+        segment_finished_requests={session.request_id},
+    )
+
+    sched._resume_downstream_chunk_receiver(session)
+    sched._remove_stopped_requests_from_queues(set(), {session})
+
+    assert session.status == RequestStatus.WAITING
+    assert session in sched.waiting.requests
+    assert session not in sched.skipped_waiting.requests
+    assert sched.num_waiting_for_streaming_input == 0
+    assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
+
+
+@pytest.mark.parametrize(
+    ("receives_chunks", "session_mode"),
+    [
+        (False, "duplex"),
+        (True, "turn"),
+    ],
+)
+def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
+    receives_chunks: bool,
+    session_mode: str,
+) -> None:
+    """Only duplex connector-driven receivers resume without an update."""
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
+    sched.num_waiting_for_streaming_input = 1
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=receives_chunks,
+        segment_finished_requests={session.request_id},
+    )
+    sched.skipped_waiting = MagicMock()
+    sched._enqueue_waiting_request = MagicMock()
+
+    sched._resume_downstream_chunk_receiver(session)
+
+    assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+    assert sched.num_waiting_for_streaming_input == 1
+    assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
+    sched.skipped_waiting.remove_requests.assert_not_called()
+    sched._enqueue_waiting_request.assert_not_called()

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     sched.recompute_kv_load_failures = False
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
+
+Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
+non-empty scheduled step and, at the top of ``update_from_output()``,
+advances ``processed_step_seq`` and drains the deferred frees. The Omni
+schedulers each carried only half of that contract:
+
+- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
+  advances) but its reimplemented ``update_from_output()`` never advanced
+  ``processed_step_seq`` nor drained -> deferred blocks accumulate without
+  bound under async/PP KV-consumer configurations.
+- ``OmniGenerationScheduler`` reimplements both: the fast-path
+  ``schedule()`` never advanced ``sched_step_seq`` (so finished requests
+  freed blocks an in-flight GPU step could still be writing) and its
+  ``update_from_output()`` never drained.
+
+These tests pin the restored contract at the call boundary; the drain and
+free logic itself is upstream-tested.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import SchedulingPolicy
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_update_sched(scheduler_cls) -> MagicMock:
+    """MagicMock scheduler driving the real ``update_from_output`` with an
+    empty batch; only the fence state carries real values."""
+    sched = MagicMock()
+    sched.requests = {}
+    sched.perf_metrics = None
+    sched.chunk_transfer_adapter = None
+    sched.defer_block_free = True
+    sched.processed_step_seq = 0
+    return sched
+
+
+def _empty_frame(*, total_num_scheduled_tokens: int):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {}
+    scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = []
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    return scheduler_output, model_runner_output
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
+    sched = _make_update_sched(scheduler_cls)
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 1
+    sched._drain_deferred_frees.assert_called_once_with()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
+    """0-token steps do not advance the fence upstream (their seq was never
+    advanced in schedule()); draining on them would run ahead of the GPU."""
+    sched = _make_update_sched(scheduler_cls)
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 0
+    sched._drain_deferred_frees.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
+    sched = _make_update_sched(scheduler_cls)
+    sched.defer_block_free = False
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 0
+    sched._drain_deferred_frees.assert_not_called()
+
+
+def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
+    """MagicMock scheduler that drives the real generation fast-path
+    ``schedule()`` over one running one-shot request."""
+    request = Request(
+        request_id="req-fence",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+
+    sched = MagicMock()
+    sched.max_num_scheduled_tokens = 8
+    sched._pause_state = PauseState.UNPAUSED
+    sched.requests = {request.request_id: request}
+    sched.running = [request]
+    sched.waiting = []  # falsy: the waiting loop never runs
+    sched.policy = SchedulingPolicy.FCFS
+    sched.chunk_transfer_adapter = None
+    sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
+    sched.num_lookahead_tokens = 0
+    sched.log_stats = False
+    sched.kv_cache_config = MagicMock(kv_cache_groups=[])
+    sched.use_v2_model_runner = False
+    sched.prev_step_scheduled_req_ids = set()
+    sched.needs_kv_cache_zeroing = False
+    sched.connector = None
+    sched.ec_connector = None
+    sched.defer_block_free = defer_block_free
+    sched.sched_step_seq = 0
+    # last_sched_seq stamping order: _update_after_schedule must observe the
+    # already-advanced fence value, as upstream schedule() guarantees.
+    sched.seq_at_update_after_schedule = None
+
+    def _capture(_scheduler_output):
+        sched.seq_at_update_after_schedule = sched.sched_step_seq
+
+    sched._update_after_schedule.side_effect = _capture
+    return sched, request
+
+
+def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
+    sched, request = _make_fast_path_sched(defer_block_free=True)
+
+    OmniGenerationScheduler.schedule(sched)
+
+    assert sched.sched_step_seq == 1
+    assert sched.seq_at_update_after_schedule == 1
+
+
+def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
+    sched, request = _make_fast_path_sched(defer_block_free=False)
+
+    OmniGenerationScheduler.schedule(sched)
+
+    assert sched.sched_step_seq == 0

tests/core/sched/test_omni_sched_ec_request_finish.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
+
+Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
+BEFORE the encoder cache is freed (so the connector can inspect per-request
+state), honors its delayed-free result, and returns the generated
+``ec_transfer_params``; ``update_from_output()`` then carries them into the
+``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
+always returned ``None``; both Omni schedulers discarded the second tuple
+item, so a remote encoder-cache handle could never reach the frontend.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
+
+
+def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
+    """Minimal real-`_free_request` harness; `order` records the EC hook
+    firing relative to the encoder-cache free."""
+    order = MagicMock()
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._omits_kv_transfer_cache = {}
+    sched._connector_finished = lambda request: (False, None)
+    sched.ec_connector = ec_connector
+    if ec_connector is not None:
+        order.attach_mock(ec_connector.request_finished, "request_finished")
+    sched.encoder_cache_manager = MagicMock()
+    order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
+    sched.finished_req_ids = set()
+    sched._new_prompt_len_snapshot = {}
+    sched.finished_req_ids_dict = None
+    sched._should_transfer_kv_for_request = lambda req_id: False
+    sched._free_blocks = MagicMock()
+    sched._free_input_coordinator_request = MagicMock()
+    sched.chunk_transfer_adapter = None
+    return sched, order
+
+
+class _FakeFinishedRequest:
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+
+    def is_finished(self) -> bool:
+        return True
+
+
+def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
+    ec_connector = MagicMock()
+    ec_connector.request_finished.return_value = (False, EC_PARAMS)
+    sched, order = _make_free_request_sched(ec_connector=ec_connector)
+    request = _FakeFinishedRequest("req-ec")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params == EC_PARAMS
+    ec_connector.request_finished.assert_called_once_with(request)
+    hook_calls = [name for name, *_ in order.mock_calls]
+    assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
+    sched._free_blocks.assert_called_once_with(request)
+
+
+def test_free_request_honors_ec_delayed_free() -> None:
+    ec_connector = MagicMock()
+    ec_connector.request_finished.return_value = (True, EC_PARAMS)
+    sched, _ = _make_free_request_sched(ec_connector=ec_connector)
+    request = _FakeFinishedRequest("req-ec-delay")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params == EC_PARAMS
+    sched._free_blocks.assert_not_called()
+
+
+def test_free_request_without_ec_connector_returns_none() -> None:
+    sched, _ = _make_free_request_sched(ec_connector=None)
+    request = _FakeFinishedRequest("req-no-ec")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params is None
+    sched._free_blocks.assert_called_once_with(request)
+
+
+def _make_finishing_request() -> Request:
+    request = Request(
+        request_id="req-ec-finish",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+    return request
+
+
+def _make_finish_sched(request: Request) -> MagicMock:
+    """MagicMock scheduler driving the real ``update_from_output`` through a
+    finishing request; ``_free_request`` reports EC transfer params."""
+    sched = MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._maybe_decode_pooling_output.return_value = None
+    sched._handle_stopped_request.return_value = True
+    sched._free_request.return_value = (None, EC_PARAMS)
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.chunk_transfer_adapter = None
+    sched.finished_req_ids_dict = None
+    sched._new_prompt_len_snapshot = {}
+    return sched
+
+
+def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
+    request = _make_finishing_request()
+    sched = _make_finish_sched(request)
+    sched._update_request_with_output.return_value = ([42], True)
+
+    engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.ec_transfer_params == EC_PARAMS
+    sched._free_request.assert_called_once_with(request)
+
+
+def _finish_frame_parts(request: Request):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    return scheduler_output, model_runner_output
+
+
+def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
+    request = _make_finishing_request()
+    # One-shot generation finish: the whole prompt is computed.
+    request.num_computed_tokens = len(request.prompt_token_ids)
+    sched = _make_finish_sched(request)
+
+    engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.ec_transfer_params == EC_PARAMS
+    sched._free_request.assert_called_once_with(request)

tests/core/sched/test_omni_sched_prefill_stats_finalize.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
+
+Upstream v0.28 finalizes prefill statistics with
+``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
+before emitting them, which computes ``num_cache_creation_tokens`` (tokens
+computed and written to the prefix cache). Both Omni schedulers previously
+emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_request() -> Request:
+    request = Request(
+        request_id="req-prefill-stats",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+    # A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
+    request.prefill_stats.set(
+        num_prompt_tokens=3,
+        num_local_cached_tokens=1,
+        num_external_cached_tokens=0,
+    )
+    return request
+
+
+def _make_sched(request: Request) -> MagicMock:
+    sched = MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._maybe_decode_pooling_output.return_value = None
+    sched._handle_stopped_request.return_value = True
+    # The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
+
+    def _free_and_invalidate(req):
+        # Lifecycle guard: freeing releases the KV blocks, after which the
+        # manager reports 0 cached tokens. If finalize ran after
+        # _free_request (the reviewed bug), creation count would compute
+        # against 0 and the assertions below would see 0, not 2.
+        sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+        return (None, None)
+
+    sched._free_request.side_effect = _free_and_invalidate
+    sched.chunk_transfer_adapter = None
+    sched.finished_req_ids_dict = None
+    sched._new_prompt_len_snapshot = {}
+    return sched
+
+
+def _frame(request: Request):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    return scheduler_output, model_runner_output
+
+
+def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+    request = _make_request()
+    sched = _make_sched(request)
+    sched._update_request_with_output.return_value = ([42], True)
+
+    engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.prefill_stats is not None
+    assert eco.prefill_stats.num_cached_tokens == 1
+    assert eco.prefill_stats.num_cache_creation_tokens == 2
+    sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
+
+
+def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+    request = _make_request()
+    request.num_computed_tokens = len(request.prompt_token_ids)
+    sched = _make_sched(request)
+
+    engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.prefill_stats is not None
+    assert eco.prefill_stats.num_cached_tokens == 1
+    assert eco.prefill_stats.num_cache_creation_tokens == 2
+    sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)

tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json

@@ -0,0 +1,106 @@
+[
+  {
+    "test_name": "test_hunyuanvideo15_i2v_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
+      "serve_args": {
+        "flow-shift": 5.0,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames33_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 33,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 5.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
+      "serve_args": {
+        "flow-shift": 5.0,
+        "cache-backend": "cache_dit",
+        "tensor-parallel-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames33_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 33,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 5.0
+        }
+      }
+    ]
+  }
+]

tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json

@@ -0,0 +1,104 @@
+[
+  {
+    "test_name": "test_hunyuanvideo15_t2v_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+      "serve_args": {
+        "flow-shift": 5.0,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames33_steps4",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 33,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 5.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+      "serve_args": {
+        "flow-shift": 5.0,
+        "cache-backend": "cache_dit",
+        "tensor-parallel-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames33_steps4",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 33,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 5.0
+        }
+      }
+    ]
+  }
+]

tests/diffusion/attention/test_flash_attn.py

@@ -308,6 +308,91 @@ def test_flash_attn_func_preferred_over_varlen():
     print("✓ flash_attn_func forward works correctly!")
 
 
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@pytest.mark.parametrize("k_len", [40, 256])
+def test_cross_attn_key_padding_vs_sdpa(k_len):
+    """
+    Case 3: Cross-attention with a key-padding mask.
+
+    The mask covers only the key sequence; every query position must attend
+    to the valid keys. FA and SDPA outputs should be very close. k_len=256
+    covers equal Q/K lengths, which must not be mistaken for self-attention,
+    and valid key lengths differ per batch element.
+    """
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+
+    batch_size = 2
+    q_len = 256
+    valid_lens = (17, 29)
+    num_heads = 8
+    head_dim = 64
+
+    fa_impl = FlashAttentionImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
+    )
+    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+
+    torch.manual_seed(7)
+    query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
+    key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
+    value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
+
+    attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
+    for i, valid_len in enumerate(valid_lens):
+        attn_mask[i, :valid_len] = True
+    attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+
+    output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
+    output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
+
+    max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
+    mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
+
+    print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
+    print(f"Max absolute difference: {max_diff:.6f}")
+    print(f"Mean absolute difference: {mean_diff:.6f}")
+
+    assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+    assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+
+    print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
+
+
+def test_varlen_masked_routing_by_role(monkeypatch):
+    """The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
+    calls = []
+
+    def fake_varlen_func(q, k, v, **kwargs):
+        calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
+        return q
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+
+    batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
+    query = torch.randn(batch_size, seq_len, num_heads, head_dim)
+    metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
+
+    cross_impl = FlashAttentionImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
+    )
+    output = cross_impl.forward_cuda(query, query, query, metadata)
+    assert output.shape == query.shape
+    num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
+    assert num_q_rows == batch_size * seq_len
+    assert cu_seqlens_q.tolist() == [0, 4, 8]
+    assert cu_seqlens_k.tolist() == [0, 2, 5]
+
+    self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
+    output = self_impl.forward_cuda(query, query, query, metadata)
+    assert output.shape == query.shape
+    num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
+    assert num_q_rows == 5
+    assert cu_seqlens_q.tolist() == [0, 2, 5]
+    assert cu_seqlens_k.tolist() == [0, 2, 5]
+
+
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
 

tests/diffusion/distributed/test_expert_parallel_layout.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
 
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 
+@pytest.mark.cpu
+@pytest.mark.core_model
+@pytest.mark.parametrize(
+    ("fully_shard_degree", "error_message"),
+    [(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
+)
+def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
+    """Reject invalid HSDP shard sizes before allocating any process groups."""
+    world_size = 4
+    created_groups: list[object] = []
+
+    def fail_if_group_created(*args, **kwargs):
+        del args, kwargs
+        created_groups.append(object())
+        raise AssertionError("a process group was created before HSDP validation")
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=0,
+        local_rank=0,
+        device_group=object(),
+    )
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    with pytest.raises(ValueError, match=error_message):
+        parallel_state.initialize_model_parallel(
+            fully_shard_degree=fully_shard_degree,
+            use_hsdp=True,
+            backend="gloo",
+        )
+
+    assert created_groups == []
+
+
 @pytest.mark.cpu
 @pytest.mark.core_model
 def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):

tests/diffusion/distributed/test_tensor_parallel.py

@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
 
 from __future__ import annotations
 
+import gc
 import os
 import time
 from pathlib import Path
@@ -75,8 +76,13 @@ def _run_generate(
     if num_requests < 2:
         raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
 
+    gc.collect()
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
+    with current_omni_platform.device(device_index):
+        free_bytes, total_bytes = current_omni_platform.mem_get_info()
+    initial_used_mb = (total_bytes - free_bytes) / (1024**2)
+
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     try:
@@ -113,7 +119,10 @@ def _run_generate(
                 pass
 
             median_time_s = float(np.median(per_request_times_s))
-            peak_memory_mb = monitor.peak_used_mb
+            # DeviceMemoryMonitor samples absolute device usage. Normalize by
+            # this run's starting usage so compiler and allocator caches retained
+            # by the preceding TP=1 run are not charged to the TP=2 run.
+            peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
             return _extract_single_image([last_output]), median_time_s, peak_memory_mb
     finally:
         monitor.stop()

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Layer-level correctness & performance test for MoT parallel linear layers.
 
@@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
     return output
 
 
+def _assert_gen_expert_parallel_metadata(layer):
+    """The secondary expert must satisfy vLLM's parallel-linear contract."""
+    for attr in (
+        "tp_rank",
+        "tp_size",
+        "input_size",
+        "input_size_per_partition",
+        "output_size",
+        "output_size_per_partition",
+    ):
+        assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
+
+
+def test_mot_gen_expert_parallel_metadata():
+    """Keep secondary experts compatible with vLLM online quantizers."""
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        _is_tp_sharded,
+    )
+
+    with set_current_vllm_config(VllmConfig()):
+        layers = (
+            MoTQKVParallelLinear(
+                hidden_size=16,
+                head_size=4,
+                total_num_heads=4,
+                total_num_kv_heads=2,
+                bias=False,
+                vae_bias=False,
+                params_dtype=torch.bfloat16,
+                disable_tp=True,
+            ),
+            MoTRowParallelLinear(
+                16,
+                8,
+                bias=False,
+                vae_bias=False,
+                input_is_parallel=True,
+                params_dtype=torch.bfloat16,
+                disable_tp=True,
+            ),
+        )
+
+    for layer in layers:
+        _assert_gen_expert_parallel_metadata(layer)
+        assert _is_tp_sharded(layer.gen_exp) is False
+
+
 def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
     """Compare outputs, print metrics, and assert correctness.
 
@@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
             disable_tp=True,
         ).cuda()
 
+        _assert_gen_expert_parallel_metadata(mot_linear)
+
         assert text_linear.output_size_per_partition == N, (
             f"Expected output_size_per_partition={N}, "
             f"got {text_linear.output_size_per_partition}. "
@@ -423,6 +475,8 @@ def test_mot_o_proj(
             disable_tp=True,
         ).cuda()
 
+        _assert_gen_expert_parallel_metadata(mot_linear)
+
         _sync_weights(text_linear, vae_linear, mot_linear)
 
         text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
@@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
         ).cuda()
 
         with torch.no_grad():
+            # vLLM linear layers leave weights uninitialized (torch.empty), so
+            # a freshly zeroed GPU page would make both outputs zero and the
+            # cosine-similarity gate degenerate. Seed deterministic values so
+            # und-mode is compared against a meaningful reference.
+            ref_linear.weight.normal_(mean=0.0, std=0.02)
+            if ref_linear.bias is not None:
+                ref_linear.bias.normal_(mean=0.0, std=0.02)
             mot_linear.weight.copy_(ref_linear.weight)
             if bias and ref_linear.bias is not None:
                 mot_linear.bias.copy_(ref_linear.bias)
@@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
         ).cuda()
 
         with torch.no_grad():
+            # vLLM linear layers leave weights uninitialized (torch.empty), so
+            # a freshly zeroed GPU page would make both outputs zero and the
+            # cosine-similarity gate degenerate. Seed deterministic values so
+            # und-mode is compared against a meaningful reference.
+            ref_linear.weight.normal_(mean=0.0, std=0.02)
+            if ref_linear.bias is not None:
+                ref_linear.bias.normal_(mean=0.0, std=0.02)
             mot_linear.weight.copy_(ref_linear.weight)
             if bias and ref_linear.bias is not None:
                 mot_linear.bias.copy_(ref_linear.bias)

tests/diffusion/lora/test_base_linear.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
     assert torch.allclose(out, expected)
 
 
+def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
+    from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    torch.nn.Module.__init__(layer)
+    layer.base_layer = torch.nn.Module()
+    layer.n_slices = 1
+    layer.lora_a_stacked = (torch.zeros(2),)
+    layer.lora_b_stacked = (torch.zeros(3),)
+
+    layer._set_diffusion_lora_buffer_device(torch.device("meta"))
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+    def recreate_on_cpu(self, max_loras, lora_config, model_config):
+        self.lora_a_stacked = (torch.zeros(4),)
+        self.lora_b_stacked = (torch.zeros(5),)
+
+    monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
+    layer.create_lora_weights(max_loras=1, lora_config=object())
+
+    assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+
 def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
     # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
     # LoRA tensors are still allocated and non-empty.

tests/diffusion/lora/test_lora_manager.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
@@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
     assert replace_calls == ["foo"]
 
 
+def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DLOLoRALayer(DummyBaseLayerWithLoRA):
+        lora_buffer_device = None
+
+        def _set_diffusion_lora_buffer_device(self, device):
+            self.lora_buffer_device = device
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
+    monkeypatch.setattr(
+        manager_mod,
+        "from_layer_diffusion",
+        lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
+    )
+
+    pipeline = torch.nn.Module()
+    pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _FakeLinearBase()
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
+
+    assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
+
+
 def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
     import vllm_omni.diffusion.lora.manager as manager_mod
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -249,6 +249,30 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
     assert loader.take_host_weight_plan() is None
 
 
+def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
+    tmp_path: Path,
+    monkeypatch,
+):
+    canonical_root = tmp_path / "canonical"
+    canonical_root.mkdir()
+    save_file(
+        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+        str(canonical_root / "model.safetensors"),
+    )
+    loader = DiffusersPipelineLoader(
+        LoadConfig(),
+        _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
+    )
+    pipeline = _HWRPipeline(canonical_root)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+
+    with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
+        loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    assert pipeline.load_count == 0
+    assert loader.take_host_weight_plan() is None
+
+
 def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
@@ -261,10 +285,11 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
 
 
 @pytest.mark.parametrize(
-    ("dist_offload", "use_allgather"),
+    ("dist_offload", "use_allgather", "mode"),
     [
-        (False, False),
-        (True, True),
+        (False, False, "preferred"),
+        (True, True, "preferred"),
+        (True, False, "disabled"),
     ],
 )
 def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
@@ -272,12 +297,13 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
     tmp_path,
     dist_offload,
     use_allgather,
+    mode,
 ):
     """Disabled and AllGather paths must never construct or probe HWR."""
     from vllm_omni.host_weight_runtime import HostWeightRuntime
 
     root = tmp_path / "must-not-be-touched"
-    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
+    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
     model = _DummyPipelineModel(source_prefix="transformer.")
     modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
 

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -303,6 +303,83 @@ class _FakeH3Attention(nn.Module):
         super().__init__()
 
 
+class _FakeFluxLinear(_FakeH3Linear):
+    def __init__(
+        self,
+        *args: object,
+        bias: bool = False,
+        params_dtype: torch.dtype | None = None,
+        total_num_heads: int | None = None,
+        total_num_kv_heads: int | None = None,
+        **kwargs: object,
+    ) -> None:
+        super().__init__(
+            *args,
+            bias=bias,
+            params_dtype=params_dtype,
+            total_num_heads=total_num_heads,
+            total_num_kv_heads=total_num_kv_heads,
+            **kwargs,
+        )
+        self.weight.weight_loader = self._load_weight
+
+    @staticmethod
+    def _load_weight(
+        parameter: nn.Parameter,
+        loaded_weight: torch.Tensor,
+        shard_id: str | None = None,
+    ) -> None:
+        del shard_id
+        parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
+
+
+class _FakeFluxAttention(nn.Module):
+    def __init__(self, **kwargs: object) -> None:
+        del kwargs
+        super().__init__()
+        self.register_buffer("beta", torch.tensor([0.5]))
+        self.register_buffer("eps", torch.tensor([-1e-6]))
+
+
+class _Flux2KleinPipeline(nn.Module):
+    def __init__(self, transformer: nn.Module) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
+        self.vae = nn.Module()
+        self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
+
+
+def _reduced_flux2_klein(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    device: torch.device | str = "cpu",
+) -> nn.Module:
+    from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
+
+    monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
+
+    transformer = flux2.Flux2Transformer2DModel(
+        patch_size=1,
+        in_channels=8,
+        out_channels=8,
+        num_layers=1,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=1,
+        joint_attention_dim=8,
+        timestep_guidance_channels=8,
+        mlp_ratio=1.0,
+        axes_dims_rope=(2, 2, 2, 2),
+        guidance_embeds=False,
+    )
+    return transformer.to(device=device, dtype=torch.bfloat16)
+
+
 class _H3Pipeline(nn.Module):
     def __init__(self, transformer: nn.Module) -> None:
         super().__init__()
@@ -1200,3 +1277,152 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
     assert contract.implementation_id == "minimax-h3-dit"
     assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
     assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
+
+
+def test_reduced_flux2_klein_declares_complete_final_layout_contract(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
+        Flux2Transformer2DModel,
+    )
+
+    transformer = _reduced_flux2_klein(monkeypatch)
+    pipeline = _Flux2KleinPipeline(transformer)
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    context = _identity(pipeline, source)
+    targets = collect_final_layout_targets(
+        pipeline,
+        (("transformer", transformer),),
+        policy=FINAL_LAYOUT_BF16_POLICY,
+        require_materialized=True,
+    )
+    target_names = {target.name for target in targets}
+
+    contract = Flux2Transformer2DModel.host_weight_restore_contract
+    assert isinstance(contract, FinalLayoutModelContract)
+    assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+    assert contract.implementation_id == "flux2-klein-dit"
+    assert contract.version == "1"
+    assert context.model_contracts == (contract,)
+    assert transformer.stacked_params_mapping == [
+        (".to_qkv.", ".to_q.", "q"),
+        (".to_qkv.", ".to_k.", "k"),
+        (".to_qkv.", ".to_v.", "v"),
+        (".add_kv_proj", ".add_q_proj", "q"),
+        (".add_kv_proj", ".add_k_proj", "k"),
+        (".add_kv_proj", ".add_v_proj", "v"),
+    ]
+    loaded = transformer.load_weights(
+        [
+            ("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
+            ("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
+            (
+                "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+                torch.tensor([3.0], dtype=torch.bfloat16),
+            ),
+        ]
+    )
+    assert loaded == {
+        "transformer_blocks.0.attn.to_qkv.weight",
+        "transformer_blocks.0.attn.add_kv_proj.weight",
+        "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+    }
+    assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
+    assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
+    assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
+    assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
+    assert all("text_encoder" not in name and "vae" not in name for name in target_names)
+    transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_restore_validator_rejects_non_ready_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    transformer = _reduced_flux2_klein(monkeypatch)
+
+    mapping = transformer.stacked_params_mapping
+    transformer.stacked_params_mapping = None
+    with pytest.raises(ValueError, match="packed QKV mapping"):
+        transformer.validate_restored_host_weights()
+    transformer.stacked_params_mapping = mapping
+
+    weight = transformer.x_embedder.weight
+    weight.data = weight.data.float()
+    with pytest.raises(ValueError, match="must stay bf16"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.to(torch.bfloat16)
+
+    weight.data = weight.data.T
+    with pytest.raises(ValueError, match="contiguous strided layout"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.contiguous()
+
+    attention = transformer.transformer_blocks[0].attn.attn
+    attention._non_persistent_buffers_set.add("beta")
+    with pytest.raises(ValueError, match="must be persistent"):
+        transformer.validate_restored_host_weights()
+    attention._non_persistent_buffers_set.remove("beta")
+
+    attention.beta.data = attention.beta.data.float()
+    with pytest.raises(ValueError, match="scalar bf16"):
+        transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_publication_restore_fallback_and_teardown(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    cold_transformer = _reduced_flux2_klein(monkeypatch)
+    cold_model = _Flux2KleinPipeline(cold_transformer)
+    with torch.no_grad():
+        for index, parameter in enumerate(cold_transformer.parameters(), start=1):
+            parameter.fill_(index)
+        for name, buffer in cold_transformer.named_buffers():
+            if name.endswith((".beta", ".eps")):
+                buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
+
+    context = _identity(cold_model, source)
+    runtime = _runtime(tmp_path / "flux2-klein-store")
+    miss = runtime.resolve(context.identity)
+    assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+    producer = FinalLayoutBF16Producer(
+        context,
+        cold_model,
+        (("transformer", cold_transformer),),
+        max_shard_bytes=64,
+    )
+    publication = runtime.publish_after_load(context.identity, producer=producer)
+    assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+
+    warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
+    warm_model = _Flux2KleinPipeline(warm_transformer)
+    warm_context = _identity(warm_model, source)
+    assert warm_context.identity == context.identity
+    hit = runtime.resolve(warm_context.identity)
+    assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert hit.lease is not None
+
+    plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+    assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
+    plan.commit()
+    warm_transformer.validate_restored_host_weights()
+
+    cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
+    warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
+    for name in warm_context.tensor_names:
+        assert torch.equal(warm_tensors[name], cold_tensors[name])
+        assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
+    assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
+
+    del warm_tensors, warm_model, warm_transformer, plan
+    gc.collect()
+    hit.lease.close()
+    assert hit.lease.closed
+    assert isinstance(runtime.store, FilesystemHostWeightStore)
+    assert runtime.store.cleanup(context.identity) is None

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -245,7 +245,7 @@ else:
         [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
         capture_output=True,
         text=True,
-        timeout=3,
+        timeout=30,
         check=False,
     )
 

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Focused LTX-2.5 pipeline correctness tests."""
 
 import json
 from dataclasses import replace
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
 
 
 def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
-    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
+    pipe = SimpleNamespace(
+        component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
+        use_diffusion_decoder=True,
+    )
     od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
 
     monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
 
 def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     revision = "pinned-revision"
-    calls = {"components": []}
+    calls: dict[str, Any] = {"components": []}
+    natten_processor = object()
+
+    class FakeDiffusionDecoder:
+        def set_attn_processor(self, processor):
+            calls["diffusion_decoder_processor"] = processor
+
+        def enable_tiling(self):
+            calls["diffusion_decoder_tiling"] = True
+
+        def set_parallel_size(self, parallel_size, mode):
+            calls["diffusion_decoder_parallel"] = (parallel_size, mode)
+
     profile = replace(
         LTX25_DISTILLED_COMPONENT_PROFILE,
         text_encoder_cls=object,
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
         vocoder_fallback_cls=None,
         scheduler_use_dynamic_shifting=False,
     )
-    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
+    pipeline = SimpleNamespace(
+        component_profile=profile,
+        pipeline_kind="distilled_two_stage",
+        use_diffusion_decoder=True,
+    )
     od_config = SimpleNamespace(
         model="org/converted-ltx25",
         revision=revision,
         dtype=torch.bfloat16,
         quantization_config=None,
+        extras={},
+        vae_use_tiling=False,
+        parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
     )
 
     monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
             )
         return object()
 
+    def fake_native_diffusion_decoder(model, **kwargs):
+        calls["native_diffusion_decoder"] = (model, kwargs)
+        return FakeDiffusionDecoder()
+
     def fake_transformer_config(model, subfolder, local_files_only, *, revision):
         calls["transformer_config"] = (model, subfolder, local_files_only, revision)
         return {"component": "transformer"}
@@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
     monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
     monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+    monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
     monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+    monkeypatch.setattr(
+        ltx2_components,
+        "LTX2VideoVaeNeighborhoodNattenProcessor",
+        lambda: natten_processor,
+    )
     monkeypatch.setattr(
         ltx2_components,
         "create_transformer_from_config",
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
         "latent_upsampler",
     }
     assert all(call[3]["revision"] == revision for call in calls["components"])
+    assert "diffusion_decoder" not in calls["prefetch"][1]
+    assert calls["native_diffusion_decoder"] == (
+        od_config.model,
+        {"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
+    )
+    assert calls["diffusion_decoder_processor"] is natten_processor
+    assert calls["diffusion_decoder_tiling"] is True
+    assert calls["diffusion_decoder_parallel"] == (2, "tile")
     assert calls["transformer_config"] == (
         od_config.model,
         profile.transformer_subfolder,
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     assert calls["scheduler"][1]["revision"] == revision
 
 
+def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
+    revision = "diffusers-revision"
+    config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
+    calls: dict[str, Any] = {}
+
+    class FakeDecoder:
+        def init_distributed(self):
+            calls["init_distributed"] = True
+
+    expected = FakeDecoder()
+
+    def fake_load_config(model, **kwargs):
+        calls["config"] = (model, kwargs)
+        return config
+
+    def fake_resolve(model, repo_id, filename, **kwargs):
+        calls["artifact"] = (model, repo_id, filename, kwargs)
+        return "/models/native-diffvae.safetensors"
+
+    def fake_load(cls, path, passed_config, dtype):
+        calls["load"] = (path, passed_config, dtype)
+        return expected
+
+    monkeypatch.setattr(
+        ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+        "load_config",
+        staticmethod(fake_load_config),
+    )
+    monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
+    monkeypatch.setattr(
+        ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+        "from_ltx25_native_checkpoint",
+        classmethod(fake_load),
+    )
+
+    actual = ltx2_components._load_ltx25_native_diffusion_decoder(
+        "Lightricks/LTX-2.5-Diffusers",
+        local_files_only=False,
+        dtype=torch.bfloat16,
+        revision=revision,
+    )
+
+    assert actual is expected
+    assert calls["config"] == (
+        "Lightricks/LTX-2.5-Diffusers",
+        {
+            "subfolder": "diffusion_decoder",
+            "local_files_only": False,
+            "revision": revision,
+        },
+    )
+    assert calls["artifact"] == (
+        "Lightricks/LTX-2.5-Diffusers",
+        "Lightricks/LTX-2.5",
+        "vae/ltx-2.5-video-vae-bf16.safetensors",
+        {
+            "model_revision": revision,
+            "artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
+        },
+    )
+    assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
+    assert calls["init_distributed"] is True
+
+
+def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
+    class MissingNattenProcessor:
+        def __init__(self):
+            raise ImportError("kernels is unavailable")
+
+    monkeypatch.setattr(
+        ltx2_components,
+        "LTX2VideoVaeNeighborhoodNattenProcessor",
+        MissingNattenProcessor,
+    )
+
+    with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
+        ltx2_components._create_ltx25_natten_processor()
+
+
 def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
     (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
     (tmp_path / "transformer").mkdir()
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
     assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
 
 
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+    ],
+)
+@pytest.mark.parametrize(
+    ("extras", "use_diffusion_decoder"),
+    [({}, True), ({"ltx2_use_conv_vae": True}, False)],
+)
+def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
+    tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
+):
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = pipeline_cls(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            extras=extras,
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
+
+    assert pipe.use_diffusion_decoder is use_diffusion_decoder
+    assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
+
+
 @pytest.mark.parametrize(
     ("pipeline_kind", "expected"),
     [
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
         max_sequence_length=16,
     )
     prompt_context = object()
-    phase_calls = []
+    phase_calls: list[Any] = []
 
     def resolve_request_inputs(req, **kwargs):
         return request_inputs

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for LTX image-to-video input and conditioning behavior."""
 
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
 
         assert LTX2Pipeline.support_image_input is True
 
+    def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = SimpleNamespace(
+            encode=lambda image: image,
+            latents_mean=torch.full((2,), 100.0),
+            latents_std=torch.full((2,), 10.0),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.diffusion_decoder = SimpleNamespace(
+            latents_mean=torch.tensor([1.0, 3.0]),
+            latents_std=torch.tensor([2.0, 4.0]),
+            config=SimpleNamespace(scaling_factor=2.0),
+        )
+        monkeypatch.setattr(
+            ltx2_conditioning,
+            "retrieve_latents",
+            lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+        )
+
+        actual = pipe._encode_i2v_image_latents(
+            torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            generator=None,
+            dtype=torch.float32,
+        )
+
+        torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
+
     def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
@@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
     assert actual.stride()[1:] == (1, actual.shape[1])
 
 
+def test_prepare_video_latents_uses_diffusion_decoder_statistics():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    pipeline.use_diffusion_decoder = True
+    pipeline.vae_spatial_compression_ratio = 1
+    pipeline.vae_temporal_compression_ratio = 1
+    pipeline.transformer_spatial_patch_size = 1
+    pipeline.transformer_temporal_patch_size = 1
+    pipeline.vae = SimpleNamespace(
+        latents_mean=torch.full((2,), 100.0),
+        latents_std=torch.full((2,), 10.0),
+        config=SimpleNamespace(scaling_factor=1.0),
+    )
+    pipeline.diffusion_decoder = SimpleNamespace(
+        latents_mean=torch.tensor([1.0, 3.0]),
+        latents_std=torch.tensor([2.0, 4.0]),
+        config=SimpleNamespace(scaling_factor=2.0),
+    )
+
+    actual = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        height=1,
+        width=1,
+        num_frames=1,
+        noise_scale=0.0,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+    )
+
+    torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
+
+
 def test_prepare_audio_latents_samples_directly_in_packed_token_space():
     pipeline = _make_pipeline(LTX2Pipeline)
 

tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for the fixed LTX refinement-phase adapter."""
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
 ):
     pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
     manifest = object()
-    captured = {}
+    captured: dict[str, Any] = {}
 
     def resolve_artifact(*_args, **kwargs):
-        assert not kwargs
+        assert kwargs == {
+            "model_revision": None,
+            "artifact_revision": None,
+        }
         return "adapter"
 
     def parse(transformer, path):

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for the shared LTX pipeline runtime and public contracts."""
 
@@ -8,6 +8,7 @@ import os
 import tempfile
 from dataclasses import replace
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -165,7 +166,19 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
     assert detect_ltx_model_version(str(tmp_path)) == "2"
 
 
-def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
+@pytest.mark.parametrize(
+    ("model", "repo_id", "expected_revision"),
+    [
+        ("Lightricks/test", "Lightricks/test", "model-revision"),
+        ("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
+    ],
+)
+def test_ltx_artifact_uses_source_revision_and_hub_fallback(
+    monkeypatch,
+    model,
+    repo_id,
+    expected_revision,
+):
     filename = "ltx-sidecar.safetensors"
     calls = []
 
@@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch)
     monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
 
     assert (
-        resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
+        resolve_ltx_artifact(
+            model,
+            repo_id,
+            filename,
+            model_revision="model-revision",
+            artifact_revision="artifact-revision",
+        )
+        == "/cache/ltx-sidecar.safetensors"
     )
-    assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
+    assert calls == [
+        {
+            "repo_id": repo_id,
+            "filename": filename,
+            "revision": expected_revision,
+        }
+    ]
 
 
 def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
     expected.write_bytes(b"sidecar")
     monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
 
-    assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
+    assert resolve_ltx_artifact(
+        str(tmp_path),
+        "Lightricks/test",
+        filename,
+        model_revision="unused-model-revision",
+        artifact_revision="unused-artifact-revision",
+    ) == str(expected)
+
+
+def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
+    calls = []
+
+    def fake_download(**kwargs):
+        calls.append(kwargs)
+        return "/cache/ltx-sidecar.safetensors"
+
+    monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+
+    assert (
+        resolve_ltx_artifact(
+            str(tmp_path),
+            "Lightricks/test",
+            "ltx-sidecar.safetensors",
+            model_revision="local-model-revision",
+            artifact_revision="artifact-revision",
+        )
+        == "/cache/ltx-sidecar.safetensors"
+    )
+    assert calls == [
+        {
+            "repo_id": "Lightricks/test",
+            "filename": "ltx-sidecar.safetensors",
+            "revision": "artifact-revision",
+        }
+    ]
 
 
 @pytest.mark.parametrize(
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
         max_sequence_length=16,
     )
     prompt_context = object()
-    phase_calls = []
+    phase_calls: list[Any] = []
 
     def resolve_request_inputs(req, **kwargs):
         return request_inputs
@@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
     pipeline.vae_spatial_compression_ratio = 32
     pipeline.vae_temporal_compression_ratio = 8
     pipeline.latent_upsampler = FakeUpsampler()
-    activated_slots = []
+    activated_slots: list[Any] = []
     if pipeline_cls is LTX2TwoStagePipeline:
         pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
     object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
@@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
         )
 
 
+def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
+    request = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="",
+        height=128,
+        width=128,
+        num_frames=9,
+        frame_rate=24.0,
+        num_inference_steps=30,
+        guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=16,
+    )
+
+    with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
+        validate_pipeline_request(
+            request,
+            pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+            vae_spatial_compression_ratio=32,
+            vae_temporal_compression_ratio=8,
+            pipeline_name="LTX2Pipeline",
+            min_video_latent_spatial_size=(7, 7),
+        )
+
+    # ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
+    validate_pipeline_request(
+        request,
+        pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+        vae_spatial_compression_ratio=32,
+        vae_temporal_compression_ratio=8,
+        pipeline_name="LTX2Pipeline",
+    )
+
+    validate_pipeline_request(
+        replace(request, height=256, width=256),
+        pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+        vae_spatial_compression_ratio=32,
+        vae_temporal_compression_ratio=8,
+        pipeline_name="LTX2Pipeline",
+        min_video_latent_spatial_size=(7, 7),
+    )
+
+
 @pytest.mark.parametrize(
     ("direct_kwargs", "sampling_kwargs", "error"),
     [
@@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step
 
 
 def test_ltx_two_stage_entries_select_the_official_adapter_slots():
-    phase_switches = []
+    phase_switches: list[Any] = []
     ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
     ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
     for phase in LTX2_TWO_STAGE_RECIPE.phases:

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+"""Unit tests for LTX video VAE tiling and distributed decode behavior."""
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -11,6 +12,466 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
+    import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+
+    class FakeBaseVocoder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.input_dtype = None
+
+        def forward(self, value):
+            self.input_dtype = value.dtype
+            return value
+
+    monkeypatch.setattr(
+        ltx_runtime.torch,
+        "autocast",
+        lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
+    )
+    vocoder = FakeBaseVocoder()
+
+    output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
+
+    assert vocoder.input_dtype == torch.bfloat16
+    assert output.dtype == torch.bfloat16
+
+
+class TestLTXDiffusionDecoder:
+    @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
+    def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None)
+
+        with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"):
+            model.set_parallel_size(2, mode=mode)
+
+    def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        calls = []
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.distributed_executor = SimpleNamespace(
+            set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode))
+        )
+
+        model.set_parallel_size(4, mode="tile")
+
+        assert calls == [(4, "tile")]
+
+    def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            LTX2VideoDiffusionDecoder3d,
+            LTX2VideoDiffusionDecoderModel,
+        )
+
+        class IdentityAttention:
+            def build_block_mask(self, _hidden_states):
+                return None
+
+        class IdentityBlock(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.attn = IdentityAttention()
+
+            def forward(self, hidden_states, _block_mask):
+                return hidden_states
+
+        class IdentityUpsample(torch.nn.Module):
+            def forward(self, hidden_states, *, drop_leading_frame):
+                assert drop_leading_frame
+                return hidden_states
+
+        decoder = object.__new__(LTX2VideoDiffusionDecoder3d)
+        torch.nn.Module.__init__(decoder)
+        decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])])
+        decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()])
+        decoder.trailing_pad_latent_frames = 2
+        decoder.temporal_compression_ratio = 8
+        decoder.stage5_kernel = (11, 3, 3)
+
+        # 17 frames contain one real frame followed by 16 propagated ghost
+        # frames. The reference keeps 11 for stage-5 NATTEN instead of
+        # cropping the tensor immediately back to one.
+        context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4))
+        assert context.shape == (1, 11, 3, 3, 4)
+
+        class FakeDecoder(torch.nn.Module):
+            def forward(self, z, *, generator, num_inference_steps):
+                assert generator is None
+                assert num_inference_steps is None
+                return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2)
+
+        model = object.__new__(LTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.decoder = FakeDecoder()
+        model.use_tiling = False
+        model.tile_sample_min_num_frames = 80
+        model.tile_sample_min_height = 768
+        model.tile_sample_min_width = 768
+        model.temporal_compression_ratio = 8
+        model.spatial_compression_ratio = 2
+
+        output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0]
+        assert output.shape == (1, 3, 1, 4, 6)
+
+    @pytest.mark.parametrize(
+        ("model_version", "extras", "expected"),
+        [
+            ("2.5", {}, True),
+            ("2.5", {"ltx2_use_conv_vae": True}, False),
+            ("2.3", {}, False),
+        ],
+    )
+    def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+        assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected
+
+    def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+        with pytest.raises(TypeError, match="ltx2_use_conv_vae"):
+            _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
+
+    def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            convert_ltx25_native_diffusion_decoder_state_dict,
+        )
+
+        qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
+        qkv_bias = torch.arange(6, dtype=torch.float32)
+        projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        projection_bias = torch.ones(2)
+        gate = torch.tensor([2.0, 3.0])
+        mean = torch.tensor([0.1, 0.2])
+        std = torch.tensor([0.3, 0.4])
+        native = {
+            "encoder.ignored": torch.ones(1),
+            "per_channel_statistics.mean-of-means": mean,
+            "vae.per_channel_statistics.std-of-means": std,
+            "decoder.type_emb": torch.ones(2),
+            "decoder.coarse_head.weight": torch.ones(1),
+            "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight,
+            "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias,
+            "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]),
+            "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]),
+            "decoder.diff_blocks.0.attn.proj.weight": projection_weight,
+            "decoder.diff_blocks.0.attn.proj.bias": projection_bias,
+            "decoder.diff_blocks.0.gate_msa": gate,
+            "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2),
+        }
+
+        converted = convert_ltx25_native_diffusion_decoder_state_dict(native)
+
+        torch.testing.assert_close(converted["latents_mean"], mean)
+        torch.testing.assert_close(converted["latents_std"], std)
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0]))
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0]))
+        torch.testing.assert_close(
+            converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight
+        )
+        torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias)
+        assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted
+        assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
+
+    def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            convert_ltx25_native_diffusion_decoder_state_dict,
+        )
+
+        with pytest.raises(ValueError, match="not divisible by 3"):
+            convert_ltx25_native_diffusion_decoder_state_dict(
+                {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)}
+            )
+
+    def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        seen = {}
+
+        class ConvVae:
+            dtype = torch.float32
+            config = SimpleNamespace(timestep_conditioning=True)
+
+            def decode(self, *_args, **_kwargs):
+                raise AssertionError("the convolutional VAE decoder must not run")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def decode(self, latents, *, generator, return_dict):
+                seen["latents"] = latents
+                seen["generator"] = generator
+                seen["return_dict"] = return_dict
+                return (latents + 1,)
+
+        class AudioVae:
+            dtype = torch.float32
+
+            def decode(self, audio_latents, *, return_dict):
+                assert return_dict is False
+                return (audio_latents + 2,)
+
+        class VideoProcessor:
+            def postprocess_video(self, video, *, output_type):
+                seen["output_type"] = output_type
+                return video
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+        pipe.audio_vae = AudioVae()
+        pipe.vocoder = torch.nn.Identity()
+        pipe.video_processor = VideoProcessor()
+
+        generator = torch.Generator().manual_seed(123)
+        latents = torch.ones(1, 1)
+        output = pipe._decode_output(
+            latents=latents,
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=generator,
+            device=torch.device("cpu"),
+            decode_timestep=0.7,
+            decode_noise_scale=0.5,
+            prompt_batch_size=1,
+        )
+
+        torch.testing.assert_close(seen["latents"], latents)
+        assert seen["generator"] is generator
+        assert seen["return_dict"] is False
+        assert seen["output_type"] == "pt"
+        torch.testing.assert_close(output.output[0], latents + 1)
+        torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0))
+
+    def test_diffusion_decode_denormalizes_with_decoder_statistics(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.full((2,), 100.0),
+            latents_std=torch.full((2,), 10.0),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.diffusion_decoder = SimpleNamespace(
+            latents_mean=torch.tensor([10.0, 20.0]),
+            latents_std=torch.tensor([2.0, 4.0]),
+            config=SimpleNamespace(scaling_factor=2.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.tensor(0.0),
+            latents_std=torch.tensor(1.0),
+        )
+        forward_ctx = SimpleNamespace(
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            original_audio_num_frames=1,
+            latent_mel_bins=2,
+        )
+
+        video, audio = pipe._unpack_and_denormalize_stage(
+            forward_ctx,
+            torch.tensor([[[1.0, 2.0]]]),
+            torch.zeros(1, 1, 2),
+        )
+
+        torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0]))
+        torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2))
+
+    def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        class ConvVae:
+            config = SimpleNamespace(timestep_conditioning=False)
+
+            def is_distributed_enabled(self):
+                raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def decode(self, *_args, **_kwargs):
+                raise AssertionError("only the output rank may run DiffVAE decode")
+
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+
+        output = pipe._decode_output(
+            latents=torch.ones(1, 1),
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=None,
+            device=torch.device("cpu"),
+            decode_timestep=0.0,
+            decode_noise_scale=None,
+            prompt_batch_size=1,
+        )
+
+        assert output.output[0].numel() == 0
+        assert output.output[1].numel() == 0
+
+    def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        class ConvVae:
+            config = SimpleNamespace(timestep_conditioning=False)
+
+            def is_distributed_enabled(self):
+                raise AssertionError("DiffVAE decode must query the active decoder")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def __init__(self):
+                self.decode_calls = 0
+
+            def is_distributed_enabled(self):
+                return True
+
+            def decode(self, *_args, **_kwargs):
+                self.decode_calls += 1
+                return (torch.empty(0, 3, 0, 0, 0),)
+
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+
+        output = pipe._decode_output(
+            latents=torch.ones(1, 1),
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=None,
+            device=torch.device("cpu"),
+            decode_timestep=0.0,
+            decode_noise_scale=None,
+            prompt_batch_size=1,
+        )
+
+        assert pipe.diffusion_decoder.decode_calls == 1
+        assert output.output[0].numel() == 0
+        assert output.output[1].numel() == 0
+
+    def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        expected = torch.ones(1, 3, 1, 2, 2)
+        seen = {}
+
+        def native_tiled_decode(self, z, generator=None, num_inference_steps=None):
+            seen["z"] = z
+            seen["generator"] = generator
+            seen["num_inference_steps"] = num_inference_steps
+            return expected
+
+        monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode)
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.is_distributed_enabled = lambda: False
+        generator = torch.Generator().manual_seed(17)
+        z = torch.zeros(1, 1, 1, 1, 1)
+
+        result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode(
+            model,
+            z,
+            generator=generator,
+            num_inference_steps=1,
+        )
+
+        assert result is expected
+        assert seen["z"] is z
+        assert seen["generator"] is generator
+        assert seen["num_inference_steps"] == 1
+
+    def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
+        from diffusers.utils.torch_utils import randn_tensor
+
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        model = DistributedLTX2VideoDiffusionDecoderModel(
+            out_channels=1,
+            latent_channels=1,
+            patch_size=1,
+            decoder_head_dim=8,
+            decoder_stage_channels=(8, 8, 8, 8, 8),
+            decoder_stage_depths=(1, 1, 1, 1, 1),
+            decoder_stage_kernels=((1, 1, 1),) * 4,
+            decoder_upsample_strides=((1, 1, 1),) * 4,
+            decoder_upsample_channel_reductions=(1, 1, 1, 1),
+            decoder_stage5_kernel=(1, 1, 1),
+            spatial_compression_ratio=1,
+            temporal_compression_ratio=1,
+        )
+        model.tile_sample_min_num_frames = 2
+        model.tile_sample_stride_num_frames = 1
+        model.tile_sample_min_height = 2
+        model.tile_sample_stride_height = 1
+        model.tile_sample_min_width = 2
+        model.tile_sample_stride_width = 1
+        model.distributed_executor = SimpleNamespace(group=None)
+        model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1)
+        monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None)
+
+        z = torch.zeros(1, 1, 2, 3, 3)
+        distributed_generator = torch.Generator().manual_seed(123)
+        reference_generator = torch.Generator().manual_seed(123)
+        tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1)
+
+        assert grid_spec.grid_shape == (1, 3, 3)
+        assert grid_spec.split_dims == (2, 3, 4)
+        for task in tasks:
+            tile_shape = model._tiled_pixel_shape_from_features(
+                task.tensor,
+                drop_leading_frame=task.drop_leading_frame,
+                crop_trailing_ghost=task.crop_trailing_ghost,
+            )
+            expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype)
+            actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype)
+            torch.testing.assert_close(actual, expected)
+
+        assert torch.equal(distributed_generator.get_state(), reference_generator.get_state())
+
+
 class TestLTXOutputRank:
     @pytest.mark.parametrize(
         ("distributed_vae_state", "expected_decode_calls"),
@@ -113,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
             },
             output_dtype=torch.float32,
         )
-        seen = {}
+        seen: dict[str, Any] = {}
 
         def exec_tile(task):
             return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))

tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CUDA regression tests for LTX-2 BWE vocoder precision."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class _BWEConvVocoder(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(
+            1,
+            1,
+            kernel_size=1,
+            bias=False,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        self.bwe_generator = torch.nn.Identity()
+        self.input_dtype = None
+        self.conv_output_dtype = None
+
+    def forward(self, value):
+        self.input_dtype = value.dtype
+        output = self.conv(value)
+        self.conv_output_dtype = output.dtype
+        return output
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
+    from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
+
+    vocoder = _BWEConvVocoder()
+    generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
+
+    output = _run_ltx_vocoder(vocoder, generated_mel)
+
+    assert next(vocoder.parameters()).dtype == torch.bfloat16
+    assert vocoder.input_dtype == torch.float32
+    assert vocoder.conv_output_dtype == torch.float32
+    assert output.dtype == torch.bfloat16

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -38,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
     pipeline.default_audio_shift = 3.0
     pipeline.device = torch.device("cpu")
     pipeline.od_config = SimpleNamespace()
+    pipeline.text_encoder = object()
     cache_spec = object()
     quality_plan = SimpleNamespace(cache_dit=cache_spec)
     pipeline._quality_policy = Mock()
@@ -110,6 +111,40 @@ def test_pipeline_import_registry_and_component_discovery():
     assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
 
 
+def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(1, 1)
+    pipeline.video_vae = torch.nn.Linear(1, 1)
+    pipeline.audio_vae = torch.nn.Linear(1, 1)
+    pipeline.text_encoder = None
+    pipeline._dit_modules = ["transformer"]
+    pipeline._encoder_modules = []
+    pipeline._vae_modules = ["video_vae", "audio_vae"]
+
+    discovered = ModuleDiscovery.discover(pipeline)
+
+    assert discovered.dit_names == ["transformer"]
+    assert discovered.encoder_names == []
+    assert discovered.vae_names == ["video_vae", "audio_vae"]
+
+
+def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+    cache = Mock()
+    video_vae = Mock()
+    audio_vae = Mock()
+
+    pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae)
+
+    video_vae.set_omni_component_cache.assert_called_once_with(cache)
+    audio_vae.set_omni_component_cache.assert_called_once_with(cache)
+
+
 def _write_partition_index(path, *, partition, tasks):
     path.mkdir(parents=True)
     (path / "model_index.json").write_text(
@@ -575,6 +610,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
     pipeline.default_audio_shift = 3.0
     pipeline.device = torch.device("cpu")
     pipeline.od_config = SimpleNamespace()
+    pipeline.text_encoder = Mock()
     pipeline._base_schedule_by_partition = schedules
     pipeline._quality_policy = Mock()
     pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -998,7 +1034,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
         config.hidden_size,
         bias=False,
     )
-    attn_call = {}
+    attn_call: dict[str, object] = {}
 
     def fake_attention(query, key, value):
         attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
@@ -1020,7 +1056,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
 def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
 
-    sdpa_call = {}
+    sdpa_call: dict[str, object] = {}
 
     def fake_sdpa(query, key, value, **kwargs):
         sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
@@ -1095,7 +1131,7 @@ def test_encoder_forward_forwards_video_inputs():
 
 
 def test_reference_video_shape_uses_h3_adapt_shape_policy():
-    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+    from vllm_omni.model_executor.models.minimax_h3.reference_video import (
         _reference_video_shape,
     )
 
@@ -1681,6 +1717,85 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
     assert transcode_calls[0][1]["target_frame_count"] == 209
 
 
+def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+    commands: list[list[str]] = []
+    monkeypatch.setattr(
+        reference_video_module.subprocess,
+        "run",
+        lambda command, **_kwargs: commands.append(command),
+    )
+
+    reference_video_module._transcode_reference_video(
+        "input.mp4",
+        target_width=768,
+        target_height=448,
+        target_frame_count=0,
+        workdir=str(tmp_path),
+        duration_seconds=4.0,
+    )
+
+    assert "-frames:v" not in commands[0]
+
+
+def test_prepared_reference_video_descriptor_round_trip(tmp_path):
+    from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+        deserialize_prepared_reference_videos,
+        serialize_prepared_reference_videos,
+    )
+
+    item = {
+        "original_path": str(tmp_path / "source.mp4"),
+        "prepared_path": str(tmp_path / "prepared.mp4"),
+        "input_has_audio": True,
+        "width": 1344,
+        "height": 768,
+        "start_time_seconds": 0.0,
+        "duration_seconds": 5.2,
+        "audio_duration_seconds": 5.2,
+    }
+
+    descriptor = serialize_prepared_reference_videos([item], str(tmp_path))
+
+    assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item])
+
+
+def test_prepared_reference_video_descriptor_rejects_unknown_fields():
+    from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+        deserialize_prepared_reference_videos,
+    )
+
+    with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"):
+        deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
+
+
+def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _reuse_prepared_reference_videos,
+    )
+
+    prepared_path = tmp_path / "prepared.mp4"
+    prepared_path.touch()
+    prepared = [{"prepared_path": str(prepared_path)}]
+
+    assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
+
+
+def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _reuse_prepared_reference_videos,
+    )
+
+    with pytest.raises(ValueError, match="count does not match"):
+        _reuse_prepared_reference_videos([], expected_count=1)
+    with pytest.raises(ValueError, match="prepared reference video is unavailable"):
+        _reuse_prepared_reference_videos(
+            [{"prepared_path": str(tmp_path / "missing.mp4")}],
+            expected_count=1,
+        )
+
+
 def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
@@ -2120,7 +2235,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path):
 
 
 def test_g4_standalone_audio_duration_and_total_duration_contract():
-    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+    from vllm_omni.model_executor.models.minimax_h3.reference_video import (
         validate_reference_audio_waveforms,
     )
 
@@ -2195,7 +2310,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
     ],
 )
 def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
-    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+    from vllm_omni.model_executor.models.minimax_h3.reference_video import (
         _validate_reference_video_metadata,
     )
 

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
     [
         "model-level CPU offload (--enable-cpu-offload)",
         "layerwise offload (--enable-layerwise-offload)",
-        "distributed layerwise offload (--enable-distributed-layerwise-offload)",
     ],
 )
 def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
@@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
         )
 
 
+def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+    loaded = pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+
+    assert loaded is not None
+    assert captured["unsupported_offload_mode"] is None
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -254,18 +254,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
     assert transformer_config.ignored_layers == ignored_layers
     assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
     assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
-
-
-def test_pipeline_strips_prequantized_text_encoder_config():
-    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
-    from vllm_omni.quantization import ComponentQuantizationConfig
-
-    prequantized = Mock()
-    prequantized.get_name.return_value = "modelopt"
-    online_fp8 = Mock()
-    online_fp8.get_name.return_value = "fp8"
-    component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
-
-    assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
-    assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
-    assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -187,7 +187,6 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
         "model": model_ref,
         "enforce_eager": True,
         "tensor_parallel_size": 2,
-        "text_encoder_tp_size": 2,
         "vae_use_tiling": True,
     }
 

tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py

@@ -1,13 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import pytest
 import torch
 import torch.nn.functional as F
 
+from tests.helpers.mark import hardware_marks
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    *hardware_marks(res={"npu": "A3"}, num_cards=1),
+]
 
 
 @pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py

@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion]
+
+
+def _selected_operators():
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import (
+        resolve_h3_vae_operators,
+    )
+
+    if not current_omni_platform.is_available():
+        pytest.skip("No accelerator is available")
+    device = current_omni_platform.get_torch_device()
+    operators = resolve_h3_vae_operators(device)
+    if operators is None:
+        pytest.skip("No H3 VAE operator implementation is registered for this target")
+    return device, operators
+
+
+def _operator_set(*, supports=lambda _device: True):
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet
+
+    return H3VAEOperatorSet(
+        supports=supports,
+        qk_norm_rope=lambda *_args: None,
+        scaled_residual=lambda *_args: None,
+    )
+
+
+def _qk_reference(x, cos, sin):
+    normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype)
+    rotary = normalized[..., :48]
+    first, second = rotary.chunk(2, dim=-1)
+    rotated = torch.cat((-second, first), dim=-1)
+    return torch.cat(
+        (rotary * cos + rotated * sin, normalized[..., 48:]),
+        dim=-1,
+    )
+
+
+def _vit_norm_input(_module, hidden_states):
+    return hidden_states.float()
+
+
+def _same_dtype_norm_input(_module, hidden_states):
+    return hidden_states
+
+
+def _failing_norm_input(_module, _hidden_states):
+    raise RuntimeError("unsupported remote normalization semantics")
+
+
+@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
+def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
+    device, operators = _selected_operators()
+
+    torch.manual_seed(17)
+    qkv = torch.randn(
+        batch,
+        sequence,
+        32,
+        192,
+        device=device,
+        dtype=torch.float16,
+    )
+    q, k, _ = qkv.chunk(3, dim=-1)
+    cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16)
+    sin = torch.randn_like(cos)
+
+    expected_q = _qk_reference(q, cos, sin)
+    expected_k = _qk_reference(k, cos, sin)
+    with torch.inference_mode():
+        actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5)
+
+    assert actual is not None
+    actual_q, actual_k = actual
+    assert torch.equal(actual_q, expected_q)
+    assert torch.equal(actual_k, expected_k)
+
+
+def test_h3_vae_scaled_residual_is_bit_exact():
+    device, operators = _selected_operators()
+
+    torch.manual_seed(29)
+    residual = torch.randn(195, 2048, device=device, dtype=torch.float32)
+    branch = torch.randn(195, 2048, device=device, dtype=torch.float16)
+    scale = torch.randn(2048, device=device, dtype=torch.float32)
+    expected = residual + branch * scale
+
+    with torch.inference_mode():
+        actual = operators.scaled_residual(residual, branch, scale)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_h3_vae_exact_ops_reject_unsupported_inputs():
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
+        try_qk_norm_rope_exact,
+    )
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact
+
+    q = torch.randn(1, 2, 32, 64, dtype=torch.float32)
+    cos = torch.randn(1, 2, 1, 48, dtype=torch.float32)
+    residual = torch.randn(2, 2048)
+    branch = torch.randn(2, 2048, dtype=torch.float16)
+    scale = torch.randn(2048)
+    with torch.inference_mode():
+        assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None
+        assert try_scaled_residual_exact(residual, branch, scale) is None
+
+
+def _make_decoder():
+    class Attention(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.to_qkv = nn.Linear(8, 24)
+            self.to_out = nn.Linear(8, 8)
+            self.norm_q = nn.RMSNorm(8, elementwise_affine=False)
+            self.norm_k = nn.RMSNorm(8, elementwise_affine=False)
+            self.spatial_parallel = False
+            self.dim_head = 8
+
+        def perform_attention(self, query, _key, _value, _pack_info):
+            return query
+
+        def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
+            return hidden_states
+
+    class FeedForward(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.w1 = nn.Linear(8, 32)
+            self.w2 = nn.Linear(16, 8)
+            self.use_gated = True
+            self.act_fn = nn.SiLU()
+            self._compile_forward_enabled = False
+            self._compile_forward_fatal = False
+            self._compiled_forward = None
+
+        def forward(self, hidden_states):
+            hidden_states = self.w1(hidden_states)
+            gate, hidden_states = hidden_states.chunk(2, dim=-1)
+            return self.w2(self.act_fn(gate) * hidden_states)
+
+    class Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attn = Attention()
+            self.ff = FeedForward()
+            self.norm1 = nn.RMSNorm(8)
+            self.norm2 = nn.RMSNorm(8)
+            self.scale1 = nn.Parameter(torch.zeros(8))
+            self.scale2 = nn.Parameter(torch.zeros(8))
+            self.use_scale = True
+
+        def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
+            return hidden_states
+
+    decoder = nn.Module()
+    decoder.transformer_blocks = nn.ModuleList([Block(), Block()])
+    decoder.proj_out = nn.Linear(8, 8)
+    return decoder
+
+
+def test_video_vae_installs_exact_optimizations(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.decoder = torch.nn.Module()
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = FakeModel()
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+    install = Mock(return_value=True)
+    monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install)
+    monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock())
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("cuda"),
+        load_device=torch.device("cpu"),
+    )
+
+    install.assert_called_once_with(
+        video_vae.model.decoder,
+        device=torch.device("cuda"),
+    )
+
+
+def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    operators = _operator_set()
+    monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
+    decoder = _make_decoder()
+
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+
+    for block in decoder.transformer_blocks:
+        assert block.attn.to_qkv.weight.dtype == torch.float16
+        assert block.attn.to_out.weight.dtype == torch.float16
+        assert block.ff.w1.weight.dtype == torch.float16
+        assert block.ff.w2.weight.dtype == torch.float16
+        assert block.forward.__func__.__name__ == "_optimized_transformer_block"
+        assert block.attn.forward.__func__.__name__ == "_optimized_attention"
+        assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward"
+        assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope
+        assert block._omni_scaled_residual is operators.scaled_residual
+    assert decoder.proj_out.weight.dtype == torch.float32
+
+    # Repeated installation is idempotent.
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+
+
+def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(
+        vae_ops,
+        "resolve_h3_vae_operators",
+        lambda _device: _operator_set(),
+    )
+    decoder = _make_decoder()
+    for block in decoder.transformer_blocks:
+        # The official checkpoint loads this JSON boolean-like field as int 0.
+        block.attn.spatial_parallel = 0
+
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+    assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False)
+
+
+def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    device, operators = _selected_operators()
+    monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
+    decoder = _make_decoder().to(device)
+    feed_forward = decoder.transformer_blocks[0].ff
+    reference_forward = type(feed_forward).forward
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=device,
+    )
+
+    hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32)
+    with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16):
+        expected = reference_forward(feed_forward, hidden_states)
+        actual = feed_forward(hidden_states)
+
+    assert expected.dtype == torch.float16
+    assert torch.equal(actual, expected)
+
+
+def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None)
+    decoder = _make_decoder()
+
+    assert not vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+    assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
+    for block in decoder.transformer_blocks:
+        assert block.attn.to_qkv.weight.dtype == torch.float32
+        assert block.ff.w1.weight.dtype == torch.float32
+        assert block.forward.__func__.__name__ == "forward"
+        assert block.attn.forward.__func__.__name__ == "forward"
+
+
+@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input])
+def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(
+        vae_ops,
+        "resolve_h3_vae_operators",
+        lambda _device: _operator_set(),
+    )
+    decoder = _make_decoder()
+    forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__
+    monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input)
+
+    assert not vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+    assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
+    assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
+
+
+@pytest.mark.parametrize(
+    "break_contract",
+    [
+        lambda block: delattr(block.attn, "spatial_parallel"),
+        lambda block: setattr(block.attn, "spatial_parallel", 2),
+        lambda block: setattr(block.attn, "perform_attention", None),
+        lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))),
+        lambda block: setattr(block.ff, "_compile_forward_enabled", True),
+    ],
+)
+def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(
+        vae_ops,
+        "resolve_h3_vae_operators",
+        lambda _device: _operator_set(),
+    )
+    decoder = _make_decoder()
+    break_contract(decoder.transformer_blocks[0])
+
+    assert not vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("meta"),
+    )
+    assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
+
+
+def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
+
+    first = _operator_set(supports=lambda _device: False)
+    added = _operator_set(supports=lambda device: device.type == "meta")
+    monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added))
+
+    assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added
+    assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
+
+
+def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
+
+    platform = Mock()
+    platform.is_cuda.return_value = True
+    platform.is_available.return_value = True
+    monkeypatch.setattr(dispatch, "HAS_TRITON", True)
+    monkeypatch.setattr(dispatch, "current_omni_platform", platform)
+
+    for capability in (90, 100, 103):
+        platform.get_device_capability.return_value.to_int.return_value = capability
+        assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None
+
+    for capability in (89, 101, 110):
+        platform.get_device_capability.return_value.to_int.return_value = capability
+        assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import gc
 from typing import Any
 
 import numpy as np
@@ -15,7 +16,11 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 AUDIO_MODEL: dict[str, dict[str, int | None]] = {
-    "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
+    # The inference peak includes backend workspaces as well as model weights.
+    # On ROCm, large MIOpen workspaces and allocator fragmentation can mask the
+    # resident-weight reduction, so use a conservative floor that still catches
+    # a disabled/no-op layerwise offloader.
+    "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 512},
 }
 
 IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
@@ -63,8 +68,9 @@ def run_inference(
 ) -> tuple[float, Any]:
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
-    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-    monitor.start()
+    with current_omni_platform.device(device_index):
+        free_bytes, total_bytes = current_omni_platform.mem_get_info()
+    initial_used_mb = (total_bytes - free_bytes) / (1024**2)
 
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
@@ -78,22 +84,37 @@ def run_inference(
         # cache_backend="cache_dit",
         **params["runner_params"],
     ) as runner:
+        # Measure steady-state inference memory, not model construction. Enabling
+        # layerwise offload first loads the model and then replaces each block's
+        # device storage with CPU-backed weights.  Monitoring that transition
+        # captures both the original model and temporary staging allocations,
+        # which is not representative of layerwise-offloaded inference.
+        monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
         current_omni_platform.reset_peak_memory_stats()
+        monitor.start()
 
-        # Refer to tests/e2e/offline_inference/test_wan22.py
-        # Use minimal settings for testing
-        output = runner.omni.generate(
-            "A cat sitting on a table",
-            OmniDiffusionSamplingParams(
-                generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-                guidance_scale=1.0,
-                num_inference_steps=num_inference_steps,
-                **params["sampler_params"],
-            ),
-        )
+        try:
+            # Refer to tests/e2e/offline_inference/test_wan22.py
+            # Use minimal settings for testing
+            output = runner.omni.generate(
+                "A cat sitting on a table",
+                OmniDiffusionSamplingParams(
+                    generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+                    guidance_scale=1.0,
+                    num_inference_steps=num_inference_steps,
+                    **params["sampler_params"],
+                ),
+            )
+        finally:
+            monitor.stop()
 
-    peak = monitor.peak_used_mb
-    monitor.stop()
+    # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
+    # starting usage so process-wide compiler/workspace caches retained from a
+    # previous run do not make the second measurement order-dependent.
+    peak = max(0.0, monitor.peak_used_mb - initial_used_mb)
+
+    gc.collect()
+    current_omni_platform.empty_cache()
 
     return peak, output
 
@@ -148,8 +169,10 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     assert expected_saved_memory is not None
 
     # Verify that layerwise offloading significantly reduces memory usage
-    # Passes only if the actual savings exceeds the expected savings
-    assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, (
+    # Passes only if the actual savings meets the expected savings
+    actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
+    assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
         f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
-        f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
+        f"should be at least {expected_saved_memory} MB less than no offload peak memory "
+        f"{no_offload_peak_memory} MB (actual savings: {actual_saved_memory} MB)"
     )

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -35,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadHook,
     PinnedResidentLayerGroup,
 )
-from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+)
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
@@ -807,7 +810,10 @@ def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_ru
     assert torch.equal(block.weight, expected)
 
 
-def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
+def test_hwr_registration_uses_transport_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
     dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
     plan, _, lease = _fake_hwr_plan()
     backend = DistributedLayerwiseOffloadBackend(
@@ -1125,6 +1131,51 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+    def test_hwr_registration_failure_falls_back_to_bounded_staging(
+        self,
+        monkeypatch,
+        patched_offload_runtime,
+    ):
+        plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        pipeline = _HWRPipeline()
+        original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
+
+        def fail_registration(*args, **kwargs):
+            del args, kwargs
+            raise HostRegistrationError("registration unavailable in CPU test")
+
+        def allocate_unpinned_staging(hooks, resident_group=None):
+            # Registration selection still observes pin_cpu_memory=True. The
+            # CPU-only test avoids asking the host for CUDA-pinned allocations.
+            for hook in hooks:
+                hook.pin_memory = False
+            return original_staging_allocator(hooks, resident_group)
+
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
+        monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
+
+        backend.enable(pipeline)
+
+        assert carrier.taken
+        assert backend._using_rank_local_mmap
+        assert not backend._using_registered_mmap
+        assert backend._host_registration is None
+        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+        assert not lease.closed
+
+        backend.disable()
+        assert lease.closed
+
     def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
         self,
         monkeypatch,
@@ -2132,8 +2183,9 @@ class TestDynamicSlotTracking:
                 prev_idx = (i - 1) % num_blocks
 
                 # Dynamic slot tracking: read from prev's prefetched slot
-                if prefetched_slots[prev_idx] is not None:
-                    current_slots[i] = prefetched_slots[prev_idx]
+                prefetched_slot = prefetched_slots[prev_idx]
+                if prefetched_slot is not None:
+                    current_slots[i] = prefetched_slot
 
                 read_slot = current_slots[i]
 

tests/diffusion/quantization/test_svdquant_config.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    UnquantizedLinearMethod,
+)
+
+from vllm_omni.quantization import svdquant_config
+from vllm_omni.quantization.factory import (
+    SUPPORTED_QUANTIZATION_METHODS,
+    build_quant_config,
+)
+from vllm_omni.quantization.svdquant_config import (
+    DiffusionSVDQuantConfig,
+    DiffusionSVDQuantLinearMethod,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_config_from_checkpoint_dict() -> None:
+    config = DiffusionSVDQuantConfig.from_config(
+        {
+            "rank": 32,
+            "precision": "nvfp4",
+            "act_unsigned": False,
+            "modules_to_not_convert": ["blocks.0.adaln_proj.linear"],
+        }
+    )
+
+    assert config.rank == 32
+    assert config.precision == "nvfp4"
+    assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"]
+    assert config.get_supported_act_dtypes() == [torch.bfloat16]
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"rank": 0}, "rank must be positive"),
+        ({"precision": "int4"}, "NVFP4 checkpoints only"),
+        ({"act_unsigned": True}, "unsigned activations"),
+    ],
+)
+def test_config_rejects_unsupported_phase1_options(
+    kwargs: dict[str, object],
+    message: str,
+) -> None:
+    with pytest.raises(ValueError, match=message):
+        DiffusionSVDQuantConfig(**kwargs)
+
+
+def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None:
+    config = build_quant_config(
+        {
+            "quant_method": "svdquant",
+            "rank": 16,
+            "precision": "nvfp4",
+            "w4a16_modules": ["adaln_proj.linear"],
+            "w4a16_group_size": 64,
+        }
+    )
+
+    assert isinstance(config, DiffusionSVDQuantConfig)
+    assert config.rank == 16
+    assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_skipped_and_non_linear_modules_remain_unquantized() -> None:
+    config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"])
+    linear = Mock(spec=LinearBase)
+
+    assert isinstance(
+        config.get_quant_method(linear, "blocks.0.adaln_proj.linear"),
+        UnquantizedLinearMethod,
+    )
+    assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None
+
+
+def test_linear_method_creates_canonical_partitioned_weights(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32))
+    layer = torch.nn.Module()
+
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[64, 64],
+        input_size=128,
+        output_size=128,
+        params_dtype=torch.bfloat16,
+    )
+
+    assert layer.qweight.shape == (128, 64)
+    assert layer.qweight.dtype == torch.int8
+    assert layer.wscales.shape == (8, 128)
+    assert layer.wscales.dtype == torch.float8_e4m3fn
+    assert layer.proj_down.shape == (128, 32)
+    assert layer.proj_up.shape == (128, 32)
+    assert layer.smooth_factor.shape == (128,)
+    assert layer.wcscales.shape == (128,)
+    assert layer.wtscale.shape == (1,)
+
+    assert layer.qweight.input_dim == 1
+    assert layer.qweight.output_dim == 0
+    assert layer.wscales.input_dim == 0
+    assert layer.wscales.output_dim == 1
+    assert layer.proj_down.input_dim == 0
+    assert layer.proj_up.output_dim == 0
+    assert layer.wcscales.output_dim == 0
+
+
+def test_linear_method_rejects_misaligned_input_partition(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig())
+
+    with pytest.raises(ValueError, match="divisible by the block size 16"):
+        method.create_weights(
+            torch.nn.Module(),
+            input_size_per_partition=24,
+            output_partition_sizes=[32],
+            input_size=24,
+            output_size=32,
+            params_dtype=torch.bfloat16,
+        )
+
+
+def test_active_linear_uses_svdquant_method(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+    config = DiffusionSVDQuantConfig()
+    linear = Mock(spec=LinearBase)
+
+    method = config.get_quant_method(linear, "blocks.0.attn.out_proj")
+
+    assert isinstance(method, DiffusionSVDQuantLinearMethod)

tests/diffusion/quantization/test_svdquant_linear.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeNvfp4Kernel:
+    def __init__(self, base_weight: torch.Tensor | None = None) -> None:
+        self.base_weight = base_weight
+        self.processed = False
+        self.last_input: torch.Tensor | None = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        del layer
+        self.processed = True
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        del layer
+        assert bias is None
+        assert self.base_weight is not None
+        self.last_input = x.clone()
+        return torch.mm(x, self.base_weight)
+
+
+def _register_parameter(
+    layer: torch.nn.Module,
+    name: str,
+    value: torch.Tensor,
+) -> None:
+    layer.register_parameter(
+        name,
+        torch.nn.Parameter(value, requires_grad=False),
+    )
+
+
+def test_supports_only_validated_datacenter_blackwell() -> None:
+    assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0))
+    assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3))
+    assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0))
+    assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0))
+    assert not svdquant._supports_capability(None)
+
+
+def test_rejects_nvfp4_backend_with_incompatible_layer_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    svdquant._nvfp4_kernel.cache_clear()
+    monkeypatch.setattr(
+        svdquant,
+        "init_nvfp4_linear_kernel",
+        lambda: _FakeNvfp4Kernel(),
+    )
+
+    with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"):
+        svdquant._nvfp4_kernel()
+
+    svdquant._nvfp4_kernel.cache_clear()
+
+
+def test_prepare_weights_uses_existing_nvfp4_layout_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeNvfp4Kernel()
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4)
+    wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4)
+    _register_parameter(layer, "qweight", qweight)
+    _register_parameter(layer, "wscales", wscales)
+    _register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16))
+    _register_parameter(
+        layer,
+        "wcscales",
+        torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16),
+    )
+
+    method.process_weights_after_loading(layer)
+
+    assert kernel.processed
+    assert not hasattr(layer, "qweight")
+    assert not hasattr(layer, "wscales")
+    assert layer.weight.dtype == torch.uint8
+    assert torch.equal(layer.weight.view(torch.int8), qweight)
+    assert torch.equal(layer.weight_scale, wscales.transpose(0, 1))
+    torch.testing.assert_close(layer.alpha, torch.tensor([2.0]))
+    assert layer.input_global_scale_inv.dtype == torch.float32
+    assert layer.output_channel_scale.shape == (4,)
+
+
+def test_prepare_weights_drops_identity_output_scale(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeNvfp4Kernel()
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    _register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8))
+    _register_parameter(
+        layer,
+        "wscales",
+        torch.ones(2, 4, dtype=torch.float8_e4m3fn),
+    )
+    _register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16))
+    _register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16))
+
+    method.process_weights_after_loading(layer)
+
+    assert layer.output_channel_scale is None
+
+
+def test_compatibility_path_applies_qkv_scale_and_rank_correction(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    base_weight = torch.tensor(
+        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+        dtype=torch.bfloat16,
+    )
+    kernel = _FakeNvfp4Kernel(base_weight)
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+
+    layer = torch.nn.Module()
+    layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16)
+    layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16)
+    layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)
+    layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16)
+    layer.output_size_per_partition = 3
+    x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16)
+    bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16)
+
+    actual = method.apply(layer, x, bias)
+
+    smoothed = x / layer.smooth_factor
+    expected = torch.mm(smoothed, base_weight)
+    expected = expected * layer.output_channel_scale
+    expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T)
+    expected = expected + bias
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(kernel.last_input, smoothed)
+
+
+def test_compatibility_path_rejects_fp16_activations(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+    layer = torch.nn.Module()
+    with pytest.raises(ValueError, match="requires BF16 activations"):
+        method.apply(
+            layer,
+            torch.zeros(1, 16, dtype=torch.float16),
+        )

tests/diffusion/quantization/test_svdquant_tp_loading.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
+    layer = object.__new__(linear_type)
+    torch.nn.Module.__init__(layer)
+    layer.tp_rank = tp_rank
+    return layer
+
+
+def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
+    values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
+    return values.reshape(shape).to(dtype)
+
+
+def _method(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
+
+
+def test_column_parallel_loader_shards_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[32],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = loaded.narrow(parameter.output_dim, 32, 32)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_row_parallel_loader_shards_svdquant_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(RowParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=64,
+        output_partition_sizes=[64],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_down": _values((128, 8), torch.bfloat16),
+        "smooth_factor": _values((128,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        shard_size = parameter.shape[parameter.input_dim]
+        expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
+    layer.total_num_heads = 4
+    layer.total_num_kv_heads = 2
+    layer.num_heads = 2
+    layer.num_kv_heads = 1
+    layer.num_kv_head_replicas = 1
+    layer.head_size = 8
+    layer.v_head_size = 8
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[16, 8, 8],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = torch.cat(
+            [
+                loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
+                for part in expected_global_outputs
+            ],
+            dim=parameter.output_dim,
+        )
+        torch.testing.assert_close(parameter, expected)

tests/diffusion/test_diffusion_engine_metrics.py

@@ -172,12 +172,29 @@ class TestMetricKeys:
         assert ".cancel(" not in abort_branch
 
     def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
+        """A metrics-only sentinel contributes its queue depth and is not routed.
+
+        Absorption lives in ``_absorb_diffusion_metrics`` rather than inline in a
+        loop, so both orchestration loops share one implementation. The ordering
+        is asserted where it now lives: the snapshot inside the helper, and the
+        absorb-before-route guard inside every loop that polls diffusion output.
+        """
         source = _read_source(_ORCHESTRATOR_PATH)
-        loop_source = _get_function_source(source, "Orchestrator", "_orchestration_loop")
-        snapshot_pos = loop_source.index("_update_stage_replica_waiting(")
-        sentinel_pos = loop_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
-        route_pos = loop_source.index("pool.record_output_timestamps([diffusion_output])")
-        assert snapshot_pos < sentinel_pos < route_pos
+
+        # Snapshot before the verdict, so a sentinel still reports its waiting
+        # depth on the way out instead of being dropped unaccounted.
+        absorb_source = _get_function_source(source, "Orchestrator", "_absorb_diffusion_metrics")
+        snapshot_pos = absorb_source.index("_update_stage_replica_waiting(")
+        sentinel_pos = absorb_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
+        assert snapshot_pos < sentinel_pos
+
+        # Absorb before routing, or a sentinel reaches the downstream consumer
+        # as if it were a real request output.
+        for loop_name in ("_orchestration_loop", "_orchestration_loop_event_driven"):
+            loop_source = _get_function_source(source, "Orchestrator", loop_name)
+            absorb_pos = loop_source.index("self._absorb_diffusion_metrics(")
+            route_pos = loop_source.index("record_output_timestamps(")
+            assert absorb_pos < route_pos, loop_name
 
 
 class TestVaeDecodeEmit:

tests/diffusion/test_diffusion_output_formatter.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
     result = results[0]
     assert result.images == ["frame-0"]
     assert result.prompt == "prompt-0"
-    assert result.final_output_type == "image"
+    assert result.final_output_type == "video"
     assert result.multimodal_output == {
         "metadata": {
             "audio": {"sample_rate": 48000},
@@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
 
     assert result.images == ["frame-0"]
     assert result.prompt == "watch and listen"
-    assert result.final_output_type == "image"
+    assert result.final_output_type == "video"
     assert result.multimodal_output == {
         "metadata": {
             "audio": {"sample_rate": 16000},

tests/diffusion/test_diffusion_streaming_output.py

@@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint:
             lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
         )
         mocker.patch(
-            "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
-            return_value="diffusion",
+            "vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
+            return_value=True,
         )
         mocker.patch(
             "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",

tests/diffusion/test_diffusion_worker_process_title.py

@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+import shutil
+import subprocess
+import sys
+import textwrap
+import time
+from dataclasses import dataclass
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@dataclass(frozen=True)
+class _FakeGroup:
+    world_size: int = 1
+    rank_in_group: int = 0
+
+
+_GROUP_GETTERS = (
+    "get_dp_group",
+    "get_pp_group",
+    "get_sp_group",
+    "get_cfg_group",
+    "get_tp_group",
+    "get_fs_group",
+    "get_hsdp_replicate_group",
+    "get_ep_group",
+)
+
+
+def _patch_groups(
+    mocker: MockerFixture,
+    groups: dict[str, _FakeGroup],
+) -> dict[str, Mock]:
+    patched_getters = {}
+    for getter_name in _GROUP_GETTERS:
+        patched_getters[getter_name] = mocker.patch.object(
+            worker_module,
+            getter_name,
+            return_value=groups.get(getter_name, _FakeGroup()),
+        )
+    return patched_getters
+
+
+def test_setup_uses_default_name_before_model_parallel_init(
+    mocker: MockerFixture,
+) -> None:
+    mocker.patch.object(
+        worker_module,
+        "model_parallel_is_initialized",
+        return_value=False,
+    )
+    group_getters = _patch_groups(mocker, {})
+    set_process_title = mocker.patch.object(worker_module, "set_process_title")
+    decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+    worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True)
+
+    set_process_title.assert_called_once_with(
+        name="DiffusionWorker",
+        prefix="vLLM-Omni",
+    )
+    decorate_logs.assert_called_once_with("DiffusionWorker")
+    for group_getter in group_getters.values():
+        group_getter.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
+    [
+        ({}, False, False, 1, "DiffusionWorker"),
+        ({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"),
+        ({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"),
+        ({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"),
+        (
+            {
+                "get_cfg_group": _FakeGroup(2, 1),
+                "get_tp_group": _FakeGroup(2, 1),
+            },
+            False,
+            False,
+            1,
+            "DiffusionWorker_CFG1_TP1",
+        ),
+        (
+            {
+                "get_dp_group": _FakeGroup(2, 1),
+                "get_sp_group": _FakeGroup(4, 2),
+                "get_tp_group": _FakeGroup(2, 1),
+            },
+            False,
+            False,
+            1,
+            "DiffusionWorker_DP1_SP2_TP1",
+        ),
+        (
+            {"get_fs_group": _FakeGroup(4, 3)},
+            False,
+            True,
+            1,
+            "DiffusionWorker_FS3",
+        ),
+        (
+            {
+                "get_sp_group": _FakeGroup(2, 1),
+                "get_fs_group": _FakeGroup(4, 3),
+            },
+            False,
+            True,
+            1,
+            "DiffusionWorker_SP1_FS3",
+        ),
+        (
+            {
+                "get_fs_group": _FakeGroup(1, 0),
+                "get_hsdp_replicate_group": _FakeGroup(2, 1),
+            },
+            False,
+            True,
+            2,
+            "DiffusionWorker_RP1",
+        ),
+        ({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"),
+        ({}, True, False, 1, "DiffusionWorker"),
+    ],
+    ids=[
+        "all-singleton",
+        "tp-only",
+        "dp-only",
+        "pp-only",
+        "tp-cfg",
+        "dp-sp-tp",
+        "hsdp-only",
+        "hsdp-sp",
+        "hsdp-replicated",
+        "expert-parallel",
+        "singleton-ep",
+    ],
+)
+def test_setup_uses_initialized_parallel_groups(
+    mocker: MockerFixture,
+    groups: dict[str, _FakeGroup],
+    enable_ep: bool,
+    use_hsdp: bool,
+    hsdp_replicate_size: int,
+    expected_name: str,
+) -> None:
+    mocker.patch.object(
+        worker_module,
+        "model_parallel_is_initialized",
+        return_value=True,
+    )
+    group_getters = _patch_groups(mocker, groups)
+    set_process_title = mocker.patch.object(worker_module, "set_process_title")
+    decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+    worker_module._setup_diffusion_worker_proc_title_and_log_prefix(
+        enable_ep=enable_ep,
+        use_hsdp=use_hsdp,
+        hsdp_replicate_size=hsdp_replicate_size,
+    )
+
+    set_process_title.assert_called_once_with(
+        name=expected_name,
+        prefix="vLLM-Omni",
+    )
+    decorate_logs.assert_called_once_with(expected_name)
+    if use_hsdp:
+        group_getters["get_fs_group"].assert_called_once_with()
+        if hsdp_replicate_size > 1:
+            group_getters["get_hsdp_replicate_group"].assert_called_once_with()
+        else:
+            group_getters["get_hsdp_replicate_group"].assert_not_called()
+    else:
+        group_getters["get_fs_group"].assert_not_called()
+        group_getters["get_hsdp_replicate_group"].assert_not_called()
+    if enable_ep:
+        group_getters["get_ep_group"].assert_called_once_with()
+    else:
+        group_getters["get_ep_group"].assert_not_called()
+
+
+def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
+    mocker.patch.object(
+        worker_module,
+        "model_parallel_is_initialized",
+        return_value=False,
+    )
+    mocker.patch.dict(sys.modules, {"setproctitle": None})
+    decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+    worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
+
+    decorate_logs.assert_called_once_with("DiffusionWorker")
+
+
+@pytest.mark.skipif(
+    sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None,
+    reason="requires Linux ps and setproctitle",
+)
+def test_process_title_is_visible_through_ps() -> None:
+    expected_name = "vLLM-Omni::DiffusionWorker_TP1"
+    child_script = textwrap.dedent(
+        """
+        import time
+
+        from vllm.utils.system_utils import set_process_title
+
+
+        set_process_title(
+            name="DiffusionWorker_TP1",
+            prefix="vLLM-Omni",
+        )
+        time.sleep(30)
+        """
+    )
+    process = subprocess.Popen(
+        [sys.executable, "-c", child_script],
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    observed_title = ""
+    try:
+        for _ in range(100):
+            if process.poll() is not None:
+                stderr = process.stderr.read() if process.stderr is not None else ""
+                pytest.fail(f"child process exited before title check: {stderr}")
+            result = subprocess.run(
+                ["ps", "-o", "args=", "-p", str(process.pid)],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+            observed_title = result.stdout.strip()
+            if expected_name in observed_title:
+                break
+            time.sleep(0.1)
+        else:
+            pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}")
+    finally:
+        process.terminate()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait(timeout=5)

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -52,6 +52,11 @@ def client(mock_engine):
         c.shutdown()
 
 
+def test_inline_client_implements_shared_prompt_hook_protocol(client):
+    assert client.prompt_transform_func is None
+    assert client.prompt_expand_func is None
+
+
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_success(client, mock_engine):
     mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -1209,6 +1209,67 @@ def test_cleanup_preserves_pending_save(build_adapter):
     assert len(adapter._pending_save_reqs) == 1
 
 
+def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
+    )
+
+    adapter.save_async(multimodal_output=None, request=request)
+    stale_task = adapter._pending_save_reqs.popleft()
+    adapter.cleanup(request.request_id, request.external_req_id)
+    adapter._send_single_request(stale_task)
+
+    connector.put.assert_not_called()
+    assert request.external_req_id not in adapter.put_req_chunk
+
+    replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
+    adapter.save_async(multimodal_output=None, request=replacement)
+    current_task = adapter._pending_save_reqs.popleft()
+    adapter._send_single_request(current_task)
+
+    connector.put.assert_called_once()
+    assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
+
+
+def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
+    adapter, connector = build_adapter(stage_id=1)
+    first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+    put_started = threading.Event()
+    release_put = threading.Event()
+
+    def blocking_put(**kwargs):
+        put_started.set()
+        release_put.wait(timeout=2)
+        return True, 1, {}
+
+    connector.put.side_effect = blocking_put
+    adapter.save_async(multimodal_output=None, request=first)
+    task = adapter._pending_save_reqs.popleft()
+    sender = threading.Thread(target=adapter._send_single_request, args=(task,))
+    sender.start()
+    assert put_started.wait(timeout=1)
+    cleaner = threading.Thread(
+        target=adapter.finish_requests,
+        args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
+    )
+    cleaner.start()
+    try:
+        cleaner.join(timeout=0.5)
+        assert not cleaner.is_alive()
+        assert task["sender_token"].cancelled
+        assert first.external_req_id in adapter._sender_tokens
+    finally:
+        release_put.set()
+        sender.join(timeout=1)
+        cleaner.join(timeout=1)
+    assert not sender.is_alive()
+    assert first.external_req_id not in adapter._sender_tokens
+    assert first.external_req_id not in adapter.put_req_chunk
+
+
 def test_cleanup_only_affects_target_request(build_adapter):
     """Cleanup for one request must not affect another request's state."""
     adapter, _ = build_adapter(stage_id=1)
@@ -1507,6 +1568,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     scheduler.structured_output_manager.should_advance.return_value = False
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
 
     request = _HashableRequest(
@@ -1587,6 +1649,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager = mocker.MagicMock()
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
     scheduler.waiting_for_transfer_free = set()
     scheduler.transfer_triggered_requests = set()
@@ -1742,6 +1805,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     scheduler.structured_output_manager.should_advance.return_value = False
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
     scheduler._pending_finish_reqs = list(pending_finish_reqs)
 
@@ -2216,3 +2280,43 @@ def test_expiry_is_per_request(build_adapter):
 
     assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
     assert "healthy" in adapter._waiting_since
+
+
+def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
+    from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+        talker2code2wav_async_chunk,
+    )
+
+    adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
+    external_req_id = "ext-native-abort"
+    first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+    first.resumable = True
+    first.model_intermediate_buffer = None
+    first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
+    first_frame = torch.full((1, 31), 101, dtype=torch.long)
+
+    talker2code2wav_async_chunk(
+        adapter,
+        {"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
+        first,
+        is_finished=False,
+    )
+
+    adapter.requests_num_chunks_sent[external_req_id] = 1
+    adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
+
+    assert external_req_id not in adapter.request_payload
+    assert external_req_id not in adapter.requests_num_chunks_sent
+
+    replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+    replacement.resumable = True
+    replacement.model_intermediate_buffer = None
+    replacement.additional_information = first.additional_information
+    replacement_payload = talker2code2wav_async_chunk(
+        adapter,
+        {"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
+        replacement,
+        is_finished=False,
+    )
+
+    assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """E2E accuracy guards against pinned official Lightricks LTX references.
 
@@ -22,6 +22,7 @@ import shutil
 import subprocess
 import sys
 import time
+from dataclasses import asdict, dataclass
 from pathlib import Path
 
 import numpy as np
@@ -42,7 +43,7 @@ LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
 LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
 LTX25_OFFICIAL_COMMON_FILES = {
     "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
-    "video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
+    "video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
     "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
 }
 LTX25_OFFICIAL_VARIANT_FILES = {
@@ -83,35 +84,27 @@ NEGATIVE_PROMPT = (
 # Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
 # Degenerate one-token attention falls back to PyTorch SDPA MATH.
 ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
-VIDEO_SSIM_MEAN_THRESHOLD = 0.95
-VIDEO_SSIM_MIN_THRESHOLD = 0.90
-VIDEO_PSNR_MEAN_THRESHOLD = 30.0
-AUDIO_RELATIVE_L2_THRESHOLD = 0.2
-AUDIO_COSINE_THRESHOLD = 0.95
-
-# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
-# The metrics artifact preserves the exact observed values.
-LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
-LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
-LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
-LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
-LTX25_AUDIO_COSINE_THRESHOLD = 0.95
-
-# Full/SFT uses independently converted dev weights, but the release contract
-# still requires decoded-video parity at the same 0.99 SSIM floor as the
-# distilled paths when both runtimes use BF16 cuDNN attention.
-LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
-LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
-LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
-LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
-LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
-LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
-LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
 LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
 
 
+@dataclass(frozen=True)
+class LTXAccuracyThresholds:
+    video_ssim_mean: float
+    video_ssim_min: float
+    video_psnr_mean_db: float
+    audio_relative_l2: float
+    audio_cosine_similarity: float
+
+
+STRICT_THRESHOLDS = LTXAccuracyThresholds(
+    video_ssim_mean=0.99,
+    video_ssim_min=0.99,
+    video_psnr_mean_db=40.0,
+    audio_relative_l2=0.10,
+    audio_cosine_similarity=0.99,
+)
+
+
 def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
     from vllm_omni.outputs import OmniRequestOutput
 
@@ -215,6 +208,9 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
         model = Path(configured_model).resolve()
         snapshot_revision = model.name if model.parent.name == "snapshots" else None
         assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
+        assert (model / "diffusion_decoder" / "config.json").is_file(), (
+            f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}"
+        )
         return model, configured_model, configured_revision or snapshot_revision
     model_id = configured_model or LTX25_OMNI_MODEL_ID
     revision = configured_revision
@@ -230,6 +226,7 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
                 "connectors/config.json",
                 "connectors/diffusion_pytorch_model.safetensors.index.json",
                 "connectors/diffusion_pytorch_model-*.safetensors",
+                "diffusion_decoder/config.json",
                 "scheduler/*",
                 "text_encoder/config.json",
                 "text_encoder/generation_config.json",
@@ -419,8 +416,6 @@ def test_ltx25_i2v_requests_pin_official_crf() -> None:
 
 
 def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
-    from benchmarks.diffusion.quantization_quality import compute_lpips_video
-
     assert reference.shape == prediction.shape
     assert reference.ndim == 4 and reference.shape[-1] == 3
     ssim_scores: list[float] = []
@@ -435,7 +430,6 @@ def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
         "ssim_mean": float(np.mean(ssim_scores)),
         "ssim_min": float(np.min(ssim_scores)),
         "psnr_mean_db": float(np.mean(psnr_scores)),
-        "lpips_alex_mean": compute_lpips_video(reference, prediction),
         "max_abs": float(difference.max()),
         "mean_abs": float(difference.mean()),
     }
@@ -463,13 +457,24 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
     }
 
 
+def _assert_strict_similarity(
+    video_metrics: dict[str, float],
+    audio_metrics: dict[str, float | bool],
+) -> None:
+    assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+    assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+    assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+    assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+    assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
+
+
 @pytest.mark.slow
 @pytest.mark.benchmark
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
-    """Compare Omni with the official runtime using the same checkpoint connector weights."""
+    """Compare Omni with official canonical DiffVAE and matching connector weights."""
     _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
@@ -544,6 +549,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
     omni_metadata = json.loads((omni_output / "metadata.json").read_text())
     assert official_metadata["official_revision"] == official_revision
     assert official_metadata["pipeline"] == "DistilledPipeline"
+    assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
     assert official_metadata["attention_backend"] == ATTENTION_BACKEND
     assert official_metadata["connector_model"] == str(omni_model)
     assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
@@ -569,6 +575,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
         "omni_model": omni_model_source,
         "omni_model_revision": omni_model_revision,
         "resolved_omni_model_path": str(omni_model),
+        "thresholds": asdict(STRICT_THRESHOLDS),
         "peak_memory_mb": {
             "official_allocated": official_metadata["peak_memory_allocated_mb"],
             "official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -582,13 +589,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
-    assert video_metrics["ssim_mean"] >= ssim_mean_threshold
-    assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
-    assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
-    assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
-    assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
-    assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
+    _assert_strict_similarity(video_metrics, audio_metrics)
 
 
 @pytest.mark.slow
@@ -598,7 +599,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
 def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
-    """Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
+    """Compare raw official Full/SFT and DiffVAE weights with Omni."""
     _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
@@ -701,6 +702,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
             "full_two_stage": "TI2VidTwoStagesPipeline",
         }[pipeline_mode]
     )
+    assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
     assert official_metadata["attention_backend"] == ATTENTION_BACKEND
     assert official_metadata["connector_model"] == str(omni_model)
     assert omni_metadata["model_class_name"] == model_class_name
@@ -729,6 +731,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
         "omni_model": omni_model_source,
         "omni_model_revision": omni_model_revision,
         "resolved_omni_model_path": str(omni_model),
+        "thresholds": asdict(STRICT_THRESHOLDS),
         "peak_memory_mb": {
             "official_allocated": official_metadata["peak_memory_allocated_mb"],
             "official_reserved": official_metadata["peak_memory_reserved_mb"],
@@ -742,12 +745,4 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
-    assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
-    assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
-    lpips_threshold = (
-        LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
-    )
-    assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
-    assert video_metrics["lpips_alex_mean"] <= lpips_threshold
-    assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD
+    _assert_strict_similarity(video_metrics, audio_metrics)

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -3,9 +3,8 @@
 
 """E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
 
-The original reduced one-stage guards remain unchanged. Four complementary
-default-shape cases cover both model versions, both two-stage weight families,
-and T2V/I2V without expanding every combination into a duplicate golden.
+Four complementary default-shape cases cover both model versions, both
+two-stage weight families, and T2V/I2V without duplicating one-stage coverage.
 """
 
 from __future__ import annotations
@@ -69,25 +68,23 @@ class LTXAccuracyThresholds:
     video_ssim_mean: float
     video_ssim_min: float
     video_psnr_mean_db: float
-    # Zero disables gating for waveform-sensitive audio that is not expected
-    # to align with the official output.
     audio_relative_l2: float
     audio_cosine_similarity: float
 
 
 STRICT_THRESHOLDS = LTXAccuracyThresholds(
-    video_ssim_mean=0.95,
-    video_ssim_min=0.90,
-    video_psnr_mean_db=30.0,
-    audio_relative_l2=0.20,
-    audio_cosine_similarity=0.95,
+    video_ssim_mean=0.99,
+    video_ssim_min=0.99,
+    video_psnr_mean_db=40.0,
+    audio_relative_l2=0.10,
+    audio_cosine_similarity=0.99,
 )
 
 
 @dataclass(frozen=True)
 class LTXAccuracyCase:
     name: str
-    pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
+    pipeline_kind: Literal["distilled", "two_stage"]
     model_id: str
     model_revision: str
     model_env: str
@@ -99,8 +96,6 @@ class LTXAccuracyCase:
     num_inference_steps: int
     seed: int
     stg_block: int | None
-    thresholds: LTXAccuracyThresholds
-    prompt: str = PROMPT
     gemma_model_id: str | None = None
     gemma_model_revision: str | None = None
     gemma_model_env: str | None = None
@@ -168,59 +163,6 @@ I2V_IMAGE = LTXArtifact(
 )
 
 
-LEGACY_CASES = (
-    LTXAccuracyCase(
-        name="ltx2",
-        pipeline_kind="one_stage",
-        model_id="Lightricks/LTX-2",
-        model_revision=LTX2_REVISION,
-        model_env="VLLM_TEST_LTX2_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX2_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=29,
-        thresholds=STRICT_THRESHOLDS,
-    ),
-    LTXAccuracyCase(
-        name="ltx2_3",
-        pipeline_kind="one_stage",
-        model_id="diffusers/LTX-2.3-Diffusers",
-        model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
-        model_env="VLLM_TEST_LTX23_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX23_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=28,
-        thresholds=STRICT_THRESHOLDS,
-    ),
-    LTXAccuracyCase(
-        name="ltx2_3_i2v",
-        pipeline_kind="one_stage",
-        model_id="diffusers/LTX-2.3-Diffusers",
-        model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
-        model_env="VLLM_TEST_LTX23_MODEL",
-        model_class_name="LTX2Pipeline",
-        checkpoint=LTX23_CHECKPOINT,
-        width=512,
-        height=384,
-        num_frames=25,
-        num_inference_steps=20,
-        seed=42,
-        stg_block=28,
-        thresholds=STRICT_THRESHOLDS,
-        image=I2V_IMAGE,
-    ),
-)
-
-
 TWO_STAGE_CASES = (
     LTXAccuracyCase(
         name="ltx2_distilled_t2v",
@@ -240,7 +182,6 @@ TWO_STAGE_CASES = (
         gemma_model_id="Lightricks/LTX-2",
         gemma_model_revision=LTX2_REVISION,
         gemma_model_env="VLLM_TEST_LTX2_MODEL",
-        thresholds=STRICT_THRESHOLDS,
     ),
     LTXAccuracyCase(
         name="ltx23_distilled_i2v",
@@ -260,7 +201,6 @@ TWO_STAGE_CASES = (
         gemma_model_id="diffusers/LTX-2.3-Diffusers",
         gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
         gemma_model_env="VLLM_TEST_LTX23_MODEL",
-        thresholds=STRICT_THRESHOLDS,
         image=I2V_IMAGE,
     ),
     LTXAccuracyCase(
@@ -280,7 +220,6 @@ TWO_STAGE_CASES = (
         seed=10,
         stg_block=29,
         enable_layerwise_offload=True,
-        thresholds=STRICT_THRESHOLDS,
     ),
     LTXAccuracyCase(
         name="ltx23_two_stage_layer_fused_i2v",
@@ -299,12 +238,11 @@ TWO_STAGE_CASES = (
         seed=10,
         stg_block=28,
         enable_layerwise_offload=True,
-        thresholds=STRICT_THRESHOLDS,
         image=I2V_IMAGE,
     ),
 )
 
-CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
+WEEKLY_CASES = TWO_STAGE_CASES
 
 
 def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
@@ -460,10 +398,49 @@ def _resolve_image(case: LTXAccuracyCase) -> Path | None:
     return _resolve_artifact(case.image)
 
 
+def _omni_model_with_pinned_artifacts(
+    model: Path,
+    output_root: Path,
+    artifacts: tuple[Path | None, ...],
+) -> Path:
+    """Expose the already-resolved sidecars to Omni without copying the model."""
+    resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None)
+    if not resolved_artifacts:
+        return model
+
+    overlay = output_root / "omni-model"
+    overlay.mkdir()
+    for source in model.iterdir():
+        (overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir())
+    for artifact in resolved_artifacts:
+        target = overlay / artifact.name
+        if target.exists():
+            assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}"
+        else:
+            target.symlink_to(artifact.resolve())
+    return overlay
+
+
+def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None:
+    model = tmp_path / "model"
+    model.mkdir()
+    component = model / "model_index.json"
+    component.write_text("{}\n")
+    artifact = tmp_path / "sidecar.safetensors"
+    artifact.write_bytes(b"pinned")
+    output_root = tmp_path / "output"
+    output_root.mkdir()
+
+    overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,))
+
+    assert (overlay / component.name).samefile(component)
+    assert (overlay / artifact.name).samefile(artifact)
+
+
 def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
     request: dict[str, object] = {
         "pipeline_kind": case.pipeline_kind,
-        "prompt": case.prompt,
+        "prompt": PROMPT,
         "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
         "width": case.width,
         "height": case.height,
@@ -473,14 +450,17 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
         "seed": case.seed,
     }
     if case.pipeline_kind != "distilled":
-        assert case.stg_block is not None
         request.update(
             video_cfg_scale=3.0,
             audio_cfg_scale=7.0,
-            video_stg_scale=1.0,
-            audio_stg_scale=1.0,
             video_modality_scale=3.0,
             audio_modality_scale=3.0,
+        )
+    if case.pipeline_kind == "two_stage":
+        assert case.stg_block is not None
+        request.update(
+            video_stg_scale=1.0,
+            audio_stg_scale=1.0,
             video_rescale_scale=0.7,
             audio_rescale_scale=0.7,
             video_stg_blocks=[case.stg_block],
@@ -543,7 +523,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
     }
 
 
-@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
+@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name)
 @pytest.mark.slow
 @pytest.mark.benchmark
 @pytest.mark.diffusion
@@ -561,9 +541,11 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
     checkpoint = _resolve_artifact(case.checkpoint, model)
     spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
     distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
+    omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora))
     image = _resolve_image(case)
+    request = _request(case, image)
     request_path = output_root / "request.json"
-    request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
+    request_path.write_text(json.dumps(request, indent=2) + "\n")
 
     runner = Path(__file__).with_name("run_ltx_reference.py")
     runner_args = [
@@ -619,7 +601,7 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
             "--output-dir",
             str(omni_output),
             "--model",
-            str(model),
+            str(omni_model),
             "--model-class-name",
             case.model_class_name,
         ]
@@ -651,18 +633,16 @@ def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Pat
         "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
         "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
         "enable_layerwise_offload": enable_layerwise_offload,
-        "thresholds": asdict(case.thresholds),
-        "request": _request(case, image),
+        "thresholds": asdict(STRICT_THRESHOLDS),
+        "request": request,
         "video": video_metrics,
         "audio": audio_metrics,
     }
     (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
     print(json.dumps(result, indent=2))
 
-    assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
-    assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
-    assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
-    if case.thresholds.audio_relative_l2 > 0:
-        assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
-    if case.thresholds.audio_cosine_similarity > 0:
-        assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
+    assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+    assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+    assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+    assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+    assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity

tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -119,6 +119,9 @@ def _model_name(model_env_var: str = MODEL_ENV_VAR, component: str = "FL2VA") ->
 def _server_args() -> list[str]:
     return [
         "--trust-remote-code",
+        # Both official-reference requests use fixed shapes. Keep regional
+        # compilation static so cold Inductor codegen is stable across CI nodes.
+        "--no-diffusion-compile-dynamic",
         "--num-gpus",
         "4",
         "--usp",

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -537,7 +537,7 @@ def api_server(server_case: ServerCase, mock_async_omni, tmp_path):
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)
 
-    port = get_open_port()
+    port = get_open_port(host="0.0.0.0")
     args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
 
     def run_server():
@@ -814,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
             ServerCase(
                 served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
                 outputs=[_build_diffusion_video_output()],
             ),
             "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
@@ -825,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
             ServerCase(
                 served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
                 outputs=[_build_diffusion_video_output()],
             ),
             "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
@@ -886,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -897,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),

tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py

@@ -8,6 +8,7 @@ from dataclasses import FrozenInstanceError
 
 import pytest
 
+from tests.helpers.stage_config import get_deploy_duplex_max_sessions
 from vllm_omni.config.stage_config import (
     DuplexSessionRuntimeConfig,
     load_deploy_config,
@@ -66,3 +67,18 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str,
 
     with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
         load_deploy_config(deploy_path)
+
+
+@pytest.mark.parametrize(
+    ("deploy_yaml", "expected"),
+    [
+        ("minicpmo_4_5.yaml", 4),
+        ("minicpmo_4_5_8x4090.yaml", 1),
+        ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
+    ],
+)
+def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None:
+    # Guards the admission probe. A capacity edit, a config that declares none,
+    # or an overlay inheriting one from its base must surface here rather than
+    # as another nightly duplex admission timeout.
+    assert get_deploy_duplex_max_sessions(deploy_yaml) == expected

tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
+    NemotronVoiceChatDataPlaneContext,
+    NemotronVoiceChatDataPlaneSession,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_RUNTIME = {
+    "nvc_text_bos_id": 0,
+    "nvc_text_eos_id": 1,
+    "nvc_text_pad_id": 12,
+    "nvc_function_sotc_id": 20,
+    "nvc_function_eotc_id": 21,
+    "nvc_tokenizer_ref": "test-tokenizer",
+}
+
+
+def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
+    projector = NemotronVoiceChatDataPlaneSession(encode_audio)
+    projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
+    return projector
+
+
+def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
+    metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
+    if function_token is not None:
+        metadata["nvc_function_token"] = [function_token]
+    completion = SimpleNamespace(multimodal_output=metadata)
+    return SimpleNamespace(
+        stage_id=0,
+        request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
+    )
+
+
+def test_function_channel_projects_completed_call_without_ending_speech() -> None:
+    projector = _projector()
+    projector._decode = lambda token_ids: (
+        '[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
+    )
+    context = NemotronVoiceChatDataPlaneContext(epoch=0)
+
+    events = []
+    for function_token in (20, 99, 21):
+        events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
+
+    listen = [event for event in events if event.get("is_listen") is True]
+    function = [event for event in events if event.get("function_call") is True]
+    assert len(listen) == 3
+    assert len(function) == 1
+    assert function[0]["name"] == "weather"
+    assert function[0]["arguments"] == '{"city":"Shanghai"}'
+
+    projector._decode = lambda _token_ids: "not-json"
+    bad_events = []
+    for function_token in (20, 98, 21):
+        bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
+    assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
+
+
+def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
+    projector = _projector(lambda *_: "audio")
+    old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
+    new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
+
+    projector.begin_request(old_request)
+    list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
+    projector.begin_request(new_request)
+
+    # EOTC in the new epoch must not close the old epoch's partial function.
+    new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
+    assert not [event for event in new_events if event.get("function_call") is True]
+    assert not [event for event in new_events if event.get("end_of_turn") is True]
+
+    audio = SimpleNamespace(
+        stage_id=2,
+        request_output=SimpleNamespace(
+            request_id=new_request,
+            outputs=[
+                SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
+            ],
+        ),
+    )
+    audio_events = list(projector.project_output(audio))
+    assert sum(event.get("end_of_turn") is True for event in audio_events) == 1