[Cleanup][Full Duplex] Defer handoff recovery to PR 6170

Signed-off-by: Sy03 <1370724210@qq.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/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/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**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

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

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),
@@ -1661,6 +1659,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 +2127,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 +2224,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 +2854,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

@@ -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,7 @@ 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_rejects_reserved_diffusion_kv_mode(tmp_path):
@@ -1183,6 +1192,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_aborted_queue_sweep.py

@@ -0,0 +1,89 @@
+"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
+
+``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
+handing control to upstream, because upstream ``Scheduler.schedule()`` raises
+``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
+admits in a finished state -- which kills the stage's engine core, not just
+the request.
+
+Regression for the sweep missing ``skipped_waiting``, the third queue
+upstream admits from. An aborted duplex session parked there was re-selected
+by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
+stage.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_POLICIES = [
+    pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
+    pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
+]
+
+
+class _StubRequest:
+    """Minimal ``Request`` stub with the surface the sweep exercises.
+
+    Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
+    relies on so the same stub works under both scheduling policies.
+    """
+
+    def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.priority = 0
+        self.arrival_time = arrival_time
+
+    def __lt__(self, other: _StubRequest) -> bool:
+        return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
+
+
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
+    """Bare scheduler with just the surface the sweep reads."""
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(policy)
+    for req in waiting:
+        scheduler.waiting.add_request(req)
+    scheduler.skipped_waiting = create_request_queue(policy)
+    for req in skipped_waiting:
+        scheduler.skipped_waiting.add_request(req)
+    scheduler.running = list(running)
+    return scheduler
+
+
+@pytest.mark.parametrize("policy", _POLICIES)
+def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
+    monkeypatch: pytest.MonkeyPatch,
+    policy: SchedulingPolicy,
+) -> None:
+    """The public schedule path must remove aborted skipped requests first."""
+    aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
+    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.num_waiting_for_streaming_input = 0
+    scheduler._process_pending_omni_inputs = lambda model_mode: None
+    scheduler._resync_streaming_input_counter = lambda: None
+    scheduler._should_defer_waiting_admission = lambda: False
+    scheduler._restore_omni_wait_queues = lambda: None
+    scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
+    scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
+    scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
+
+    sentinel = object()
+
+    def upstream_schedule(_self, _throttle_prefills: bool = False):
+        assert list(scheduler.skipped_waiting) == []
+        return sentinel
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
+
+    assert OmniARScheduler.schedule(scheduler) is sentinel

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -28,11 +28,16 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_id=stage_id, max_model_len=None),
+    )
     sched.num_waiting_for_streaming_input = 0
     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
 
 
@@ -207,6 +212,81 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.chunk_transfer_adapter.save_async.assert_not_called()
 
 
+def test_stale_async_frame_is_dropped_before_output_processing() -> None:
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.num_in_flight_tokens = 2
+    session.num_computed_tokens = session.num_tokens + 1
+    session.num_output_placeholders = 1
+    session.async_tokens_to_discard = 1
+    session.sampling_params = SimpleNamespace(num_logprobs=1)
+    num_computed_tokens = session.num_computed_tokens
+    num_output_placeholders = session.num_output_placeholders
+
+    sched = MagicMock()
+    sched.requests = {session.request_id: session}
+    sched.perf_metrics = None
+    sched.structured_output_manager.should_advance.return_value = False
+
+    def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
+        request.async_tokens_to_discard = 0
+        return token_ids, False
+
+    sched._update_request_with_output.side_effect = discard_stale_output
+    sched._process_kv_transfer_trigger.return_value = False
+    sched.chunk_transfer_adapter = MagicMock()
+    sched.running = [session]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.multimodal_outputs = None
+    model_runner_output.inter_stage_outputs = [{"hidden": object()}]
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {session.request_id: 0}
+    model_runner_output.routed_experts = None
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert session.async_tokens_to_discard == 0
+    assert session.status == RequestStatus.RUNNING
+    assert session.num_computed_tokens == num_computed_tokens
+    assert session.num_output_placeholders == num_output_placeholders
+    sched.chunk_transfer_adapter.save_async.assert_not_called()
+
+    session.sampling_params = SamplingParams(max_tokens=8)
+    next_payload = {"hidden": object()}
+    sched._update_request_with_output.side_effect = None
+    sched._update_request_with_output.return_value = ([43], False)
+    model_runner_output.sampled_token_ids = [[43]]
+    model_runner_output.inter_stage_outputs = [next_payload]
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+        next_payload,
+        session,
+        False,
+    )
+
+
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()
@@ -276,29 +356,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 +392,73 @@ 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 == {}
+
+
+def test_native_duplex_prompt_reset_prevents_stage1_context_overflow() -> None:
+    sched = _make_scheduler(stage_id=1)
+    sched.vllm_config.model_config.max_model_len = 4096
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=False,
+        segment_finished_requests=set(),
+    )
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.resumable = True
+    session.prompt_token_ids = [0] * 3
+    session._all_token_ids.clear()
+    session._all_token_ids.extend([0] * 4095)
+    session._output_token_ids.extend([1] * 4080)
+    session.num_prompt_tokens = 3
+    session.num_computed_tokens = 4095
+
+    update = _make_update([0] * 16)
+    update.model_intermediate_buffer = {"native_duplex": True}
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 16
+    assert list(session._all_token_ids) == [0] * 16
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 16
+    assert session.num_computed_tokens == 0
+    assert session.model_intermediate_buffer == update.model_intermediate_buffer
+    assert session.status == RequestStatus.WAITING

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -76,6 +76,9 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
+    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+    # abort sweep at the top of ``schedule()`` reads it.
+    sched.skipped_waiting = _MockQueue()
     return sched
 
 

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -0,0 +1,151 @@
+"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
+``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
+and in ``OmniARScheduler.schedule``.
+
+Upstream keeps that counter incrementally (``+1`` when a resumable request
+parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
+next update arrives or the request is finished) and
+``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
+queues. Omni rewrites ``request.status`` outside both hooks -- the
+chunk-transfer adapter's park/restore and
+``_realign_request_status_to_queues`` -- so a counted request can leave
+``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
+counter for the lifetime of the process.
+
+An inflated counter does not degrade throughput, it hangs the stage:
+``EngineCore.has_work()`` reads false while a live request sits in
+``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
+``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
+server answered the first conversation after a boot and was silent from the
+second on while still reporting healthy.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StubRequest:
+    def __init__(self, request_id: str, status: RequestStatus) -> None:
+        self.request_id = request_id
+        self.status = status
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+_SCHEDULER_PARAMS = [
+    pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
+    pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
+]
+
+
+def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = requests
+    scheduler.running = running
+    scheduler.waiting = waiting
+    scheduler.skipped_waiting = skipped if skipped is not None else []
+    scheduler.num_waiting_for_streaming_input = counter
+    return scheduler
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_clears_counter_left_by_a_stomped_status(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """The abort case that wedged the talker.
+
+    The request was counted while parked for streaming input, then had its
+    status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
+    ``finish_requests`` took the running branch and never decremented. With
+    no parked requests left, the counter must land at 0.
+    """
+    aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={},  # upstream already dropped it
+        running=[],
+        waiting=[],
+        counter=1,
+    )
+
+    monkeypatch.setattr(
+        scheduler_mod.VLLMScheduler,
+        "finish_requests",
+        lambda self, request_ids, finished_status: [aborted],
+    )
+
+    scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
+
+    assert scheduler.num_waiting_for_streaming_input == 0
+
+
+def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
+    so the resync has to look at both or it under-counts and reintroduces the
+    spin the counter exists to prevent."""
+    parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
+    scheduler = _make_scheduler(
+        OmniARScheduler,
+        requests={parked.request_id: parked},
+        running=[],
+        waiting=[],
+        counter=0,
+        skipped=[parked],
+    )
+
+    scheduler._resync_streaming_input_counter()
+
+    assert scheduler.num_waiting_for_streaming_input == 1
+
+
+def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pins placement in ``schedule``: the chunk-transfer adapter rewrites
+    status during ``process_pending_chunks``, so the resync has to run after
+    that and before upstream schedules, while ``self.waiting`` is still the
+    real queue."""
+    parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = {parked.request_id: parked}
+    scheduler.running = [parked]
+    scheduler.waiting = []
+    scheduler.skipped_waiting = []
+    scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
+    scheduler.policy = "fcfs"
+
+    seen: dict[str, int] = {}
+
+    monkeypatch.setattr(
+        ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
+    )
+    monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
+    monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
+
+    def fake_schedule(self, throttle_prefills=False):
+        seen["counter"] = self.num_waiting_for_streaming_input
+        raise _StopScheduleError
+
+    class _StopScheduleError(Exception):
+        pass
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
+
+    with pytest.raises(_StopScheduleError):
+        OmniARScheduler.schedule(scheduler)
+
+    assert seen["counter"] == 0

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_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/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

@@ -1038,6 +1038,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 +1402,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",
@@ -1964,7 +1965,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 +1986,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_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/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:

tests/diffusion/offloader/test_sequential_backend.py

@@ -8,7 +8,13 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
+from vllm_omni.diffusion.offloader.sequential_backend import (
+    ModelLevelOffloadBackend,
+    SequentialOffloadHook,
+    apply_sequential_offload,
+    remove_sequential_offload,
+    sequential_offload_component,
+)
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    offloaded: list[nn.Module] = []
+    monkeypatch.setattr(
+        SequentialOffloadHook,
+        "_to_cpu",
+        lambda self, module: offloaded.append(module),
+    )
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=torch.device("cpu"),
+        offload_initial_dits=True,
+    )
+
+    assert offloaded == [dit]
+    remove_sequential_offload([dit, encoder])
+
+
+def test_direct_component_activation_failure_still_offloads(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    component = _create_simple_module()
+    apply_sequential_offload(
+        dit_modules=[_create_simple_module()],
+        encoder_modules=[component],
+        device=torch.device("cpu"),
+    )
+    hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    offloaded: list[nn.Module] = []
+
+    def fail_activation(module: nn.Module) -> None:
+        raise RuntimeError("activation failed")
+
+    monkeypatch.setattr(hook, "pre_forward", fail_activation)
+    monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module))
+
+    with pytest.raises(RuntimeError, match="activation failed"):
+        with sequential_offload_component(component):
+            pass
+
+    assert offloaded == [component]
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""
@@ -85,7 +137,12 @@ class TestMoveParamsPinMemory:
         original_isinstance = isinstance
 
         def fake_isinstance(obj, cls):
-            if cls.__name__ == "DTensor":
+            # torchada's Tensor.to() checks tuples such as
+            # ``(str, torch.device)``.  Keep the DTensor probe compatible with
+            # the normal isinstance contract instead of breaking those calls.
+            if type(cls) is tuple:
+                return any(fake_isinstance(obj, candidate) for candidate in cls)
+            if getattr(cls, "__name__", None) == "DTensor":
                 return True
             return original_isinstance(obj, cls)
 

tests/diffusion/quantization/test_component_routing.py

@@ -9,6 +9,9 @@ import torch
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
 )
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+    CompressedTensorsConfig,
+)
 from vllm.model_executor.models.utils import WeightsMapper
 
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
@@ -344,7 +347,41 @@ class TestTalkerVisualRouting:
 
 
 # ===================================================================
-# 4. ComponentQuantizationConfig.resolve
+# 4. ComponentQuantizationConfig.apply_vllm_mapper
+# ===================================================================
+
+
+class TestComponentApplyVllmMapper:
+    def test_maps_compressed_tensors_component_and_default(self):
+        thinker = CompressedTensorsConfig(
+            target_scheme_map={},
+            ignore=["thinker.model.layers.0.mlp.gate"],
+            quant_format="pack-quantized",
+        )
+        default = CompressedTensorsConfig(
+            target_scheme_map={},
+            ignore=["talker.model.layers.0.mlp.gate"],
+            quant_format="pack-quantized",
+        )
+        config = ComponentQuantizationConfig(
+            component_configs={"thinker": thinker, "code2wav": None},
+            default_config=default,
+        )
+        mapper = WeightsMapper(
+            orig_to_new_prefix={
+                "thinker.model": "thinker.language_model.model",
+                "talker.model": "talker.language_model.model",
+            }
+        ).get_unstacked_mapper()
+
+        config.apply_vllm_mapper(mapper)
+
+        assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"]
+        assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"]
+
+
+# ===================================================================
+# 5. ComponentQuantizationConfig.resolve
 # ===================================================================
 
 

tests/diffusion/quantization/test_quantization_quality.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
 """
 Quantization quality gate for diffusion models.
 
@@ -53,6 +53,7 @@ _BENCH_MODULE_NAME = "benchmarks.diffusion.quantization_quality"
 
 if _BENCH_MODULE_NAME not in sys.modules:
     _spec = importlib.util.spec_from_file_location(_BENCH_MODULE_NAME, _BENCH_MODULE_PATH)
+    assert _spec is not None and _spec.loader is not None
     _mod = importlib.util.module_from_spec(_spec)
     sys.modules[_BENCH_MODULE_NAME] = _mod
     _spec.loader.exec_module(_mod)
@@ -93,7 +94,7 @@ class QualityTestConfig:
             return self.quantized_model
         return self.model or ""
 
-    def quantization_ref(self) -> str | None:
+    def quantization_ref(self) -> str | dict[str, object] | None:
         if self.quantized_model is not None:
             return None
         return self.quantization
@@ -500,6 +501,8 @@ def _quality_param(c: QualityTestConfig):
         marks.append(
             pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
         )
+    if c.id == "fp8_ltx2":
+        marks.append(pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/6245"))
     return pytest.param(c, id=c.id, marks=marks)
 
 

tests/diffusion/test_async_output_timeout.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the async-output wait bound (_async_output_timeout)."""
+
+import logging
+
+import pytest
+
+from vllm_omni.diffusion.diffusion_engine import (
+    _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+    _ASYNC_OUTPUT_TIMEOUT_ENV,
+    _async_output_timeout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestAsyncOutputTimeoutIsConfigurable:
+    def test_default_is_generous(self, monkeypatch):
+        """The bound covers a copy queued behind a full denoise step, so the
+        default must not be tight enough to abort a slow but healthy render.
+        """
+        monkeypatch.delenv(_ASYNC_OUTPUT_TIMEOUT_ENV, raising=False)
+        assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT == 600.0
+
+    def test_env_var_overrides_the_default(self, monkeypatch):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "1800")
+        assert _async_output_timeout() == 1800.0
+
+    def test_env_var_accepts_a_float(self, monkeypatch):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "45.5")
+        assert _async_output_timeout() == 45.5
+
+    def test_value_is_not_frozen_at_import(self, monkeypatch):
+        """Resolved per call rather than captured in a module constant, so the
+        value never goes stale relative to os.environ and no module reload is
+        needed to exercise it.
+        """
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "60")
+        assert _async_output_timeout() == 60.0
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "120")
+        assert _async_output_timeout() == 120.0
+
+
+class TestAsyncOutputTimeoutRejectsBadValues:
+    """This runs on the request path, so a typo in the environment must degrade
+    to the default rather than fail the generation.
+    """
+
+    @pytest.mark.parametrize("value", ["", "abc", "30s", "None"])
+    def test_non_numeric_falls_back_to_the_default(self, monkeypatch, caplog, value):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+        with caplog.at_level(logging.WARNING):
+            assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+
+    @pytest.mark.parametrize("value", ["0", "-1", "-0.5"])
+    def test_non_positive_falls_back_to_the_default(self, monkeypatch, caplog, value):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+        with caplog.at_level(logging.WARNING):
+            assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT

tests/diffusion/test_compile.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 
 import vllm_omni.diffusion.compile as compile_module
 from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -76,6 +77,37 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
 
 
+def test_regionally_compile_keeps_hook_dispatch_outside_compiled_graph(monkeypatch):
+    model = _ModelWithWrappedRepeatedBlocks()
+    block = model.transformer_blocks[0]
+    registry = HookRegistry.get_or_create(block)
+    hook = ModelHook()
+    registry.register_hook("test", hook)
+
+    wrapped_forward = block.forward
+    original_forward = block._omni_original_forward
+    compile_calls = []
+
+    def _compile(fn, *args, **kwargs):
+        compile_calls.append(fn)
+
+        def _compiled(*fn_args, **fn_kwargs):
+            return f"compiled:{fn(*fn_args, **fn_kwargs)}"
+
+        return _compiled
+
+    monkeypatch.setattr(compile_module.torch, "compile", _compile)
+
+    regionally_compile(model)
+
+    assert compile_calls[0] is original_forward
+    assert compile_calls[0] is not wrapped_forward
+    assert block.forward is wrapped_forward
+    assert block._omni_original_forward is not original_forward
+    assert hook.fn_ref.original_forward is block._omni_original_forward
+    assert block("ok") == "compiled:ok"
+
+
 def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch):
     """cache-dit matches blocks via inspect.signature(block.forward).
 

tests/diffusion/test_diffusion_config_propagation.py

@@ -149,10 +149,18 @@ class TestCreateDefaultDiffusion:
         assert ea["enforce_eager"] is True
         assert ea["lora_path"] == "/tmp/lora"
 
-    def test_diffusion_kv_mode_roundtrip(self):
-        od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
+    def test_diffusion_kv_mode_roundtrip(self, monkeypatch):
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+        od = _roundtrip_diffusion_config(
+            model="x",
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+        )
 
         assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+        assert od.diffusion_kv_max_rows_per_request == 2
 
     def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
         monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -12,6 +12,7 @@ import pytest
 from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
     assert engine.scheduler.get_request_state(request_id) is None
 
 
+def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]:
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+    cleanup = Mock()
+    engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup)
+    return engine, cleanup
+
+
+def test_paged_terminal_output_clears_worker_rows() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    engine._finalize_finished_request = lambda request_id, *_args: request_id
+    engine._put_output = lambda *_args: None
+
+    engine._emit_finished_outputs({"req-0", "req-1"})
+
+    cleanup.assert_called_once()
+    assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"}
+
+
+def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    engine.scheduler = SimpleNamespace(
+        get_request_state=lambda _request_id: None,
+        finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")),
+    )
+
+    engine._abort_requests(["req-idle", "req-idle"])
+
+    cleanup.assert_called_once_with(["req-idle"])
+    engine.scheduler.finish_requests.assert_not_called()
+
+
+def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY)
+    engine._finalize_finished_request = lambda request_id, *_args: request_id
+    engine._put_output = lambda *_args: None
+    engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None)
+
+    engine._emit_finished_outputs({"req-0"})
+    engine._abort_requests(["req-idle"])
+
+    cleanup.assert_not_called()
+
+
 def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
     od_config = SimpleNamespace(
         custom_pipeline_args=None,

tests/diffusion/test_diffusion_model_runner.py

@@ -11,6 +11,7 @@ import torch
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 
@@ -193,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         parallel_config=SimpleNamespace(use_hsdp=False),
         streaming_output=False,
     )
+    runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
+        vllm_config=runner.vllm_config,
+        od_config=runner.od_config,
+        device=runner.device,
+    )
     runner.kv_transfer_manager = SimpleNamespace(
         receive_kv_cache=lambda req, target_device=None: None,
         receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,