[Perf][TTS] Add deadline-aware audio scheduling

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

@@ -17,7 +17,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
+        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
 
     - label: "Simple · Diffusion Test · Shard %N/%t"
       agent_pool: mi300_1
@@ -29,7 +29,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-diffusion-$$BUILDKITE_PARALLEL_JOB.xml"
+        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -38,7 +38,7 @@ steps:
       grade: Blocking
       commands:
         - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'"
 
     - label: "Simple · Other Test"
       agent_pool: mi300_1
@@ -49,7 +49,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py"
 
 - group: ":card_index_dividers: Diffusion Test"
   steps:

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

@@ -17,7 +17,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
+        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
 
     - label: "Simple · Diffusion Test"
       agent_pool: mi300_1
@@ -28,7 +28,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu'  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "pytest -sv tests/diffusion -m 'core_model and cpu'  --ignore=tests/diffusion/cache/test_teacache_extractors.py"
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -37,7 +37,7 @@ steps:
       grade: Blocking
       commands:
         - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'"
 
     - label: "Simple · Other Test"
       agent_pool: mi300_1
@@ -48,7 +48,7 @@ steps:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py"
 
 - label: "Custom Pipeline Test"
   agent_pool: mi300_1

.buildkite/common/scripts/run_cov_split.sh

@@ -16,6 +16,11 @@
 # At least one of --offline/--online is required, so a job with only one mode runs
 # just that half. Total runtime is bounded by the step's timeout_in_minutes.
 #
+# Coverage + artifact upload run only on scheduled weekly E2E sweeps
+# (WEEKLY=1 and BUILDKITE_BRANCH=main — same gate as bootstrap --e2e).
+# All other invocations (PR merge-test, regular main merge, local) run the
+# same combined pytest as the job used before coverage was added.
+#
 # Exits non-zero if any half or the upload failed.
 set -uo pipefail
 
@@ -44,6 +49,17 @@ if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then
     exit 2
 fi
 
+_should_run_coverage() {
+    [[ "${WEEKLY:-}" == "1" && "${BUILDKITE_BRANCH:-}" == "main" ]]
+}
+
+if ! _should_run_coverage; then
+    PATHS=("${OFFLINE[@]}" "${ONLINE[@]}")
+    echo "--- pytest (no coverage): ${MODEL_ID}"
+    pytest -s -v "${PATHS[@]}" "${PYTEST_ARGS[@]}"
+    exit $?
+fi
+
 # Two steps can test the same model, and artifacts are scoped to the build (a
 # nightly one carries the ready, merge and nightly tiers at once), so unsuffixed
 # reports would share a path -- a glob download then races to rename onto one file.

.buildkite/common/scripts/upload_pipeline.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
 
 Bootstrap mode (``bootstrap-upload-steps.yml``):
@@ -58,11 +59,12 @@ BOOTSTRAP_UPLOAD_IF_KEYS = {
 E2E_GROUP_MARKER = "E2E Test"
 CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml"
 
-CUDA_NIGHTLY_ONLY = (
-    '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
-)
-NPU_NIGHTLY_ONLY = (
-    '(build.branch == "main" && build.env("NIGHTLY") == "1") || '
+# Bootstrap Buildkite ``if`` expressions.
+# ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
+# ``*_UPLOAD_IF``: full gate for uploading that child pipeline.
+NIGHTLY_MAIN_IF = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
+NIGHTLY_LABEL_IF = (
+    f"({NIGHTLY_MAIN_IF}) || "
     '(build.branch != "main" && ('
     'build.pull_request.labels includes "nightly-test" || '
     'build.pull_request.labels includes "omni-test" || '
@@ -71,6 +73,19 @@ NPU_NIGHTLY_ONLY = (
     'build.pull_request.labels includes "diffusion-x2v-test"'
     "))"
 )
+WEEKLY_E2E_IF = 'build.branch == "main" && build.env("WEEKLY") == "1"'
+WEEKLY_MAIN_IF = 'build.branch == "main" && (build.env("WEEKLY") == "1" || build.env("NON_CRITICAL") == "1")'
+WEEKLY_LABEL_IF = f'({WEEKLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "weekly-test")'
+READY_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "ready"'
+MERGE_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
+MERGE_MAIN_IF = (
+    'build.branch == "main" && build.env("NIGHTLY") != "1" && '
+    'build.env("WEEKLY") != "1" && build.env("NON_CRITICAL") != "1"'
+)
+READY_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || ({READY_LABEL_IF})"
+MERGE_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || (({MERGE_MAIN_IF}) || ({MERGE_LABEL_IF}))"
+BOOTSTRAP_DISABLED_IF = "false"
+BOOTSTRAP_ENABLED_IF = "true"
 
 
 # --- Logging ---
@@ -102,74 +117,47 @@ def _format_bootstrap_if(expr: str) -> str:
 
 
 def _compute_bootstrap_if_exprs(*, decision, platform: str) -> dict[str, str]:
-    disabled = "false"
-    nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
-    nightly_only = NPU_NIGHTLY_ONLY if platform == "npu" else CUDA_NIGHTLY_ONLY
-
     if platform == "npu":
-        ready_pr = (
-            'build.branch != "main" && ('
-            'build.pull_request.labels includes "npu-test" || '
-            'build.pull_request.labels includes "ready"'
-            ")"
-        )
-        nightly_label_if = nightly_only
-        weekly_label_if = disabled
-        merge_base = disabled
+        ready_upload = READY_LABEL_IF
+        merge_upload = BOOTSTRAP_DISABLED_IF
+        weekly_label_if = BOOTSTRAP_DISABLED_IF
     else:
-        ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"'
-        merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"'
-        merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
-        nightly_label_if = (
-            '(build.branch == "main" && build.env("NIGHTLY") == "1") || '
-            '(build.branch != "main" && ('
-            'build.pull_request.labels includes "nightly-test" || '
-            'build.pull_request.labels includes "omni-test" || '
-            'build.pull_request.labels includes "tts-test" || '
-            'build.pull_request.labels includes "diffusion-x2iat-test" || '
-            'build.pull_request.labels includes "diffusion-x2v-test"'
-            "))"
-        )
-        weekly_label_if = (
-            '(build.branch == "main" && build.env("WEEKLY") == "1") || '
-            '(build.branch != "main" && build.pull_request.labels includes "weekly-test")'
-        )
-        merge_base = f"({nightly_main}) || (({merge_main}) || ({merge_pr}))"
-
-    ready_base = f"({nightly_main}) || ({ready_pr})"
-    weekly_main = 'build.branch == "main" && build.env("WEEKLY") == "1"'
+        ready_upload = READY_UPLOAD_IF
+        merge_upload = MERGE_UPLOAD_IF
+        weekly_label_if = WEEKLY_LABEL_IF
 
     if decision.skip_all:
         # Docs / skip-mark only: no PR-label escape hatch. Main scheduled
-        # NIGHTLY=1 still runs L4 + L2/L3 (--e2e); WEEKLY=1 still runs L5.
-        image_expr = f"({nightly_main}) || ({weekly_main})" if platform == "cuda" else nightly_main
-        ready_expr = nightly_main
-        merge_expr = nightly_main if platform == "cuda" else disabled
-        nightly_expr = nightly_main
-        weekly_expr = weekly_main if platform == "cuda" else disabled
+        # NIGHTLY=1 still runs L4; WEEKLY=1 / NON_CRITICAL=1 still run L5.
+        # main+WEEKLY=1 also uploads L2/L3 (those steps then pass --e2e).
+        image_expr = f"({NIGHTLY_MAIN_IF}) || ({WEEKLY_MAIN_IF})" if platform == "cuda" else NIGHTLY_MAIN_IF
+        ready_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
+        merge_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
+        nightly_expr = NIGHTLY_MAIN_IF
+        weekly_expr = WEEKLY_MAIN_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
     elif decision.skip_l2_l3:
         l2_enabled = decision.is_run("npu", "l2") if platform == "npu" else decision.is_run("cuda", "l2")
         l3_enabled = platform == "cuda" and decision.is_run("cuda", "l3")
 
-        ready_expr = ready_base if l2_enabled else disabled
-        merge_expr = merge_base if l3_enabled else disabled
-        nightly_expr = nightly_label_if
-        weekly_expr = weekly_label_if if platform == "cuda" else disabled
+        ready_expr = ready_upload if l2_enabled else BOOTSTRAP_DISABLED_IF
+        merge_expr = merge_upload if l3_enabled else BOOTSTRAP_DISABLED_IF
+        nightly_expr = NIGHTLY_LABEL_IF
+        weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF
 
-        image_parts = [f"({nightly_label_if})"]
+        image_parts = [f"({NIGHTLY_LABEL_IF})"]
         if platform == "cuda":
             image_parts.append(f"({weekly_label_if})")
         if l2_enabled:
-            image_parts.insert(0, f"({ready_base})")
+            image_parts.insert(0, f"({ready_upload})")
         if l3_enabled:
-            image_parts.insert(1 if l2_enabled else 0, f"({merge_base})")
+            image_parts.insert(1 if l2_enabled else 0, f"({merge_upload})")
         image_expr = " || ".join(image_parts)
     else:
-        image_expr = "true"
-        ready_expr = ready_base
-        merge_expr = merge_base if platform == "cuda" else disabled
-        nightly_expr = nightly_label_if
-        weekly_expr = weekly_label_if if platform == "cuda" else disabled
+        image_expr = BOOTSTRAP_ENABLED_IF
+        ready_expr = ready_upload
+        merge_expr = merge_upload if platform == "cuda" else BOOTSTRAP_DISABLED_IF
+        nightly_expr = NIGHTLY_LABEL_IF
+        weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF
 
     return {
         "image": _format_bootstrap_if(image_expr),

.buildkite/cuda/bootstrap-upload-steps.yml

@@ -8,28 +8,28 @@ steps:
     agents:
       queue: "cpu_queue_premerge"
 
-  # L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
+  # L2 Test — PR ready label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group)
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
     commands:
       - |
         EXTRA=""
-        if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
+        if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
           EXTRA="--e2e"
         fi
         python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-ready.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  # L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
+  # L3 Test — PR merge-test label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group)
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
     commands:
       - |
         EXTRA=""
-        if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
+        if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
           EXTRA="--e2e"
         fi
         python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml
@@ -45,7 +45,7 @@ steps:
     agents:
       queue: "cpu_queue_premerge"
 
-  # L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
+  # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (scheduled), or PR with weekly-test label
   - label: "Upload Weekly Pipeline"
     depends_on: image-build
     key: upload-weekly-pipeline

.buildkite/cuda/test-merge.yml

@@ -9,23 +9,27 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Simple · Model Executor Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/model_executor -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Diffusion Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/diffusion -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Engine&Entrypoints Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Diffusion Test"

.buildkite/cuda/test-nightly.yml

@@ -89,32 +89,6 @@ steps:
             exit $$EXIT
         mirror_hardwares: h100_2
 
-      - label: ":full_moon: Omni · Perf Test · vLLM Text"
-        key: nightly-omni-performance-vllm-text
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        mirror_hardwares: h100_2
-
-      - label: ":full_moon: Omni · Perf Test · Multi-Replica"
-        key: nightly-omni-performance-multi-replicas
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        mirror_hardwares: h100_3
-
       - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test"
         key: nightly-omni-performance-minicpmo-4-5
         timeout_in_minutes: 180
@@ -255,6 +229,12 @@ steps:
           - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
         mirror_hardwares: h100_4
 
+      - label: ":full_moon: HiDream-O1-Image · Accuracy Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
         key: nightly-diffusion-x2iat-performance-qwen-image
         timeout_in_minutes: 180
@@ -406,8 +386,6 @@ steps:
     depends_on:
       - nightly-omni-performance-no-async-chunk
       - nightly-omni-performance-async-chunk
-      - nightly-omni-performance-vllm-text
-      - nightly-omni-performance-multi-replicas
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2v-performance
@@ -419,8 +397,6 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
       - python tools/nightly/generate_nightly_perf_excel.py

.buildkite/cuda/test-ready.yml

@@ -9,23 +9,27 @@ steps:
     depends_on: upload-ready-pipeline
     steps:
       - label: "Simple · Model Executor Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/model_executor -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Diffusion Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/diffusion -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Engine&Entrypoints Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
+        timeout_in_minutes: 20
         commands:
-          - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Diffusion Test"

.buildkite/cuda/test-weekly.yml

@@ -5,51 +5,91 @@ env:
   HF_HUB_ETAG_TIMEOUT: 60
 
 steps:
-
-  - label: "Reliability Test - qwen3-omni"
-    timeout_in_minutes: 180
+  - label: "Simple · CPU Coverage Test"
+    key: weekly-simple-cpu-coverage
     depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    if: build.env("WEEKLY") == "1"
+    timeout_in_minutes: 60
     commands:
-      - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
-    mirror_hardwares: h100_2
+      - |
+        set +e
+        REPORT="coverage-core-model-cpu-$${BUILDKITE_STEP_ID:-local}.xml"
+        pytest -sv tests/ -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:$${REPORT}
+        EXIT=$$?
+        gzip -9 -f "$${REPORT}"
+        buildkite-agent artifact upload "$${REPORT}.gz"
+        exit $$EXIT
+    mirror_hardwares: l4_1
 
-  - label: "Reliability Test - wan22"
-    timeout_in_minutes: 180
-    depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
-    commands:
-      - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
-    mirror_hardwares: h100_2
-
-  - label: "Reliability Test - hunyuan-image"
-    timeout_in_minutes: 180
-    depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
-    commands:
-      - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
-    mirror_hardwares: h100_4
-
-  - group: "Reliability Test - Invalid parameters Test"
+  - group: "Reliability Test"
     depends_on: upload-weekly-pipeline
     if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
     steps:
 
-      - label: "Invalid parameters Test · H100"
+      - label: "Reliability Test · qwen3-omni"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+        mirror_hardwares: h100_2
+
+      - label: "Reliability Test · wan22"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+        mirror_hardwares: h100_2
+
+      - label: "Reliability Test · hunyuan-image"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
+        mirror_hardwares: h100_4
+
+      - label: "Reliability Test · Invalid parameters · H100"
         timeout_in_minutes: 120
         commands:
           - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
         mirror_hardwares: h100_2
 
-      - label: "Invalid parameters Test · L4"
+      - label: "Reliability Test · Invalid parameters · L4"
         timeout_in_minutes: 120
         commands:
           - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
         mirror_hardwares: l4_1
 
-  - group: ":card_index_dividers: E2E Tests"
+  - group: "Perf Test"
     depends_on: upload-weekly-pipeline
     if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    steps:
+
+      - label: ":full_moon: Omni · Perf Test · vLLM Text"
+        key: weekly-omni-performance-vllm-text
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        mirror_hardwares: h100_2
+
+      - label: ":full_moon: Omni · Perf Test · Multi-Replica"
+        key: weekly-omni-performance-multi-replicas
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        mirror_hardwares: h100_3
+
+  - group: ":card_index_dividers: E2E Tests"
+    depends_on: upload-weekly-pipeline
+    if: build.env("NON_CRITICAL") == "1"
     steps:
 
       - label: "Omni · L4"
@@ -75,7 +115,7 @@ steps:
         mirror_hardwares: h100_2
 
       - label: "TTS · L4"
-        timeout_in_minutes: 90
+        timeout_in_minutes: 120
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"
@@ -126,7 +166,7 @@ steps:
         mirror_hardwares: l4_4
 
       - label: "Diffusion · H100 · Single-GPU"
-        timeout_in_minutes: 120
+        timeout_in_minutes: 150
         commands:
           - >-
             pytest -sv

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

@@ -43,7 +43,7 @@ steps:
     depends_on: upload-ready-pipeline
     steps:
       - label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
-        mirror_hardwares: a2b3_npu_1
+        mirror_hardwares: a3_npu_2
         env:
           VLLM_WORKER_MULTIPROC_METHOD: spawn
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"

.claude/skills/find-simplifications/SKILL.md

@@ -0,0 +1,135 @@
+---
+name: find-simplifications
+description: Find evidence-backed simplification candidates in vLLM-Omni and, when requested, turn them into focused proposals or code changes. Use for audits of dead, duplicated, speculative, over-generalized, unnecessarily defensive, or hand-rolled code. Use review-pr for ordinary correctness review and diffusion-perf-opt for performance-first optimization.
+---
+
+# Find vLLM-Omni Simplifications
+
+Find a small number of well-proven ways to reduce code, state, APIs, or maintenance cost without weakening supported behavior. Treat this as an architecture and ownership audit, not a dead-code search or style cleanup.
+
+The default deliverable is a read-only report. Do not add TODOs, design documents, issues, or code unless the user asks for those changes.
+
+## Establish the Contract
+
+Before judging a candidate:
+
+1. Read the repository instructions supplied by the host.
+2. Read [`docs/design/architecture_overview.md`](../../../docs/design/architecture_overview.md).
+3. Select only the relevant active module and feature documents from [`docs/design/index.md`](../../../docs/design/index.md). Archived design pages are historical context, not current contracts.
+4. Identify the comparison base with `git merge-base HEAD origin/main`; do not assume the current branch represents the intended baseline.
+5. If tests or CI coverage may change, use [`vllm-omni-test`](../vllm-omni-test/SKILL.md) for test placement, markers, and commands.
+
+Separate an implementation accident from an intentional boundary. In particular:
+
+- vLLM-Omni extends and sometimes overrides upstream vLLM. A similar upstream implementation is evidence for reuse only after comparing API, lifecycle, and multi-stage semantics.
+- Registries, import strings, model configuration, deployment YAML, plugins, and optional dependencies create dynamic consumers that a static reference count can miss.
+- Hardware backends, model families, execution modes, and connector implementations may be intentionally parallel. Do not propose collapsing one only because it is inactive in the local environment.
+- Stage, replica, request, cache, and transport state can look redundant while representing different owners. Prove that two fields or barriers encode the same fact before merging them.
+- A shorter hot path is not a simplification if it adds synchronization, tensor copies, host materialization, recompilation, or a measurable quality/performance regression.
+
+## Strong Candidates
+
+A strong candidate removes a meaningful surface and has a clear replacement or deletion boundary. Examples include:
+
+- A public method, configuration option, registry entry, route, message field, or event has no production consumer and no compatibility obligation.
+- Tests or documentation are the only consumers of behavior that is not part of a supported contract.
+- Multiple layers mirror the same request, lifecycle, cache-validity, readiness, or topology state.
+- AR, diffusion, platform, or model-specific implementations duplicate a stable common operation and can share it without hiding performance-critical differences.
+- An inherited upstream vLLM primitive is reimplemented locally without an Omni-specific semantic difference.
+- Compatibility or fallback branches protect environments the project no longer supports.
+- Serialization, IPC, output formatting, or connector code materializes the same payload more than once or converts between equivalent representations.
+- A helper, abstraction, package split, or extension point has only one real implementation and adds more indirection than policy.
+- Hand-written infrastructure can be replaced by an existing project dependency, Python standard-library feature, PyTorch primitive, or upstream vLLM facility with net deletion and equal observability.
+- Tests maintain large mocks, snapshots, or expected data solely for an unused API or representation.
+
+Do not elevate typo fixes, formatting, isolated renames, or vague complexity complaints into simplification candidates. Bundle small cleanups only when they support one substantive removal.
+
+## Survey by Ownership Boundary
+
+For broad audits, cover the largest or riskiest production deltas first. Useful domains are:
+
+- `vllm_omni/engine/`: orchestration, stage pools, request state, output convergence, cancellation, and replica lifecycle.
+- `vllm_omni/worker/` and AR model execution: upstream vLLM overlays, scheduler/cache state, collective RPC, adapters, and weight loading.
+- `vllm_omni/diffusion/`: request/step execution, batching, output materialization, IPC, offload, parallelism, and model pipelines.
+- `vllm_omni/entrypoints/`: duplicated validation/defaulting, offline/online parity, streaming, and route compatibility.
+- Connectors and distributed code: ownership, copies, barriers, handles, cleanup, and failure propagation across process/device/node boundaries.
+- Platform backends: genuine vendor differences versus copied control flow that has drifted.
+- Configuration, registries, recipes, and examples: duplicate sources of truth and options that never reach a runtime consumer.
+- Tests and CI: redundant fixtures, mocks of obsolete contracts, and expensive coverage that can move to a cheaper level without losing the guarded behavior.
+
+When the user explicitly requests a repository-wide or many-candidate audit and parallel agent work is available and authorized, divide these domains among agents and require the same evidence format from each. Otherwise survey them sequentially.
+
+## Prove or Reject Each Candidate
+
+Use `rg` first, then read every relevant caller and implementation. Search exact Python symbols plus serialized names, route paths, configuration keys, registry strings, YAML values, and CLI spellings.
+
+Classify evidence:
+
+- **Production:** `vllm_omni/`, runtime configuration, registries, serving entrypoints, connectors, recipes used for supported deployments, and executable loader paths.
+- **Validation/support:** `tests/`, benchmarks, examples, docs, CI, and developer tooling.
+- **External/dynamic:** upstream vLLM contracts, third-party plugins, model repositories, optional backends, and string-loaded objects. State what was checked and what remains unknown.
+
+For each candidate, answer:
+
+1. What exact surface would disappear or merge?
+2. Who owns the state or behavior today?
+3. Which production consumers exist, including dynamic consumers?
+4. What supported behavior, compatibility, performance, memory, or observability could change?
+5. What is the smallest validation that would catch a bad simplification?
+6. Is the result net deletion after tests, adapters, migration glue, and documentation are counted?
+
+Reject or downgrade the candidate when a live consumer exists and removal is actually a feature decision; when an active design contract explains the separation; when the change only moves complexity; or when hardware/model evidence is unavailable for a performance-sensitive claim.
+
+## Lifecycle and Data-Movement Audits
+
+For asynchronous or distributed code, map the ownership graph before recommending deletion:
+
+```text
+request admission -> orchestrator -> stage/replica -> worker/device
+  -> connector or IPC boundary -> output aggregation -> terminal cleanup
+```
+
+Associate every lock, event, sentinel, queue, readiness flag, timeout, abort path, and cache marker with an owner and transition. Coalesce mechanisms only when they represent the same state and have the same failure boundary. Preserve separate mechanisms when they protect publication versus rollback, local versus remote ownership, first-terminal-outcome arbitration, or cleanup after partial failure.
+
+For tensors and media, also trace device, process, and ownership transitions. Count GPU-to-CPU copies, shared-memory/object-store handles, serialization, decode/encode steps, and unlink/release responsibility. Prefer eliminating a transfer or representation over merely hiding it behind a helper.
+
+## Dependency Replacements
+
+Prefer facilities already available in the repository before adding a dependency. For a proposed replacement:
+
+- Name the exact local behavior it covers and any semantics that remain in glue code.
+- Check maintenance, license, optional-dependency boundaries, platform support, and transitive footprint.
+- Compare net deletion, test burden, import/startup cost, and runtime behavior.
+- Benchmark when the replaced code is on a model, scheduler, connector, serialization, or serving hot path.
+
+A wrapper that retains the same state machine or conversion logic is not a simplification.
+
+## Deliver the Result
+
+Report only candidates supported by concrete evidence. For each one include:
+
+- **Surface:** files and symbols involved.
+- **Evidence:** production consumers found, dynamic paths checked, and relevant design contract.
+- **Simplification:** exactly what to remove, merge, inline, or reuse.
+- **Benefit:** deleted API/state/code, reduced data movement, or reduced maintenance/test burden.
+- **Risk and validation:** behavior that could change and the narrowest useful tests or measurements.
+- **Confidence:** high or medium, with the unresolved evidence named for medium-confidence items.
+
+Also list representative ideas that were rejected when that prevents repeated investigation. Prefer three high-confidence candidates over a long speculative inventory.
+
+If the user asks for a durable design proposal, place it under `docs/design/feature/` and register it in `docs/design/index.md`; use current design-document conventions and keep implementation details proportional to the decision. If the user asks for implementation, make the smallest coherent change, update affected tests/docs, and preserve public compatibility unless the user explicitly accepts a breaking change.
+
+## Folding Work from Another Branch
+
+Diff the source branch against its own merge base with `origin/main`, not against the current feature branch. Port only independently justified candidates. Consolidate overlaps into the candidate with stronger evidence, and do not preserve a candidate count by carrying weaker duplicates.
+
+## Validation
+
+Match validation to the outgoing diff:
+
+- Documentation or skill changes: run the repository skill validator when available, check links/paths, search for stale source-specific terminology, and run `git diff --check`.
+- Python changes: run the narrowest relevant pytest coverage through the repository environment, then lint/format the changed files.
+- Distributed, model, hardware, performance, or accuracy changes: run available CPU/static checks and state the exact GPU/NPU/model gap; never present simulated evidence as device validation.
+- Before submission, use [`precheck-pr`](../precheck-pr/SKILL.md) for the branch-level self-review.
+
+When summarizing, distinguish checks that passed from checks that were not runnable, and state the snapshot or branch against which the audit was performed.

.claude/skills/precheck-pr/SKILL.md

@@ -67,6 +67,8 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type
 
 **Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt.
 
+**Also run the Simplification pass on every PR** using [find-simplifications](../find-simplifications/SKILL.md). Keep it diff-scoped: inspect added or modified public APIs, state, abstractions, fallback and compatibility paths, data movement, and only the adjacent ownership needed to prove a candidate. Do not expand the pass into a repository-wide audit. Zero candidates is a valid pass; report only evidence-backed opportunities. Mark a candidate as blocking only when it proves a correctness or merge-readiness problem, such as newly unreachable code. Otherwise report it as a warning.
+
 **Also confirm local pre-commit gates** from [docs/contributing/README.md](../../../docs/contributing/README.md#linting). GitHub Actions `SKIP` does not prove they passed. Full new-hook list: Omni SPDX (`vLLM-Omni project`); forbidden imports (stdlib `re`/`base64`, pickle, Hugging Face Hub API, Triton/TileLang); no new `torch.cuda` call sites; shellcheck (macOS/Windows must install the binary); mypy-3.10; test files have CI level + hardware marks; TTS `_tts_model_type` branches in `serving_speech.py` must not increase; markdownlint on `docs/`; Buildkite YAML schema. ✗ for growing `allowed_files` / `ALLOWED_FILES` or raising `MAX_MODEL_TYPE_BRANCHES` without review.
 
 ### Step 5: Print Report
@@ -82,6 +84,7 @@ Pre-check report for <branch>
   PR title format    ✓
   Code quality       ⚠ 1 broad except, 2 Any hints
   Examples policy    ✗ new model-specific Python example
+  Simplification     ⚠ helper duplicates an existing owner
   PR desc integrity  ✓
   Registry/config    ✓
   Dead code          ⚠ 2 warnings

.claude/skills/readme.md

@@ -31,9 +31,12 @@ self-check.
 - [`diffusion-perf-opt`](diffusion-perf-opt/SKILL.md): guides diffusion model
   performance optimization, including profiling traces, parallel strategies,
   stage timing analysis, and benchmark-driven tuning
+- [`find-simplifications`](find-simplifications/SKILL.md): finds evidence-backed
+  opportunities to remove or merge dead, duplicated, speculative,
+  over-generalized, or unnecessarily defensive vLLM-Omni code
 - [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a
-  PR by validating title format, dead code, accuracy and performance claims,
-  and merge readiness
+  PR by validating title format, dead code, simplification opportunities,
+  accuracy and performance claims, and merge readiness
 - [`quantization`](quantization/SKILL.md): guides quantization method selection,
   model integration, checkpoint loading, and quality/performance validation
   for vLLM-Omni

.claude/skills/review-pr/SKILL.md

@@ -101,6 +101,7 @@ If docs and live code disagree, verify the code/tests and report the drift.
 | [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
 | [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
 | [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. |
+| [find-simplifications](../find-simplifications/SKILL.md) | Every review; run a diff-scoped subtraction and simplification pass after correctness blockers. |
 
 ### Delivery and reviewer coordination
 
@@ -183,9 +184,13 @@ promotion gate. Candidate or draft rules are questions, not blockers, unless
 current code, tests, or policy enforce them. Inspect both sides of any config,
 registry, serialization, connector, cache, or stage boundary.
 
-When a diff adds or expands a helper, class, fallback, compatibility branch, or
-public behavior, run a subtraction pass: remove out-of-scope behavior and check
-whether each new abstraction can be deleted, merged, moved, or inlined.
+Read and apply [find-simplifications](../find-simplifications/SKILL.md) on every
+review. Constrain it to the diff and the adjacent ownership, callers, or
+consumers needed to prove a candidate. Check whether added or expanded helpers,
+classes, state, fallback and compatibility branches, data movement, or public
+behavior can be deleted, merged, moved, or inlined. Zero candidates is a valid
+result. Do not widen the review into repository backlog or report speculative
+style preferences as simplification findings.
 
 ### 6. Verify the changed path
 

.gitignore

@@ -245,7 +245,6 @@ datasets/
 !recipes/cosmos3/negative_prompt.json
 !tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
-!tests/assets/soulxsinger/*.json
 *.jsonl
 *.parquet
 

README.md

@@ -12,10 +12,10 @@ Easy, fast, and cheap omni-modality model serving for everyone
 | <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
 </p>
 
-
 ---
 
 *Latest News* 🔥
+
 - [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
 - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
 - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
@@ -89,11 +89,18 @@ If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.or
 ```
 
 ## Join the Community
+
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
 
 ## Star History
 
-[![Star History Chart](https://api.star-history.com/svg?repos=vllm-project/vllm-omni&type=date&legend=top-left)](https://www.star-history.com/#vllm-project/vllm-omni&type=date&legend=top-left)
+<a href="https://www.star-history.com/?repos=vllm-project%2Fvllm-omni&type=date&legend=top-left">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&theme=dark&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
+    <source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
+    <img alt="Star History Chart" src="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
+  </picture>
+</a>
 
 ## License
 

docs/api/README.md

@@ -114,10 +114,6 @@ 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/stage_configs.md

@@ -74,8 +74,10 @@ For fields whose deploy default is `null`, the deploy layer contributes no
 override. The effective value may still come from a platform section, an
 explicit CLI or stage override, or the downstream vLLM engine default.
 
-Note: for diffusion path, `distributed_executor_backend` currently defaults to
-`mp`, and `ray` / `external_launcher` are not fully supported yet.
+Note: for the diffusion path, an omitted `distributed_executor_backend` selects
+`uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output
+segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker
+subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
 
 ### Stage fields
 

docs/contributing/README.md

@@ -334,7 +334,7 @@ These PR shapes all trigger skip-ci:
 Skip-ci does **not** apply when the diff also touches product code (for example `vllm_omni/`), or when test files change assertions, imports, fixtures, or other non-skip logic. If the diff cannot be resolved (non-PR branches outside `main`), CI runs as usual.
 
 !!! note
-    Skipping L2/L3 does **not** disable the Docker image build step. Nightly (L4) upload can still run when the PR has a `nightly-test` label or on scheduled `main` builds with `NIGHTLY=1`. Bootstrap child steps live in `bootstrap-upload-steps.yml`; the hook entry step runs `upload_pipeline.py --upload <platform>/bootstrap-upload-steps.yml`, which injects `if` by step `key` from skip-ci before upload. See [CI Settings — Diff-aware CI](./ci/ci_settings.md#diff-aware-ci) and [Test System Overview](./ci/test_system_overview.md).
+    Skipping L2/L3 does **not** disable the Docker image build step. Nightly (L4) upload can still run when the PR has a `nightly-test` label or on scheduled `main` builds with `NIGHTLY=1`. Weekly (L5) upload can still run with `weekly-test`, `WEEKLY=1`, or `NON_CRITICAL=1`. Bootstrap child steps live in `bootstrap-upload-steps.yml`; the hook entry step runs `upload_pipeline.py --upload <platform>/bootstrap-upload-steps.yml`, which injects `if` by step `key` from skip-ci before upload. See [CI Settings — Diff-aware CI](./ci/ci_settings.md#diff-aware-ci) and [Test System Overview](./ci/test_system_overview.md).
 
 ### Code Quality
 

docs/contributing/ci/ci_settings.md

@@ -8,7 +8,7 @@ For CI levels (L1–L5) and triggers, see [Test System Overview](./test_system_o
 
 Canonical layout (prefer these paths for new changes):
 
-```
+```text
 .buildkite/
 ├── common/                          # Shared across platforms
 │   ├── scripts/
@@ -48,7 +48,7 @@ Canonical layout (prefer these paths for new changes):
     └── scripts/
 ```
 
-**Placement rules**
+### Placement rules
 
 | Rule | Detail |
 | ---- | ------ |
@@ -81,7 +81,14 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr
     | L2 | `cuda/test-ready.yml` | `ready` label |
     | L3 | `cuda/test-merge.yml` | `merge-test` label / main merge |
     | L4 | `cuda/test-nightly.yml` | `nightly-test` label or `NIGHTLY=1` |
-    | L5 | `cuda/test-weekly.yml` | `weekly-test` label or `WEEKLY=1` |
+    | L5 | `cuda/test-weekly.yml` | `weekly-test` label, `WEEKLY=1`, or `NON_CRITICAL=1` (see step `if`s below) |
+
+    On scheduled `main` builds, `test-weekly.yml` is uploaded when **`WEEKLY=1` or `NON_CRITICAL=1`**. Inside that file:
+
+    - **Reliability** / **Perf Test** / **Simple · CPU Coverage Test** → `WEEKLY=1` (or PR label `weekly-test` for Reliability/Perf). Pipeline upload for scheduled env vars still requires `main`.
+    - **E2E Tests** group (slow Omni/TTS/Diffusion sweeps) → `NON_CRITICAL=1`
+
+    **`--e2e`:** when uploading ready/merge on `main` with **`WEEKLY=1`**, bootstrap passes `--e2e` so only the **E2E Test** group is kept (same flag `run_cov_split.sh` uses to enable per-model coverage).
 
     **Upload:** `upload_pipeline.py --upload` expands uploader-only keys before Buildkite upload.
 
@@ -112,6 +119,8 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr
 
     **Test YAML:** `npu/test-npu-ready.yml` (L2), `test-npu-nightly.yml` (L4).
 
+    **L2 trigger:** PR label `ready` (CUDA/AMD L2). PR labels `ready` + `npu-test` trigger NPU CI only.
+
     **Upload:** same as CUDA—`upload_pipeline.py --upload`.
 
     **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. Presets: `a2b3_npu_1`, `a2b3_npu_4`, `a3_npu_2` in `common/ci_mirror_hardwares.yml`.
@@ -204,15 +213,15 @@ Label triggers (`ready`, `merge-test`) are unchanged—diff-aware logic only red
 | Diff shape | Path | Default L2/L3 |
 | --- | --- | --- |
 | Product code / non-whitelisted paths | normal CI | all on |
-| Docs and/or qualifying skip-mark only | `skip_all` | all off (scheduled `main` NIGHTLY/WEEKLY exceptions below) |
+| Docs and/or qualifying skip-mark only | `skip_all` | all off (scheduled `main` NIGHTLY/WEEKLY/`NON_CRITICAL` exceptions below) |
 | Whitelisted CI YAML only | yaml-gated (`skip_l2_l3`) | per platform/level matrix |
 | Docs/skip-mark **+** whitelisted CI YAML | yaml-gated (`skip_l2_l3`) | same as CI-YAML-only (does **not** widen to normal CI) |
 | Non-qualifying skip-mark and no CI YAML | normal CI | all on |
 | Diff unavailable | normal CI | all on |
 
-**`skip_all` exceptions (CUDA/NPU bootstrap only):** PR labels (`nightly-test`, `merge-test`, `npu-test`, `weekly-test`, …) do **not** revive jobs. On `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4 plus L2/L3 (with `--e2e`); `WEEKLY=1` still builds the image and uploads L5 (CUDA).
+**`skip_all` exceptions (CUDA/NPU bootstrap only):** PR labels (`nightly-test`, `merge-test`, `weekly-test`, `ready`, …) do **not** revive jobs under `skip_all`. On `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4; `WEEKLY=1` or `NON_CRITICAL=1` still builds the image and uploads L5 (CUDA). `WEEKLY=1` on `main` also uploads L2/L3 with `--e2e`.
 
-**Yaml-gated nightly/weekly:** L4/L5 upload steps keep their normal label / `NIGHTLY` / `WEEKLY` conditions (for example PR `nightly-test` still uploads nightly). Only L2/L3 upload steps are matrix-gated.
+**Yaml-gated nightly/weekly:** L4/L5 upload steps keep their normal label / `NIGHTLY` / `WEEKLY` / `NON_CRITICAL` conditions (for example PR `nightly-test` still uploads nightly). Only L2/L3 upload steps are matrix-gated.
 
 ##### Whitelisted CI YAML → platform
 
@@ -285,12 +294,12 @@ CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test
 | `source_file_dependencies` | List of path **prefixes**. If any changed file equals a prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. |
 | `mirror_hardwares` | Expand to `agents` + `plugins` (+ optional `image`) from `ci_mirror_hardwares.yml`. |
 
-**Policy**
+### Policy
 
 - **Always uploaded** (no key): groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3).
 - **Diff-gated**: every **E2E Test** leaf job. List the smallest prefix set per step—pytest file(s), model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable. A **group** may define the key instead; the whole group drops if no prefix matches.
 
-**YAML examples**
+### YAML examples
 
 ```yaml
       - label: "Diffusion · Qwen Image Test"
@@ -312,7 +321,7 @@ CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test
         mirror_hardwares: l4_4
 ```
 
-**Local dry-run**
+### Local dry-run
 
 ```bash
 python3 .buildkite/common/scripts/upload_pipeline.py .buildkite/cuda/test-ready.yml
@@ -324,18 +333,19 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
 
 ### Labels and grouping
 
-- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main.
+- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` weekly runs on main (`WEEKLY=1`).
 - Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps.
 
 ### Per-model coverage
 
 Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
 per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
-Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
-diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
-another diffusion coverage job if a second pilot is needed again.) This is a
-pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
-is gated on coverage.
+Test". Coverage + artifact upload run only when bootstrap uploaded that job with
+`--e2e` (`WEEKLY=1` and `BUILDKITE_BRANCH=main`, via `run_cov_split.sh`); otherwise
+the same job runs plain pytest without `--cov`. It runs on a single GPU so the
+pilot is cheap to reproduce.
+
+**L1 package coverage (CPU):** ready/merge Simple Test jobs run without `--cov`. Scheduled weekly on `main` with `WEEKLY=1` runs **Simple · CPU Coverage Test** in `.buildkite/cuda/test-weekly.yml` (`pytest -m 'core_model and cpu'` + Cobertura XML artifact).
 
 **Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
 `<model_id>` is the model's directory name under
@@ -354,8 +364,8 @@ imported. A run covering zero lines produces the same 7 MB. gzip takes it to
 roughly 410 KB and changes nothing about the data.
 
 `<step_id>` is what keeps two steps that cover the same model apart. Artifacts are
-scoped to the build, and a nightly build carries the ready, merge and nightly tiers
-at once, so several same-model steps land in one namespace — `qwen3_tts` already has
+scoped to the build, and a weekly build can carry ready/merge (with `--e2e`) and
+weekly tiers at once, so several same-model steps land in one namespace — `qwen3_tts` already has
 both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two
 uploads on one path do not fail loudly: an exact-name download can report the path
 as ambiguous, and a glob download fetches every match concurrently and renames each
@@ -425,7 +435,7 @@ just confirming both files exist.
 | Check | Command / location |
 | ----- | ------------------ |
 | CUDA render | `python3 .buildkite/common/scripts/upload_pipeline.py .buildkite/cuda/test-<level>.yml` |
-| No leaked uploader keys | `… \| grep -E 'mirror_hardwares|source_file_dependencies'` → empty |
+| No leaked uploader keys | Pipe the render into `grep -E 'mirror_hardwares\|source_file_dependencies'` — expect empty |
 | Skip-ci / upload unit tests | `pytest tests/buildkite/` |
 | Local job replay (CUDA L2+) | `tools/run_ready_jobs.sh`, `tools/run_merge_jobs.sh`, `tools/nightly/run_nightly_jobs.sh` (read YAML from `cuda/`) |
 

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -1,14 +1,15 @@
 When you want to add L4-level ***performance test*** cases, add entries to JSON files under `tests/dfx/perf/tests/` and run them via `tests/dfx/perf/scripts/run_benchmark.py` (omni / TTS) or `run_diffusion_benchmark.py` (diffusion).
 
-**Config file layout (in-tree examples)**
+#### Config file layout (in-tree examples)
 
 | Model type | Runner | Example JSON files |
 | ---------- | ------ | ------------------ |
-| Omni | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` |
+| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` in `mark`) |
+| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) |
 | TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
 | Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
 
-**How runners pick cases**
+#### How runners pick cases
 
 Without **`--test-config-file`**, each runner scans all `*.json` under `tests/dfx/perf/tests/` but only keeps its own model type:
 
@@ -17,9 +18,9 @@ Without **`--test-config-file`**, each runner scans all `*.json` under `tests/df
 
 Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-omni"`) or `"diffusion"` in the `mark` array. Omni / TTS JSON has neither.
 
-**Running perf cases**
+#### Running perf cases
 
-Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"`).
+Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"` or `-m "slow and omni and H100"` for weekly Omni configs).
 
 ```JSON
 {
@@ -55,20 +56,22 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul
 }
 ```
 
-**Parameter Explanation**
+#### Parameter Explanation
 
-*Overview*
+##### Overview
 
-| Field            | Required | Description                                                     |
-| ---------------- | -------- | --------------------------------------------------------------- |
-| test_name        | Yes      | Unique identifier for the test case                             |
-| mark             | No       | Pytest marks for this case (see **`mark` field** below). Omit only for configs not meant to be filtered by `-m`. |
-| server_params    | Yes      | Server-side configuration parameters                            |
-| benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
-| server_type      | Diffusion only | When set (e.g. `"vllm-omni"`), routes the case to `run_diffusion_benchmark.py`. |
-| benchmark_endpoint | Diffusion only | API path for the benchmark client (e.g. `/v1/videos`, `/v1/images/generations`). |
+| Field              | Required       | Description                                  |
+| ------------------ | -------------- | -------------------------------------------- |
+| test_name          | Yes            | Unique identifier for the test case          |
+| mark               | No             | Pytest marks; see **`mark` field** below     |
+| server_params      | Yes            | Server-side configuration parameters         |
+| benchmark_params   | Yes            | Benchmark running parameters                 |
+| server_type        | Diffusion only | Routes case to run_diffusion_benchmark.py    |
+| benchmark_endpoint | Diffusion only | Benchmark API path                           |
 
-**`mark` field**
+Omit `mark` only for configs not meant to be filtered by `-m`. `server_type` is typically `"vllm-omni"`. `benchmark_endpoint` examples: `/v1/videos`, `/v1/images/generations`.
+
+#### `mark` field
 
 Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
 
@@ -119,11 +122,11 @@ Multi-GPU diffusion (example: Cosmos3 with `cfg-parallel-size=2`):
 - Diffusion perf: include `"diffusion"` in the `mark` array
 - HunyuanImage local-weight cases: add `"local_model"` to the `mark` array
 
-**Parametrization IDs**
+#### Parametrization IDs
 
 Each `(server, benchmark index)` pair becomes one `pytest.param` with id `{test_name}-{suffix}`. The suffix comes from `benchmark_params[].name` when set (for example `test_tts-p0`, `test_omni-p1`); otherwise it is derived from fields like `task` / `eval_phase`.
 
-**Benchmark result filenames**
+#### Benchmark result filenames
 
 Result files use the **runtime** hardware label from `get_runtime_resource_label()` (detected GPU/NPU name on the machine that ran the job), **not** `mark.hardware_marks.res`. On the default H100 CI pool, `H100` is **omitted** from filenames (`resource_label_for_filename`).
 
@@ -132,7 +135,7 @@ Examples:
 - Omni/TTS: `result_{test_name}_{optional_hw}_{dataset}_....json` under `BENCHMARK_DIR`
 - Diffusion: one aggregate `diffusion_result_{config_stem}_{optional_hw}_{timestamp}.json` per source JSON file (array of all runs from that file)
 
-**Local commands**
+#### Local commands
 
 ```bash
 # Bulk load + filter by JSON mark
@@ -146,18 +149,19 @@ pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
   --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
 ```
 
+<!-- markdownlint-disable-next-line MD051 -->
 See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
 
-**`server_params` Configuration**
+#### `server_params` Configuration
 
-*Basic Parameters*
+##### Basic Parameters
 
-| Parameter         | Required | Example                            | Description                   |
-| ----------------- | -------- | ---------------------------------- | ----------------------------- |
-| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
-| stage_config_name | Yes      | "qwen3_omni_moe.yaml"              | Deploy configuration file name |
+| Parameter         | Required | Example                   | Description                    |
+| ----------------- | -------- | ------------------------- | ------------------------------ |
+| model             | Yes      | Qwen/Qwen3-Omni-30B-A3B…  | Model name or path             |
+| stage_config_name | Yes      | qwen3_omni_moe.yaml       | Deploy configuration file name |
 
-*Dynamic Configuration (update/delete)*
+##### Dynamic Configuration (update/delete)
 
 Supports incremental modifications based on the basic configuration:
 
@@ -166,9 +170,9 @@ Supports incremental modifications based on the basic configuration:
 | update    | Update or add configuration items    |
 | delete    | Delete specified configuration items |
 
-**Example**:
+##### Example
 
-```
+```json
 "update": {
     "async_chunk": true,  // Enable asynchronous chunk processing
     "stages": {
@@ -184,20 +188,22 @@ Supports incremental modifications based on the basic configuration:
 }
 ```
 
-**`benchmark_params` Configuration**
+#### `benchmark_params` Configuration
 
 You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
 
-1.  Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
-2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
-3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
-4.  Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys.
-5.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
+1. Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
+2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
+3. Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
+4. Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys.
+5. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
 
-| Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
-| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| name            | string      | No       | `"1024x1024_steps4"` | Stable suffix for pytest param id and logs. |
-| num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
-| request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-| max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
-| baseline        | object      | No       | see below       | Optional thresholds. **Hardware-nested only**: top-level keys must be `[hardware-resource]` pytest markers from `pyproject.toml` (`get_hardware_mark_list()`). Under each hardware key, metric names are free-form; values are scalars or sweep-aligned lists. Example: `{"H100": {"throughput_qps": […], "custom_stage_ms": 1.0}, "A3": {…}}`. Runtime filename aliases (`_RUNTIME_DEVICE_ALIASES`) are separate and do not expand this allowlist. |
+| Parameter       | Type          | Required | Example/Values       | Description                          |
+| --------------- | ------------- | -------- | -------------------- | ------------------------------------ |
+| name            | string        | No       | `"1024x1024_steps4"` | Stable suffix for pytest param id    |
+| num_prompts     | int / array   | Yes      | 10, [10, 20, 30]     | Number of requests (scalar or sweep) |
+| request_rate    | float / array | No       | 0.5, [0.5, 1, inf]   | Queries per second (QPS mode)        |
+| max_concurrency | int / array   | No       | 1, [1, 2, 3]         | Max in-flight requests               |
+| baseline        | object        | No       | see below            | Optional hardware-nested thresholds  |
+
+`num_prompts`, `request_rate`, and `max_concurrency` accept scalars or arrays; single values expand to match sweep length. QPS and concurrency modes are mutually exclusive. `baseline` top-level keys must be hardware pytest markers from `pyproject.toml`; metric names and values are free-form (scalars or sweep-aligned lists). Omit `baseline` to run without threshold assertions.

docs/contributing/ci/test_execution_guide.md

@@ -1,13 +1,20 @@
 # Test Execution Guide
+
 ## Setting Up the Test Environment
+
 ### Creating a Container
+
 vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
 For a local test environment, you can follow the steps below to create a container:
+
 ## Installing Dependencies
+
 ### vLLM & vLLM-Omni
+
 vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
 
 ### Test Case Dependencies
+
 When running test cases, you may need to install the following dependencies:
 
 ```bash
@@ -16,6 +23,7 @@ apt-get install -y espeak-ng jq
 ```
 
 ## Running Tests
+
 Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
 
 ### CI job runners (L2–L4): logs, timing, and timeouts
@@ -25,7 +33,7 @@ Our test scripts use the pytest framework. First, please use `git clone https://
 **Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`):
 
 | Path | Description |
-|------|-------------|
+| ------ | ------------- |
 | `jobs/*.sh` | Generated per-step bash wrappers |
 | `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) |
 | `<job_key>.log` | Tee output for each job |
@@ -33,7 +41,7 @@ Our test scripts use the pytest framework. First, please use `git clone https://
 
 **Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example:
 
-```
+```text
 === Job timing summary ===
   TTS_VoxCPM2_Test  25m 40s  OK
   Omni_Qwen3-Omni_Test  8m 05s  TIMED OUT
@@ -114,7 +122,6 @@ Failed jobs: 1/2
     pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
     ```
 
-
 === "L4 level"
 
     **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/cuda/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
@@ -160,7 +167,7 @@ Failed jobs: 1/2
     pytest -sv tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json
     pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
     ```
-    Nightly **Perf Test** jobs in ``test-nightly.yml`` use ``--test-config-file`` only (no ``-m``). E2e L4 function tests still use ``full_model`` + ``--run-level full_model`` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml)). Example:
+    Nightly **Perf Test** jobs in [``test-nightly.yml``](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml) use ``--test-config-file`` only (no ``-m``). Weekly **Perf Test** in [``test-weekly.yml``](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml) runs ``test_qwen3_omni_vllm_text.json`` and ``test_qwen3_omni_multi_replicas.json`` (JSON ``mark`` includes ``slow``). E2e L4 function tests use ``full_model`` + ``--run-level full_model``. Example:
 
 === "L5 level"
 
@@ -186,9 +193,13 @@ Failed jobs: 1/2
 
     ```
 
-    The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).
+    The latest L5 CI jobs are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml):
+
+    - **Reliability** (+ invalid-parameter), **Perf Test** (Omni vLLM Text / Multi-Replica), and **Simple · CPU Coverage Test** → `WEEKLY=1` (or PR label `weekly-test` for Reliability/Perf). Scheduled upload of those env-gated steps still requires `main`.
+    - **E2E Tests** (slow Omni / TTS / Diffusion sweeps) → `NON_CRITICAL=1` (pipeline upload also accepts `WEEKLY=1` / `weekly-test`; the group `if` still needs `NON_CRITICAL`)
 
 You can find more information about markers in [Test Writing Guide — Markers for Tests](./test_writing_guide.md#markers-for-tests).
 
 ## Adding New Test Cases
+
 Please refer to the [Test System Overview](./test_system_overview.md) and [Test Writing Guide](./test_writing_guide.md).

docs/contributing/ci/test_system_overview.md

@@ -12,8 +12,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <th>Level</th>
       <th>Scope & Focus</th>
       <th>Model Coverage Strategy</th>
-      <th>Feature Coverage Strategy</th>
-      <th>Interface Coverage Strategy</th>
+      <th>Module Strategy</th>
       <th>Tags</th>
       <th>Time Cost</th>
       <th>Test Dir</th>
@@ -31,7 +30,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>/</td>
       <td>/</td>
       <td>/</td>
-      <td>/</td>
       <td><a href="https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md">PR Checklist</a></td>
       <td>/</td>
       <td>/</td>
@@ -43,7 +41,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>/</td>
       <td>/</td>
       <td>/</td>
-      <td>/</td>
       <td><a href="./failures.md">CI Failures</a></td>
       <td>/</td>
       <td>/</td>
@@ -53,7 +50,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>Unit tests for components like entrypoints, models</td>
       <td>/</td>
       <td>/</td>
-      <td>/</td>
       <td><code>core_model and cpu</code></td>
       <td rowspan="2">&lt;15min</td>
       <td>
@@ -71,9 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
     <tr>
       <td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
       <td>Online (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
-      <td>High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy)</td>
-      <td>High-priority features and component tests that require launching instances</td>
-      <td>High-priority interfaces (using random lightweight models)</td>
+      <td>Key models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy)</td>
+      <td>Key features and modules that require launching instances (prefer random weights)</td>
       <td><code>core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>
         <strong>Model E2E:</strong><br>
@@ -95,9 +90,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
     <tr>
       <td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
       <td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
-      <td>High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy</td>
-      <td>Medium-priority features and component tests that require launching instances</td>
-      <td>Medium-priority interfaces (using random lightweight models)</td>
+      <td>Key models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy</td>
+      <td>Key features and modules that require launching instances (using real weights)</td>
       <td><code>advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>&lt;30min</td>
       <td>
@@ -120,9 +114,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
     <tr>
       <td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
       <td>Online: full functional scenarios + performance test + doc test + accuracy test</td>
-      <td>High-priority models: function, performance, accuracy, and doc testing<br>Medium-priority models: function and doc testing</td>
-      <td>Low-priority features and component tests that require launching instances</td>
-      <td>Low-priority interfaces (using real weights)</td>
+      <td>Key models: function, performance, accuracy, and doc testing</td>
+      <td>Other features and modules that require launching instances (using real weights)</td>
       <td><code>full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>&lt;3 hour</td>
       <td>
@@ -147,11 +140,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>GPU</td>
     </tr>
     <tr>
-      <td><strong>L5</strong><br>(Stability & Reliability)</td>
-      <td>Online: long-term stability test + reliability test</td>
-      <td>Long-term stability and reliability testing for high-priority models<br>Low-priority models: function and doc testing</td>
+      <td><strong>L5</strong><br>(Stability & Reliability & selected Perf)</td>
+      <td>Online: long-term stability + reliability; coverage test</td>
+      <td>Long-term stability and reliability testing for key models<br>Non-critical scenario performance and accuracy testing for key models<br>Other models: function and doc testing</td>
       <td>/</td>
-      <td>Invalid-parameter validation for high-priority interfaces</td>
       <td><code>slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td> Depends on reality </td>
       <td>
@@ -178,8 +170,8 @@ For per-level test authoring (markers, examples), see [Test Writing Guide](./tes
 
 Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues.
 
-1.  ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
-2.  ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
+1. ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
+2. ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
 
 ## Notes
 

docs/contributing/ci/test_writing_guide.md

@@ -509,7 +509,7 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 #### 3.2 Testing Content and Scope
 
 - ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
-- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below.
+- ***Performance Testing***: JSON configs under `tests/dfx/perf/tests/` drive Omni, TTS, and diffusion serving benchmarks via `run_benchmark.py` or `run_diffusion_benchmark.py`, measuring throughput, latency, and related metrics under configured workloads.
 - ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
 
 #### 3.3 Test Directory and Execution Files
@@ -567,7 +567,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
 - ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
 - ***Run Command***:
     - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`). Weekly CI (`.buildkite/cuda/test-weekly.yml`) runs one step per suite (`WEEKLY=1` or PR label `weekly-test`)
+    - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`).
 - ***Script Example***:
 
 <details>

docs/design/feature/async_diffusion_output.md

@@ -124,8 +124,8 @@ When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_
 1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
    When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
 
-2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
-   The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
+2. **Timeout Diagnostics (`_async_output_timeout()`)**:
+   The engine waits for `wait_output_ready()` with a timeout of `_async_output_timeout()` — 600.0s by default, overridable per run with the `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` environment variable (a non-numeric or non-positive value is ignored with a warning and the default applies, since this resolves on the request path). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter. The packing itself takes milliseconds, but the wait is queued behind the step's GPU work, so its wall-clock duration tracks step time; the bound is generous because a hung engine is surfaced by the worker monitor and `check_health()`, not by this timeout.
 
 3. **Queue Serialization & Atomic Future Resolution**:
    - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).

docs/design/feature/attention_backend_selection.md

@@ -61,7 +61,8 @@ unstructured keyword dictionaries:
 - `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
   validation;
 - `skip_softmax` is consumed by TRTLLM; and
-- `block_sparse` is consumed by block-sparse backends such as RainFusion.
+- `block_sparse` is consumed by block-sparse backends such as RainFusion; and
+- `fastvideo_vsa_topk` is consumed by FastVideo VSA.
 
 A backend reads only the fields it owns and rejects incompatible values. New
 options should be added to a shared typed spec only when more than one backend
@@ -87,3 +88,19 @@ An implementation change is complete only when it:
    corresponding user guide; and
 6. updates the overview matrix without moving algorithm details back into the
    landing page.
+
+## FastVideo VSA model metadata
+
+FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
+the original latent video grid. Wan integrations therefore attach the
+post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
+backend validates that `T * H * W` equals the sequence length, derives the
+runtime block count, and selects the configured top-k key/value blocks for
+each query block.
+
+Checkpoint-specific compensation remains model-owned. Wan layers pass an
+optional learned `gate_compress` projection through attention metadata when
+the checkpoint contains those weights. Checkpoints without the projection do
+not allocate or execute it. When top-k selects every block, native checkpoints
+route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
+to retain checkpoint semantics.

docs/design/feature/host_weight_runtime.md

@@ -15,8 +15,11 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
 ## Status
 
 The first implementation provides contracts and a CPU local-filesystem store.
-It does not enable cached loading for a model by itself. A consumer must still
-provide an exact identity, a representation producer, and a model restorer.
+The initial diffusion consumer contract 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.
 
 V1 includes:
 
@@ -24,16 +27,17 @@ V1 includes:
 - coordinated local lookup and one-producer publication;
 - descriptor-backed safetensors mmap leases;
 - preferred and required resolution policy;
+- explicit, separately reported post-load publication;
 - validation, deny, quarantine, cleanup, and capacity controls; and
 - typed reports for every terminal resolution outcome.
 
 V1 does not include:
 
-- a public CLI or default-on model integration;
-- a generic BF16, FP8, quantized, or model-specific producer;
+- 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;
 - a remote artifact provider or cross-node coordination;
-- post-load publication (V1 accepts only pre-load-safe producers);
 - automatic eviction; or
 - a change to DLO AllGather or no-AllGather behavior.
 
@@ -84,6 +88,12 @@ flowchart TD
     F -->|"preferred"| C
     F -->|"required"| E["Fail startup"]
     P -->|"nonretryable failure"| E
+    C --> W{"Post-load publication enabled?"}
+    W -->|"yes"| B2["Publish final model through POST_LOAD_ONLY producer"]
+    B2 --> CL["Close validated publication lease"]
+    CL --> R2["Return separate publication report"]
+    R2 --> D
+    W -->|"no"| D["Keep canonical model"]
 ```
 
 The modes are:
@@ -94,10 +104,12 @@ The modes are:
 | `preferred` | Return an exact lease when available; otherwise use canonical fallback only for a miss, invalid cache entry, or typed retryable failure. |
 | `required` | Fail when an exact lease cannot be acquired. |
 
-An unsupported capability, semantic identity collision, producer failure, or
-publication failure is nonretryable and remains visible even in preferred
-mode. Storage policy must never disguise a semantic or configuration error as
-a cache miss.
+During pre-load resolution, an unsupported capability, semantic identity
+collision, producer failure, or publication failure is nonretryable and remains
+visible even in preferred mode. A post-load publication failure is likewise
+visible in its own report, but cannot revise the canonical-fallback outcome.
+Storage policy must never disguise a semantic or configuration error as a cache
+miss.
 
 ## Loader and restoration sequence
 
@@ -118,24 +130,35 @@ sequenceDiagram
     R->>S: lookup(exact identity)
     alt validated hit
         S-->>R: HostWeightLease
+        R-->>L: lease and LOCAL_HIT report
+        L->>X: plan_restore(model, lease)
+        X-->>L: validation-only restore plan
+        L->>X: commit() once
     else miss with allowed producer
         R->>S: get_or_build(identity, producer)
         S->>P: produce(store-scoped writer)
         P-->>S: final-layout tensors and metadata
         S-->>R: validated HostWeightLease
+        R-->>L: lease and LOCAL_PRODUCTION report
+        L->>X: plan_restore(model, lease)
+        X-->>L: validation-only restore plan
+        L->>X: commit() once
     else policy permits canonical fallback
         R-->>L: CANONICAL_FALLBACK
         L->>L: Run canonical loader
+        opt explicit post-load publication enabled
+            L->>R: publish_after_load(identity, POST_LOAD_ONLY producer)
+            R->>S: get_or_build(identity, producer)
+            S-->>R: validated HostWeightLease or typed failure
+            R->>R: close publication lease
+            R-->>L: separate publication report
+        end
     else required or nonretryable failure
         R-->>L: FAILED report
     end
-    R-->>L: lease and terminal report
-    L->>X: plan_restore(model, lease)
-    X-->>L: validation-only restore plan
-    L->>X: commit() once
-    L->>T: consume restored CPU tensors
+    L->>T: consume final CPU tensors and any lease
     T-->>L: transfer teardown complete
-    L->>L: close lease
+    L->>L: close lease when one was acquired
 ```
 
 `plan_restore()` must not mutate the model or lease. `commit() -> None` is the
@@ -143,6 +166,15 @@ sole one-shot model mutation. If planning fails, canonical fallback may reuse
 the untouched model. If commit begins and fails, the partially hydrated model
 must be discarded and canonical fallback must construct a fresh model.
 
+`publish_after_load()` is synchronous in V1 and accepts only a
+`POST_LOAD_ONLY` producer. Its report is separate from the terminal resolution,
+so a publication failure cannot rewrite a successful canonical fallback. A
+successful publication closes the store-returned lease inside the runtime and
+warms only future startups. It does not restore, rebind, or otherwise mutate the
+canonically loaded model serving the current startup.
+`allow_local_build` gates producers during pre-load resolution, while
+`allow_post_load_publish` independently gates this explicit post-load path.
+
 ## Exact representation identity
 
 The store performs exact lookup. It never silently converts one representation
@@ -188,6 +220,56 @@ Dynamic LoRA overlays are not part of a reusable base-weight artifact. A
 statically merged adapter is cacheable only as a separate identity containing
 the adapter fingerprint and merge semantics.
 
+### Initial diffusion final-layout contract
+
+The shared diffusion contract covers complete final-layout DiT parameters and
+persistent buffers. Text encoders, VAEs, non-persistent derived state, and other
+pipeline components remain outside the artifact. One explicit representation
+policy selects allowed dtypes, tensor roles, physical layout identity, producer
+ABI, manifest schema, and restoration schema.
+
+The contract is intentionally separate from loader activation:
+
+- `FinalLayoutRequest` contains typed loader identity/configuration fingerprints,
+  TP coordinate, and conservative SP semantics. It has no open metadata bag,
+  DP coordinate, SP rank, device identity, DLO transfer mode, registration
+  policy, or store path.
+- `FinalLayoutArtifactSpec` binds one `WeightRepresentation` and runtime-layout
+  name to explicit producer/restorer schemas and a canonical, versioned
+  implementation ABI descriptor. Compatibility never depends on reflective
+  source inspection.
+- `PreparedWeightSource` snapshots immutable revisions or exact local file
+  content plus a typed checkpoint-adapter identity before ordinary
+  materialization. Source replacement before or during production fails
+  publication. A hash-looking symlink basename is trusted only for an explicit
+  Hugging Face Hub source whose repository ID and
+  `models--.../snapshots/<revision> -> blobs/<hash>` topology validate; every
+  local or otherwise unverified symlink target is content-hashed.
+- the tensor ownership digest records exact runtime names, kinds, shapes,
+  semantic roles, dtypes, and strides from a CPU or meta model skeleton;
+- `FinalLayoutTensorRestorer` accepts only an exact lease identity, validates
+  complete policy-defined coverage without mutation, and returns a one-shot
+  commit plan;
+- each model declares one dtype-neutral `FinalLayoutModelContract` with an
+  explicit implementation version and a post-commit validator; and
+- `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.
+
+Other representations reuse source identity, typed parallel identity, tensor
+ownership, and exact restoration only when their policy proves those semantics.
+For example, runtime FP8 needs a separate policy/producer for generated scales,
+quantization metadata, and Cutlass physical layouts; it is not enabled by
+changing a dtype string on the BF16 producer.
+
+This stage makes no startup, sharing, or DLO performance claim. A following
+consumer PR owns disabled/preferred/required precedence, mixed-component loader
+transactions, warm-hit restoration, and transactional lease handoff. A TP2
+prewarm deployment will require a matching TP2 producer cohort to populate both
+TP-coordinate identities even though the store coordinates each artifact
+independently.
+
 ## Host sharing and GPU transport
 
 Every process receives its own virtual mappings and `HostWeightLease`, but
@@ -252,7 +334,8 @@ removing live mappings.
 | Retryable lock, domain, or capacity failure | Canonical fallback | Fail |
 | Unsupported producer or backend | Fail | Fail |
 | Identity collision | Fail | Fail |
-| Producer or publication failure | Fail | Fail |
+| Pre-load producer or atomic publication failure | Fail | Fail |
+| Post-load publication failure after canonical fallback | Keep canonical model; report publication failure | Not reached through required resolution |
 | Restore planning failure | Canonical fallback may reuse untouched model | Fail |
 | Restore commit failure | Discard model; fallback requires a fresh instance | Fail |
 
@@ -275,7 +358,8 @@ A consumer PR must:
    a failed restore commit.
 7. Retain the lease until every transport operation that may access mapped
    memory has completed.
-8. Emit and test the terminal resolution report.
+8. Emit and test the terminal resolution report and any separate post-load
+   publication report.
 
 Consumer validation must prove:
 

docs/design/module/host_weight_runtime.md

@@ -55,7 +55,10 @@ flowchart LR
 The contracts have deliberately narrow ownership:
 
 - `HostWeightRuntime` applies resolution and fallback policy and emits one
-  terminal report. It does not own a model or call the canonical loader.
+  terminal resolution report. Explicit post-load publication emits a separate
+  report and closes its store-returned lease before returning; it does not
+  revise the completed resolution.
+  The runtime does not own a model or call the canonical loader.
 - `HostWeightStore` owns artifact lookup, coordinated construction,
   publication, validation, quarantine, and lifecycle.
 - `HostWeightLease` owns process-local tensor views, mapped-file resources, and
@@ -135,9 +138,23 @@ producer cancellation requires a future process-isolated producer contract. A
 waiter timing out never cancels another process's valid build.
 
 Remote providers are represented by protocols but are explicitly unsupported
-in the first implementation. V1 also invokes only `PRE_LOAD_SAFE` producers;
-enabling post-load publication is rejected during configuration until a
-separate post-load entry point and model-lifetime contract exist.
+in the first implementation. Pre-load `resolve()` invokes only
+`PRE_LOAD_SAFE` producers. A loader that completed canonical materialization
+may explicitly call `publish_after_load()` with a `POST_LOAD_ONLY` producer
+when `allow_post_load_publish` is enabled.
+
+`allow_local_build` controls producer use during pre-load resolution;
+`allow_post_load_publish` independently controls the explicit post-load entry
+point. This permits lookup-only warm starts that populate a missing artifact
+only after the canonical loader has established the current model.
+
+Post-load publication is synchronous in V1, but its outcome is independent of
+the already completed resolution. A store failure is returned in a separate
+publication report and cannot turn a valid canonical model into a failed
+resolution. On success, the runtime closes the validated store lease before
+returning the report. Post-load publication only warms future startups; it does
+not restore, rebind, or otherwise mutate the canonically loaded model serving
+the current startup.
 
 ## Restoration transaction boundary
 

docs/models/supported_models.md

@@ -68,17 +68,12 @@ 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` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
-| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
-| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
 | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -109,6 +104,7 @@ th {
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
 
 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the

docs/user_guide/diffusion/attention_backends.md

@@ -20,6 +20,7 @@ For the internal selector, registry, and platform contract, see
 | Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
 | Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Backend options
 
@@ -35,6 +36,7 @@ For the internal selector, registry, and platform contract, see
 | `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Configuration
 
@@ -105,6 +107,7 @@ Backend-specific typed blocks are documented with their consumers:
   and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
 - `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
 - `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
+- `fastvideo_vsa_topk`: [FastVideo VSA](attention_backends/fastvideo_vsa.md#choose-top-k).
 
 ## Platform defaults
 

docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md

@@ -0,0 +1,125 @@
+# FastVideo VSA
+
+FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
+self-attention by partitioning the post-patch latent grid into spatiotemporal
+blocks. For every query block, VSA scores the key/value blocks and computes
+attention only against the selected top-k blocks.
+
+The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
+
+VSA is a CUDA-only, explicitly selected backend. It requires the
+`fastvideo-kernel` package and currently supports non-causal self-attention
+with equal query and key/value sequence lengths. Unsupported shapes, masks,
+dtypes, sequence-parallel execution, or kernel failures fall back to
+`TORCH_SDPA` and emit a warning with the reason.
+
+## Enable the backend
+
+For online serving, select the backend with the existing attention backend
+flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
+for every query block:
+
+```bash
+vllm-omni serve <model> \
+  --diffusion-attention-backend FASTVIDEO_VSA \
+  --fastvideo-vsa-topk 64
+```
+
+The backwards-compatible environment variable selects the backend with the
+default `topk=64`:
+
+```bash
+export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
+vllm-omni serve <model>
+```
+
+To tune top-k, pass the CLI backend and top-k flags together as shown above,
+or use the equivalent structured configuration:
+
+```bash
+vllm-omni serve <model> \
+  --diffusion-attention-config \
+  '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
+```
+
+Do not combine `--diffusion-attention-backend` with an explicit
+`diffusion_attention_config.default.backend`. The top-k value must be positive
+and is valid only when the default backend is `FASTVIDEO_VSA`.
+
+For a deploy YAML stage, use either the shorthand fields:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_attention_backend: FASTVIDEO_VSA
+    fastvideo_vsa_topk: 64
+```
+
+or the structured configuration:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_attention_config:
+      default:
+        backend: FASTVIDEO_VSA
+        fastvideo_vsa_topk: 64
+```
+
+## Choose top-k
+
+At runtime the backend logs the sequence shape and derived block count:
+
+```text
+FASTVIDEO_VSA routing: seq_len=27280, dit_seq_shape=(31, 22, 40),
+block_size=(4, 8, 8), num_blocks=120, topk=64, keep_ratio=53.3%,
+checkpoint_mode=native, route=VSA
+```
+
+Use `num_blocks` as the upper bound when tuning top-k. Top-k is the number of
+key/value blocks retained for each query block, not the number of tokens or
+the total number of blocks processed by the layer.
+
+- A smaller top-k increases sparsity and may reduce attention computation, but
+  it can remove relevant blocks and reduce visual quality or temporal
+  consistency. Routing and padding overhead can also make a smaller value
+  slower for some shapes.
+- A larger top-k retains more context and generally approaches dense-attention
+  quality, but reduces the potential speedup and increases memory traffic.
+- `topk > num_blocks` is invalid for the runtime shape and falls back to SDPA.
+- For a native checkpoint, `topk == num_blocks` routes to SDPA because scoring
+  every block provides no sparsity benefit.
+- For a FastVideo DMD checkpoint, `topk == num_blocks` stays on the VSA
+  all-block path so the checkpoint keeps its trained compensation semantics.
+
+There is no universally optimal value. Resolution, frame count, GPU, kernel
+version, and checkpoint all affect both quality and latency. Start from the
+logged `num_blocks`, test several keep ratios on the target workload, and
+compare output quality against the same checkpoint running dense attention.
+
+## Checkpoint behavior
+
+The backend does not expose a user-selectable gate mode. Wan checkpoints that
+contain learned `to_gate_compress` weights use that projection automatically.
+When those weights are absent, the unused projection is removed and VSA uses
+the sparse branch without learned compensation.
+
+FastVideo DMD checkpoints use their fixed distilled timestep schedule. Native
+Wan checkpoints keep their normal scheduler and inference-step configuration;
+selecting VSA does not turn a native checkpoint into a distilled model.
+
+## Verify routing and fallback
+
+Check the startup and first-forward logs instead of assuming that selecting
+the backend guarantees sparse execution:
+
+- `route=VSA` means top-k block selection is active.
+- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
+  blocks through the VSA kernel.
+- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
+  executed; the warning includes the reason.
+
+The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
+standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
+and no active sequence-parallel context. NPU and XPU paths do not execute the
+FastVideo VSA CUDA kernel.

docs/user_guide/diffusion/cpu_offload.md

@@ -49,7 +49,6 @@ higher-priority strategy is selected.
 | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
 | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
 | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
-| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
 | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
 | Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
 

docs/user_guide/diffusion_features.md

@@ -124,12 +124,12 @@ The following tables show which models support each feature:
 | **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ✅           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **HiDream-O1-Image**     |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ✅*        |
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -457,6 +457,7 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
 - Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
+- For a bounded active-stream deployment, `active_stream_policy: edf` can prioritize streams whose server-estimated playback buffer is closest to exhaustion. Configure it under the connector's `extra` block and use a `codec_chunk_ramp`; this is opt-in, estimates playback from server-side PCM emission (no client playback acknowledgements), and targets first-audio latency and continuity rather than single-request completion latency.
 
 ---
 

examples/offline_inference/audex/README.md

@@ -17,7 +17,7 @@ Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
 pair contract) lives in `common.py`; each task script keeps only its own
 prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
 two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
-Open, AudioX) use the generic
+Open) use the generic
 `examples/offline_inference/text_to_audio/text_to_audio.py` instead.
 
 ## text_to_speech.py — text → speech

examples/offline_inference/audex/text_to_audio.py

@@ -7,7 +7,7 @@ thinker generates interleaved 4-codebook <audiocodec_N> RVQ tokens under an
 RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
 mono WAV.
 
-(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
+(For DIFFUSION text-to-audio models — Stable Audio Open — use
 ``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
 TTA is a two-stage autoregressive pipeline with a different interface.)
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -1,117 +0,0 @@
-"""
-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_audio/README.md

@@ -1,11 +1,10 @@
 # Text-To-Audio
 
-A unified script for text/video-to-audio generation. Supported models:
+A unified script for text-to-audio generation. Supported models:
 
 | Model | Tasks | Notes |
 |-------|-------|-------|
 | `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` |
-| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` |
 
 The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
 
@@ -50,42 +49,12 @@ python text_to_audio.py \
   --output stable_audio_output.wav
 ```
 
-### AudioX
-
-AudioX supports six tasks. Sampler and reference knobs (declared in
-`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body`
-JSON flag, routed into sampling `extra_args`.
-
-Text tasks (`t2a` / `t2m`):
-
-```bash
-python text_to_audio.py \
-  --model zhangj1an/AudioX --task t2a \
-  --prompt "Fireworks burst twice, followed by a clock ticking." \
-  --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \
-  --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
-  --output t2a.wav
-```
-
-Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`:
-
-```bash
-python text_to_audio.py \
-  --model zhangj1an/AudioX --task tv2a \
-  --prompt "drum beating sound and human talking" \
-  --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
-  --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
-  --output tv2a.wav
-```
-
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`.
-- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`).
-- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX).
-- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX).
-- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`.
+- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s`).
+- `--audio-length`: audio duration in seconds (audio length for Stable Audio).
+- `--extra-body`: JSON dict of model-specific knobs, merged into sampling `extra_args`.
 - `--negative-prompt`: negative prompt for classifier-free guidance.
 - `--seed`: integer seed for deterministic generation.
 - `--guidance-scale`: classifier-free guidance scale.

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -2,15 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Example script for text/video-to-audio generation with diffusion audio models.
+Example script for text-to-audio generation with diffusion audio models.
 
 This script supports:
   * Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio.
-  * AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m.
 
-Model-specific generation knobs are declared in `vllm_omni/model_extras` and
-routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so
-the same script serves both models without per-model branching in the engine.
+Model-specific generation knobs can be passed via `--extra-body` and are
+merged into `sampling_params.extra_args`.
 
 Usage:
     # Stable Audio Open
@@ -18,15 +16,6 @@ Usage:
     python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
     python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
     python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache
-
-    # AudioX (text-to-audio / text-to-music)
-    python text_to_audio.py --model zhangj1an/AudioX --task t2a \
-        --prompt "Fireworks burst twice, then a clock ticking." \
-        --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0
-    # AudioX video-conditioned (v2*/tv2*) — pass a video file/URL
-    python text_to_audio.py --model zhangj1an/AudioX --task tv2a \
-        --prompt "drum beating sound and human talking" --video clip.mp4 \
-        --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}'
 """
 
 import argparse
@@ -38,16 +27,11 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
-AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
-AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
-
 
 def parse_extra_body(value: str) -> dict[str, Any]:
     """Parse a JSON object string for --extra-body."""
@@ -65,7 +49,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--model",
         default="stabilityai/stable-audio-open-1.0",
-        help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.",
+        help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0.",
     )
     parser.add_argument(
         "--prompt",
@@ -96,19 +80,13 @@ def parse_args() -> argparse.Namespace:
         "--audio-start",
         type=float,
         default=0.0,
-        help=(
-            "Audio start offset in seconds. Applies to both models "
-            "(audio_start_in_s for Stable Audio, seconds_start for AudioX)."
-        ),
+        help="Audio start offset in seconds (maps to audio_start_in_s).",
     )
     parser.add_argument(
         "--audio-length",
         type=float,
         default=10.0,
-        help=(
-            "Audio duration in seconds. Maps to audio length for Stable Audio "
-            "(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX."
-        ),
+        help=("Audio duration in seconds. Maps to audio length for Stable Audio (max ~47s for stable-audio-open-1.0)."),
     )
     parser.add_argument(
         "--num-inference-steps",
@@ -122,41 +100,20 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of audio waveforms to generate for the given prompt.",
     )
-    parser.add_argument(
-        "--task",
-        type=str,
-        default=None,
-        choices=list(AUDIOX_TASKS),
-        help=(
-            "[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. "
-            "v2*/tv2* require --video. Ignored for Stable Audio Open."
-        ),
-    )
-    parser.add_argument(
-        "--video",
-        type=str,
-        default=None,
-        help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.",
-    )
     parser.add_argument(
         "--extra-body",
         type=parse_extra_body,
         default=None,
         help=(
-            "Model-specific generation params as a JSON object, e.g. "
-            '\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered '
-            "against the model's declared extra_body_params (see vllm_omni/model_extras), "
-            "so unknown keys are silently dropped. Values here override the equivalent flags."
+            "Model-specific generation params as a JSON object, merged into "
+            "sampling extra_args. Values here override the equivalent flags."
         ),
     )
     parser.add_argument(
         "--output",
         type=str,
         default=None,
-        help=(
-            "Path to save the generated audio (WAV format). Defaults to "
-            "audiox_output.wav for AudioX and stable_audio_output.wav otherwise."
-        ),
+        help=("Path to save the generated audio (WAV format). Defaults to stable_audio_output.wav."),
     )
     parser.add_argument(
         "--sample-rate",
@@ -275,8 +232,6 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
 
 def main():
     args = parse_args()
-    model_name = str(args.model).lower() if args.model else ""
-    is_audiox = "audiox" in model_name
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     cache_config = None
     if args.cache_backend == "tea_cache":
@@ -285,11 +240,9 @@ def main():
         }
 
     print(f"\n{'=' * 60}")
-    print("Text/Video-to-Audio Generation")
+    print("Text-to-Audio Generation")
     print(f"{'=' * 60}")
     print(f"  Model: {args.model}")
-    if args.task is not None:
-        print(f"  Task: {args.task}")
     print(f"  Prompt: {args.prompt}")
     print(f"  Negative prompt: {args.negative_prompt}")
     print(f"  Duration: {args.audio_length}s")
@@ -305,8 +258,6 @@ def main():
     print(f"  Seed: {args.seed}")
     print(f"{'=' * 60}\n")
 
-    # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
-    # is resolved from its config.
     omni_kwargs: dict[str, Any] = dict(
         model=args.model,
         use_hsdp=args.use_hsdp,
@@ -318,11 +269,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
     )
-    if is_audiox:
-        omni_kwargs["model_class_name"] = "AudioXPipeline"
     omni = Omni(**omni_kwargs)
-    model_class_name = get_model_class_name(omni)
-    declared_extra_body_params = get_extra_body_params(model_class_name)
 
     diffusion_params = OmniDiffusionSamplingParams(
         generator=generator,
@@ -332,38 +279,19 @@ def main():
         seed=args.seed,
     )
 
-    # Negative prompt is opt-in for both models (off by default). Including a
+    # Negative prompt is opt-in (off by default). Including a
     # non-empty negative prompt enables the separate CFG conditioning branch.
     prompt: dict[str, Any] = {"prompt": args.prompt}
     if args.negative_prompt:
         prompt["negative_prompt"] = args.negative_prompt
 
-    if model_class_name == "AudioXPipeline":
-        if args.task is None:
-            raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).")
-        if args.task in AUDIOX_VIDEO_TASKS and not args.video:
-            raise ValueError(f"AudioX task {args.task!r} requires --video.")
-        # AudioX core knobs are surfaced as shared flags; sampler/reference knobs
-        # (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and
-        # are filtered against the declared extra_body contract.
-        user_extra: dict[str, Any] = {
-            "audiox_task": args.task,
-            "seconds_start": args.audio_start,
-            "seconds_total": args.audio_length,
-            "video_path": args.video,
-        }
-        if args.extra_body:
-            user_extra.update(args.extra_body)
-        apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
-    else:
-        # Stable Audio Open path (backward-compatible).
-        audio_end_in_s = args.audio_start + args.audio_length
-        diffusion_params.extra_args = {
-            "audio_start_in_s": args.audio_start,
-            "audio_end_in_s": audio_end_in_s,
-        }
-        if args.extra_body:
-            diffusion_params.extra_args.update(args.extra_body)
+    audio_end_in_s = args.audio_start + args.audio_length
+    diffusion_params.extra_args = {
+        "audio_start_in_s": args.audio_start,
+        "audio_end_in_s": audio_end_in_s,
+    }
+    if args.extra_body:
+        diffusion_params.extra_args.update(args.extra_body)
 
     generation_start = time.perf_counter()
 
@@ -374,7 +302,7 @@ def main():
 
     print(f"Total generation time: {generation_time:.2f} seconds")
 
-    output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav")
+    output = args.output or "stable_audio_output.wav"
     output_path = Path(output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     suffix = output_path.suffix or ".wav"

examples/offline_inference/text_to_image/text_to_image.py

@@ -89,7 +89,7 @@ def parse_args() -> argparse.Namespace:
         "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
         "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
         "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
-        "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
+        "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
         "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
     )
     parser.add_argument(
@@ -599,17 +599,6 @@ 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/text_to_speech/README.md

@@ -420,8 +420,7 @@ against transformers 5.x, so nothing is lost by holding it there.
 
 (The same clash is why this is documented here rather than declared in
 `pyproject.toml`: an extra resolves together with the base dependencies, so
-`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
-SVS documents `nemo_toolkit[asr]` the same way.)
+`vllm-omni[gepard]` would be unsatisfiable rather than opt-in.)
 
 On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
 consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
@@ -617,100 +616,3 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
 - `--num-prompts N` replicates the prompt for performance measurement.
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.
-
----
-
-## SoulX-Singer
-
-Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
-
-### Prerequisites
-
-Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
-
-```bash
-# 1. DiT weights
-export BASE=path/to/SoulX-Singer
-export PREPROCESS=path/to/SoulX-Singer-Preprocess
-export SVC_DIR=path/to/SoulX-Singer-svc
-
-huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
-
-# 2. Preprocess weights (required)
-huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
-export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
-
-# 3. SVS / SVC view directories
-mkdir -p "$SVC_DIR"
-cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
-mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
-
-cat > "$BASE/config.json" <<'EOF'
-{
-  "model_type": "soulxsinger",
-  "architectures": ["SoulXSingerPipeline"],
-  "max_num_seqs": 1
-}
-EOF
-
-cat > "$SVC_DIR/config.json" <<'EOF'
-{
-  "model_type": "soulxsinger",
-  "architectures": ["SoulXSingerSVCPipeline"],
-  "max_num_seqs": 1
-}
-EOF
-```
-
-`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
-
-**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
-
-```bash
-pip install "BS-RoFormer"   # vocal sep + F0 on GPU — SVS and SVC
-```
-
-Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
-
-```bash
-# install optional dependencies:
-pip install -e ".[soulx-svs]"
-```
-
-English SVS adds NeMo ASR and NLTK data; pass `--language English`:
-
-```bash
-pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
-python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
-```
-
-**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
-
-### Quick start
-
-```bash
-# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
-python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
-    --model "$BASE" \
-    --preprocess-weights-dir "$PREPROCESS" \
-    --control score \
-    --num-inference-steps 32 \
-    -o output.wav
-
-python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
-    --model "$SVC_DIR" \
-    --preprocess-weights-dir "$PREPROCESS" \
-    --svc \
-    --num-inference-steps 32 \
-    -o output_svc.wav
-```
-
-`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
-
-### Notes
-
-- Output: 24 kHz mono WAV; batch only.
-- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
-- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
-
----

examples/offline_inference/text_to_speech/soulxsinger/benchmark.py

@@ -1,162 +0,0 @@
-"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
-
-Usage:
-    python benchmark.py --model /path/to/SoulX-Singer --svs \\
-        --prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
-        --target-audio ../SoulX-Singer/example/audio/music.mp3 \\
-        --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-        -o benchmark.wav
-
-    python benchmark.py ... --enable-diffusion-pipeline-profiler
-"""
-
-import argparse
-import statistics
-import time
-from pathlib import Path
-
-import soundfile as sf
-from end2end import (
-    SVC_DEPLOY_CONFIG,
-    SVS_DEPLOY_CONFIG,
-    add_inference_args,
-    build_sampling,
-    extract_audio,
-    resolve_preprocess_weights_dir,
-)
-
-from vllm_omni.entrypoints.omni import Omni
-
-
-def _audio_duration_sec(outputs) -> float:
-    audio_np, sr = extract_audio(outputs)
-    return audio_np.size / sr
-
-
-def _stage_durations_from_output(outputs) -> dict[str, float]:
-    omni_out = outputs[0]
-    durations = getattr(omni_out, "stage_durations", None) or {}
-    if durations:
-        return dict(durations)
-    ro = omni_out
-    if ro is not None:
-        return dict(getattr(ro, "stage_durations", None) or {})
-    return {}
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
-    add_inference_args(parser)
-    parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
-    parser.add_argument("--runs", type=int, default=3, help="Measured runs")
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
-    )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        default=None,
-        choices=["bfloat16", "bf16", "float16", "fp16", "half"],
-        help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
-    )
-    parser.add_argument(
-        "--output",
-        "-o",
-        type=str,
-        default=None,
-        help="Save WAV from the last measured run (after warmup)",
-    )
-    parser.add_argument(
-        "--save-each-run",
-        action="store_true",
-        help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...",
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
-    deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
-    mode = "SVS" if args.svs else "SVC"
-
-    omni_kwargs: dict = {
-        "model": args.model,
-        "deploy_config": deploy_config,
-        "async_chunk": False,  # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
-    }
-    if args.enable_diffusion_pipeline_profiler:
-        omni_kwargs["enable_diffusion_pipeline_profiler"] = True
-    if args.enforce_eager:
-        omni_kwargs["enforce_eager"] = True
-    if args.dtype is not None:
-        omni_kwargs["dtype"] = args.dtype
-
-    compile_mode = "eager" if args.enforce_eager else "torch.compile"
-    dtype_label = args.dtype or "from deploy YAML"
-    print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
-
-    omni = Omni(**omni_kwargs)
-    kind = "svs" if args.svs else "svc"
-    sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
-    prompt = {"prompt_token_ids": [0]}
-
-    latencies_ms: list[float] = []
-    rtfs: list[float] = []
-    last_stages: dict[str, float] = {}
-    last_measured_outputs = None
-
-    total_runs = args.warmup + args.runs
-    for run_idx in range(total_runs):
-        is_warmup = run_idx < args.warmup
-        label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
-        t0 = time.perf_counter()
-        outputs = list(omni.generate([prompt], sampling))
-        elapsed_ms = (time.perf_counter() - t0) * 1000.0
-        audio_sec = _audio_duration_sec(outputs)
-        rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
-        print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
-        if not is_warmup:
-            latencies_ms.append(elapsed_ms)
-            rtfs.append(rtf)
-            last_stages = _stage_durations_from_output(outputs)
-            last_measured_outputs = outputs
-            if args.output and args.save_each_run:
-                measured_idx = run_idx - args.warmup + 1
-                out_path = Path(args.output)
-                run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
-                audio_np, sr = extract_audio(outputs)
-                sf.write(str(run_path), audio_np, sr)
-                print(f"  saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
-
-    omni.close()
-
-    if args.output and last_measured_outputs is not None:
-        audio_np, sr = extract_audio(last_measured_outputs)
-        out_path = Path(args.output)
-        out_path.parent.mkdir(parents=True, exist_ok=True)
-        sf.write(str(out_path), audio_np, sr)
-        print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
-
-    if latencies_ms:
-        print("\n=== Summary (measured runs) ===")
-        print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
-        print(f"RTF:       mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
-        if last_stages:
-            print("\nStage durations (last measured run):")
-            for name, value in sorted(last_stages.items()):
-                if name.endswith("_ms"):
-                    print(f"  {name}: {value:.1f} ms")
-                else:
-                    print(f"  {name}: {value * 1000.0:.1f} ms")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/text_to_speech/soulxsinger/end2end.py

@@ -1,285 +0,0 @@
-"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
-
-Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
-
-Usage:
-    # SVS (default)
-    python end2end.py --model /path/to/SoulX-Singer \\
-        --prompt-audio path/to/prompt/audio.mp3 \\
-        --target-audio path/to/target/audio.mp3 \\
-        --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-        -o output.wav
-
-    # SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
-    python end2end.py --model /path/to/SoulX-Singer-svc \\
-        --prompt-audio path/to/prompt/audio.mp3 \\
-        --target-audio path/to/target/audio.mp3 \\
-        --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
-        -o output_svc.wav
-"""
-
-import argparse
-import os
-from pathlib import Path
-
-import numpy as np
-import soundfile as sf
-
-from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-REPO_ROOT = Path(__file__).resolve().parents[4]
-SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
-SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
-DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
-DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
-_SAMPLE_RATE = 24000
-
-
-def _require_paths(paths: dict[str, str | None]) -> None:
-    missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
-    if missing:
-        raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
-
-
-def resolve_preprocess_weights_dir(explicit: str | None) -> str:
-    if explicit:
-        path = Path(explicit).expanduser().resolve()
-        if not (path / "rmvpe" / "rmvpe.pt").is_file():
-            raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
-        return str(path)
-    env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
-    if env:
-        path = Path(env).expanduser().resolve()
-        if (path / "rmvpe" / "rmvpe.pt").is_file():
-            return str(path)
-    raise FileNotFoundError(
-        "SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
-        "SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
-        "(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
-    )
-
-
-def build_sampling(
-    args: argparse.Namespace,
-    *,
-    preprocess_weights_dir: str,
-    kind: str,
-) -> OmniDiffusionSamplingParams:
-    extra: dict = {
-        "prompt_audio": os.path.abspath(args.prompt_audio),
-        "target_audio": os.path.abspath(args.target_audio),
-        "preprocess_weights_dir": preprocess_weights_dir,
-        "vocal_sep": args.vocal_sep,
-        "auto_shift": args.auto_shift,
-        "pitch_shift": args.pitch_shift,
-    }
-    if kind == "svs":
-        extra["language"] = args.language
-        extra["control"] = args.control
-        if args.prompt_metadata_path:
-            extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
-        if args.target_metadata_path:
-            extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
-        if args.prompt_metadata_path or args.target_metadata_path:
-            extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
-
-    extra = validate_soulx_extra_args(kind, extra)
-
-    return OmniDiffusionSamplingParams(
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        seed=args.seed,
-        extra_args=extra,
-    )
-
-
-def extract_audio(outputs) -> tuple[np.ndarray, int]:
-    out = outputs[0]
-    mm = getattr(out, "multimodal_output", None)
-    if mm is None:
-        ro = out
-        if ro is not None:
-            mm = getattr(ro, "multimodal_output", None)
-            if mm is None and ro.outputs:
-                mm = getattr(ro.outputs[0], "multimodal_output", None)
-    if not mm or "audio" not in mm:
-        raise RuntimeError("No audio in multimodal_output")
-
-    audio_val = mm["audio"]
-    if isinstance(audio_val, list):
-        chunks = [
-            np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
-            for chunk in audio_val
-            if chunk is not None
-        ]
-        audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
-    elif hasattr(audio_val, "cpu"):
-        audio_np = audio_val.detach().cpu().numpy().squeeze()
-    else:
-        audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
-
-    sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
-    if isinstance(sr_val, list) and sr_val:
-        sr_val = sr_val[-1]
-    if hasattr(sr_val, "item"):
-        sr_val = sr_val.item()
-    return audio_np.astype(np.float32, copy=False), int(sr_val)
-
-
-def add_inference_args(parser: argparse.ArgumentParser) -> None:
-    mode = parser.add_mutually_exclusive_group()
-    mode.add_argument(
-        "--svc",
-        action="store_true",
-        help="Assert SVC mode (must match config.json architectures)",
-    )
-    parser.add_argument(
-        "--model",
-        type=str,
-        required=True,
-        help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
-    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
-    )
-    parser.add_argument(
-        "--prompt-audio",
-        type=str,
-        default=str(DEFAULT_PROMPT_AUDIO),
-        help="Prompt / reference audio (mp3 or wav)",
-    )
-    parser.add_argument(
-        "--target-audio",
-        type=str,
-        default=str(DEFAULT_TARGET_AUDIO),
-        help="Target content audio (mp3 or wav)",
-    )
-    parser.add_argument(
-        "--prompt-metadata-path",
-        type=str,
-        default=None,
-        help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
-    )
-    parser.add_argument(
-        "--target-metadata-path",
-        type=str,
-        default=None,
-        help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
-    )
-    parser.add_argument(
-        "--prompt-wav-path",
-        type=str,
-        default=None,
-        help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
-    )
-    parser.add_argument(
-        "--preprocess-weights-dir",
-        type=str,
-        default=None,
-        help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
-    )
-    parser.add_argument(
-        "--language",
-        type=str,
-        default="Mandarin",
-        help="SVS lyrics language (Mandarin, Cantonese, English)",
-    )
-    parser.add_argument(
-        "--control",
-        type=str,
-        default="melody",
-        choices=["score", "melody"],
-        help="SVS control mode",
-    )
-    parser.add_argument(
-        "--vocal-sep",
-        action=argparse.BooleanOptionalAction,
-        default=False,
-        help="Run vocal separation during preprocess",
-    )
-    parser.add_argument(
-        "--auto-shift",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Auto pitch shift to align target pitch with prompt",
-    )
-    parser.add_argument("--pitch-shift", type=int, default=0)
-    parser.add_argument("--num-inference-steps", type=int, default=32)
-    parser.add_argument("--guidance-scale", type=float, default=3.0)
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
-    )
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
-    add_inference_args(parser)
-    parser.add_argument("--output", "-o", type=str, default="output.wav")
-    args = parser.parse_args()
-    return args
-
-
-def main() -> None:
-    args = parse_args()
-    kind = resolve_soulx_kind(args.model)
-
-    if not (args.svc ^ (kind == "svs")):
-        raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
-
-    required = {
-        "prompt_audio": args.prompt_audio,
-        "target_audio": args.target_audio,
-    }
-    if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
-        raise ValueError(
-            "SVS precomputed metadata requires both --prompt-metadata-path and "
-            "--target-metadata-path (or neither, for online preprocess)."
-        )
-    if args.prompt_metadata_path:
-        required["prompt_metadata_path"] = args.prompt_metadata_path
-        required["target_metadata_path"] = args.target_metadata_path
-    if args.prompt_wav_path:
-        required["prompt_wav_path"] = args.prompt_wav_path
-    _require_paths(required)
-    preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
-    deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
-    mode = kind.upper()
-
-    print(f"Loading SoulX-Singer {mode} from {args.model}")
-    print(f"  deploy: {deploy_config}")
-    print(f"  prompt: {args.prompt_audio}")
-    print(f"  target: {args.target_audio}")
-
-    omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False)  # SoulX-Singer is batch-only
-    sampling = build_sampling(
-        args,
-        preprocess_weights_dir=preprocess_dir,
-        kind=kind,
-    )
-    prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
-
-    print("Running generate (inline preprocess + DiT)...")
-    outputs = list(omni.generate([prompt], sampling))
-    if not outputs:
-        raise RuntimeError("No output from omni.generate")
-
-    err = getattr(outputs[0], "error", None)
-    if err is not None:
-        raise RuntimeError(f"Pipeline error: {err}")
-
-    audio_np, sr = extract_audio(outputs)
-    sf.write(args.output, audio_np, sr)
-    print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
-    omni.close()
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/x_to_text/README.md

@@ -10,12 +10,10 @@ 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
-and MammothModa2. BAGEL uses its registered default deploy.
+model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
+BAGEL uses its registered default deploy.
 
 ## Text-To-Text
 
@@ -35,22 +33,6 @@ 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
@@ -72,16 +54,6 @@ 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 |
@@ -102,4 +74,4 @@ the Dev checkpoint.
 
 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 and MammothModa2 use their repository deploy defaults unless overridden.
+BAGEL uses its repository deploy default unless overridden.

examples/offline_inference/x_to_text/x_to_text.py

@@ -15,7 +15,6 @@ 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:
@@ -62,7 +61,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 in {"hunyuan_image3", "mammoth_moda2"},
+        "trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
         "log_stats": args.log_stats,
     }
     if args.enforce_eager is not None:
@@ -71,8 +70,6 @@ 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/audiox/README.md

@@ -1,64 +0,0 @@
-# AudioX online serving
-
-Launches the `AudioXPipeline` behind vLLM-Omni's OpenAI-compatible chat endpoint and provides a
-minimal Python client that covers all six tasks (`t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`).
-
-## Start the server
-
-```bash
-cd examples/online_serving/audiox
-bash run_server.sh                 # defaults: MODEL=zhangj1an/AudioX, PORT=8099
-```
-
-Environment overrides: `MODEL`, `PORT`, `DIFFUSION_ATTENTION_BACKEND`.
-
-## Call from Python
-
-```bash
-# text-to-audio
-python openai_chat_client.py --task t2a \
-    --prompt "Fireworks burst twice, followed by a period of silence before a clock begins ticking." \
-    --output t2a.wav
-
-# text-to-music
-python openai_chat_client.py --task t2m \
-    --prompt "Uplifting ukulele tune for a travel vlog" \
-    --output t2m.wav
-
-# video-to-audio (no text)
-python openai_chat_client.py --task v2a --video path/to/clip.mp4 --output v2a.wav
-
-# text+video-to-audio
-python openai_chat_client.py --task tv2a \
-    --prompt "drum beating sound and human talking" \
-    --video path/to/clip.mp4 \
-    --output tv2a.wav
-```
-
-The client sends:
-
-- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`,
-  `seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server
-  routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's
-  `sampling_params.extra_args`.
-- For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files)
-
-## curl
-
-```bash
-curl -sS -X POST http://localhost:8099/v1/chat/completions \
-  -H 'Content-Type: application/json' \
-  -d '{
-    "model": "zhangj1an/AudioX",
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
-    "extra_body": {
-      "num_inference_steps": 250,
-      "guidance_scale": 7.0,
-      "seed": 42,
-      "audiox_task": "t2m",
-      "seconds_total": 10.0,
-      "sigma_min": 0.3,
-      "sigma_max": 500.0
-    }
-  }' > t2m.json
-```

examples/online_serving/audiox/openai_chat_client.py

@@ -1,119 +0,0 @@
-#!/usr/bin/env python3
-"""AudioX OpenAI-compatible chat client.
-
-AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the
-chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content
-item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed)
-are sent as flattened keys under ``extra_body``; the server routes the declared knobs into
-``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``).
-
-Usage:
-  python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav
-  python openai_chat_client.py --task tv2a --prompt "drum beating" --video clip.mp4 -o tv2a.wav
-"""
-
-from __future__ import annotations
-
-import argparse
-import base64
-import io
-import mimetypes
-import sys
-from pathlib import Path
-
-import requests
-import soundfile
-import torch
-
-VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
-TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"})
-
-
-def _to_data_url(path: str) -> str:
-    mime, _ = mimetypes.guess_type(path)
-    mime = mime or "video/mp4"
-    with open(path, "rb") as f:
-        data = base64.b64encode(f.read()).decode("ascii")
-    return f"data:{mime};base64,{data}"
-
-
-def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
-    audio = audio.to(torch.float32)
-    audio = audio / audio.abs().max().clamp(min=1e-8)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    # soundfile expects channels-last (T, C); project convention is (C, T).
-    soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
-
-
-def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
-    for choice in body.get("choices", []):
-        audio_obj = choice.get("message", {}).get("audio")
-        if not (isinstance(audio_obj, dict) and audio_obj.get("data")):
-            continue
-        data, sr = soundfile.read(io.BytesIO(base64.b64decode(audio_obj["data"])), dtype="float32", always_2d=True)
-        return torch.from_numpy(data).transpose(0, 1), sr
-    brief = {k: v for k, v in body.items() if k != "choices"}
-    raise RuntimeError(f"no audio in response message.audio: {brief}")
-
-
-def main() -> int:
-    p = argparse.ArgumentParser(description="AudioX OpenAI chat client")
-    p.add_argument("--task", required=True, choices=["t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m"])
-    p.add_argument("--prompt", "-p", default="", help="Text prompt (required for t2*/tv2*).")
-    p.add_argument("--video", help="Video path or URL (required for v2*/tv2*).")
-    p.add_argument("--output", "-o", default="audiox_out.wav")
-    p.add_argument("--server", "-s", default="http://localhost:8099")
-    p.add_argument("--model", default="zhangj1an/AudioX")
-    p.add_argument("--steps", type=int, default=250)
-    p.add_argument("--guidance-scale", type=float, default=7.0)
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument("--seconds-total", type=float, default=10.0)
-    p.add_argument("--seconds-start", type=float, default=0.0)
-    p.add_argument("--sigma-min", type=float, default=0.03)
-    p.add_argument("--sigma-max", type=float, default=1000.0)
-    args = p.parse_args()
-
-    if args.task in VIDEO_TASKS and not args.video:
-        print(f"ERROR: task {args.task!r} requires --video", file=sys.stderr)
-        return 2
-    if args.task in TEXT_TASKS and not args.prompt.strip() and args.task not in {"v2a", "v2m"}:
-        print(f"ERROR: task {args.task!r} requires --prompt", file=sys.stderr)
-        return 2
-
-    content: list[dict] = [{"type": "text", "text": args.prompt}]
-    if args.task in VIDEO_TASKS:
-        vurl = args.video if args.video.startswith(("http://", "https://")) else _to_data_url(args.video)
-        content.append({"type": "video_url", "video_url": {"url": vurl}})
-
-    payload = {
-        "model": args.model,
-        "messages": [{"role": "user", "content": content}],
-        "extra_body": {
-            "num_inference_steps": args.steps,
-            "guidance_scale": args.guidance_scale,
-            "seed": args.seed,
-            "audiox_task": args.task,
-            "seconds_start": args.seconds_start,
-            "seconds_total": args.seconds_total,
-            "sigma_min": args.sigma_min,
-            "sigma_max": args.sigma_max,
-        },
-    }
-
-    print(f"POST {args.server}/v1/chat/completions  task={args.task} steps={args.steps}")
-    r = requests.post(
-        f"{args.server}/v1/chat/completions",
-        headers={"Content-Type": "application/json"},
-        json=payload,
-        timeout=600,
-    )
-    r.raise_for_status()
-    audio, sr = _decode_audio_from_response(r.json())
-    _save_wav(audio, Path(args.output), sr)
-    dur = audio.shape[-1] / sr
-    print(f"saved {args.output}  sr={sr}Hz  duration={dur:.2f}s  channels={audio.shape[0]}")
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

examples/online_serving/audiox/run_server.sh

@@ -1,16 +0,0 @@
-#!/bin/bash
-# AudioX online serving startup script.
-
-MODEL="${MODEL:-zhangj1an/AudioX}"
-PORT="${PORT:-8099}"
-DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-FLASH_ATTN}"
-
-echo "Starting AudioX server..."
-echo "Model: $MODEL"
-echo "Port: $PORT"
-echo "Diffusion attention backend: $DIFFUSION_ATTENTION_BACKEND"
-
-DIFFUSION_ATTENTION_BACKEND="$DIFFUSION_ATTENTION_BACKEND" \
-    vllm serve "$MODEL" --omni \
-        --model-class-name AudioXPipeline \
-        --port "$PORT"

examples/online_serving/text_to_speech/README.md

@@ -25,7 +25,6 @@ For the full list of supported architectures across all modalities, see
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
-| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
 
 CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
 
@@ -715,65 +714,3 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
 - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
 - Voice cloning is gated upstream and may require a recent `mistral_common`.
 - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
-
----
-
-## SoulX-Singer
-
-Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
-
-### Prerequisites
-
-Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
-
-### Launch
-
-```bash
-# SVS (default)
-export MODEL=/path/to/SoulX-Singer
-export PREPROCESS=/path/to/SoulX-Singer-Preprocess
-bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
-
-# SVC
-export MODE=svc
-export MODEL=/path/to/SoulX-Singer-svc
-bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
-```
-
-Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
-```bash
-export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
-vllm serve $MODEL --omni \
-    --deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
-    --port 8192 --trust-remote-code --enforce-eager
-```
-
-### Sending requests
-
-Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
-
-```bash
-# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
-python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
-    --prompt-audio /path/on/server/zh_prompt.mp3 \
-    --target-audio /path/on/server/music.mp3 \
-    --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
-    -o output.wav
-```
-
-Use precomputed metadata to skip online preprocess with following command:
-```bash
-python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
-    --prompt-metadata-path /path/on/server/zh_prompt.json \
-    --target-metadata-path /path/on/server/music.json \
-    --audio-path /path/on/server/zh_prompt.mp3 \
-    -o output.wav
-```
-
-`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
-
-### Notes
-
-- Output: 24 kHz mono WAV; batch only.
-- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
-- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.

examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py

@@ -1,186 +0,0 @@
-#!/usr/bin/env python3
-"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
-
-Sends prompt audio via ``input_audio`` and target accompaniment via
-``extra_args['target_audio']`` (server-local path). For integrated preprocess,
-also pass ``preprocess_weights_dir`` in ``extra_args``.
-
-Usage:
-  python openai_chat_client.py \\
-      --prompt-audio /path/on/server/zh_prompt.mp3 \\
-      --target-audio /path/on/server/music.mp3 \\
-      --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
-      -o output.wav
-"""
-
-from __future__ import annotations
-
-import argparse
-import base64
-import io
-import sys
-from pathlib import Path
-
-import requests
-import soundfile
-import torch
-
-
-def _audio_to_data_url(path: Path) -> str:
-    with path.open("rb") as handle:
-        data = base64.b64encode(handle.read()).decode("ascii")
-    return f"data:audio/mpeg;base64,{data}"
-
-
-def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
-    audio = audio.to(torch.float32)
-    peak = audio.abs().max().clamp(min=1e-8)
-    audio = audio / peak
-    path.parent.mkdir(parents=True, exist_ok=True)
-    soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
-
-
-def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
-    for choice in body.get("choices", []):
-        audio_obj = choice.get("message", {}).get("audio")
-        if isinstance(audio_obj, dict) and audio_obj.get("data"):
-            data, sr = soundfile.read(
-                io.BytesIO(base64.b64decode(audio_obj["data"])),
-                dtype="float32",
-                always_2d=True,
-            )
-            return torch.from_numpy(data).transpose(0, 1), sr
-    brief = {k: v for k, v in body.items() if k != "choices"}
-    raise RuntimeError(f"no audio in response message.audio: {brief}")
-
-
-def main() -> int:
-    repo_root = Path(__file__).resolve().parents[4]
-    default_assets = repo_root / "tests" / "assets" / "soulxsinger"
-
-    parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
-    parser.add_argument("--port", type=int, default=8192)
-    parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
-    parser.add_argument(
-        "--prompt-audio",
-        default=str(default_assets / "zh_prompt.mp3"),
-        help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
-    )
-    parser.add_argument(
-        "--target-audio",
-        default=str(default_assets / "music.mp3"),
-        help="Target accompaniment path on the server (extra_args['target_audio'])",
-    )
-    parser.add_argument(
-        "--prompt-metadata-path",
-        default=None,
-        help="SVS precomputed prompt metadata.json",
-    )
-    parser.add_argument(
-        "--target-metadata-path",
-        default=None,
-        help="SVS precomputed target metadata.json",
-    )
-    parser.add_argument(
-        "--audio-path",
-        default=None,
-        help="SVS prompt vocal wav for precomputed metadata",
-    )
-    parser.add_argument("--preprocess-weights-dir", default=None)
-    parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
-    parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
-    parser.add_argument("--language", default="Mandarin")
-    parser.add_argument("--num-inference-steps", type=int, default=32)
-    parser.add_argument("--guidance-scale", type=float, default=3.0)
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Optional CFM seed. Omit for non-deterministic sampling.",
-    )
-    parser.add_argument(
-        "--auto-shift",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Auto pitch shift (default: on, original upstream infer.sh)",
-    )
-    parser.add_argument(
-        "--control",
-        default="melody",
-        choices=["melody", "score"],
-        help="SVS control mode",
-    )
-    parser.add_argument("--vocal-sep", action="store_true")
-    args = parser.parse_args()
-
-    meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
-    if any(meta_paths) and not all(meta_paths):
-        print(
-            "ERROR: precomputed metadata requires --prompt-metadata-path, "
-            "--target-metadata-path, and --audio-path together.",
-            file=sys.stderr,
-        )
-        return 2
-
-    extra_args: dict = {
-        "vocal_sep": args.vocal_sep,
-        "auto_shift": args.auto_shift,
-        "pitch_shift": 0,
-    }
-    if all(meta_paths):
-        extra_args.update(
-            {
-                "prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
-                "target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
-                "audio_path": str(Path(args.audio_path).expanduser().resolve()),
-            }
-        )
-        content = [{"type": "text", "text": "soulx-singer"}]
-    else:
-        prompt_path = Path(args.prompt_audio).expanduser().resolve()
-        if not prompt_path.is_file():
-            print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
-            return 2
-        extra_args["prompt_audio"] = str(prompt_path)
-        extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
-        if args.preprocess_weights_dir:
-            extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
-        content = [
-            {"type": "text", "text": "soulx-singer"},
-            {
-                "type": "input_audio",
-                "input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
-            },
-        ]
-    if not args.svc:
-        extra_args["language"] = args.language
-        extra_args["control"] = args.control
-
-    payload = {
-        "model": args.model,
-        "modalities": ["audio"],
-        "messages": [{"role": "user", "content": content}],
-        "num_inference_steps": args.num_inference_steps,
-        "guidance_scale": args.guidance_scale,
-        "extra_args": extra_args,
-    }
-    if args.seed is not None:
-        payload["seed"] = args.seed
-
-    print(f"POST http://localhost:{args.port}/v1/chat/completions")
-    response = requests.post(
-        f"http://localhost:{args.port}/v1/chat/completions",
-        headers={"Content-Type": "application/json"},
-        json=payload,
-        timeout=1800,
-    )
-    response.raise_for_status()
-    audio, sample_rate = _decode_audio_from_response(response.json())
-    _save_wav(audio, Path(args.output), sample_rate)
-    duration = audio.shape[-1] / sample_rate
-    print(f"saved {args.output}  sr={sample_rate}Hz  duration={duration:.2f}s")
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

examples/online_serving/text_to_speech/soulxsinger/run_server.sh

@@ -1,39 +0,0 @@
-#!/bin/bash
-# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
-#
-# Usage:
-#   MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
-#
-# Audio paths in client extra_args must be readable on the server host.
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
-
-MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
-MODE="${MODE:-svs}"
-PORT="${PORT:-8192}"
-GPUS="${GPUS:-0}"
-
-if [[ "$MODE" == "svc" ]]; then
-  DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
-else
-  DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
-fi
-
-echo "Starting SoulX-Singer server"
-echo "  MODEL=$MODEL"
-echo "  MODE=$MODE"
-echo "  PORT=$PORT"
-echo "  DEPLOY_CONFIG=$DEPLOY_CONFIG"
-echo "  CUDA_VISIBLE_DEVICES=$GPUS"
-
-CUDA_VISIBLE_DEVICES="$GPUS" \
-vllm serve "$MODEL" \
-    --omni \
-    --deploy-config "$DEPLOY_CONFIG" \
-    --host 0.0.0.0 \
-    --port "$PORT" \
-    --trust-remote-code \
-    --enforce-eager

pyproject.toml

@@ -135,13 +135,6 @@ docs = [
     "pydantic",
 ]
 
-soulx-svs = [
-    "funasr>=1.0.0",
-    "g2pM>=0.1.2.5",
-    "g2p-en>=2.1.0",
-    "ToJyutping>=3.2.0",
-]
-
 [project.urls]
 Homepage = "https://github.com/vllm-project/vllm-omni"
 Repository = "https://github.com/vllm-project/vllm-omni"
@@ -192,6 +185,7 @@ select = [
     "N",      # pep8-naming
     "UP",     # pyupgrade
     "TID251", # flake8-tidy-imports.banned-api
+    "TC004",  # flake8-type-checking: runtime-used imports
 ]
 ignore = [
     "E203",  # whitespace before ':' (conflicts with black)

recipes/HiDream-ai/HiDream-O1-Image.md

@@ -0,0 +1,113 @@
+# HiDream-O1-Image
+
+> HiDream-O1 text-to-image generation through the shared offline example
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-O1-Image`
+- Task: Text-to-image
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to generate native-resolution images with HiDream-O1 on one
+or two CUDA GPUs. The two-GPU command shards the Qwen3-VL attention and MLP
+weights, as well as the timestep and pixel projections.
+
+## References
+
+- Upstream model:
+  [`HiDream-ai/HiDream-O1-Image`](https://huggingface.co/HiDream-ai/HiDream-O1-Image)
+- Related offline example:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- Related issue:
+  [`vllm-project/vllm-omni#3479`](https://github.com/vllm-project/vllm-omni/issues/3479)
+
+## Hardware Support
+
+## GPU
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-O1-Image \
+  --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+  --height 2048 \
+  --width 2048 \
+  --num-inference-steps 50 \
+  --guidance-scale 5.0 \
+  --seed 42 \
+  --output hidream_o1_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- The pipeline supports BF16 text-to-image generation.
+- Requested sizes are matched to the closest resolution supported by the checkpoint.
+
+### 2x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-O1-Image \
+  --tensor-parallel-size 2 \
+  --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+  --height 2048 \
+  --width 2048 \
+  --num-inference-steps 50 \
+  --guidance-scale 5.0 \
+  --seed 42 \
+  --output hidream_o1_tp2_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_tp2_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- Tensor parallelism shards the attention, MLP, timestep, and pixel projection weights across both GPUs.
+- Image editing and image-conditioned generation are not supported.

recipes/MammothModa2/MammothModa2.md

@@ -1,212 +0,0 @@
-# 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-4090.md

@@ -110,9 +110,10 @@ outputs of each pair agreed to within 115 bytes.
 Four-GPU T2VA repeats: 273.8 s and 269.5 s, peaking at 15,560 MiB and
 15,612 MiB. Four-GPU Ref2VA is a single successful end-to-end run (545.2 s,
 16,448 MiB). Follow-up Ref2VA requests on this build completed diffusion but
-then hit the engine's hardcoded 30 s async output wait
-(`_ASYNC_OUTPUT_TIMEOUT` in `diffusion_engine.py`) during post-compute D2H,
-so they are not reported as repeats.
+then hit the engine's async output wait during post-compute D2H, so they are
+not reported as repeats. That wait was a hardcoded 30 s when these numbers were
+taken; it is now `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` (default 600 s), so a rerun on
+current main should not lose these repeats to the bound.
 
 Relative to the two-GPU baseline on the same shape, four-GPU TP2×USP2 was about
 1.6× faster for both tasks, while per-GPU peak HBM stayed in the mid-15 GiB

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -118,6 +118,21 @@ On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
 above) or HSDP — see
 [§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
 
+### CPU MP4 response encoding (Atlas A2)
+
+Non-streaming MP4 responses use one public automatic encoder. It checks the
+runtime frame shape, common dtype, and RGB channel-plane contiguity for every
+request; compatible inputs use direct planar PyAV frames, while unsupported
+inputs fall back to the legacy muxer before opening the PyAV container. No CLI
+flag, model declaration, or user configuration is required. Streaming fMP4
+output remains on its existing incremental path.
+
+Current performance and correctness validation is limited to one Atlas A2 host
+with 8x Ascend 910B4-1 NPUs, the `FL2VA`/`t2va` partition, one request at a
+time, 1344x768 at 24 fps, and 5, 8.7, and 15 second requests. This optimization
+only changes CPU MP4 response encoding; it does not change DiT execution or
+stage 0.
+
 ### Optional optimizations
 
 Two independent optimizations may be enabled on top of the configuration

recipes/MiniMaxAI/MiniMax-H3.md

@@ -65,6 +65,31 @@ On AMD ROCm, install without the `[fa4]` extra (FA4 is CUDA-only) and use
 `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
 reference-video preparation and MP4 output.
 
+## CPU MP4 response encoding
+
+For CUDA and ROCm deployments, non-streaming MP4 responses are encoded on the
+host CPU through PyAV/libx264 after generation. The response encoder selects the
+path automatically at runtime: inputs with a supported frame shape, common dtype,
+and RGB per-channel contiguous layout use direct planar frames; other inputs
+fall back to the legacy path before the PyAV container is opened. No CLI flag,
+model declaration, or user configuration is required. Streaming fMP4 output is
+unchanged.
+
+A community benchmark on 2x Xeon 8480C reported the following comparison
+between the legacy and direct planar paths
+([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)):
+
+| Metric | Legacy | Direct planar | Change |
+| --- | ---: | ---: | ---: |
+| Median wall time | 1.805 s | 1.394 s | -22.8% |
+| Median process CPU time | 3.613 s | 3.207 s | -11.2% |
+| Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) |
+
+Across the 1.0-8.7 s sweep, wall-time improvement was approximately 21.8-22.6%;
+outputs were byte-identical and full decode passed. This is evidence for host
+CPU response encoding. Actual gains depend on the CPU and runtime, and should
+not be interpreted as GPU, DiT, or stage 0 speedups.
+
 ## Start a server
 
 Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery
@@ -158,7 +183,7 @@ vllm serve "${MODEL}" \
 Use the profile that matches the per-GPU memory capacity:
 
 | Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
-|---|---:|---:|---:|---|---|---|
+| --- | ---: | ---: | ---: | --- | --- | --- |
 | `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
 | `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
 
@@ -177,8 +202,8 @@ increasing it on a different request shape.
 At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
 2 x RTX 5090 without OOM:
 
-| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
-|---:|---:|---:|---:|---|
+| Shape    | Frames        | Client E2E | Sampled peak/GPU       | Output validation                                            |
+| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: |
 | 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
 
 This is a single end-to-end validation run, not a warmed multi-run latency
@@ -247,7 +272,7 @@ vllm serve "${MODEL}" \
 Do not add `--enforce-eager` to this performance configuration. The first
 request includes regional compilation; warm the server once before measuring
 steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
-1. The H3 VAE supports its native `tile` mode, not
+`1`. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
 ### Attention Backends
@@ -505,7 +530,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
 `0.26.0+rocm723` wheel (HIP 7.2).
 
 | Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
+| ---------- | --------------- | ----------------- |
 | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
@@ -687,7 +712,7 @@ at most 15 seconds combined.
 ## Official input matrix and limits
 
 | Task | Supported references | Limits |
-|------|----------------------|--------|
+| ------ | ---------------------- | -------- |
 | T2VA | text only | prompt must be non-empty |
 | FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
 | Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
@@ -718,13 +743,17 @@ Omitting `quality` preserves the startup default: it uses the reference path
 normally, or the server-configured profile when the server was started with
 `--cache-backend cache_dit`.
 
+Turbo is independent of this quality switch: `quality` selects a Cache-DiT
+policy, while Turbo changes the active LoRA weights and sampling schedule. See
+[LoRA](#lora) below.
+
 The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
 TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
 `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
 balanced switch order.
 
 | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
-|---|---:|---:|---:|---:|---|
+| --- | ---: | ---: | ---: | ---: | --- |
 | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
 | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
 
@@ -734,10 +763,49 @@ balanced switch order.
 > workload. The values above apply to this deployment and are not universal
 > guarantees. `lossless` remains the exact reference path.
 
+## LoRA
+
+### Turbo LoRA
+
+Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
+
+```text
+lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+```
+
+Download only that file:
+
+```bash
+export TURBO_DIR=/path/to/minimax-h3-turbo
+export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+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
+`--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:
+
+```bash
+-F 'num_inference_steps=5' \
+-F 'flow_shift=6' \
+-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
+-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
+```
+
+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.
+The five requested sigma points produce the four denoiser evaluations expected
+by the Turbo artifact.
+
 ## Key parameters
 
 | Parameter | Recommended value | Notes |
-|-----------|-------------------|-------|
+| ----------- | ------------------- | ------- |
 | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
 | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
 | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
@@ -762,9 +830,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi
 The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
 excluded warmup followed by three requests.
 
-| Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
-| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
+| Workload                               | Configuration                               | Observed result                      |
+| -------------------------------------- | ------------------------------------------- | ------------------------------------ |
+| FL2VA, 209 frames, 1248x768            | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency    |
 | Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
 
 These measurements describe the validated shapes rather than a general

recipes/README.md

@@ -39,7 +39,6 @@ recipes/
 
 | Recipe | Task | Hardware |
 |--------|------|----------|
-| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
 | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
 | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ |
@@ -47,6 +46,7 @@ recipes/
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
+| [`HiDream-ai/HiDream-O1-Image.md`](./HiDream-ai/HiDream-O1-Image.md) | Text-to-image with shared offline example | 1x or 2x H100 80GB |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
@@ -57,7 +57,6 @@ recipes/
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
 | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
-| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |

recipes/audiox/AudioX.md

@@ -1,130 +0,0 @@
-# AudioX
-
-> AudioX MMDiT for unified audio + music generation: t2a / t2m / v2a / v2m / tv2a / tv2m
-
-## Summary
-
-- Vendor: HKUSTAudio (project), `zhangj1an/AudioX` weight bundle
-- Model: `zhangj1an/AudioX`
-- Task: Text/video → audio or music. Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`.
-- Mode: Offline inference + online serving (pure diffusion)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe to run AudioX for sound-effect (`*2a`) or music (`*2m`) generation
-from a text prompt and/or video clip. AudioX is a unified diffusion transformer
-that produces stereo 44.1 kHz audio up to ~10 s per call.
-
-## References
-
-- Project page: <https://zeyuet.github.io/AudioX/>
-- vLLM-Omni weight bundle: <https://huggingface.co/zhangj1an/AudioX>
-- Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline`
-- Input transforms: `vllm_omni.transformers_utils.processors.audiox`
-- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs)
-- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py)
-- Online: standard OpenAI chat-completions endpoint (see commands below)
-
-## Hardware Support
-
-## GPU
-
-### 1x L4 24GB
-
-#### Environment
-
-- OS: Ubuntu 22.04
-- Python: 3.10+
-- Driver / runtime: CUDA 12.4
-- vLLM version: 0.20.0
-- vLLM-Omni version: 0.1.x
-
-#### Command
-
-AudioX uses the **standard** `text_to_audio` example offline and the standard
-OpenAI chat-completions endpoint online. All model-specific knobs
-(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`,
-`cfg_rescale`, `video_path`, `audio_path`) are declared in
-`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) /
-`--extra-body` (offline). The offline example also exposes `--task`, `--video`,
-`--audio-start`, and `--audio-length` shortcuts for the most common knobs.
-
-Offline — text tasks (`t2a` / `t2m`):
-
-```bash
-huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
-python examples/offline_inference/text_to_audio/text_to_audio.py \
-  --model ./audiox_weights --task t2a \
-  --prompt "Fireworks burst twice, followed by a clock ticking." \
-  --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
-  --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
-  --output t2a.wav
-```
-
-Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`:
-
-```bash
-python examples/offline_inference/text_to_audio/text_to_audio.py \
-  --model ./audiox_weights --task tv2a \
-  --prompt "drum beating sound and human talking" \
-  --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
-  --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
-  --output tv2a.wav
-```
-
-Online:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \
-  vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099
-
-curl -sS -X POST http://localhost:8099/v1/chat/completions \
-  -H 'Content-Type: application/json' \
-  -d '{
-    "model": "zhangj1an/AudioX",
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
-    "extra_body": {
-      "num_inference_steps": 250,
-      "guidance_scale": 7.0,
-      "seed": 42,
-      "audiox_task": "t2m",
-      "seconds_total": 10.0,
-      "sigma_min": 0.3,
-      "sigma_max": 500.0
-    }
-  }' > t2m.json
-```
-
-For `v2*` / `tv2*` online, attach the video as a `video_url` content item
-(data URI for local files) in the message `content`.
-
-#### Verification
-
-```bash
-# Health check
-curl http://localhost:8099/health
-
-# Listen to the saved file (stereo, 44.1 kHz, sigma_min=0.03, sigma_max=1000 — upstream defaults)
-ffprobe t2a.wav
-```
-
-#### Notes
-
-- Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio.
-- Output rate: 44.1 kHz stereo regardless of requested rate.
-- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the
-  `audiox_task` knob (in `extra_body` online, or `--task` offline).
-- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video`
-  (→ `video_path` knob), online via a `video_url` content item.
-- Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`).
-- Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with
-  `QKVParallelLinear`); cross-attention K/V is also TP-sharded.
-
-### Known limitations
-
-- Inference uses an inlined DPM-Solver++(3M) SDE sampler (k-diffusion port). Replacing it with
-  diffusers' `EDMDPMSolverMultistepScheduler` introduces a fixed ~861 Hz resonance and is not
-  recommended.
-- Generation is fixed at 10 s (configured by the bundle's `sample_size`); longer outputs require
-  a different bundle.

requirements/common.txt

@@ -12,7 +12,7 @@
 transformers >= 5.10.1, < 5.15
 av>=14.0.0
 omegaconf>=2.3.0
-diffusers==0.38.0
+diffusers==0.40.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1

tests/assets/soulxsinger/music.json

@@ -1,16 +0,0 @@
-[
-  {
-    "index": "vocal_240_51240",
-    "language": "Mandarin",
-    "time": [
-      240,
-      51240
-    ],
-    "duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
-    "text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>",
-    "phoneme": "<SP> zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>",
-    "note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0",
-    "note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1",
-    "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.Line truncated
-  }
-]

tests/assets/soulxsinger/zh_prompt.json

@@ -1,16 +0,0 @@
-[
-  {
-    "index": "vocal_320_10687",
-    "language": "Mandarin",
-    "time": [
-      320,
-      10687
-    ],
-    "duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
-    "text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意",
-    "phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4",
-    "note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0",
-    "note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2",
-    "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
-  }
-]

tests/benchmarks/benchmark_video_response_encoding.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark legacy and automatic MP4 response encoding on synthetic input."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import statistics
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from vllm_omni.entrypoints.openai.video_api_utils import (
+    _encode_video_bytes,
+    _encode_video_bytes_legacy,
+)
+
+
+@dataclass(frozen=True)
+class BenchmarkVariant:
+    label: str
+    encoder: Callable[..., bytes]
+
+
+def _positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed < 1:
+        raise argparse.ArgumentTypeError("value must be at least 1")
+    return parsed
+
+
+def _build_inputs(
+    *,
+    frames: int,
+    height: int,
+    width: int,
+    fps: int,
+    audio_sample_rate: int,
+    seed: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    planar = rng.random((3, frames, height, width), dtype=np.float32)
+    video = planar.transpose(1, 2, 3, 0)
+    if not all(video[..., channel].flags.c_contiguous for channel in range(3)):
+        raise RuntimeError("synthetic video does not expose contiguous channel planes")
+
+    sample_count = max(1, round(frames / fps * audio_sample_rate))
+    timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
+    left = np.sin(2 * np.pi * 220 * timeline).astype(np.float32, copy=False)
+    right = np.sin(2 * np.pi * 330 * timeline).astype(np.float32, copy=False)
+    return video, np.stack((left, right))
+
+
+def _measure(
+    variant: BenchmarkVariant,
+    *,
+    video: np.ndarray,
+    audio: np.ndarray,
+    fps: int,
+    audio_sample_rate: int,
+) -> tuple[bytes, dict[str, object]]:
+    cpu_start = time.process_time_ns()
+    wall_start = time.perf_counter_ns()
+    output = variant.encoder(
+        video,
+        fps=fps,
+        audio=audio,
+        audio_sample_rate=audio_sample_rate,
+        video_codec_options={"preset": "ultrafast", "threads": "0"},
+    )
+    wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000
+    process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000
+    return output, {
+        "label": variant.label,
+        "wall_ms": wall_ms,
+        "process_cpu_ms": process_cpu_ms,
+        "output_bytes": len(output),
+        "output_sha256": hashlib.sha256(output).hexdigest(),
+    }
+
+
+def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]:
+    selected = [record for record in records if record["label"] == label]
+    wall_values = [float(record["wall_ms"]) for record in selected]
+    cpu_values = [float(record["process_cpu_ms"]) for record in selected]
+    return {
+        "runs": len(selected),
+        "wall_ms": {
+            "median": statistics.median(wall_values),
+            "min": min(wall_values),
+            "max": max(wall_values),
+        },
+        "process_cpu_ms": {
+            "median": statistics.median(cpu_values),
+            "min": min(cpu_values),
+            "max": max(cpu_values),
+        },
+    }
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--frames", type=_positive_int, default=12)
+    parser.add_argument("--height", type=_positive_int, default=96)
+    parser.add_argument("--width", type=_positive_int, default=160)
+    parser.add_argument("--fps", type=_positive_int, default=24)
+    parser.add_argument("--audio-sample-rate", type=_positive_int, default=32000)
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--rounds", type=_positive_int, default=3)
+    parser.add_argument("--seed", type=int, default=20260817)
+    parser.add_argument("--output-json", type=Path)
+    args = parser.parse_args()
+    if args.warmup < 0:
+        parser.error("--warmup must be non-negative")
+    return args
+
+
+def main() -> None:
+    args = _parse_args()
+    video, audio = _build_inputs(
+        frames=args.frames,
+        height=args.height,
+        width=args.width,
+        fps=args.fps,
+        audio_sample_rate=args.audio_sample_rate,
+        seed=args.seed,
+    )
+    baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
+    candidate = BenchmarkVariant("automatic", _encode_video_bytes)
+
+    for _ in range(args.warmup):
+        for variant in (baseline, candidate):
+            _measure(
+                variant,
+                video=video,
+                audio=audio,
+                fps=args.fps,
+                audio_sample_rate=args.audio_sample_rate,
+            )
+
+    records: list[dict[str, object]] = []
+    for round_index in range(args.rounds):
+        order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
+        for variant in order:
+            _, record = _measure(
+                variant,
+                video=video,
+                audio=audio,
+                fps=args.fps,
+                audio_sample_rate=args.audio_sample_rate,
+            )
+            record["round"] = round_index + 1
+            records.append(record)
+
+    output_hashes = {str(record["output_sha256"]) for record in records}
+    if len(output_hashes) != 1:
+        raise RuntimeError(f"benchmark variants produced different outputs: {sorted(output_hashes)}")
+
+    paired_savings_ms = []
+    paired_savings_percent = []
+    for round_index in range(args.rounds):
+        pair = [record for record in records if record["round"] == round_index + 1]
+        baseline_ms = float(next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
+        candidate_ms = float(next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
+        paired_savings_ms.append(baseline_ms - candidate_ms)
+        paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
+
+    result = {
+        "config": {
+            "frames": args.frames,
+            "height": args.height,
+            "width": args.width,
+            "fps": args.fps,
+            "audio_sample_rate": args.audio_sample_rate,
+            "warmup": args.warmup,
+            "rounds": args.rounds,
+            "seed": args.seed,
+            "video_shape": list(video.shape),
+            "video_strides": list(video.strides),
+            "codec_options": {"preset": "ultrafast", "threads": "0"},
+        },
+        "records": records,
+        "summary": {
+            baseline.label: _summarize(records, baseline.label),
+            candidate.label: _summarize(records, candidate.label),
+            "paired_wall_saving_ms": {
+                "median": statistics.median(paired_savings_ms),
+                "min": min(paired_savings_ms),
+                "max": max(paired_savings_ms),
+            },
+            "paired_wall_saving_percent": {
+                "median": statistics.median(paired_savings_percent),
+                "min": min(paired_savings_percent),
+                "max": max(paired_savings_percent),
+            },
+            "output_sha256": output_hashes.pop(),
+        },
+    }
+    output_json = json.dumps(result, indent=2, sort_keys=True)
+    if args.output_json is not None:
+        args.output_json.write_text(output_json + "\n", encoding="utf-8")
+    print(output_json)
+
+
+if __name__ == "__main__":
+    main()

tests/buildkite/test_upload_pipeline.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations
 
 import sys
@@ -61,16 +62,23 @@ def test_bootstrap_steps_loaded_from_file() -> None:
 
 
 def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
-    """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 still gates scheduled CI."""
+    """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 / NON_CRITICAL=1 still gates scheduled CI."""
     rendered = _render(["docs/foo.md"])
     assert "key: image-build" in rendered
     assert "key: upload-nightly-pipeline" in rendered
     assert "key: upload-weekly-pipeline" in rendered
-    # Scheduled nightly also uploads L2/L3 with --e2e
-    assert "key: upload-ready-pipeline" in rendered
-    assert "key: upload-merge-pipeline" in rendered
-    assert 'build.env("NIGHTLY") == "1"' in rendered
+    # Scheduled main+WEEKLY=1 uploads L2/L3 with --e2e; NIGHTLY still gates L4 only.
+    doc = yaml.safe_load(rendered)
+    by_key = {step["key"]: step for step in doc["steps"]}
+    assert "NIGHTLY" not in by_key["upload-ready-pipeline"]["if"]
+    assert 'build.branch == "main"' in by_key["upload-ready-pipeline"]["if"]
+    assert 'build.env("WEEKLY") == "1"' in by_key["upload-ready-pipeline"]["if"]
+    assert "NIGHTLY" not in by_key["upload-merge-pipeline"]["if"]
+    assert 'build.branch == "main"' in by_key["upload-merge-pipeline"]["if"]
+    assert 'build.env("WEEKLY") == "1"' in by_key["upload-merge-pipeline"]["if"]
+    assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
     assert 'build.env("WEEKLY") == "1"' in rendered
+    assert 'build.env("NON_CRITICAL") == "1"' in rendered
     assert "nightly-test" not in rendered
     assert "weekly-test" not in rendered
     assert "merge-test" not in rendered
@@ -78,6 +86,19 @@ def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
     assert "if: false" not in rendered
 
 
+def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None:
+    """NPU skip_all: scheduled NIGHTLY still uploads L4, not L2 ready."""
+    rendered = _render_bootstrap_pipeline(
+        BOOTSTRAP_STEPS_TEMPLATE,
+        decision=resolve_ci_decision(["docs/foo.md"]),
+        path=Path(".buildkite/npu/bootstrap-upload-steps.yml"),
+    )
+    doc = yaml.safe_load(rendered)
+    by_key = {step["key"]: step for step in doc["steps"]}
+    assert "upload-ready-pipeline" not in by_key
+    assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
+
+
 def test_yaml_gated_l45_only_does_not_unconditionally_build_image() -> None:
     rendered = _render([".buildkite/cuda/test-nightly.yml"])
     assert "if: true" not in rendered

tests/config/test_config_factory.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 StageConfigFactory and related classes.
 """
@@ -1267,8 +1267,6 @@ 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),
@@ -1306,6 +1304,7 @@ stages:
             ({"model_type": "step_audio_2"}, None, "step_audio_2"),
             (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
             (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
+            (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"),
         ],
     )
     def test_migrated_models_are_discovered_without_explicit_deploy(
@@ -1661,6 +1660,7 @@ stages:
             ]
         }
         original_foo = fake_config["stages"][0]["foo"]
+        assert isinstance(original_foo, dict)
         original_b = original_foo["b"]
 
         with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
@@ -2128,7 +2128,12 @@ class TestBaseConfigInheritance:
             sampling = deploy.stages[1].default_sampling_params
             assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
             assert sampling["temperature"] == 0.8, filename
-            assert sampling["top_k"] == 100, filename
+            # Upstream utils.TTSSamplingParams. min_tokens is not checked here:
+            # the duplex overlay drops it to 0 because the model budgets each
+            # generate_chunk itself.
+            assert sampling["top_k"] == 25, filename
+            assert sampling["top_p"] == 0.85, filename
+            assert sampling["repetition_penalty"] == 1.05, filename
 
     def test_ci_inherits_from_main(self):
         ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
@@ -2220,6 +2225,28 @@ class TestPlatformOverrides:
         assert rocm.stages[0].enforce_eager is None
         assert rocm.stages[1].enforce_eager is True
 
+    def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
+        pipeline = resolve_pipeline_config(
+            "qwen3_omni_moe",
+            Q3_OMNI_ALL_STAGES_HF_CONFIG,
+        )
+        assert isinstance(pipeline, PipelineConfig)
+
+        cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
+        cuda_stages = merge_pipeline_deploy(pipeline, cuda)
+        expected = {"custom_ops": ["+rotary_embedding"]}
+        assert cuda_stages[0].yaml_engine_args["compilation_config"] == expected
+        assert all("compilation_config" not in stage.yaml_engine_args for stage in cuda_stages[1:])
+
+        for platform in ("cpu", "musa", "npu", "rocm", "xpu"):
+            deploy = _apply_platform_overrides(
+                load_deploy_config(deploy_path),
+                platform=platform,
+            )
+            config = deploy.stages[0].compilation_config or {}
+            assert "+rotary_embedding" not in config.get("custom_ops", [])
+
     def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
         pipeline = resolve_pipeline_config("minicpmo_4_5")
         assert isinstance(pipeline, PipelineConfig)
@@ -2828,13 +2855,14 @@ class TestObjectStorageConfigResolution:
         monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
 
         class Recorder:
+            def __init__(self) -> None:
+                self.pulls = pulls
+
             def set_files(self, files: list[tuple[str, str]]) -> None:
                 materialized_files.clear()
                 materialized_files.extend(files)
 
-        recorder = Recorder()
-        recorder.pulls = pulls
-        return recorder
+        return Recorder()
 
     def test_passthrough_for_non_uri(self, fake_object_storage):
         assert _materialize_object_storage_configs("org/model") == "org/model"

tests/config/test_omni_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for the additive structured Omni config."""
 
 from __future__ import annotations
@@ -1090,7 +1090,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
 
 
 def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
     monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
 
     cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
         diffusion_attention_backend="flash_attn",
@@ -1100,6 +1103,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
         kv_cache_skip_layers=[2],
         static_lora_scale=0.25,
         diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
         diffusers_load_kwargs=None,
         diffusers_call_kwargs=None,
     )
@@ -1116,7 +1120,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     assert cfg.diffusers_call_kwargs == {}
 
 
-def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
+def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
     deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
     deploy_path.write_text(
         "\n".join(
@@ -1128,6 +1135,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                 "    diffusion_attention_backend: flash_attn",
                 "    fa_deterministic: true",
                 "    diffusion_kv_mode: paged_scheduler",
+                "    diffusion_kv_max_rows_per_request: 2",
             ]
         )
     )
@@ -1141,6 +1149,36 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
     assert stage.diffusion_config.fa_deterministic is True
     assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+    assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
+
+
+def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_attention_backend: FASTVIDEO_VSA",
+                "    fastvideo_vsa_topk: 96",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path=str(deploy_path),
+    ).stage_by_id(0)
+
+    attention_config = stage.diffusion_config.diffusion_attention_config
+    assert attention_config.default is not None
+    assert attention_config.default.backend == "FASTVIDEO_VSA"
+    assert attention_config.default.fastvideo_vsa_topk == 96
 
 
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
@@ -1183,6 +1221,7 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "prompt_embed_cache_size",
         "diffusion_kv_cache_dtype",
         "diffusion_kv_mode",
+        "diffusion_kv_max_rows_per_request",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
     assert {
         "revision",

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -33,6 +33,9 @@ def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
     sched.skipped_waiting = set()
+    sched._free_request_blocks = MagicMock()
+    sched.encoder_cache_manager = MagicMock()
+    sched._inflight_prefills = set()
     return sched
 
 
@@ -276,29 +279,33 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
         "meta": {"turn_eos_token_id": 99},
     }
     assert session.status == RequestStatus.WAITING
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
 
 
-def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
+def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_watermark() -> None:
     sched = _make_scheduler(stage_id=1)
+    session = _make_request()
     sched.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=False,
         segment_finished_requests=set(),
+        requests_num_chunks_sent={session.external_req_id: 59},
     )
-    session = _make_request()
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
     session.prompt_token_ids = [0] * 59
     session._all_token_ids.clear()
     session._all_token_ids.extend(session.prompt_token_ids)
     session.num_prompt_tokens = 59
     session.num_computed_tokens = 59
+    session.num_in_flight_tokens = 2
     update = _make_update([0] * 10)
     update.additional_information = None
     update.model_intermediate_buffer = {
         "ids": {"tts": list(range(8))},
         "hidden_states": {"tts": [[0.0]] * 8},
         "meta": {
-            "replace_streaming_prompt": True,
             "next_stage_prompt_len": 10,
+            "replace_streaming_prompt": True,
         },
     }
 
@@ -308,6 +315,42 @@ def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None
     assert list(session._all_token_ids) == [0] * 10
     assert session.num_prompt_tokens == 10
     assert session.num_computed_tokens == 0
+    assert session.num_stale_output_tokens == 2
     assert session.additional_information is None
     assert session.model_intermediate_buffer == update.model_intermediate_buffer
     assert session.status == RequestStatus.WAITING
+    assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
+
+
+def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
+    sched = _make_scheduler(stage_id=1)
+    session = _make_request()
+    session.external_req_id = "external-ar-streaming-test"
+    session.num_in_flight_tokens = 2
+    # _update_request_as_session() may have already fenced this frame before
+    # the connector marks the explicit replacement ready.
+    session.num_stale_output_tokens = 2
+    session.num_output_placeholders = 2
+    session.spec_token_ids = [-1, -1]
+    sched.requests = {session.request_id: session}
+    sched._inflight_prefills.add(session)
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        replaced_streaming_prompt_ids={session.request_id},
+        requests_with_ready_chunks={session.request_id},
+        requests_num_chunks_sent={session.external_req_id: 4090},
+    )
+
+    sched._reset_ready_async_chunk_replacements()
+    sched._reset_ready_async_chunk_replacements()
+
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
+    assert session not in sched._inflight_prefills
+    assert session.num_stale_output_tokens == 2
+    assert session.num_output_placeholders == 0
+    assert session.spec_token_ids == []
+    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 == {}

tests/core/sched/test_omni_generation_scheduler_playback_deadline.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_generation_scheduler_records_audio_output_for_edf(mocker):
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler.chunk_transfer_adapter = SimpleNamespace(
+        playback_deadline_enabled=True,
+        record_audio_output=mocker.Mock(),
+    )
+    multimodal_output = {
+        "model_outputs": torch.zeros(24000),
+        "sr": torch.tensor(24000),
+    }
+
+    scheduler._record_playback_deadline_output("req", multimodal_output)
+
+    scheduler.chunk_transfer_adapter.record_audio_output.assert_called_once_with("req", 24000, 24000)
+
+
+def test_generation_scheduler_skips_audio_accounting_when_edf_is_disabled(mocker):
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler.chunk_transfer_adapter = SimpleNamespace(
+        playback_deadline_enabled=False,
+        record_audio_output=mocker.Mock(),
+    )
+
+    scheduler._record_playback_deadline_output("req", {"model_outputs": torch.zeros(24000)})
+
+    scheduler.chunk_transfer_adapter.record_audio_output.assert_not_called()

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

@@ -10,7 +10,7 @@
           "num_cards": 3
         }
       },
-      "full_model",
+      "slow",
       "omni"
     ],
     "server_params": {

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -10,7 +10,7 @@
           "num_cards": 2
         }
       },
-      "full_model",
+      "slow",
       "omni"
     ],
     "server_params": {

tests/diffusion/attention/test_attention_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
 
@@ -90,6 +90,14 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
+    def test_fastvideo_vsa_topk_serialized(self):
+        spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96)
+        assert spec.backend_kwargs() == {"topk": 96}
+
+    def test_fastvideo_vsa_topk_rejected_for_other_backend(self):
+        with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"):
+            AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96)
+
     def test_block_sparse_defaults_applied_when_backend_selected(self):
         spec = AttentionSpec(backend="RAINFUSION_ATTN")
         assert spec.block_sparse.sparsity == 0.8

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+    FastVideoVSAImpl,
+)
+from vllm_omni.diffusion.attention.backends.registry import (
+    DiffusionAttentionBackendEnum,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_fastvideo_vsa_backend_is_registered():
+    assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
+
+
+def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    fake_module = types.ModuleType("fastvideo_kernel")
+
+    def fake_video_sparse_attn_bshd(
+        q,
+        k,
+        v,
+        variable_block_sizes,
+        q_variable_block_sizes,
+        compress_attn_weight,
+        topk,
+        block_size,
+    ):
+        calls["q_shape"] = tuple(q.shape)
+        calls["vbs"] = variable_block_sizes.detach().cpu().tolist()
+        calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist()
+        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+        calls["compress_shape"] = tuple(compress_attn_weight.shape)
+        calls["topk"] = topk
+        calls["block_size"] = block_size
+        return q + k + v
+
+    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+    query = torch.randn(1, 300, 2, 8)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)})
+    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+    output = impl.forward_cuda(query, key, value, attn_metadata)
+
+    assert output.shape == query.shape
+    assert calls["q_shape"] == (1, 1024, 2, 8)
+    assert calls["vbs"] == [192, 48, 48, 12]
+    assert calls["q_vbs"] == [192, 48, 48, 12]
+    assert calls["compress_shape"] == (1, 1024, 2, 8)
+    assert calls["compress_sum"] == 0.0
+    assert calls["topk"] == 1
+    assert calls["block_size"] == (4, 8, 8)
+
+
+def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch):
+    calls: dict[str, Any] = {}
+    fake_module = types.ModuleType("fastvideo_kernel")
+
+    def fake_video_sparse_attn_bshd(
+        q,
+        k,
+        v,
+        variable_block_sizes,
+        q_variable_block_sizes,
+        compress_attn_weight,
+        topk,
+        block_size,
+    ):
+        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+        calls["compress_shape"] = tuple(compress_attn_weight.shape)
+        return q + k + v
+
+    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+    query = torch.randn(1, 300, 2, 8)
+    gate = torch.ones_like(query)
+    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate})
+    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert output.shape == query.shape
+    assert calls["compress_shape"] == (1, 1024, 2, 8)
+    # Only the 300 real tokens carry ones; padded gate slots stay zero.
+    assert calls["compress_sum"] == gate.numel()
+
+
+def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch):
+    calls: dict[str, Any] = {}
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+
+    def fake_fallback(query, key, value, attn_metadata, reason):
+        calls["reason"] = reason
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+    query = torch.randn(1, 512, 2, 8)
+    output = impl.forward_cuda(query, query, query, AttentionMetadata())
+
+    assert output.shape == query.shape
+    assert calls["reason"] == "vsa_dit_seq_shape metadata is required"
+
+
+def test_fastvideo_vsa_falls_back_for_mask(monkeypatch):
+    calls: dict[str, Any] = {}
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"min_seq_len": 1},
+    )
+
+    def fake_fallback(query, key, value, attn_metadata, reason):
+        calls["reason"] = reason
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+    query = torch.randn(1, 512, 2, 8)
+    mask = torch.ones(1, 512, dtype=torch.bool)
+    output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask))
+
+    assert output.shape == query.shape
+    assert calls["reason"]
+
+
+def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
+    query = torch.randn(1, 512, 2, 8)
+    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)})
+
+    all_blocks = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1},
+    )
+    too_many = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1},
+    )
+
+    # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
+    assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
+    assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"

tests/diffusion/attention/test_piecewise_attn.py

@@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_
 from vllm_omni.diffusion.attention.backends.abstract import QueryRange
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     Segment,
+    build_paged_piecewise_plan,
     build_segments,
     piecewise_attn,
+    run_paged_piecewise_plan,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -84,6 +86,72 @@ def test_build_segments_full_span_retains_global_kv_end():
     assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
 
 
+def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
+    full_attn_spans = [[(2, 5)], [(4, 7)]]
+    plan = build_paged_piecewise_plan(
+        full_attn_spans,
+        query_offsets=[0, 2],
+        query_lens=[6, 6],
+        seq_lens=[6, 8],
+        device=DEVICE,
+    )
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    key = query + 100
+    value = query + 200
+    calls = []
+
+    def run_segment(segment_query, segment_key, segment_value, segment_batch):
+        calls.append(
+            (
+                [segment.mode for segment in segment_batch.row_segments],
+                [segment.kv_end for segment in segment_batch.row_segments],
+                segment_batch.query_indices.tolist(),
+            )
+        )
+        torch.testing.assert_close(segment_key, segment_query + 100)
+        torch.testing.assert_close(segment_value, segment_query + 200)
+        return segment_query
+
+    output = run_paged_piecewise_plan(
+        query,
+        key,
+        value,
+        plan,
+        plan.segments,
+        run_segment,
+    )
+
+    torch.testing.assert_close(output, query)
+    assert calls == [
+        (["causal", "causal"], [2, 4], [0, 1, 6, 7]),
+        (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
+        (["causal", "causal"], [6, 8], [5, 11]),
+    ]
+
+
+def test_paged_piecewise_runner_uses_runner_output_shape():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        lambda segment_query, *_args: segment_query[..., :3],
+    )
+
+    assert output.shape == (6, 2, 3)
+    torch.testing.assert_close(output, query[..., :3])
+
+
 @pytest.mark.parametrize("global_spans", SPAN_CASES)
 @pytest.mark.parametrize("q_range", Q_RANGE_CASES)
 @pytest.mark.parametrize("batch_size", BATCH_CASES)

tests/diffusion/attention/test_ulysses_uaa.py

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import os
-import socket
 import tempfile
 from types import SimpleNamespace
 
@@ -12,6 +11,7 @@ import numpy as np
 import pytest
 import torch
 
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.config import set_current_diffusion_config
@@ -27,20 +27,6 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
-    os.environ["RANK"] = str(rank)
-    os.environ["LOCAL_RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(master_port)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
@@ -71,7 +57,7 @@ def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
 def _run_attention_case(
     local_rank: int,
     world_size: int,
-    master_port: int,
+    init_method: str,
     input_file: str,
     output_file: str,
     num_heads: int,
@@ -85,10 +71,9 @@ def _run_attention_case(
     device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
     current_omni_platform.set_device(device)
 
-    _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port)
     os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
-    init_distributed_environment(world_size=world_size, rank=local_rank)
+    init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
     initialize_model_parallel(
         data_parallel_size=1,
         cfg_parallel_size=1,
@@ -202,10 +187,8 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
     batch_size = 2
     head_size = 8
 
-    base_port = _find_free_port()
-    sp_port = _find_free_port()
-    while sp_port == base_port:
-        sp_port = _find_free_port()
+    base_init_method = get_distributed_init_method()
+    sp_init_method = get_distributed_init_method()
 
     with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
         input_file = f_in.name
@@ -224,7 +207,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
         # Baseline (no SP)
         torch.multiprocessing.spawn(
             _run_attention_case,
-            args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"),
+            args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
             nprocs=1,
         )
 
@@ -233,7 +216,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
             _run_attention_case,
             args=(
                 sp_world_size,
-                sp_port,
+                sp_init_method,
                 input_file,
                 sp_file,
                 num_heads,
@@ -278,10 +261,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
     # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
     split_sizes = [3, 2, 3, 2]
 
-    base_port = _find_free_port()
-    sp_port = _find_free_port()
-    while sp_port == base_port:
-        sp_port = _find_free_port()
+    base_init_method = get_distributed_init_method()
+    sp_init_method = get_distributed_init_method()
 
     with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
         input_file = f_in.name
@@ -300,7 +281,19 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
         # Baseline (no SP)
         torch.multiprocessing.spawn(
             _run_attention_case,
-            args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"),
+            args=(
+                1,
+                base_init_method,
+                input_file,
+                baseline_file,
+                num_heads,
+                head_size,
+                1,
+                "strict",
+                1,
+                None,
+                "mem_efficient",
+            ),
             nprocs=1,
         )
 
@@ -309,7 +302,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
             _run_attention_case,
             args=(
                 sp_world_size,
-                sp_port,
+                sp_init_method,
                 input_file,
                 sp_file,
                 num_heads,

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -0,0 +1,772 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from contextlib import nullcontext
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.worker import block_table as native_block_table
+
+from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.metadata import (
+    DiffusionKVContextMetadata,
+    DiffusionKVMetadata,
+    DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _registration_backend(parallel_config: object):
+    model_config = SimpleNamespace(set_attention_geometry=Mock())
+    vllm_config = SimpleNamespace(
+        attention_config=SimpleNamespace(use_non_causal=False),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+        model_config=model_config,
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=SimpleNamespace(parallel_config=parallel_config),
+        device=torch.device("cpu"),
+    )
+    return backend, vllm_config, model_config
+
+
+def _parallel_config(**overrides):
+    values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict")
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2))
+    layer = SimpleNamespace(
+        num_heads=8,
+        softmax_scale=0.125,
+        skip_sequence_parallel=False,
+        attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"),
+    )
+    unrelated = object()
+    vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated})
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=8,
+        dtype=torch.float16,
+        non_causal=True,
+    )
+    constructor_calls = []
+
+    class FakeLayerAdapter:
+        def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None:
+            constructor_calls.append((layer_name, layer, spec, ulysses_degree))
+            self.num_heads = layer.num_heads // ulysses_degree
+            self.num_kv_heads = spec.num_kv_heads // ulysses_degree
+            self.head_size = spec.head_size
+            self.spec = replace(spec, num_kv_heads=self.num_kv_heads)
+
+    monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter)
+    monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+    specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+    installed = vllm_config.compilation_config.static_forward_context["layer-0"]
+    assert isinstance(installed, FakeLayerAdapter)
+    assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated
+    assert specs["layer-0"].num_kv_heads == 2
+    assert constructor_calls == [("layer-0", layer, spec, 2)]
+    model_config.set_attention_geometry.assert_called_once_with(
+        num_heads=4,
+        num_kv_heads=2,
+        head_size=8,
+    )
+    assert vllm_config.attention_config.use_non_causal is True
+
+
+def test_cache_layer_registration_rejects_backend_without_paged_support() -> None:
+    backend, _, _ = _registration_backend(_parallel_config())
+    layer = SimpleNamespace(
+        attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"),
+    )
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=8,
+        dtype=torch.float16,
+    )
+
+    with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"):
+        backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+
+def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend, _, _ = _registration_backend(_parallel_config())
+
+    def unsupported_block_tables():
+        raise NotImplementedError("platform has no paged BlockTables")
+
+    monkeypatch.setattr(
+        model_runner_backend_module.current_omni_platform,
+        "get_diffusion_kv_block_tables_cls",
+        unsupported_block_tables,
+    )
+
+    with pytest.raises(NotImplementedError, match="no paged BlockTables"):
+        backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))})
+
+
+@pytest.mark.parametrize(
+    ("parallel_config", "message"),
+    [
+        (_parallel_config(ring_degree=2), "Ring"),
+        (_parallel_config(allgather_degree=2), "AllGather-KV"),
+        (_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"),
+    ],
+)
+def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None:
+    backend, _, _ = _registration_backend(parallel_config)
+
+    with pytest.raises(NotImplementedError, match=message):
+        backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False))
+
+
+def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None:
+    backend, _, _ = _registration_backend(
+        _parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa")
+    )
+
+    assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1
+
+
+class _FakeSpec:
+    def __init__(self, block_size: int) -> None:
+        self.block_size = block_size
+        self.capacity_calls: list[tuple[object, int]] = []
+
+    def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int:
+        self.capacity_calls.append((vllm_config, max_len))
+        return math.ceil(max_len / self.block_size)
+
+
+class _FakeStagedTable:
+    def __init__(self) -> None:
+        self.staged: list[tuple[int, tuple[int, ...]]] = []
+        self.clear_calls = 0
+
+    def clear_staged_writes(self) -> None:
+        self.staged.clear()
+        self.clear_calls += 1
+
+
+class _FakeNumBlocks:
+    def __init__(self, num_groups: int, max_num_reqs: int) -> None:
+        self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32)
+        self.copy_calls = 0
+
+    def copy_to_uva(self) -> None:
+        self.copy_calls += 1
+
+
+class _FakeBlockTables:
+    instances: list[_FakeBlockTables] = []
+
+    def __init__(self, **kwargs: object) -> None:
+        self.kwargs = kwargs
+        self.block_sizes = list(kwargs["block_sizes"])
+        self.max_num_reqs = int(kwargs["max_num_reqs"])
+        self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"])
+        self.cp_size = int(kwargs["cp_size"])
+        self.cp_rank = int(kwargs["cp_rank"])
+        self.cp_interleave = int(kwargs["cp_interleave"])
+        self.blocks_per_kv_block = [1] * len(self.block_sizes)
+        self.block_tables = [_FakeStagedTable() for _ in self.block_sizes]
+        self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs)
+        self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = []
+        self.apply_calls = 0
+        self.layout_refreshes = 0
+        self.fail_on_append_call: int | None = None
+        self.fail_apply = False
+        self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes]
+        self.__class__.instances.append(self)
+
+    def append_block_ids(
+        self,
+        req_index: int,
+        new_block_ids: tuple[list[int], ...],
+        overwrite: bool,
+    ) -> None:
+        call_number = len(self.append_calls) + 1
+        if self.fail_on_append_call == call_number:
+            raise RuntimeError("fake append failure")
+        copied_ids = tuple(list(group_ids) for group_ids in new_block_ids)
+        self.append_calls.append((req_index, copied_ids, overwrite))
+        for group_index, group_ids in enumerate(copied_ids):
+            self.num_blocks.np[group_index, req_index] = len(group_ids)
+            if group_ids:
+                self.block_tables[group_index].staged.append((req_index, tuple(group_ids)))
+
+    def apply_staged_writes(self) -> None:
+        self.apply_calls += 1
+        if self.fail_apply:
+            raise RuntimeError("fake apply failure")
+        for group_index, group_table in enumerate(self.block_tables):
+            for row, group_ids in group_table.staged:
+                self.committed[group_index][row] = group_ids
+            group_table.staged.clear()
+        self.num_blocks.copy_to_uva()
+
+    def init_block_table_layout_tensors(self) -> None:
+        self.layout_refreshes += 1
+
+
+def _metadata(
+    request_id: str = "req-0",
+    generation: int = 1,
+    *,
+    seq_len: int = 9,
+    sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]),
+    contexts: tuple[DiffusionKVContextMetadata, ...] | None = None,
+) -> DiffusionKVMetadata:
+    if contexts is None:
+        contexts = (
+            DiffusionKVContextMetadata(
+                context_id="text",
+                cache_role="cross_attention",
+                num_tokens=5,
+                block_ids=([6, 7], [8]),
+            ),
+        )
+    return DiffusionKVMetadata(
+        request_id=request_id,
+        allocation_generation=generation,
+        sequences=(
+            DiffusionKVSequenceMetadata(
+                sequence_id=0,
+                prefix_len=4,
+                target_len=5,
+                seq_len=seq_len,
+                block_ids=sequence_block_ids,
+                context_ids=tuple(context.context_id for context in contexts),
+            ),
+        ),
+        contexts=contexts,
+    )
+
+
+def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None:
+    _FakeBlockTables.instances.clear()
+
+    def get_block_table_width(max_num_blocks, block_size):
+        events.append(("width", max_num_blocks, block_size))
+        return max_num_blocks + 1
+
+    def init_backend(kv_cache_config, vllm_config, device):
+        events.append(("backend", kv_cache_config, vllm_config, device))
+        group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
+        return [[group], [group]], "cg-support", [4, 8]
+
+    def init_cache(
+        runner_kv_caches,
+        forward_context,
+        kv_cache_config,
+        attn_groups,
+        device,
+        cache_dtype,
+        kernel_block_sizes,
+        vllm_config,
+    ):
+        events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+        runner_kv_caches.extend(["cache-0", "cache-1"])
+        return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
+
+    monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend)
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache)
+    monkeypatch.setattr(
+        model_runner_backend_module.current_omni_platform,
+        "get_diffusion_kv_block_tables_cls",
+        lambda: _FakeBlockTables,
+    )
+    monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False)
+
+
+def _runner(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    max_rows_per_request: int = 3,
+    max_num_seqs: int = 2,
+    max_len: int = 16,
+    max_num_batched_tokens: int = 32,
+    cp_size: int = 1,
+    cp_rank: int = 0,
+    cp_interleave: int = 1,
+) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]:
+    events: list[tuple] = []
+    _install_native_modules(monkeypatch, events)
+    od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        diffusion_kv_max_rows_per_request=max_rows_per_request,
+        max_num_seqs=max_num_seqs,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=max_len),
+        scheduler_config=SimpleNamespace(
+            max_num_batched_tokens=max_num_batched_tokens,
+            max_num_seqs=max_num_seqs,
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=cp_size,
+            cp_kv_cache_interleave_size=cp_interleave,
+        ),
+        compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+        cache_config=SimpleNamespace(cache_dtype="auto"),
+    )
+    monkeypatch.setattr(
+        model_runner_backend_module,
+        "get_dcp_group",
+        lambda: SimpleNamespace(rank_in_group=cp_rank),
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=od_config,
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+    }
+
+    config = SimpleNamespace(
+        num_blocks=32,
+        kv_cache_groups=[
+            SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+            SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+        ],
+    )
+    backend.initialize_kv_cache(config)
+    return backend, _FakeBlockTables.instances[-1], events
+
+
+def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, events = _runner(monkeypatch)
+
+    assert events[:2] == [("width", 4, 4), ("width", 2, 8)]
+    assert [event[0] for event in events[2:]] == ["backend", "cache"]
+    assert block_tables.kwargs == {
+        "block_sizes": [4, 8],
+        "max_num_reqs": 6,
+        "max_num_batched_tokens": 32,
+        "max_num_blocks_per_group": [5, 3],
+        "device": "cuda:1",
+        "kernel_block_sizes": [4, 8],
+        "cp_size": 1,
+        "cp_rank": 0,
+        "cp_interleave": 1,
+    }
+    assert runner.block_tables is block_tables
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+    assert runner.kv_caches == ["cache-0", "cache-1"]
+    assert runner.vllm_config.scheduler_config.max_num_seqs == 6
+
+
+def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    _, block_tables, _ = _runner(
+        monkeypatch,
+        cp_size=2,
+        cp_rank=1,
+        cp_interleave=4,
+    )
+
+    assert block_tables.kwargs["cp_size"] == 2
+    assert block_tables.kwargs["cp_rank"] == 1
+    assert block_tables.kwargs["cp_interleave"] == 4
+
+
+def test_initialize_failure_does_not_publish_partial_cache_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    events: list[tuple] = []
+    _install_native_modules(monkeypatch, events)
+    successful_init_cache = model_runner_backend_module.init_kv_cache
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=16),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+        compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+        cache_config=SimpleNamespace(cache_dtype="auto"),
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=SimpleNamespace(
+            diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+            diffusion_kv_max_rows_per_request=2,
+        ),
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"),
+        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"),
+    }
+    config = SimpleNamespace(
+        num_blocks=32,
+        kv_cache_groups=[
+            SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+            SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+        ],
+    )
+
+    def fail_init_cache(runner_kv_caches, *_args, **_kwargs):
+        runner_kv_caches.append("partial-cache")
+        backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding"
+        raise RuntimeError("cache binding failed")
+
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache)
+    with pytest.raises(RuntimeError, match="cache binding failed"):
+        backend.initialize_kv_cache(config)
+
+    assert backend.kv_cache_config is None
+    assert backend.kv_caches == []
+    assert backend.block_tables is None
+    assert backend.paged_attention_adapter is None
+    assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 1
+
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
+    backend.initialize_kv_cache(config)
+    assert backend.kv_cache_config is not None
+    assert backend.kv_caches == ["cache-0", "cache-1"]
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 2
+
+
+def test_initialize_rejects_config_for_different_registered_layers() -> None:
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=SimpleNamespace(),
+        od_config=SimpleNamespace(),
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None)
+    }
+    config = SimpleNamespace(
+        kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])],
+    )
+
+    with pytest.raises(ValueError, match="layer mismatch"):
+        backend.initialize_kv_cache(config)
+
+
+def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    assert runner.install_diffusion_kv_metadata(_metadata()) is True
+
+    assert block_tables.append_calls == [
+        (0, ([1, 2, 3], [4, 5]), True),
+        (1, ([6, 7], [8]), True),
+    ]
+    assert block_tables.apply_calls == 1
+    assert runner.get_diffusion_kv_row("req-0", 0) == 0
+    assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
+    sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
+    context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
+    assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
+    assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
+
+
+def test_block_table_mutations_invalidate_prepared_attention_batches(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runner, _, _ = _runner(monkeypatch)
+    assert runner.paged_attention_adapter is not None
+    runner.paged_attention_adapter.invalidate_prepared_batches = Mock()
+
+    runner.install_diffusion_kv_metadata(_metadata())
+    runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with()
+
+    runner.remove_diffusion_kv_requests(["req-0"])
+    assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2
+
+    runner.refresh_block_table_layout()
+    assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3
+
+
+def test_cache_layer_registration_is_frozen_after_physical_initialization(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runner, _, _ = _runner(monkeypatch)
+
+    with pytest.raises(RuntimeError, match="cannot be changed"):
+        runner.register_kv_cache_layers({})
+
+
+def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    first_sequence = _metadata().sequences[0]
+    metadata = replace(
+        _metadata(),
+        sequences=(
+            first_sequence,
+            replace(
+                first_sequence,
+                sequence_id=1,
+                block_ids=([9, 10, 11], [12, 13]),
+            ),
+        ),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+
+    assert block_tables.append_calls == [
+        (0, ([1, 2, 3], [4, 5]), True),
+        (1, ([9, 10, 11], [12, 13]), True),
+        (2, ([6, 7], [8]), True),
+    ]
+    assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2
+    assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2
+
+
+def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+    spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8)
+
+    metadata = _metadata(
+        sequence_block_ids=([1, 2], [4, 5]),
+        contexts=(),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+    assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)]
+
+
+def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    metadata = _metadata(
+        sequence_block_ids=([1, 1, 3], [4, 4]),
+        contexts=(),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+    assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)]
+
+
+def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    metadata = _metadata()
+    runner.install_diffusion_kv_metadata(metadata)
+
+    assert runner.install_diffusion_kv_metadata(metadata) is False
+    assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False
+    assert block_tables.apply_calls == 1
+
+    with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"):
+        runner.install_diffusion_kv_metadata(metadata)
+    with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"):
+        runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5])))
+    assert block_tables.apply_calls == 1
+
+
+@pytest.mark.parametrize(
+    ("metadata", "message"),
+    [
+        (_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"),
+        (_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"),
+        (_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"),
+        (
+            _metadata(
+                seq_len=17,
+                sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]),
+                contexts=(),
+            ),
+            "exceeds row capacity",
+        ),
+    ],
+)
+def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic(
+    monkeypatch: pytest.MonkeyPatch,
+    metadata: DiffusionKVMetadata,
+    message: str,
+) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    with pytest.raises(ValueError, match=message):
+        runner.install_diffusion_kv_metadata(metadata)
+
+    assert block_tables.append_calls == []
+    assert block_tables.apply_calls == 0
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    bad_context = DiffusionKVContextMetadata(
+        context_id="text",
+        cache_role="cross_attention",
+        num_tokens=5,
+        block_ids=([6, 32], [8]),
+    )
+
+    with pytest.raises(ValueError, match="outside"):
+        runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,)))
+
+    assert block_tables.append_calls == []
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+
+
+def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    block_tables.fail_on_append_call = 2
+
+    with pytest.raises(RuntimeError, match="fake append failure"):
+        runner.install_diffusion_kv_metadata(_metadata())
+
+    assert np.count_nonzero(block_tables.num_blocks.np) == 0
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    block_tables.fail_apply = True
+
+    with pytest.raises(RuntimeError, match="fake apply failure"):
+        runner.install_diffusion_kv_metadata(_metadata())
+
+    assert block_tables.apply_calls == 1
+    assert np.count_nonzero(block_tables.num_blocks.np) == 0
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+    assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables)
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+    second_context = DiffusionKVContextMetadata(
+        context_id="image",
+        cache_role="cross_attention",
+        num_tokens=1,
+        block_ids=([9], [10]),
+    )
+
+    with pytest.raises(ValueError, match="adapter limit is 2"):
+        runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context)))
+    assert block_tables.append_calls == []
+
+    runner.install_diffusion_kv_metadata(_metadata())
+    append_count = len(block_tables.append_calls)
+    with pytest.raises(ValueError, match="only 0 rows are free"):
+        runner.install_diffusion_kv_metadata(_metadata("req-1"))
+    assert len(block_tables.append_calls) == append_count
+
+
+def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=())
+    duplicate_sequence = _metadata().sequences[0]
+    duplicate = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=0,
+        sequences=(duplicate_sequence, duplicate_sequence),
+    )
+
+    with pytest.raises(ValueError, match="at least one sequence"):
+        runner.install_diffusion_kv_metadata(empty)
+    with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"):
+        runner.install_diffusion_kv_metadata(duplicate)
+    assert block_tables.append_calls == []
+
+
+@pytest.mark.parametrize(
+    ("metadata", "message"),
+    [
+        (_metadata(generation=-1), "non-negative integer"),
+        (_metadata(request_id=""), "non-empty string"),
+        (_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"),
+    ],
+)
+def test_request_generation_and_block_id_types_are_validated(
+    monkeypatch: pytest.MonkeyPatch,
+    metadata: DiffusionKVMetadata,
+    message: str,
+) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    with pytest.raises(ValueError, match=message):
+        runner.install_diffusion_kv_metadata(metadata)
+
+    assert block_tables.append_calls == []
+    assert runner._diffusion_kv_request_states == {}
+
+
+def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+    runner.install_diffusion_kv_metadata(_metadata())
+
+    assert runner.remove_diffusion_kv_requests(["missing"]) == 0
+    assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2
+    assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)]
+    assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]]
+    assert runner.remove_diffusion_kv_requests(["req-0"]) == 0
+    with pytest.raises(KeyError, match="No native diffusion KV row"):
+        runner.get_diffusion_kv_row("req-0", 0)
+
+    runner.install_diffusion_kv_metadata(_metadata("req-1"))
+    assert runner.get_diffusion_kv_row("req-1", 0) == 0
+    assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1
+
+
+def test_cached_step_does_not_reinstall_metadata() -> None:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.pipeline = object()
+    runner.install_diffusion_kv_metadata = Mock()
+    runner._supports_step_mode = Mock(return_value=False)
+    scheduler_output = DiffusionSchedulerOutput(
+        step_id=1,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]),
+        finished_req_ids=set(),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+    with pytest.raises(ValueError, match="does not support step execution"):
+        runner.execute_stepwise(scheduler_output)
+    runner.install_diffusion_kv_metadata.assert_not_called()
+
+
+def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    runner.refresh_block_table_layout()
+    assert block_tables.layout_refreshes == 1
+
+    dense_backend = DiffusionKVModelRunnerBackend(
+        vllm_config=SimpleNamespace(),
+        od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY),
+        device="cuda:1",
+    )
+    dense_backend.block_tables = Mock()
+    assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0
+    dense_backend.refresh_block_table_layout()
+    dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called()
+    with pytest.raises(ValueError, match="Dense diffusion execution"):
+        dense_backend.install_diffusion_kv_metadata(_metadata())

tests/diffusion/diffusion_kv/test_config.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from vllm_omni.config.omni_config import _DiffusionConfigProjection
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
@@ -39,6 +40,8 @@ def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
     with pytest.raises(ValueError, match="does not support imported AR KV"):
         OmniDiffusionConfig.from_kwargs(
             diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=1,
+            max_num_batched_tokens=1,
             omni_kv_config={"need_recv_cache": True},
         )
 
@@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
     config = OmniDiffusionConfig.from_kwargs(
         model_class_name="FutureDiffusionModel",
         diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        max_num_batched_tokens=1,
     )
 
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+    )
+
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+def test_paged_scheduler_requires_a_worker_row_limit() -> None:
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler")
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler")
+
+
+@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5])
+def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None:
+    with pytest.raises((TypeError, ValueError)):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
+    with pytest.raises((TypeError, ValueError)):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)

tests/diffusion/diffusion_kv/test_initialization.py

@@ -22,6 +22,7 @@ class _Executor:
         self.specs = [{"layer0": object()}, {"layer0": object()}]
         self.memory = [1024, 1024]
         self.configs = None
+        self.resolved_max_model_len = None
         self.profile_requests = None
 
     def get_kv_cache_specs(self):
@@ -31,8 +32,9 @@ class _Executor:
         self.profile_requests = profile_requests
         return self.memory
 
-    def set_kv_cache_configs(self, configs) -> None:
+    def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None:
         self.configs = configs
+        self.resolved_max_model_len = resolved_max_model_len
 
 
 def _od_config(**overrides):
@@ -55,6 +57,7 @@ def _od_config(**overrides):
         ),
         gpu_memory_utilization=0.9,
         kv_cache_memory_bytes=None,
+        max_num_seqs=1,
         max_num_batched_tokens=64,
         num_gpus=1,
     )
@@ -64,9 +67,9 @@ def _od_config(**overrides):
 
 def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
     executor = _Executor()
-    od_config = _od_config(num_gpus=2)
+    od_config = _od_config(num_gpus=2, max_model_len=-1)
     vllm_config = SimpleNamespace(
-        model_config=SimpleNamespace(max_model_len=4096),
+        model_config=SimpleNamespace(max_model_len=256),
         max_in_flight_tokens=128,
     )
     worker_configs = [object(), object()]
@@ -91,6 +94,7 @@ def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
     build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
     assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
     assert executor.configs == worker_configs
+    assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len
     assert executor.profile_requests is profile_requests
 
 

tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from importlib import import_module, util
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+
+def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test")
+    module_name = "vllm.v1.worker.gpu.block_table"
+    try:
+        module_spec = util.find_spec(module_name)
+    except ModuleNotFoundError as exc:
+        if exc.name is None or not module_name.startswith(f"{exc.name}."):
+            raise
+        pytest.skip("installed vLLM does not provide GPU BlockTables")
+    if module_spec is None:
+        pytest.skip("installed vLLM does not provide GPU BlockTables")
+    block_table_module = import_module(module_name)
+    block_tables_cls = getattr(block_table_module, "BlockTables", None)
+    if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"):
+        pytest.skip("installed vLLM does not provide the required BlockTables API")
+
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    block_tables = block_tables_cls(
+        block_sizes=[4],
+        max_num_reqs=1,
+        max_num_batched_tokens=16,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[4],
+    )
+    block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True)
+    block_tables.apply_staged_writes()
+    idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device)
+    positions = torch.arange(9, dtype=torch.int64, device=device)
+
+    slot_mappings = block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        positions,
+        num_tokens_padded=9,
+    )
+    torch.accelerator.synchronize(device)
+
+    assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]]

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F
+from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+    DiffusionPagedAttentionAdapter,
+    DiffusionPagedAttentionLayerAdapter,
+    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionRowBinding,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+_LAYER_NAME = "model.layers.0.attn"
+_NUM_HEADS = 2
+_HEAD_SIZE = 64
+_BLOCK_SIZE = 16
+
+
+class _SmokeDiffusionAttention(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.num_heads = _NUM_HEADS
+        self.num_kv_heads = _NUM_HEADS
+        self.head_size = _HEAD_SIZE
+        self.softmax_scale = _HEAD_SIZE**-0.5
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    vllm_config = VllmConfig(
+        cache_config=CacheConfig(
+            block_size=_BLOCK_SIZE,
+            cache_dtype="float16",
+            enable_prefix_caching=False,
+        ),
+        attention_config=AttentionConfig(use_non_causal=True),
+    )
+    model_config = _DiffusionVllmModelConfig(
+        model=None,
+        dtype=torch.float16,
+        max_model_len=32,
+        original_max_model_len=32,
+        hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"),
+    )
+    vllm_config.model_config = model_config
+    model_config.set_attention_geometry(
+        num_heads=_NUM_HEADS,
+        num_kv_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+    )
+    vllm_config.compilation_config.static_forward_context.clear()
+
+    diffusion_layer = _SmokeDiffusionAttention().to(device)
+    spec = FullAttentionSpec(
+        block_size=_BLOCK_SIZE,
+        num_kv_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+        dtype=torch.float16,
+        non_causal=True,
+    )
+    with set_current_vllm_config(vllm_config):
+        native_layer = DiffusionPagedAttentionLayerAdapter(
+            layer_name=_LAYER_NAME,
+            layer=diffusion_layer,
+            spec=spec,
+            vllm_config=vllm_config,
+            device=device,
+        )
+    vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer
+
+    canonical_spec = native_layer.spec
+    kv_cache_config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=canonical_spec.page_size_bytes * 4,
+                shared_by=[_LAYER_NAME],
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=[_LAYER_NAME],
+                kv_cache_spec=canonical_spec,
+            )
+        ],
+    )
+    with set_current_vllm_config(vllm_config):
+        attn_groups, _, kernel_block_sizes = init_attn_backend(
+            kv_cache_config,
+            vllm_config,
+            device,
+        )
+        runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        init_kv_cache(
+            runner_kv_caches,
+            vllm_config.compilation_config.static_forward_context,
+            kv_cache_config,
+            attn_groups,
+            device,
+            vllm_config.cache_config.cache_dtype,
+            kernel_block_sizes,
+            vllm_config,
+        )
+        block_tables = BlockTables(
+            block_sizes=[_BLOCK_SIZE],
+            max_num_reqs=1,
+            max_num_batched_tokens=32,
+            max_num_blocks_per_group=[2],
+            device=device,
+            kernel_block_sizes=kernel_block_sizes,
+        )
+
+    block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
+    block_tables.apply_staged_writes()
+    adapter = DiffusionPagedAttentionAdapter(
+        vllm_config=vllm_config,
+        device=device,
+        kv_cache_config=kv_cache_config,
+        block_tables=block_tables,
+        attn_groups=attn_groups,
+        layers={_LAYER_NAME: native_layer},
+        resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
+            row_index=0,
+            max_seq_len=19,
+        ),
+    )
+    prefix_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=17,
+                seq_len=17,
+            )
+        ]
+    )
+
+    torch.manual_seed(1)
+    query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    omni_backend = FlashAttentionImpl(
+        num_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+        softmax_scale=diffusion_layer.softmax_scale,
+        num_kv_heads=_NUM_HEADS,
+    )
+    with adapter.activate(prefix_batch):
+        prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17])
+        prefix_output = omni_backend.forward_paged(prefix_context)
+    prefix_slot_mappings = prefix_batch.slot_mappings.clone()
+
+    suffix_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=2,
+                seq_len=19,
+                kv_start_pos=17,
+            )
+        ]
+    )
+    with adapter.activate(suffix_batch):
+        suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:])
+        suffix_output = omni_backend.forward_paged(suffix_context)
+    suffix_slot_mappings = suffix_batch.slot_mappings.clone()
+
+    piecewise_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=19,
+                seq_len=19,
+            )
+        ]
+    )
+    mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril()
+    mixed_mask[5:10, 5:10] = True
+    mixed_metadata = AttentionMetadata(
+        attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+        full_attn_spans=[[(5, 10)]],
+    )
+    with adapter.activate(piecewise_batch):
+        piecewise_context = adapter.prepare_layer_context(
+            _LAYER_NAME,
+            query,
+            key,
+            value,
+            omni_attn_metadata=mixed_metadata,
+        )
+        piecewise_output = omni_backend.forward_paged(piecewise_context)
+
+    prefix_reference = (
+        F.scaled_dot_product_attention(
+            query[:17].transpose(0, 1).unsqueeze(0),
+            key[:17].transpose(0, 1).unsqueeze(0),
+            value[:17].transpose(0, 1).unsqueeze(0),
+            is_causal=False,
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    suffix_reference = (
+        F.scaled_dot_product_attention(
+            query[17:].transpose(0, 1).unsqueeze(0),
+            key.transpose(0, 1).unsqueeze(0),
+            value.transpose(0, 1).unsqueeze(0),
+            is_causal=False,
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    piecewise_reference = (
+        F.scaled_dot_product_attention(
+            query.transpose(0, 1).unsqueeze(0),
+            key.transpose(0, 1).unsqueeze(0),
+            value.transpose(0, 1).unsqueeze(0),
+            attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    torch.accelerator.synchronize(device)
+
+    assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]]
+    assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]]
+    torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -25,6 +25,23 @@ class _Backend:
         return True
 
 
+class _RunnerKVBackend:
+    def __init__(self) -> None:
+        self.kv_cache_config = None
+
+    def register_kv_cache_layers(self, layers):
+        return {layer_name: spec for layer_name, (_layer, spec) in layers.items()}
+
+    def initialize_kv_cache(self, config) -> None:
+        configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names}
+        if configured_layers != {"image_attention"}:
+            raise ValueError(
+                "Rank-local Diffusion KVCacheConfig layer mismatch: "
+                f"expected=['image_attention'], configured={sorted(configured_layers)}"
+            )
+        self.kv_cache_config = config
+
+
 def _attention(*, enabled: bool) -> Attention:
     attention = Attention.__new__(Attention)
     nn.Module.__init__(attention)
@@ -47,6 +64,7 @@ def _runner(attention: Attention) -> DiffusionModelRunner:
     pipeline = nn.Module()
     pipeline.image_attention = attention
     runner.pipeline = pipeline
+    runner.diffusion_kv_backend = _RunnerKVBackend()
     return runner
 
 
@@ -102,14 +120,21 @@ def test_worker_selects_its_rank_local_config() -> None:
     worker = object.__new__(DiffusionWorker)
     worker.rank = 1
     worker.od_config = SimpleNamespace(num_gpus=2)
+    worker.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=None),
+        cache_config=SimpleNamespace(num_gpu_blocks=None),
+    )
+    worker._maybe_get_memory_pool_context = lambda _tag: nullcontext()
     worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
-    configs = [object(), object()]
+    configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)]
 
-    worker.set_kv_cache_configs(configs)
+    worker.set_kv_cache_configs(configs, 64)
 
     assert worker.installed is configs[1]
+    assert worker.vllm_config.model_config.max_model_len == 64
+    assert worker.vllm_config.cache_config.num_gpu_blocks == 8
     with pytest.raises(ValueError, match="rank count mismatch"):
-        worker.set_kv_cache_configs(configs[:1])
+        worker.set_kv_cache_configs(configs[:1], 64)
 
 
 def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:

tests/diffusion/distributed/test_hsdp.py

@@ -1,10 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
 import gc
-import os
-import socket
 
 import pytest
 import torch
@@ -12,6 +10,7 @@ import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.tensor import DeviceMesh, DTensor
 
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.hsdp import (
     HSDPInferenceConfig,
@@ -37,35 +36,17 @@ class _PackedModel(nn.Module):
         self.root_weight = nn.Parameter(torch.ones(2))
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind(("127.0.0.1", 0))
-        return int(sock.getsockname()[1])
-
-
 @pytest.fixture(scope="module")
 def cpu_process_group():
     if dist.is_initialized():
         yield
         return
 
-    master_port = _find_free_port()
-    os.environ.update(
-        {
-            "RANK": "0",
-            "LOCAL_RANK": "0",
-            "WORLD_SIZE": "1",
-            "MASTER_ADDR": "127.0.0.1",
-            "MASTER_PORT": str(master_port),
-        }
-    )
-    dist.init_process_group("gloo", rank=0, world_size=1)
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
     try:
         yield
     finally:
         dist.destroy_process_group()
-        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
-            os.environ.pop(key, None)
         gc.collect()
 
 

tests/diffusion/distributed/test_pipeline_parallel.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
 """Tests for PipelineParallelMixin and related PP communication helpers.
 
 CPU unit tests exercise mixin logic without a process group. Distributed CPU
@@ -9,8 +9,6 @@ real multi-GPU NCCL collectives.
 
 from __future__ import annotations
 
-import os
-import socket
 from typing import Literal
 
 import pytest
@@ -21,6 +19,7 @@ from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
 from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -40,31 +39,18 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
-def _find_free_port() -> str:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return str(s.getsockname()[1])
-
-
-def update_environment_variables(envs_dict: dict[str, str]) -> None:
-    for k, v in envs_dict.items():
-        os.environ[k] = v
-
-
 def _cleanup_distributed() -> None:
-    """Destroy omni distributed state and clear test-process leftovers.
+    """Destroy omni distributed state left dangling after a test.
 
     ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
     aliased to the destroyed group. Parent-process baselines must clear that
-    dangling reference (and RANK/MASTER_* env) so later unit tests that call
+    dangling reference so later unit tests that call
     ``vllm.distributed.initialize_model_parallel`` do not see stale state.
     """
     destroy_distributed_env()
     import vllm.distributed.parallel_state as vllm_parallel_state
 
     vllm_parallel_state._PP = None
-    for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
-        os.environ.pop(key, None)
 
 
 # ---------------------------------------------------------------------------
@@ -416,24 +402,20 @@ def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
 def init_dist(
     local_rank: int,
     world_size: int,
-    master_port: str,
+    init_method: str,
     device_kind: DeviceKind,
 ) -> torch.device:
     """Initialise the distributed environment for a spawned worker."""
     device = _worker_device(local_rank, device_kind)
     if device_kind == "cuda":
         current_omni_platform.set_device(device)
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": master_port,
-        }
-    )
     backend = "gloo" if device_kind == "cpu" else None
-    init_distributed_environment(backend=backend)
+    init_distributed_environment(
+        world_size=world_size,
+        rank=local_rank,
+        distributed_init_method=init_method,
+        backend=backend,
+    )
     return device
 
 
@@ -493,11 +475,11 @@ def make_pipeline_and_inputs(
 def isend_irecv_worker(
     local_rank: int,
     world_size: int,
-    master_port: str,
+    init_method: str,
     device_kind: DeviceKind,
     result_queue,
 ):
-    device = init_dist(local_rank, world_size, master_port, device_kind)
+    device = init_dist(local_rank, world_size, init_method, device_kind)
     initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
     pp_group = get_pp_group()
 
@@ -517,13 +499,13 @@ def isend_irecv_worker(
     _cleanup_distributed()
 
 
-def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
+def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
     torch.multiprocessing.spawn(
         isend_irecv_worker,
-        args=(pp_size, master_port, device_kind, q),
+        args=(pp_size, init_method, device_kind, q),
         nprocs=pp_size,
     )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
@@ -538,7 +520,7 @@ def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) ->
 @pytest.mark.parametrize("pp_size", [2])
 def test_isend_irecv_tensor_dict(pp_size: int):
     """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
-    _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
+    _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_distributed_init_method())
 
 
 @pytest.mark.full_model
@@ -547,7 +529,7 @@ def test_isend_irecv_tensor_dict(pp_size: int):
 @pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
 def test_isend_irecv_tensor_dict_parity(pp_size: int):
     """Nightly: isend/irecv on real multi-GPU NCCL."""
-    _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
+    _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_distributed_init_method())
 
 
 # ---------------------------------------------------------------------------
@@ -578,7 +560,7 @@ def compute_single_gpu_baseline(
     if key in _baseline_cache:
         return _baseline_cache[key]
 
-    device = init_dist(0, 1, _find_free_port(), device_kind)
+    device = init_dist(0, 1, get_distributed_init_method(), device_kind)
     try:
         initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
@@ -604,7 +586,7 @@ def compute_single_gpu_baseline(
 def predict_noise_worker(
     local_rank: int,
     world_size: int,
-    master_port: str,
+    init_method: str,
     pp_size: int,
     cfg_size: int,
     do_true_cfg: bool,
@@ -614,7 +596,7 @@ def predict_noise_worker(
     result_queue,
 ):
     """Generic predict-noise worker parameterized by PP and CFG topology."""
-    device = init_dist(local_rank, world_size, master_port, device_kind)
+    device = init_dist(local_rank, world_size, init_method, device_kind)
     initialize_model_parallel(
         pipeline_parallel_size=pp_size,
         cfg_parallel_size=cfg_size,
@@ -661,7 +643,7 @@ def _run_predict_noise(
     rtol: float,
     atol: float,
     device_kind: DeviceKind,
-    master_port: str,
+    init_method: str,
 ) -> None:
     test_config = {
         "num_layers": num_layers,
@@ -681,7 +663,7 @@ def _run_predict_noise(
     world_size = pp_size * cfg_size
     torch.multiprocessing.spawn(
         predict_noise_worker,
-        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
+        args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
     )
 
@@ -786,7 +768,7 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, r
         rtol=rtol,
         atol=atol,
         device_kind="cpu",
-        master_port=_find_free_port(),
+        init_method=get_distributed_init_method(),
     )
 
 
@@ -809,7 +791,7 @@ def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers,
         rtol=rtol,
         atol=atol,
         device_kind="cuda",
-        master_port=_find_free_port(),
+        init_method=get_distributed_init_method(),
     )
 
 
@@ -824,7 +806,7 @@ def compute_scheduler_step_baseline(
     device_kind: DeviceKind,
 ) -> torch.Tensor:
     """Single-process reference: predict_noise + scheduler_step."""
-    device = init_dist(0, 1, _find_free_port(), device_kind)
+    device = init_dist(0, 1, get_distributed_init_method(), device_kind)
     try:
         initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
@@ -853,7 +835,7 @@ def compute_scheduler_step_baseline(
 def scheduler_step_worker(
     local_rank: int,
     world_size: int,
-    master_port: str,
+    init_method: str,
     pp_size: int,
     cfg_size: int,
     do_true_cfg: bool,
@@ -861,7 +843,7 @@ def scheduler_step_worker(
     device_kind: DeviceKind,
     result_queue,
 ):
-    device = init_dist(local_rank, world_size, master_port, device_kind)
+    device = init_dist(local_rank, world_size, init_method, device_kind)
     initialize_model_parallel(
         pipeline_parallel_size=pp_size,
         cfg_parallel_size=cfg_size,
@@ -905,7 +887,7 @@ def _run_scheduler_step(
     do_true_cfg: bool,
     input_seed: int,
     device_kind: DeviceKind,
-    master_port: str,
+    init_method: str,
 ) -> None:
     test_config = {
         "num_layers": 4,
@@ -925,7 +907,7 @@ def _run_scheduler_step(
     world_size = pp_size * cfg_size
     torch.multiprocessing.spawn(
         scheduler_step_worker,
-        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
+        args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
     )
 
@@ -957,7 +939,7 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
         do_true_cfg=do_true_cfg,
         input_seed=input_seed,
         device_kind="cpu",
-        master_port=_find_free_port(),
+        init_method=get_distributed_init_method(),
     )
 
 
@@ -979,5 +961,5 @@ def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
         do_true_cfg=do_true_cfg,
         input_seed=input_seed,
         device_kind="cuda",
-        master_port=_find_free_port(),
+        init_method=get_distributed_init_method(),
     )

tests/diffusion/distributed/test_ulysses_uaa_perf.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
 """Perf smoke tests for Ulysses advanced_uaa communication overhead.
 
 This test is intended for CI monitoring only:
@@ -9,8 +9,6 @@ This test is intended for CI monitoring only:
 
 from __future__ import annotations
 
-import os
-import socket
 from dataclasses import dataclass
 
 import pytest
@@ -18,6 +16,7 @@ import torch
 import torch.distributed as dist
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.attention.parallel.ulysses import (
     _all_gather_int,
     _ulysses_all_to_all_any_o,
@@ -33,20 +32,6 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.platforms import current_omni_platform
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
-    os.environ["RANK"] = str(rank)
-    os.environ["LOCAL_RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(master_port)
-
-
 def _max_all_reduce(pg: dist.ProcessGroup, value: float, *, device: torch.device) -> float:
     t = torch.tensor([value], device=device, dtype=torch.float32)
     dist.all_reduce(t, op=dist.ReduceOp.MAX, group=pg)
@@ -77,22 +62,20 @@ def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
     if available_gpus < case.world_size:
         pytest.skip(f"Requires {case.world_size} GPUs, got {available_gpus}")
 
-    master_port = _find_free_port()
+    init_method = get_distributed_init_method()
     torch.multiprocessing.spawn(
         _perf_worker,
-        args=(case.world_size, master_port, case.ulysses_degree, case.ring_degree),
+        args=(case.world_size, init_method, case.ulysses_degree, case.ring_degree),
         nprocs=case.world_size,
     )
 
 
-def _perf_worker(local_rank: int, world_size: int, master_port: int, ulysses_degree: int, ring_degree: int) -> None:
+def _perf_worker(local_rank: int, world_size: int, init_method: str, ulysses_degree: int, ring_degree: int) -> None:
     device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
     current_omni_platform.set_device(device)
 
-    _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port)
-
     try:
-        init_distributed_environment(world_size=world_size, rank=local_rank)
+        init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
         initialize_model_parallel(
             data_parallel_size=1,
             cfg_parallel_size=1,

tests/diffusion/lora/test_lora_manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -20,9 +20,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyLoRALayer:
-    def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
+    def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
         self.n_slices = n_slices
         self.output_slices = output_slices
+        self.tp_size = tp_size
         self.set_calls: list[
             tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
         ] = []
@@ -236,6 +237,14 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
 
 def test_lora_manager_activates_packed_lora_from_sublayers():
     pipeline = torch.nn.Module()
+    bound_names = None
+
+    def validate_binding(*, lora_model, bound_lora_names):
+        nonlocal bound_names
+        assert lora_model.id == 1
+        bound_names = bound_lora_names
+
+    pipeline._validate_diffusion_lora_binding = validate_binding
     pipeline.stacked_params_mapping = [
         (".to_qkv", ".to_q", "q"),
         (".to_qkv", ".to_k", "k"),
@@ -279,6 +288,70 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
     assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
     assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+    assert bound_names == frozenset(loras)
+
+
+def test_lora_manager_rolls_back_all_layers_when_activation_fails():
+    class _FailingLoRALayer(_DummyLoRALayer):
+        fail = False
+
+        def set_lora(self, index: int, lora_a, lora_b):
+            if self.fail:
+                raise RuntimeError("bind failed")
+            super().set_lora(index, lora_a, lora_b)
+
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+    first = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+    second = _FailingLoRALayer(n_slices=1, output_slices=(2,))
+    manager._lora_modules = {
+        "transformer.first": first,
+        "transformer.second": second,
+    }
+
+    def adapter(adapter_id: int):
+        loras = {
+            name: LoRALayerWeights(
+                module_name=name,
+                rank=2,
+                lora_alpha=2,
+                lora_a=torch.full((2, 2), float(adapter_id)),
+                lora_b=torch.full((2, 2), float(adapter_id)),
+            )
+            for name in manager._lora_modules
+        }
+        return type(
+            "LM",
+            (),
+            {
+                "id": adapter_id,
+                "loras": loras,
+                "get_lora": lambda self, key: self.loras.get(key),
+            },
+        )()
+
+    manager._registered_adapters = {1: adapter(1), 2: adapter(2)}
+    manager._activate_adapter(1, scale=1.0)
+    first_calls_after_success = len(first.set_calls)
+
+    second.fail = True
+    with pytest.raises(RuntimeError, match="bind failed"):
+        manager._activate_adapter(2, scale=1.0)
+
+    assert manager._active_adapter_id is None
+    assert first.reset_calls == 1
+    assert second.reset_calls == 1
+
+    # The failed activation must not leave the old adapter eligible for the
+    # fast path: reactivating it binds every layer again.
+    second.fail = False
+    manager._activate_adapter(1, scale=1.0)
+    assert manager._active_adapter_id == 1
+    assert len(first.set_calls) == first_calls_after_success + 2
 
 
 def _dummy_lora_request(adapter_id: int) -> LoRARequest:

tests/diffusion/models/flux2/test_flux2_quantization.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for FLUX.2-dev quantization routing."""
+
+import pytest
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _RecordingBlock(nn.Module):
+    def __init__(self, *, quant_config=None, prefix: str = "", **kwargs):
+        super().__init__()
+        self.quant_config = quant_config
+        self.prefix = prefix
+
+
+def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.diffusion.models.flux2 import flux2_transformer
+
+    monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
+    monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
+
+    quant_config = Fp8Config(
+        is_checkpoint_fp8_serialized=False,
+        activation_scheme="dynamic",
+    )
+    model = flux2_transformer.Flux2Transformer2DModel(
+        in_channels=4,
+        out_channels=4,
+        num_layers=2,
+        num_single_layers=3,
+        attention_head_dim=8,
+        num_attention_heads=2,
+        joint_attention_dim=16,
+        timestep_guidance_channels=8,
+        axes_dims_rope=(2, 2, 2, 2),
+        quant_config=quant_config,
+    )
+
+    assert [block.quant_config for block in model.transformer_blocks] == [quant_config, quant_config]
+    assert [block.prefix for block in model.transformer_blocks] == [
+        "transformer_blocks.0",
+        "transformer_blocks.1",
+    ]
+    assert [block.quant_config for block in model.single_transformer_blocks] == [
+        quant_config,
+        quant_config,
+        quant_config,
+    ]
+    assert [block.prefix for block in model.single_transformer_blocks] == [
+        "single_transformer_blocks.0",
+        "single_transformer_blocks.1",
+        "single_transformer_blocks.2",
+    ]
+
+    assert isinstance(model.x_embedder, nn.Linear)
+    assert isinstance(model.context_embedder, nn.Linear)
+    assert isinstance(model.proj_out, nn.Linear)

tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+from cache_dit import ForwardPattern
+from torch import nn
+
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import (
+    HiDreamO1TextModel,
+)
+from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import (
+    HiDreamO1ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _empty_text_model() -> HiDreamO1TextModel:
+    model = HiDreamO1TextModel.__new__(HiDreamO1TextModel)
+    nn.Module.__init__(model)
+    model.layers = nn.ModuleList([nn.Identity()])
+    return model
+
+
+def test_cache_dit_targets_the_decoder_module_that_owns_layers():
+    text_model = _empty_text_model()
+    pipeline = SimpleNamespace(
+        _dit_modules=HiDreamO1ImagePipeline._dit_modules,
+        model=SimpleNamespace(
+            model=SimpleNamespace(language_model=text_model),
+        ),
+    )
+
+    with (
+        patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") as block_adapter_cls,
+        patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") as cache_dit,
+    ):
+        backend = CacheDiTBackend()
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=50)
+
+    adapter_config = HiDreamO1TextModel._cache_dit_adapter_config
+    assert isinstance(adapter_config, CacheDiTAdapterConfig)
+    assert adapter_config.block_forward_patterns == {"layers": ForwardPattern.Pattern_3}
+    assert adapter_config.has_separate_cfg is True
+
+    block_adapter_cls.assert_called_once_with(
+        transformer=text_model,
+        blocks=[text_model.layers],
+        forward_pattern=[ForwardPattern.Pattern_3],
+        has_separate_cfg=True,
+        check_forward_pattern=True,
+    )
+    cache_dit.enable_cache.assert_called_once()
+    enable_call = cache_dit.enable_cache.call_args
+    assert enable_call.args == (block_adapter_cls.return_value,)
+    assert enable_call.kwargs["cache_config"].Fn_compute_blocks == 1
+    assert enable_call.kwargs["calibrator_config"] is None
+    cache_dit.refresh_context.assert_called_once_with(
+        text_model,
+        num_inference_steps=50,
+        verbose=True,
+    )

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 _TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer"
 
-NUM_HEADS = 4
-NUM_KV_HEADS = 2
+NUM_HEADS = 32
+NUM_KV_HEADS = 8
 HEAD_DIM = 16
 IMAGE_TOKEN_LEN = 8
 SUFFIX_TOKEN_LEN = 3
@@ -35,8 +35,14 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM)
 class MockAttention(nn.Module):
     def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs):
         super().__init__()
+        self.paged_kv_active = False
+        self.calls = []
+
+    def is_paged_kv_active(self):
+        return self.paged_kv_active
 
     def forward(self, query, key, value, attn_metadata=None, **kwargs):
+        self.calls.append((query, key, value))
         return query
 
 
@@ -126,6 +132,37 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+@pytest.mark.parametrize(
+    ("paged_kv_active", "expected_kv_heads"),
+    [(False, NUM_HEADS), (True, NUM_KV_HEADS)],
+)
+def test_gqa_kv_stays_compressed_for_paged_attention(
+    paged_kv_active: bool,
+    expected_kv_heads: int,
+) -> None:
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = paged_kv_active
+    q_len = 3
+    key, value = _make_known_kv(q_len)
+
+    _call_mgr(
+        mgr,
+        bs=1,
+        q_len=q_len,
+        seq_len=q_len,
+        key_flat=key,
+        value_flat=value,
+        first_step=True,
+        num_image_tokens=0,
+        gen_timestep_scatter_index=_gen_timestep_index(1, 0),
+    )
+
+    _, key_input, value_input = mgr.attn.calls[-1]
+    expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM)
+    assert key_input.shape == expected_shape
+    assert value_input.shape == expected_shape
+
+
 # ============================================================
 # Test 1: No AR KV — basic cache → reuse
 # ============================================================

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import json
 import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
+from typing import Any, TypeVar
 from unittest.mock import Mock, patch
 
 import numpy as np
@@ -14,6 +15,14 @@ import torch.nn as nn
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
+_ItemT = TypeVar("_ItemT")
+_ResultT = TypeVar("_ResultT")
+
+
+def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT:
+    items.append(item)
+    return result
+
 
 def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -185,6 +194,7 @@ def test_startup_task_rejects_unsupported_value():
 
 def test_combined_task_inference_and_transformer_routing():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.errors import OmniClientError
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
@@ -196,6 +206,9 @@ def test_combined_task_inference_and_transformer_routing():
     assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
     assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
     assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
+    assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
+    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+        pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
     pipeline.partition = "ref2va"
     pipeline.supported_tasks = frozenset({"ref2va"})
     assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -260,7 +273,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
         ):
             (tmp_path / partition_name / component).mkdir()
 
-    created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+    created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
 
     class FakeModule(torch.nn.Module):
         def __init__(self, *args, **kwargs):
@@ -278,15 +291,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
 
         return create
 
-    tokenizer_calls = []
-    processor_calls = []
+    tokenizer_calls: list[Any] = []
+    processor_calls: list[Any] = []
     tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
-    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
-        tokenizer_calls.append((args, kwargs)) or object()
+    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+        tokenizer_calls, (args, kwargs), object()
     )
     processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
-    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
-        processor_calls.append((args, kwargs)) or object()
+    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+        processor_calls, (args, kwargs), object()
     )
     monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
     monkeypatch.setattr(
@@ -626,7 +639,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule():
 
 
 def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
-    diffuse_calls = []
+    diffuse_calls: list[dict[str, Any]] = []
     base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None})
 
@@ -643,7 +656,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
 
 
 def test_distilled_forward_accepts_the_matching_explicit_step_count():
-    diffuse_calls = []
+    diffuse_calls: list[dict[str, Any]] = []
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
 
     pipeline.forward(_t2va_batch(num_inference_steps=4))
@@ -713,8 +726,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
     )
 
     assert output.shape == q.shape
-    assert attention.attention.metadata.attn_mask is None
-    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+    metadata = attention.attention.metadata
+    assert metadata is not None
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
 
 
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
@@ -754,9 +769,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
         packed_total=8,
     )
 
-    assert attention.attention.metadata.attn_mask is None
-    assert attention.attention.metadata.extra["valid_kv_length"] == 5
-    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
+    metadata = attention.attention.metadata
+    assert metadata is not None
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
+    assert metadata.extra["npu_attn_varlen"] is True
 
 
 def test_packed_attention_keeps_padding_mask_for_other_backends():
@@ -796,8 +813,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
         packed_total=8,
     )
 
+    metadata = attention.attention.metadata
+    assert metadata is not None
     assert torch.equal(
-        attention.attention.metadata.attn_mask,
+        metadata.attn_mask,
         torch.tensor([[True, True, True, True, True, False, False, False]]),
     )
 
@@ -1038,6 +1057,7 @@ def test_model_offload_uses_hooked_text_encoder_call():
     )
     expected = torch.ones(2, 3)
     pipeline.text_encoder = Mock(return_value=expected)
+    pipeline._model_cpu_offload_modules = [pipeline.text_encoder]
     input_ids = torch.tensor([1, 2])
     vision_kwargs = {"pixel_values": torch.ones(1, 4)}
 
@@ -1401,7 +1421,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
         lambda path: f"frames:{path}",
     )
 
-    rows, shapes = pipeline._encode_video_conditions(None, count=2)
+    rows, shapes = pipeline._encode_video_conditions_resident(None, count=2)
 
     assert pipeline.video_vae.calls == [
         "frames:video-1.mp4",
@@ -1483,12 +1503,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
         "audio_codecs": (),
         "file_size": 1024,
     }
-    transcode_calls = []
+    transcode_calls: list[Any] = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
+        lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"),
     )
 
     prepared = reference_video_module.prepare_reference_videos(
@@ -1509,7 +1529,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
-    decode_calls = []
+    decode_calls: list[Any] = []
     monkeypatch.setattr(
         reference_video_module,
         "_probe_video",
@@ -1522,7 +1542,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded,
+        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded),
     )
 
     sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")
@@ -1583,7 +1603,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
-    decode_calls = []
+    decode_calls: list[Any] = []
     monkeypatch.setitem(sys.modules, "decord", None)
     monkeypatch.setattr(
         reference_video_module,
@@ -1597,7 +1617,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected,
+        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected),
     )
 
     actual = reference_video_module.load_video_frames("prepared.mp4")
@@ -1637,7 +1657,7 @@ def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
 def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls = []
+    calls: list[Any] = []
     probe = {
         "streams": [
             {
@@ -1655,7 +1675,11 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)),
+        lambda command, **kwargs: _append_and_return(
+            calls,
+            (command, kwargs),
+            SimpleNamespace(stdout=json.dumps(probe)),
+        ),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1668,7 +1692,7 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
 def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls = []
+    calls: list[Any] = []
     stream = {
         "codec_type": "video",
         "width": 4,
@@ -1688,8 +1712,11 @@ def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: calls.append((command, kwargs))
-        or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)),
+        lambda command, **kwargs: _append_and_return(
+            calls,
+            (command, kwargs),
+            SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)),
+        ),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1749,12 +1776,16 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p
             "file_size": 1024,
         },
     }
-    transcode_calls = []
+    transcode_calls: list[Any] = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
+        lambda source, **kwargs: _append_and_return(
+            transcode_calls,
+            (source, kwargs),
+            f"{source}.prepared.mp4",
+        ),
     )
 
     prepared = reference_video_module.prepare_reference_videos(
@@ -1964,7 +1995,7 @@ def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch):
 
     waveform = torch.zeros(1, 10)
     with pytest.raises(ValueError, match="max_duration_seconds must be positive"):
-        pipeline._encode_audio_conditions(
+        pipeline._encode_audio_conditions_resident(
             [(waveform, 10)],
             max_duration_seconds=0,
         )
@@ -1985,7 +2016,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
     pipeline.audio_vae = FakeAudioVAE()
     waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40)
 
-    encoded, lengths = pipeline._encode_audio_conditions(
+    encoded, lengths = pipeline._encode_audio_conditions_resident(
         [(waveform, 10)],
         max_duration_seconds=2.5,
     )

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -0,0 +1,381 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from vllm.lora.lora_weights import PackedLoRALayerWeights
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module
+from vllm_omni.diffusion.models.minimax_h3.lora import (
+    load_minimax_h3_turbo_lora,
+)
+from vllm_omni.errors import OmniClientError
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_TINY_TARGET_DIMS = {
+    "attn.to_q": (2, 3),
+    "attn.to_k": (2, 3),
+    "attn.to_v": (2, 3),
+    "attn.to_out.0": (3, 2),
+    "ff.net.0.proj": (2, 4),
+    "ff.net.2": (2, 2),
+}
+
+
+@pytest.fixture(autouse=True)
+def _use_tiny_h3_dimensions(monkeypatch):
+    monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS)
+
+
+def _request(path) -> LoRARequest:
+    return LoRARequest(
+        lora_name="turbo",
+        lora_int_id=1,
+        lora_path=str(path),
+    )
+
+
+def _write_tiny_turbo(
+    path,
+    *,
+    alpha: str = "128",
+    key_format: str = "minimax-h3-diffusers",
+    omit_target: str | None = None,
+    shape_overrides: dict[str, tuple[int, int]] | None = None,
+) -> None:
+    rank = 128
+    tensors = {}
+    overrides = shape_overrides or {}
+    target_suffixes = (
+        "attn.to_q",
+        "attn.to_k",
+        "attn.to_v",
+        "attn.to_out.0",
+        "ff.net.0.proj",
+        "ff.net.2",
+    )
+    for prefix, block_count in (
+        ("transformer_blocks", 50),
+        ("token_refiner.refiner_blocks", 2),
+    ):
+        for block_index in range(block_count):
+            for suffix in target_suffixes:
+                target = f"{prefix}.{block_index}.{suffix}"
+                if target == omit_target:
+                    continue
+                input_dim, output_dim = _TINY_TARGET_DIMS[suffix]
+                a_name = f"{target}.lora_A.default.weight"
+                b_name = f"{target}.lora_B.default.weight"
+                tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
+                if suffix == "ff.net.0.proj":
+                    if b_name in overrides:
+                        tensors[b_name] = torch.ones(overrides[b_name])
+                    else:
+                        tensors[b_name] = torch.cat(
+                            (
+                                torch.ones(output_dim // 2, rank),
+                                torch.full((output_dim // 2, rank), 2.0),
+                            ),
+                            dim=0,
+                        )
+                else:
+                    tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
+    save_file(
+        tensors,
+        str(path),
+        metadata={"alpha": alpha, "key_format": key_format},
+    )
+
+
+def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    loaded = load_minimax_h3_turbo_lora(
+        partition="fl2va",
+        lora_request=_request(path),
+        lora_path=path,
+        dtype=torch.float32,
+    )
+
+    assert loaded is not None
+    lora_model, peft_helper = loaded
+    assert peft_helper.r == 128
+    assert peft_helper.lora_alpha == 128
+    assert len(lora_model.loras) == 312
+    assert "blocks.0.attn.to_q" in lora_model.loras
+    assert "blocks.0.mlp.fc1" in lora_model.loras
+    fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
+    assert isinstance(fc1, PackedLoRALayerWeights)
+    assert len(fc1.lora_a) == 2
+    assert len(fc1.lora_b) == 2
+    torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2))
+    torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2))
+    torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0))
+    torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128))
+
+
+@pytest.mark.parametrize(
+    ("tensor_suffix", "bad_shape"),
+    [
+        ("attn.to_q.lora_A.default.weight", (128, 1)),
+        ("attn.to_q.lora_A.default.weight", (128, 3)),
+        ("attn.to_q.lora_B.default.weight", (2, 128)),
+        ("attn.to_q.lora_B.default.weight", (4, 128)),
+    ],
+)
+def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    tensor_name = f"transformer_blocks.0.{tensor_suffix}"
+    _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape})
+
+    with pytest.raises(ValueError, match="invalid global shape"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+
+
+def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    class _Pipeline:
+        def _load_diffusion_lora_adapter(self, **kwargs):
+            return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
+
+    manager = object.__new__(DiffusionLoRAManager)
+    manager.pipeline = _Pipeline()
+    manager.dtype = torch.float32
+    manager._expected_lora_modules = {"to_q", "fc1"}
+
+    lora_model, peft_helper = manager._load_adapter(_request(path))
+
+    assert lora_model.id == 1
+    assert peft_helper.lora_alpha == 128
+    assert len(lora_model.loras) == 312
+    assert "blocks.0.attn.to_q" in lora_model.loras
+    assert "blocks.0.mlp.fc1" in lora_model.loras
+
+
+def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
+    wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    wrong_alpha.parent.mkdir()
+    _write_tiny_turbo(wrong_alpha, alpha="8")
+    with pytest.raises(ValueError, match="requires alpha=128"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(wrong_alpha),
+            lora_path=wrong_alpha,
+            dtype=torch.float32,
+        )
+
+    valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    valid.parent.mkdir()
+    _write_tiny_turbo(valid)
+    with pytest.raises(ValueError, match="supports FL2VA/T2VA only"):
+        load_minimax_h3_turbo_lora(
+            partition="ref2va",
+            lora_request=_request(valid),
+            lora_path=valid,
+            dtype=torch.float32,
+        )
+
+
+@pytest.mark.parametrize(
+    "offload_mode",
+    [
+        "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):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    with pytest.raises(ValueError, match="does not support"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+            unsupported_offload_mode=offload_mode,
+        )
+
+
+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)
+
+    with pytest.raises(ValueError, match="supports only"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(unsupported),
+            lora_path=unsupported,
+            dtype=torch.float32,
+        )
+
+    supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(supported)
+    # Directory resolution selects the declared artifact even when another
+    # v1.0 checkpoint is present beside it.
+    assert (
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(tmp_path),
+            lora_path=tmp_path,
+            dtype=torch.float32,
+        )
+        is not None
+    )
+
+
+def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2")
+
+    with pytest.raises(ValueError, match="target set does not match"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+
+
+def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path):
+    path = tmp_path / "other.safetensors"
+    _write_tiny_turbo(path, key_format="other")
+
+    assert (
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+        is None
+    )
+
+    declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(declared, key_format="other")
+    with pytest.raises(ValueError, match="requires safetensors metadata"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(declared),
+            lora_path=declared,
+            dtype=torch.float32,
+        )
+
+
+def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(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 = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    pipeline._turbo_lora_adapter_ids = set()
+    request = _request("generic-peft")
+
+    def load():
+        return pipeline._load_diffusion_lora_adapter(
+            lora_request=request,
+            lora_path=request.lora_path,
+            dtype=torch.float32,
+        )
+
+    def resolve(scale):
+        sampling = SimpleNamespace(lora_request=request, lora_scale=scale)
+        return pipeline._resolve_task(
+            "ref2va",
+            {},
+            has_turbo_lora=pipeline._has_active_turbo_lora(sampling),
+        )
+
+    recognized = (object(), object())
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized)
+    assert load() is recognized
+    assert resolve(0.0) == "ref2va"
+    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+        resolve(1.0)
+
+    # Simulate manager eviction followed by a generic PEFT adapter reusing the
+    # same client-supplied ID. A real reload must replace the stale kind.
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+    assert load() is None
+    assert resolve(1.0) == "ref2va"
+
+
+def test_h3_turbo_requires_all_loaded_targets_to_bind():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._turbo_lora_adapter_ids = {1}
+    lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()})
+
+    pipeline._validate_diffusion_lora_binding(
+        lora_model=lora_model,
+        bound_lora_names=frozenset(lora_model.loras),
+    )
+    with pytest.raises(ValueError, match="bound=1/2"):
+        pipeline._validate_diffusion_lora_binding(
+            lora_model=lora_model,
+            bound_lora_names=frozenset({"to_q"}),
+        )
+
+
+@pytest.mark.parametrize(
+    ("num_inference_steps", "extra_args", "error"),
+    [
+        (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"),
+        (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"),
+        (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"),
+        (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"),
+    ],
+)
+def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    sampling = SimpleNamespace(
+        num_inference_steps=num_inference_steps,
+        extra_args=extra_args,
+    )
+
+    with pytest.raises(OmniClientError, match=error):
+        pipeline._validate_turbo_sampling(sampling)
+
+
+def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    pipeline._validate_turbo_sampling(
+        SimpleNamespace(
+            num_inference_steps=5,
+            extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0},
+        )
+    )

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_h3_profiler_targets_cover_aggregate_reference_encoders():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    targets = MiniMaxH3Pipeline._PROFILER_TARGETS
+    assert "_encode_visual_conditions" in targets
+    assert "_encode_reference_audio_conditions" in targets
+    assert "_encode_video_conditions" not in targets
+    assert "_encode_video_audio_conditions" not in targets
+    assert "_encode_audio_conditions" not in targets
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=not distributed,
+        enable_distributed_layerwise_offload=distributed,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    component = Mock()
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+    with pipeline._component_on_device(component):
+        component.load_to_device.assert_called_once_with()
+        component.offload_to_cpu.assert_not_called()
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+
+
+def test_requested_model_offload_without_backend_uses_resident_components():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.encode_ids.assert_called_once()
+
+
+def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = torch.nn.Linear(2, 2)
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    remove_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+    monkeypatch.setattr(module, "remove_sequential_offload", remove_offload)
+
+    pipeline.enable_omni_model_cpu_offload(
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+    )
+
+    dits = [pipeline.transformer, pipeline.transformers_ref]
+    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    apply_offload.assert_called_once_with(
+        dit_modules=dits,
+        encoder_modules=stages,
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_initial_dits=True,
+    )
+
+    pipeline.disable_omni_model_cpu_offload()
+
+    remove_offload.assert_called_once_with([*dits, *stages])
+
+
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = Mock()
+    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_image(image):
+        events.append(("encode", image.size))
+        if encode_fails and image.width == 32:
+            raise RuntimeError("encode failed")
+        return torch.ones(image.width // 16, 4)
+
+    pipeline.video_vae.encode_image.side_effect = encode_image
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+
+    if encode_fails:
+        with pytest.raises(RuntimeError, match="encode failed"):
+            pipeline._encode_visual_conditions(images, None, video_count=0)
+    else:
+        rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
+        assert rows.shape == (3, 4)
+        assert shapes == [(1, 1, 1), (1, 1, 2)]
+
+    assert events == [
+        ("activate", pipeline.video_vae),
+        ("encode", (16, 16)),
+        ("encode", (32, 16)),
+        ("offload", pipeline.video_vae),
+    ]
+
+
+def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = Mock()
+    pipeline.video_vae.is_distributed_enabled.return_value = False
+    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_image(image):
+        events.append(("image", image.size))
+        return torch.ones(1, 4)
+
+    def encode_video(frames):
+        events.append(("video", frames))
+        return torch.ones(3, 4), (3, 2, 2)
+
+    pipeline.video_vae.encode_image.side_effect = encode_image
+    pipeline.video_vae.encode_video.side_effect = encode_video
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+    monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
+    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+    prepared_videos = [{"prepared_path": "reference.mp4"}]
+
+    rows, shapes = pipeline._encode_visual_conditions(
+        images,
+        prepared_videos,
+        video_count=1,
+    )
+
+    assert rows.shape == (5, 4)
+    assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
+    assert events == [
+        ("activate", pipeline.video_vae),
+        ("image", (16, 16)),
+        ("image", (32, 16)),
+        ("video", "frames:reference.mp4"),
+        ("offload", pipeline.video_vae),
+    ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae = Mock()
+    pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_embedded(*args, **kwargs):
+        events.append(("embedded",))
+        if encode_fails:
+            raise RuntimeError("audio encode failed")
+        return torch.ones(1, 2), [1]
+
+    def encode_standalone(*args, **kwargs):
+        events.append(("standalone",))
+        return torch.ones(1, 2), [1]
+
+    pipeline._encode_video_audio_conditions_resident = encode_embedded
+    pipeline._encode_audio_conditions_resident = encode_standalone
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+
+    def call():
+        return pipeline._encode_reference_audio_conditions(
+            [{"input_has_audio": True}],
+            has_audio=[True],
+            standalone_audios=[(torch.ones(1), 16_000)],
+            max_duration_seconds=1.0,
+        )
+
+    if encode_fails:
+        with pytest.raises(RuntimeError, match="audio encode failed"):
+            call()
+    else:
+        result = call()
+        assert result[0].shape == (1, 2)
+        assert result[2].shape == (1, 2)
+
+    assert events[0] == ("activate", pipeline.audio_vae)
+    assert events[-1] == ("offload", pipeline.audio_vae)
+    assert events.count(("activate", pipeline.audio_vae)) == 1
+    assert events.count(("offload", pipeline.audio_vae)) == 1

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import importlib
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -9,6 +10,7 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -82,6 +84,7 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline.vae_scale_factor_spatial = 8
     pipeline.boundary_ratio = 0.875
     pipeline.expand_timesteps = False
+    pipeline.is_dmd = False
     pipeline._guidance_scale = None
     pipeline._guidance_scale_2 = None
     pipeline._num_timesteps = None
@@ -94,7 +97,7 @@ def _make_pipeline() -> Wan22Pipeline:
 
 
 def _make_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": None,
         "width": None,
         "num_frames": 1,
@@ -338,3 +341,125 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
         (3.0, 3, 28.0, False),
     ]
     assert torch.equal(result, torch.full_like(latents, 10.0))
+
+
+class _StubDMDScheduler:
+    def __init__(self) -> None:
+        self.predict_clean_calls: list[tuple[float, float, float]] = []
+        self.add_noise_calls: list[tuple[float, float, float]] = []
+
+    def predict_clean(self, model_output, sample, timestep):
+        self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep)))
+        return sample - model_output
+
+    def add_noise(self, clean_sample, noise, timestep):
+        self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep)))
+        return clean_sample + 10.0
+
+
+def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None:
+    pipeline = _make_pipeline()
+    pipeline.is_dmd = True
+    pipeline.scheduler = _StubDMDScheduler()
+    latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32)
+    timesteps = torch.tensor([1000.0, 757.0, 522.0])
+
+    pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents)  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor",
+        lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]),
+    )
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=timesteps,
+        prompt_embeds=torch.zeros(1, 8),
+        negative_prompt_embeds=None,
+        guidance_low=1.0,
+        guidance_high=1.0,
+        boundary_timestep=None,
+        dtype=torch.float32,
+        attention_kwargs={},
+        generator=torch.Generator(device="cpu").manual_seed(1),
+    )
+
+    assert pipeline.scheduler.predict_clean_calls == [
+        (1.0, 0.0, 1000.0),
+        (1.0, 9.0, 757.0),
+        (1.0, 18.0, 522.0),
+    ]
+    assert pipeline.scheduler.add_noise_calls == [
+        (-1.0, 2.0, 757.0),
+        (8.0, 2.0, 522.0),
+    ]
+    torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]]))
+
+
+def _make_gate_loading_pipeline():
+    pipeline = Wan22Pipeline.__new__(Wan22Pipeline)
+    nn.Module.__init__(pipeline)
+    gate = WanSelfAttention.__new__(WanSelfAttention)
+    nn.Module.__init__(gate)
+    gate.to_gate_compress = nn.Linear(1, 1)
+    pipeline.gate_holder = gate
+    return pipeline, gate
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        ("pipeline_wan2_2", "Wan22Pipeline"),
+        ("pipeline_wan2_2_i2v", "Wan22I2VPipeline"),
+        ("pipeline_wan2_2_s2v", "Wan22S2VPipeline"),
+        ("pipeline_wan2_2_vace", "Wan22VACEPipeline"),
+    ],
+)
+def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None:
+    module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}")
+    pipeline_cls = getattr(module, class_name)
+    pipeline = pipeline_cls.__new__(pipeline_cls)
+    expected = {"loaded"}
+
+    def fake_loader(model, weights):
+        assert model is pipeline
+        assert list(weights) == [("weight", torch.ones(1))]
+        return expected
+
+    monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader)
+
+    assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected
+
+
+def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None:
+    pipeline, gate = _make_gate_loading_pipeline()
+
+    class _Loader:
+        def __init__(self, model):
+            del model
+
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+    pipeline.load_weights(iter((("other.weight", torch.ones(1)),)))
+
+    assert pipeline.has_gate_compress_weights is False
+    assert gate.to_gate_compress is None
+
+
+def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None:
+    pipeline, gate = _make_gate_loading_pipeline()
+    original_gate = gate.to_gate_compress
+
+    class _Loader:
+        def __init__(self, model):
+            del model
+
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+    pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),)))
+
+    assert pipeline.has_gate_compress_weights is True
+    assert gate.to_gate_compress is original_gate

tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_dmd_predict_clean_and_add_noise_match_flow_formulas() -> None:
+    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+    timestep = torch.tensor(757.0)
+    sigma = scheduler.sigma_for_timestep(timestep)
+    sample = torch.tensor([2.0, -1.0])
+    model_output = torch.tensor([0.5, 0.25])
+    noise = torch.tensor([-0.5, 1.5])
+
+    clean = scheduler.predict_clean(model_output, sample, timestep)
+    expected_clean = sample - sigma * model_output
+    torch.testing.assert_close(clean, expected_clean)
+
+    renoised = scheduler.add_noise(clean, noise, timestep)
+    expected_renoised = (1.0 - sigma) * expected_clean + sigma * noise
+    torch.testing.assert_close(renoised, expected_renoised)
+
+
+def test_dmd_sigma_requires_scalar_timestep() -> None:
+    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+    with pytest.raises(ValueError, match="scalar timestep"):
+        scheduler.sigma_for_timestep(torch.tensor([757.0, 522.0]))

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -1,12 +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 DistributedLayerwiseOffloadHook and backend utilities."""
 
 import gc
 import json
-import os
-import socket
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -18,6 +16,7 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     TensorBinding,
@@ -65,27 +64,10 @@ def dummy_stream(_stream):
     yield None
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
-    os.environ["RANK"] = str(rank)
-    os.environ["LOCAL_RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(master_port)
-
-
 def _cleanup_distributed() -> None:
     if dist.is_initialized():
         dist.destroy_process_group()
 
-    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-        os.environ.pop(key, None)
-
     gc.collect()
     if current_omni_platform.is_available():
         current_omni_platform.empty_cache()
@@ -94,10 +76,7 @@ def _cleanup_distributed() -> None:
 
 @pytest.fixture(scope="module")
 def dist_group():
-    master_port = _find_free_port()
-    _set_dist_env(rank=0, world_size=1, master_port=master_port)
-
-    dist.init_process_group("gloo", rank=0, world_size=1)
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
     try:
         yield
     finally:

tests/diffusion/offloader/test_layerwise_backend.py

@@ -1,11 +1,9 @@
 # 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 LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
 import gc
-import os
-import socket
 from contextlib import contextmanager
 
 import pytest
@@ -15,6 +13,7 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
+from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
@@ -39,27 +38,10 @@ def dummy_stream(_stream):
     yield None
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
-    os.environ["RANK"] = str(rank)
-    os.environ["LOCAL_RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(master_port)
-
-
 def _cleanup_distributed() -> None:
     if dist.is_initialized():
         dist.destroy_process_group()
 
-    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-        os.environ.pop(key, None)
-
     gc.collect()
     if current_omni_platform.is_available():
         current_omni_platform.empty_cache()
@@ -68,10 +50,7 @@ def _cleanup_distributed() -> None:
 
 @pytest.fixture(scope="module")
 def dist_group():
-    master_port = _find_free_port()
-    _set_dist_env(rank=0, world_size=1, master_port=master_port)
-
-    dist.init_process_group("gloo", rank=0, world_size=1)
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
     try:
         yield
     finally: