[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup.

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/cuda/test-nightly.yml

@@ -159,7 +159,19 @@ steps:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"
           - export VLLM_MOE_USE_DEEP_GEMM="0"
-          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
+        mirror_hardwares: l4_1
+
+      - label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
+            "
         mirror_hardwares: l4_1
 
       - label: ":full_moon: TTS · Perf Test"
@@ -176,6 +188,35 @@ steps:
             exit $$EXIT
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: TTS · VoxCPM2 · Perf Test"
+        key: nightly-tts-performance-voxcpm2
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        mirror_hardwares: h100_1
+
+      - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
+        key: nightly-tts-performance-higgs-audio-v3
+        timeout_in_minutes: 180
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_USE_DEEP_GEMM="0"
+          - export VLLM_MOE_USE_DEEP_GEMM="0"
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+        mirror_hardwares: h100_1
+
   - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
     key: nightly-tiny-model-test-group
     depends_on: upload-nightly-pipeline
@@ -205,7 +246,8 @@ steps:
       build.pull_request.labels includes "nightly-test" ||
       build.pull_request.labels includes "diffusion-x2iat-test"
     steps:
-      # Use explicit file shards so single-card cases do not reserve 2+ GPUs.
+      # Use explicit file shards so single-card cases do not reserve 2+ GPUs,
+      # and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
         timeout_in_minutes: 120
@@ -213,6 +255,15 @@ steps:
           - >-
             pytest -sv
             tests/e2e/online_serving/test_qwen_image_expansion.py
+            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+            tests/e2e/online_serving/test_qwen_image_layered_expansion.py
+            tests/e2e/online_serving/test_boogu_image.py
+            tests/e2e/online_serving/test_boogu_image_edit.py
+            tests/e2e/online_serving/test_boogu_image_expansion.py
+            tests/e2e/online_serving/test_boogu_image_edit_expansion.py
+            tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/online_serving/test_krea2_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and not distributed_cuda"
             --run-level "full_model"
         mirror_hardwares: h100_1
@@ -223,14 +274,47 @@ steps:
           - >-
             pytest -sv
             tests/e2e/online_serving/test_qwen_image_expansion.py
+            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
+        timeout_in_minutes: 90
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_qwen_image_layered_expansion.py
+            -m "full_model and diffusion and H100 and distributed_cuda"
+            --run-level "full_model"
+        mirror_hardwares: h100_2
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
+        timeout_in_minutes: 120
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_bagel_expansion.py
+            -m "full_model and diffusion and H100 and distributed_cuda"
+            --run-level "full_model"
+        mirror_hardwares: h100_2
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
+        timeout_in_minutes: 120
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_glm_image_expansion.py
+            tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
+            -m "full_model and diffusion and H100"
+            --run-level "full_model"
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 120
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
         mirror_hardwares: l4_4
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -243,7 +327,7 @@ steps:
         timeout_in_minutes: 180
         commands:
           - export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
-          - pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
         mirror_hardwares: h100_1
 
@@ -271,6 +355,70 @@ steps:
             exit $$EXIT
         mirror_hardwares: h100_4
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.5.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_4
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.5.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_4
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
+        key: nightly-diffusion-x2iat-performance-qwen-image-layered
+        timeout_in_minutes: 120
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.5.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_1
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
+        key: nightly-diffusion-x2iat-performance-bagel
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.5.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        mirror_hardwares: h100_1
+
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
     depends_on: upload-nightly-pipeline
@@ -292,6 +440,12 @@ steps:
           - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
+        timeout_in_minutes: 90
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Doc Test"
         timeout_in_minutes: 60
         commands:
@@ -302,8 +456,15 @@ steps:
         timeout_in_minutes: 90
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · LTX-2.5 Online Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model"
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
         timeout_in_minutes: 180
         commands:
@@ -322,8 +483,23 @@ steps:
         commands:
           - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
+          - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
+        mirror_hardwares: h100_1
+
+      - label: ":full_moon: Diffusion X2V · LTX-2.5 Official Parity Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -m full_model --run-level full_model
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Perf Test"
         key: nightly-diffusion-x2v-performance
         timeout_in_minutes: 180
@@ -399,6 +575,16 @@ steps:
           - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
         mirror_hardwares: l4_1
 
+  - label: ":bar_chart: Testcase Statistics"
+    key: nightly-testcase-statistics
+    timeout_in_minutes: 120
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+    mirror_hardwares: h100_1
+
   # NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
   # Do not add new depends_on entries or extend this step.
   - label: ":email: Nightly Collection & Email"
@@ -409,20 +595,32 @@ steps:
       - nightly-omni-performance-vllm-text
       - nightly-omni-performance-multi-replicas
       - nightly-tts-performance
+      - nightly-tts-performance-higgs-audio-v3
       - nightly-diffusion-x2iat-performance-qwen-image
+      - nightly-diffusion-x2iat-performance-qwen-image-edit
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - nightly-diffusion-x2iat-performance-qwen-image-layered
+      - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
+      - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
     commands:
       - pip install openpyxl
       - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
       - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
       - 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-x2iat-performance-qwen-image-edit
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
       - python tools/nightly/generate_nightly_perf_excel.py
       - python tools/nightly/generate_nightly_perf_html.py
       - python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"

.buildkite/cuda/test-weekly.yml

@@ -149,4 +149,4 @@ steps:
             pytest -sv
             tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
-        mirror_hardwares: h100_4
+        mirror_hardwares: h100_2

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

@@ -96,6 +96,33 @@ steps:
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
             exit $$EXIT
+  - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
+    key: nightly-diffusion-x2iat-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
+    steps:
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_8
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export CACHE_DIT_VERSION=1.5.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
   - group: ":card_index_dividers: Diffusion I2I Model Test"
     key: nightly-HunyuanImage3-x2i-test-group
     depends_on: upload-nightly-pipeline

.pre-commit-config.yaml

@@ -66,7 +66,7 @@ repos:
       entry: python tools/pre_commit/check_tts_adapter.py
       language: python
       pass_filenames: false
-      files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
       # The budget-is-now-slack reminder is printed on a PASSING run, and
       # pre-commit hides output from passing hooks unless the hook is verbose.
       verbose: true

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -19,15 +19,17 @@ DLO weight collective.
 
 Host storage is selected separately from the transfer protocol. The loader can
 produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
-otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
-the same node can share immutable checkpoint pages when direct mmap is
-selected, while the ordinary-loader fallback still keeps a private runtime
-copy per process.
+otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
+runtime cache can normalize those final ordinary-loader DiT tensors into an
+immutable node-local mmap entry. Consequently, replicas can share either
+proven-compatible checkpoint pages or equivalent final runtime layouts. An
+opt-in CUDA registration budget can make the complete final mapping directly
+H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
 
 The Phase A shared-mmap support boundary is TP1. TP greater than one is an
 ordinary-loader compatibility path: DLO can consume the resulting TP-local
-tensors, but those configurations do not use checkpoint mmap and must not be
-used to claim shared-mmap host-memory savings.
+tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
+can share each final TP coordinate among equivalent processes.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -81,6 +83,90 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
+### Post-load runtime cache
+
+The Phase B cache preserves the same ownership boundary. DLO neither derives
+cache keys nor interprets model loaders. When enabled with no-AllGather, every
+rank first completes the ordinary loader, custom callbacks, post-load
+processing, strict validation, and calibration. The loader then:
+
+1. discovers final CPU DiT parameters and persistent buffers;
+2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
+   ambiguous layouts before changing the model;
+3. hashes tensor names, metadata, and final bytes;
+4. builds or joins one content-addressed entry under a strict POSIX lock; and
+5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
+   manifest, and remapped-tensor content validation succeeds.
+
+The identity excludes DP size/rank and SP rank. It includes TP world size/rank
+and conservative SP implementation/world-size guards. DP replicas and
+equivalent SP ranks can therefore share; different TP coordinates cannot.
+Final tensor content is authoritative, so an unmodeled loader/config change
+that changes bytes creates a different entry instead of risking silent
+sharing.
+
+Publication writes sharded safetensors and a JSON manifest into a same-
+filesystem temporary directory, flushes them, and atomically renames the
+directory. The manifest includes per-file SHA-256 digests and the final tensor
+content digest. Kernel `flock` ownership distinguishes a live writer from a
+stale lock file: process death releases the lock automatically. Stale temporary
+entries are removed only by the next lock owner. A bounded wait or any cache
+error retains the already-valid ordinary tensors.
+
+DLO maps this plan without moving the DiT to `meta` and without rerunning any
+post-load hook. It changes tensor storage while preserving `Parameter` objects,
+then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
+streaming hooks. The source mappings remain open through DLO disable. With the
+default zero registration budget they feed the existing bounded two-slot host
+staging path. Every streamed block is first packed from pageable mmap storage
+into one of those pinned slots and is then copied to the device. This removes
+model-sized private pinned storage, but adds a recurrent host-to-host copy and
+can amplify TP synchronization waits when several ranks compete for host
+bandwidth.
+
+This version is deliberately opt-in and steady-state-oriented. It requires
+roughly one local-disk model copy per runtime identity, has no eviction, and
+still performs ordinary loading plus full content-hash passes in every rank.
+Skip-load-on-hit is a separate follow-up after loader side effects can be
+modeled safely.
+
+### Registered runtime-cache transfer
+
+`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
+runtime-cache registration. It also changes loader selection: no-AllGather
+uses the ordinary loader and final runtime cache even when a TP1 direct-
+checkpoint plan is available. A checkpoint binding can carry a deferred
+adapter such as grouped-QKV reordering. Registering that raw source would not
+remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
+contains the transform-complete final tensors.
+
+After mapping and validating every final tensor, DLO groups source storages by
+backing file, page-aligns their address ranges, and coalesces overlapping or
+adjacent ranges within that mapping. The complete byte count is checked
+against the per-worker budget before any CUDA call. It then registers every
+range with the read-only CUDA host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. Failure rolls back any ranges already
+registered and preserves the existing two-slot staging path. Partial direct
+transfer is deliberately unsupported. A rollback failure aborts worker
+initialization so a still-registered range is retained until process/context
+teardown rather than being unsafely unmapped.
+
+On success, the hooks copy each immutable tensor view directly into its offset
+in the existing flattened rotating device buffer on the copy stream. Parameter
+repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
+implementation currently issues more H2D copy operations than block packing,
+but avoids the much larger recurrent CPU copy. A packed cache schema should be
+considered only if launch-count measurements justify its added format and
+publication complexity.
+
+Registration is per process and CUDA context; physical file-cache pages remain
+shared and are not copied into one private host allocation per worker. The
+operator must nevertheless budget page-locked memory and satisfy OS/CUDA
+registration limits for every worker. During shutdown, DLO first synchronizes
+outstanding device work, unregisters all ranges, and then closes the safetensors
+mappings. Worker teardown invokes this cleanup before destroying the
+distributed environment and CUDA context.
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -114,11 +200,17 @@ When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
 requests must take the same execution path at every denoising step.
 
+Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
+independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
+plus shared mmap storage avoids coupling their request schedules and failure
+domains.
+
 ### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
 one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
+ordinary runtime tensors.
 
 For direct mmap, each process retains immutable safetensors views and uses two
 bounded pinned host staging slots. Processes on the same node that map the same
@@ -127,10 +219,14 @@ the persistent private full-model copy per pure-DP process, but each process
 still packs and transfers every complete block. Sharing is node-local; each
 node has its own page cache.
 
+For a runtime-cache mapping, a positive registration budget replaces those
+host slots with direct copies from the shared final tensor views. A zero budget
+or failed registration retains the bounded staging behavior.
+
 When direct mmap preflight fails, the regular model loader remains responsible
 for preparing each rank's weights, including TP-local tensors or HSDP-managed
-parameters. In that fallback, each pure-DP process keeps a private full runtime
-copy.
+parameters. Supported CPU layouts may then enter the opt-in runtime cache.
+Unsupported layouts keep one private runtime copy per process.
 
 This mode means:
 
@@ -141,17 +237,17 @@ This mode means:
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
 - Pure DP deployments share one checkpoint-backed copy per node when direct
-  mmap is selected; the ordinary-loader fallback keeps one private runtime
-  copy per rank.
+  mmap is selected. The runtime cache can share equivalent ordinary-loader
+  outputs; its fallback keeps one private runtime copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
-| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
+| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
 
 ### Combined dimensions
@@ -182,12 +278,10 @@ requirements.
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
 ordinary loader. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
-
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+use `--dlo-no-use-allgather` or disable online quantization. Runtime-cache v1
+is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
+expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
+tied/shared storage.
 
 ## Validation coverage
 
@@ -200,6 +294,13 @@ Current source-level validation includes:
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;
+- runtime-cache content identity, sharded atomic publication, concurrent
+  writer election, corruption rebuild, lock timeout, and fail-closed layout
+  rejection;
+- loader ordering through final post-load mutation and DLO remapping without
+  rerunning post-load hooks;
+- all-or-nothing page-aligned registration, budget rejection, partial-failure
+  rollback, direct-copy staging bypass, and explicit shutdown cleanup;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
@@ -259,11 +360,99 @@ from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
 because it counts a shared physical page in every process that maps it; summed
 PSS is the appropriate node-memory comparison.
 
-The highest-value missing coverage is broader end-to-end numerical and
-lifecycle comparison against ordinary layerwise offload for DP+SP,
-HSDP+SP+no-AllGather, and TP greater than one across additional models and
-target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
-change the Phase A direct-mmap TP1 support boundary.
+### L20X runtime-cache smoke matrix
+
+Phase B was also exercised on four L20X GPUs in one NVLink domain with
+MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
+outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
+VAE TP1. This is a functional, memory, and profiler smoke test rather than a
+production-quality performance benchmark. Host values below are node-wide
+maxima during request waves from descendant-process `smaps_rollup`; they do
+not represent the private startup peak. HBM is the engine-reported request
+peak because sampled `nvidia-smi` data can miss short-lived allocations.
+
+One internal DP2xTP2 engine compared the three host-storage/transfer choices.
+Each measured wave contained two requests. NVLink TX is the physical per-link
+counter delta summed across the four GPUs; RX had the same value and is not
+added again here.
+
+| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
+|---|---:|---:|---:|---:|---:|---:|
+| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
+| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
+
+The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
+footprint while reducing node PSS by about 134 GiB relative to private loader
+weights. Cold and warm labels describe cache publication or reuse during
+initialization; both request phases ran after initialization. A warm hit still
+performed ordinary loading and full content hashing, so it was not a startup
+fast path. All measured no-AllGather private/cache outputs were byte-identical
+for matching seeds.
+
+The deployment shape targeted by the cache was tested separately as two
+independent DP1xTP2 engines on the same node. Both engines mapped the same two
+TP-coordinate entries, including the same file device/inode tuples; the cache
+occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
+ahead of profiling.
+
+| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
+|---|---:|---:|---:|---:|---:|
+| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
+| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
+
+The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
+`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
+69.3% in this host-contention smoke. Across one profiled request per engine,
+aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
+payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
+rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
+3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
+summed across ranks and can overlap; they identify the mechanism, not wall
+time. Matching private/cache outputs were byte-identical.
+
+The result confirms that zero-budget runtime-cache staging is a host-capacity
+option, not a latency-neutral replacement for private pinned weights. Broader
+SP, model, platform, and production-quality validation remains tracked in
+[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
+historical B300 rows above predate the runtime cache and must not be treated as
+Phase B memory validation.
+
+### L20X registered-runtime-cache smoke
+
+The read-only registered path was then tested with the same MiniMax-H3
+workload. These are functional and mechanism-validation measurements, not a
+production benchmark. The TP1 pair used one request per wave. Both modes moved
+exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
+and audio for both measured seeds.
+
+| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
+| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
+
+Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
+fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
+unchanged, identifying recurrent host transform/packing as the removed work.
+The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
+that launch increase was much smaller than the eliminated CPU cost.
+
+Two independent TP2 engines also registered the same two TP-coordinate cache
+entries across four workers. Compared with the zero-budget cache, mean wave
+latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
+the private-loader reference was 3.61/3.72 seconds. Registered request-time
+host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
+Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
+staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
+payload were effectively unchanged. Ordinary TP communication remained, but
+the extra wait caused by staging-driven rank skew largely disappeared.
+
+All workers explicitly unregistered before releasing their mappings, the
+system's acquired-minus-released pin counter returned to its pre-run value,
+and every GPU allocation was released. Cold runtime-cache publication remained
+expensive because it still performs ordinary loading and multiple full-content
+passes; that startup problem is outside the registered transfer scope.
 
 ## Recommendations
 
@@ -272,8 +461,12 @@ change the Phase A direct-mmap TP1 support boundary.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
 - Use **no-AllGather** when independent replica execution is required. TP1
-  direct-mmap deployments can share checkpoint pages per node; other layouts
-  retain the ordinary loader's private host memory behavior and are outside
-  the Phase A shared-mmap support scope.
+  direct-mmap deployments can share checkpoint pages per node. Enable the
+  runtime cache only when independent replicas have repeated final layouts and
+  benchmark the target CPU/memory topology, especially for matching
+  coordinates across multiple TP engines. On CUDA, provide an explicit
+  registration budget large enough for the complete final layout when
+  recurrent staging cost is unacceptable; leave it at zero when page-locking
+  that layout is not operationally acceptable.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
   HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -5,7 +5,11 @@ deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
 the host weights and reconstructs each layer at runtime. Without AllGather,
 each rank streams a complete block independently. Compatible TP1 deployments
 can share checkpoint-backed host pages among processes on the same node;
-otherwise DLO streams the ordinary loader's rank-local tensors.
+otherwise DLO streams the ordinary loader's rank-local tensors. The optional
+runtime-weight cache extends node-local sharing to final ordinary-loader
+layouts, including matching TP coordinates across independent engines. On
+CUDA, an explicit registration budget can make those shared mappings directly
+H2D-copyable and avoid recurrent mmap-to-pinned host packing.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -21,8 +25,10 @@ AllGather:       [AG N+1]           [AG N+2]
 Slots:           slot 0: Layer N    slot 1: Layer N+1
 ```
 
-AllGather communicates only request-independent weight shards, so data-
-parallel ranks may process different requests concurrently.
+AllGather communicates request-independent weight shards, but every member of
+the transfer group must request the same next block in the same collective
+order. A synchronized DP wave may contain different prompts only when those
+requests follow the same denoising and block-execution path.
 
 ## Usage
 
@@ -38,6 +44,15 @@ vllm serve /path/to/model --omni \
   --data-parallel-size 4 \
   --dlo-no-use-allgather
 
+# Independently scheduled replicas sharing final runtime weights on one node
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --tensor-parallel-size 2 \
+  --dlo-no-use-allgather \
+  --dlo-enable-runtime-cache \
+  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
+  --dlo-runtime-cache-pin-limit-gib 80
+
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
@@ -50,7 +65,10 @@ from vllm_omni import Omni
 omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
+    dlo_use_allgather=False,
+    dlo_enable_runtime_cache=True,
+    dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
+    dlo_runtime_cache_pin_limit_gib=80,
 )
 ```
 
@@ -62,6 +80,10 @@ omni = Omni(
 | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
+| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
+| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
+| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
 ## Host-weight loading
@@ -78,6 +100,50 @@ DLO may still consume those TP-local tensors, but this is a compatibility path:
 it does not share checkpoint-backed runtime weights across DP replicas and
 provides no shared-mmap host-memory guarantee.
 
+When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
+gains a second storage opportunity. Every rank completes loader callbacks,
+post-load processing, validation, and calibration. It then hashes the final
+CPU DiT parameters and persistent buffers. One process per equivalent runtime
+layout publishes sharded safetensors plus a checksummed manifest; peers
+validate and mmap the same files. DLO replaces only tensor storage, preserves
+the final `Parameter` objects, and releases the private allocator-backed copies
+before installing its streaming hooks.
+
+A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
+runtime layout even when a direct TP1 checkpoint plan is available. Raw
+checkpoint mappings may require deferred loader transforms such as grouped-QKV
+reordering; registering them would still leave recurrent host work. Publishing
+the fully transformed layout once makes every streamed source directly
+copyable.
+
+The runtime layout identity excludes DP rank and SP rank. It includes TP world
+size/rank and conservative SP implementation/world-size guards. Therefore two
+independent `TP=2` engines normally create two entries: both TP0 workers map
+one entry and both TP1 workers map the other. Cache sharing does not create a
+process group or an inference-time collective.
+
+With a zero registration budget, cached tensors remain pageable mmap views.
+DLO packs each streamed block into one of two bounded pinned host slots before
+the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
+host-to-host copies; competing TP ranks or engines can turn DDR contention and
+staging skew into collective wait time.
+
+With a positive budget on CUDA, DLO instead attempts to register the complete
+page-aligned cache mapping as read-only pinned host memory. Registration is
+all-or-nothing: the total mapped range must fit the per-worker budget, and any
+CUDA registration or verification failure rolls back partial registrations and
+uses the bounded staging path. If CUDA cannot roll back a partial registration,
+worker initialization aborts instead of unmapping memory that CUDA may still
+own. On success, each immutable tensor view is copied directly into the
+existing rotating device buffers on the H2D stream, so no private host staging
+slots are allocated.
+
+CUDA registration is process-local, but it does not clone file contents. The
+same physical page-cache pages remain shared among equivalent workers; each
+worker does, however, hold its own registration reference and must satisfy the
+host's page-locking limits. Shutdown synchronizes outstanding GPU work,
+unregisters every range, and only then closes the mmap handles.
+
 The mmap plan skips only dedicated DiT weight sources. Other component sources,
 such as a text encoder loaded through the shared diffusion loader, continue to
 use their ordinary component loader. A checkpoint source that mixes DiT and
@@ -104,6 +170,34 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+## Runtime-cache lifecycle
+
+Runtime-cache publication uses a per-layout POSIX file lock, a temporary
+directory on the same filesystem, file and final-tensor SHA-256 validation,
+`fsync`, and an atomic directory rename. A crashed writer releases its kernel
+lock; the next writer removes only that layout's stale temporary directories
+while holding the lock. Lock timeout, unsupported layout, insufficient disk,
+or validation failure keeps the already-valid ordinary tensors.
+
+The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
+vLLM-Omni may isolate that variable between stage replicas. All consumers that
+should share must receive the same `--dlo-runtime-cache-dir` on a local,
+disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
+host-memory PSS, and do not use a cross-node directory in this version.
+
+There is no automatic eviction. Budget roughly one final DiT copy per distinct
+TP coordinate, dtype/layout, and model version, plus temporary space while an
+entry is published. Operators may remove entries or the complete cache root
+only after all workers using its mmap files have stopped. The feature is
+opt-in for this reason.
+
+For direct H2D, set the registration limit above the complete page-aligned
+layout size reported in the worker log. Partial registration is intentionally
+unsupported because mixing direct and staged tensors would complicate buffer
+lifetime and make performance topology-dependent. Registration is CUDA-only;
+an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
+registration error is reported and falls back to bounded staging.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -132,22 +226,37 @@ must enter each collective.
 
 ## Limitations
 
-- Direct checkpoint mmap currently requires TP1. TP greater than one is
-  outside the Phase A shared-mmap support scope and falls back before model
-  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  but it provides no shared-mmap host-memory benefit or guarantee.
+- Direct checkpoint mmap currently requires TP1. TP greater than one falls
+  back to the ordinary TP-aware loader. With the optional runtime cache,
+  equivalent processes at the same TP coordinate can then share that final
+  rank-local layout in no-AllGather mode. A positive registration budget also
+  chooses the final runtime cache at TP1 so deferred checkpoint transforms are
+  not repeated during inference.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
 - Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
+  incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also
+  rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
+  expert-parallel, CFG-parallel, and PP layouts.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+The runtime cache improves steady-state host PSS, not startup peak: every rank
+loads private weights before remapping, and validation adds full content-hash
+passes. Direct H2D removes recurrent staging but does not change that startup
+contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+
+Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
+smoke documented in the L20X runtime-cache matrix in the
+[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
+two independent TP2 engines reduced request-time host PSS by 36.6% with the
+cache, but recurrent mmap-to-pinned staging reduced combined throughput by
+69.3%. The registered path recovered near-private per-engine latency while
+retaining the shared-memory benefit; a matched TP1 run improved mean wave
+latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
+feature design for the full smoke-test caveats and metrics.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -3,26 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
 
 
-The shared runner generates video with optional audio from text, image, and audio
-inputs. It can be extended through `vllm_omni.model_extras` without adding
-another model-specific runner.
-
-## Shared request and output contract
-
-The example builds one canonical request envelope: `prompt` is text,
-`modalities` is `["video"]`, and optional `multi_modal_data` contains RGB PIL
-images plus audio `(waveform, sample_rate)` tuples. Audio files are loaded in
-full by design; the runner does not apply a time window. Prepare a trimmed input
-file explicitly when only part of a recording should condition the model.
-DreamID-Omni's audio conditioning sequence grows with the complete input
-duration, so long recordings increase denoising memory and step time.
-
-Model-specific request translation (for example, DreamID-Omni's two negative
-prompts and its 16 kHz input requirement) is declared in `model_extras`. Model
-outputs must provide `fps` and, when audio is present, `audio_sample_rate`.
-`--fps` and `--audio-sample-rate` are available as explicit metadata overrides.
-The shared encoder normalizes supported video tensor layouts, so this runner
-does not branch on a model name or assume a model-specific frame layout.
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
 
 ## Local CLI Usage
 ### Download the Model locally
@@ -110,8 +91,7 @@ python x_to_video_audio.py \
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.
-- `--model-class-name`: optional diffusion pipeline class override.
-- `--height/--width`: output resolution (defaults 704 * 1280).
+- `--height/--width`: output resolution (defaults 704 * 1024).
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -121,9 +101,6 @@ Key arguments:
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
-- `--extra-body`: JSON object containing parameters declared by the selected model.
-- `--fps`: explicit output FPS override when a model does not return that metadata.
-- `--audio-sample-rate`: explicit output audio sample-rate override.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 
 ## Example materials

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -1,25 +1,6 @@
 # X-To-Video-Audio
 
-The shared runner generates video with optional audio from text, image, and audio
-inputs. It currently documents DreamID-Omni and MagiHuman, and can be extended
-through `vllm_omni.model_extras` without adding another model-specific runner.
-
-## Shared request and output contract
-
-The example builds one canonical request envelope: `prompt` is text,
-`modalities` is `["video"]`, and optional `multi_modal_data` contains RGB PIL
-images plus audio `(waveform, sample_rate)` tuples. Audio files are loaded in
-full by design; the runner does not apply a time window. Prepare a trimmed input
-file explicitly when only part of a recording should condition the model.
-DreamID-Omni's audio conditioning sequence grows with the complete input
-duration, so long recordings increase denoising memory and step time.
-
-Model-specific request translation (for example, DreamID-Omni's two negative
-prompts and its 16 kHz input requirement) is declared in `model_extras`. Model
-outputs must provide `fps` and, when audio is present, `audio_sample_rate`.
-`--fps` and `--audio-sample-rate` are available as explicit metadata overrides.
-The shared encoder normalizes supported video tensor layouts, so this runner
-does not branch on a model name or assume a model-specific frame layout.
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
 
 ## Local CLI Usage
 ### Download the Model locally
@@ -166,8 +147,7 @@ in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.
-- `--model-class-name`: optional diffusion pipeline class override.
-- `--height/--width`: output resolution (defaults 704 * 1280).
+- `--height/--width`: output resolution (defaults 704 * 1024).
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -177,9 +157,6 @@ Key arguments:
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
-- `--extra-body`: JSON object containing parameters declared by the selected model.
-- `--fps`: explicit output FPS override when a model does not return that metadata.
-- `--audio-sample-rate`: explicit output audio sample-rate override.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 - `--cache-backend`: enable `cache_dit` for acceleration.
 - `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -3,97 +3,32 @@
 
 import argparse
 import json
-import math
 import re
 import time
-from collections.abc import Mapping
-from dataclasses import dataclass
-from pathlib import Path
 from typing import Any
 
 import numpy as np
-import torch
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.entrypoints.openai.video_api_utils import encode_video_bytes
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_extras import (
-    build_x_to_video_audio_prompt as build_model_x_to_video_audio_prompt,
-)
-from vllm_omni.model_extras import (
-    get_extra_body_params,
-    get_input_audio_sample_rate,
-    get_model_class_name,
-    get_output_tensor_range,
-)
 
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
 AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
 
 
-@dataclass(frozen=True)
-class XToVideoAudioOutput:
-    """Canonical generated video+audio payload and its required metadata."""
-
-    video: Any
-    audio: Any | None
-    fps: float
-    audio_sample_rate: int | None
-    output_tensor_range: str = "negative_one_to_one"
-
-
-def parse_json_object(value: str) -> dict[str, Any]:
-    try:
-        parsed = json.loads(value)
-    except json.JSONDecodeError as exc:
-        raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {exc}") from exc
-    if not isinstance(parsed, dict):
-        raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
-    return parsed
-
-
-def build_x_to_video_audio_prompt(
-    prompt: str,
-    media_inputs: dict[str, Any] | None = None,
-) -> dict[str, Any]:
-    """Build the canonical request envelope for image/audio-to-video models."""
-    result: dict[str, Any] = {"prompt": prompt, "modalities": ["video"]}
-    if media_inputs:
-        result["multi_modal_data"] = media_inputs
-    return result
-
-
-def _clean_official_prompt(prompt: str) -> str:
-    """Remove metadata tags used by the official DreamID prompt fixtures."""
-    prompt = re.sub(
-        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
-        "",
-        prompt,
-        flags=re.DOTALL,
-    ).strip()
-    prompt = re.sub(
-        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
-        "",
-        prompt,
-        flags=re.DOTALL,
-    ).strip()
-    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
-    return re.sub(r"\n\s*\n", "\n", prompt).strip()
-
-
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
     parser.add_argument("--model", required=True, help="Model ckpt root directory.")
     parser.add_argument(
         "--model-type",
-        default=None,
-        help="Optional legacy model type override (for example, magi-human).",
+        default="dreamid-omni",
+        choices=["dreamid-omni", "magi-human"],
+        help="Model type.",
     )
-    parser.add_argument("--model-class-name", default=None, help="Optional diffusion pipeline class override.")
     parser.add_argument("--prompt", default=None, help="Text prompt.")
 
     parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
@@ -101,9 +36,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
     parser.add_argument(
         "--extra-body",
-        type=parse_json_object,
+        type=str,
         default=None,
-        help="JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
+        help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
         'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
     )
 
@@ -136,20 +71,13 @@ def parse_args() -> argparse.Namespace:
         default="robotic, muffled, echo, distorted",
         help="Negative prompt for audio.",
     )
-    parser.add_argument("--fps", type=float, default=None, help="Override output FPS metadata.")
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=None,
-        help="Override output audio sample-rate metadata.",
-    )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
     parser.add_argument(
         "--quantization",
         type=str,
         default=None,
         choices=["fp8", "int8"],
-        help="Online (dynamic) quantization method for the model transformer.",
+        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
     )
     parser.add_argument(
         "--enable-cpu-offload",
@@ -166,7 +94,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
-        help="Enable HSDP for supported transformer blocks.",
+        help="Enable HSDP for DreamID-Omni fused transformer blocks.",
     )
     parser.add_argument(
         "--hsdp-shard-size",
@@ -183,105 +111,20 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def load_image_and_audio(
-    image_paths: list[str] | None,
-    audio_paths: list[str] | None,
-    *,
-    audio_sample_rate: int | None = None,
-) -> tuple[list[Image.Image], list[tuple[Any, int]]]:
-    """Load complete media inputs into the canonical multimodal representation."""
-    images = []
-    audios = []
+def load_image_and_audio(image_paths, audio_paths):
+    image = []
+    audio = []
 
-    for path in image_paths or []:
+    for path in image_paths:
         with Image.open(path) as img:
-            images.append(img.convert("RGB"))
+            img = img.convert("RGB")
+            image.append(img)
 
-    for path in audio_paths or []:
-        waveform, sample_rate = load_audio(path, sr=audio_sample_rate)
-        audios.append((waveform, sample_rate))
-    return images, audios
-
-
-def extract_x_to_video_audio_output(
-    outputs: Any,
-    *,
-    fps: float | None = None,
-    audio_sample_rate: int | None = None,
-    output_tensor_range: str = "negative_one_to_one",
-) -> XToVideoAudioOutput:
-    """Normalize engine output without depending on a model's tensor layout."""
-    if isinstance(outputs, list):
-        if not outputs:
-            raise RuntimeError("No output returned from the model.")
-        result = outputs[0]
-    else:
-        result = outputs
-    if result is None:
-        raise RuntimeError("No output returned from the model.")
-
-    multimodal_output = getattr(result, "multimodal_output", None) or {}
-    if not isinstance(multimodal_output, Mapping):
-        raise RuntimeError("Model multimodal output must be a mapping.")
-    images = getattr(result, "images", None)
-    video = images[0] if images else multimodal_output.get("video")
-    if video is None:
-        raise RuntimeError("No video payload found in model output.")
-
-    resolved_fps = fps if fps is not None else multimodal_output.get("fps")
-    try:
-        resolved_fps = float(resolved_fps)
-    except (TypeError, ValueError) as exc:
-        raise RuntimeError("Model output must declare a valid 'fps', or pass --fps.") from exc
-    if not math.isfinite(resolved_fps) or resolved_fps <= 0:
-        raise RuntimeError("Output FPS must be finite and positive.")
-
-    audio = multimodal_output.get("audio")
-    resolved_sample_rate = audio_sample_rate
-    if resolved_sample_rate is None:
-        resolved_sample_rate = multimodal_output.get("audio_sample_rate")
-    if audio is not None:
-        try:
-            resolved_sample_rate = int(resolved_sample_rate)
-        except (TypeError, ValueError) as exc:
-            raise RuntimeError(
-                "Model output with audio must declare a valid 'audio_sample_rate', or pass --audio-sample-rate."
-            ) from exc
-        if resolved_sample_rate <= 0:
-            raise RuntimeError("Output audio sample rate must be positive.")
-    elif resolved_sample_rate is not None:
-        resolved_sample_rate = int(resolved_sample_rate)
-
-    return XToVideoAudioOutput(video, audio, resolved_fps, resolved_sample_rate, output_tensor_range)
-
-
-def _normalize_output_tensor_range(video: Any, source_range: str) -> Any:
-    """Normalize floating-point output according to the pipeline contract."""
-    if isinstance(video, torch.Tensor):
-        if not video.is_floating_point():
-            return video
-        video = video.detach().cpu().float().numpy()
-    if isinstance(video, np.ndarray):
-        if not np.issubdtype(video.dtype, np.floating):
-            return video
-        if source_range == "negative_one_to_one":
-            return np.clip(video, -1.0, 1.0) * 0.5 + 0.5
-        if source_range == "zero_to_one":
-            return np.clip(video, 0.0, 1.0)
-        raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
-    if isinstance(video, list):
-        return [_normalize_output_tensor_range(frame, source_range) for frame in video]
-    return video
-
-
-def encode_x_to_video_audio_output(output: XToVideoAudioOutput) -> bytes:
-    """Encode any video layout supported by the shared video API utility."""
-    return encode_video_bytes(
-        _normalize_output_tensor_range(output.video, output.output_tensor_range),
-        output.fps,
-        output.audio,
-        output.audio_sample_rate,
-    )
+    for path in audio_paths:
+        audio_array, sr = load_audio(path, sr=16000)
+        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
+        audio.append(audio_array)
+    return image, audio
 
 
 def _extract_peak_memory_mb(result: Any) -> float:
@@ -314,10 +157,45 @@ def main() -> None:
     if args.prompt_file:
         with open(args.prompt_file) as f:
             text_prompt = json.load(f)
-            if isinstance(text_prompt, str):
-                text_prompt = _clean_official_prompt(text_prompt)
-    if not isinstance(text_prompt, str):
-        raise ValueError("Prompt content must be a JSON string or --prompt text.")
+            text_prompt = re.sub(
+                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(
+                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
+            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
+
+    if args.model_type == "magi-human":
+        prompt = text_prompt
+        extra_args = json.loads(args.extra_body) if args.extra_body else {}
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed,
+            extra_args=extra_args,
+        )
+    else:
+        image, audio = load_image_and_audio(args.image_path, args.audio_path)
+
+        prompt = {
+            "prompt": text_prompt,
+            "video_negative_prompt": args.video_negative_prompt,
+            "audio_negative_prompt": args.audio_negative_prompt,
+            "multi_modal_data": {"image": image, "audio": audio},
+        }
+
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed,
+            extra_args={
+                "solver_name": args.solver_name,
+                "shift": args.shift,
+            },
+        )
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
@@ -342,83 +220,64 @@ def main() -> None:
             "scm_steps_policy": "dynamic",
         }
 
-    omni_kwargs: dict[str, Any] = dict(
+    omni_kwargs = dict(
         model=args.model,
         parallel_config=parallel_config,
+        model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )
-    if args.model_type is not None:
-        omni_kwargs["model_type"] = args.model_type
-    if args.model_class_name is not None:
-        omni_kwargs["model_class_name"] = args.model_class_name
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
     omni = Omni(**omni_kwargs)
-    try:
-        model_class_name = args.model_class_name or get_model_class_name(omni)
-        input_sample_rate = get_input_audio_sample_rate(model_class_name)
-        images, audios = load_image_and_audio(
-            args.image_path,
-            args.audio_path,
-            audio_sample_rate=input_sample_rate,
-        )
-        media_inputs: dict[str, Any] = {}
-        if images:
-            media_inputs["image"] = images
-        if audios:
-            media_inputs["audio"] = audios
+    start = time.perf_counter()
+    outputs = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
 
-        canonical_prompt = build_x_to_video_audio_prompt(text_prompt, media_inputs)
-        prompt = build_model_x_to_video_audio_prompt(
-            model_class_name,
-            canonical_prompt,
-            {
-                "video_negative_prompt": args.video_negative_prompt,
-                "audio_negative_prompt": args.audio_negative_prompt,
-            },
-        )
+    if not outputs:
+        raise RuntimeError("No output returned from the model.")
+    result = outputs[0]
+    if not result.images:
+        raise RuntimeError("No video frames found in model output.")
+    generated_video = result.images[0]
+    mm = result.multimodal_output or {}
+    generated_audio = mm.get("audio")
 
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args={},
-        )
-        declared_extra_body_params = get_extra_body_params(model_class_name)
-        extra_body = dict(args.extra_body or {})
-        legacy_extra_body = {"solver_name": args.solver_name, "shift": args.shift}
-        for key, value in legacy_extra_body.items():
-            if key in declared_extra_body_params:
-                extra_body.setdefault(key, value)
-        if declared_extra_body_params:
-            apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
-        elif extra_body:
-            sampling_params.extra_args.update({key: value for key, value in extra_body.items() if value is not None})
+    if args.model_type == "magi-human":
+        # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
+        fps = float(mm.get("fps", 25))
+        sample_rate = int(mm.get("audio_sample_rate", 24000))
+        frames = generated_video
+    else:
+        fps = float(mm.get("fps", 24))
+        sample_rate = int(mm.get("audio_sample_rate", 16000))
+        # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
+        # mux_video_audio_bytes expects (F, H, W, C) uint8.
+        if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
+            raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
+        frames = generated_video.transpose(1, 2, 3, 0)
+        frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
 
-        start = time.perf_counter()
-        outputs = omni.generate(prompt, sampling_params)
-        elapsed = time.perf_counter() - start
-        output = extract_x_to_video_audio_output(
-            outputs,
-            fps=args.fps,
-            audio_sample_rate=args.audio_sample_rate,
-            output_tensor_range=get_output_tensor_range(model_class_name),
-        )
+    audio_np = None
+    if generated_audio is not None:
+        audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
 
-        output_path = Path(args.output)
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        output_path.write_bytes(encode_x_to_video_audio_output(output))
-        print(f"Saved generated video to {output_path}")
-        print(f"Total time: {elapsed:.2f}s")
-        peak_mb = _extract_peak_memory_mb(outputs)
-        if peak_mb:
-            print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-    finally:
-        omni.close()
+    output_path = args.output
+    video_bytes = mux_video_audio_bytes(
+        frames,
+        audio_np,
+        fps=float(fps),
+        audio_sample_rate=sample_rate,
+    )
+    with open(output_path, "wb") as f:
+        f.write(video_bytes)
+    print(f"Saved generated video to {output_path}")
+    print(f"Total time: {elapsed:.2f}s")
+    peak_mb = _extract_peak_memory_mb(outputs)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
 
 if __name__ == "__main__":

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -4,10 +4,7 @@ from types import SimpleNamespace
 import pytest
 from vllm.v1.engine import FinishReason
 
-from vllm_omni.core.sched.omni_scheduler_mixin import (
-    DEFAULT_INPUT_WAIT_TIMEOUT_S,
-    OmniSchedulerMixin,
-)
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniChunkRecvHandle
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -25,25 +22,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     scheduler.requests = {"request": object()}
     scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
     scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
-
-    def _collect_timed_out(timeout_s):
-        calls.append(("chunk-timeouts", timeout_s))
-        return set()
-
-    def _collect_failed_sends():
-        calls.append(("failed-sends",))
-        return {}
-
     scheduler.chunk_transfer_adapter = SimpleNamespace(
-        receives_chunks=True,
         process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
             ("process", waiting, running, scheduler_requests)
         ),
         restore_queues=lambda waiting, running, scheduler_requests: calls.append(
             ("restore-chunks", waiting, running, scheduler_requests)
         ),
-        collect_timed_out_request_ids=_collect_timed_out,
-        collect_failed_send_request_ids=_collect_failed_sends,
     )
     scheduler.input_coordinator = SimpleNamespace(
         restore_queues=lambda waiting: calls.append(("restore-full", waiting))
@@ -56,10 +41,6 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
         ("consume", "ar"),
         ("timeouts",),
         ("process", scheduler.waiting, scheduler.running, scheduler.requests),
-        # The chunk deadline runs after chunks are applied, so a chunk that
-        # arrived this cycle resets the clock before it is measured (R1.1).
-        ("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
-        ("failed-sends",),
         ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
         ("restore-full", scheduler.waiting),
     ]

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

@@ -1,10 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit coverage for the two stage-input timeout hooks.
-
-``_process_pending_input_timeouts`` covers the full-payload transport;
-``_process_pending_chunk_timeouts`` covers async-chunk. Both share
-``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
+"""Unit coverage for _process_pending_input_timeouts.
 
 Verifies that the mixin correctly *delegates* timed-out requests to the
 base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
@@ -16,10 +12,6 @@ this file only asserts the wiring from the mixin to that API.
 
 from __future__ import annotations
 
-import importlib
-import logging
-import math
-import os
 from types import SimpleNamespace
 
 import pytest
@@ -110,248 +102,5 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
     assert scheduler.finish_calls == []
 
 
-class _FakeChunkAdapter:
-    """Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
-
-    def __init__(self, timed_out_ids, *, receives_chunks=True):
-        self._timed_out_ids = set(timed_out_ids)
-        self.receives_chunks = receives_chunks
-        self.calls = []
-
-    def collect_timed_out_request_ids(self, timeout_s):
-        self.calls.append(timeout_s)
-        return set(self._timed_out_ids)
-
-
-class _FakeChunkScheduler(OmniSchedulerMixin):
-    def __init__(self, requests, adapter):
-        self.requests = requests
-        self.chunk_transfer_adapter = adapter
-        self.finish_calls = []
-
-    def finish_requests(self, req_ids, status):
-        self.finish_calls.append((set(req_ids), status))
-
-
-def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
-    """A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
-
-    This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
-    existed, the chunk path had no deadline at all and issue #3833 could park a
-    request forever.
-    """
-    req_id = "stalled-stream"
-    adapter = _FakeChunkAdapter(timed_out_ids={req_id})
-    scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
-
-    scheduler._process_pending_chunk_timeouts()
-
-    assert len(adapter.calls) == 1
-    assert adapter.calls[0] > 0
-    assert len(scheduler.finish_calls) == 1
-    finished_ids, status = scheduler.finish_calls[0]
-    assert finished_ids == {req_id}
-    assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
-
-
-def test_process_pending_chunk_timeouts_skips_already_freed_request():
-    adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
-    scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
-
-    scheduler._process_pending_chunk_timeouts()
-
-    assert scheduler.finish_calls == []
-
-
-def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
-    """Stage 0 does not receive chunks, so it has nothing to expire."""
-    adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
-    scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
-
-    scheduler._process_pending_chunk_timeouts()
-
-    assert adapter.calls == []
-    assert scheduler.finish_calls == []
-
-
-def test_process_pending_chunk_timeouts_noop_without_adapter():
-    class _NoAdapter(OmniSchedulerMixin):
-        def __init__(self):
-            self.requests = {}
-            self.chunk_transfer_adapter = None
-            self.finish_calls = []
-
-        def finish_requests(self, req_ids, status):
-            self.finish_calls.append((set(req_ids), status))
-
-    scheduler = _NoAdapter()
-    scheduler._process_pending_chunk_timeouts()
-    assert scheduler.finish_calls == []
-
-
-def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
-    """The shared knob disables both transports' nets together."""
-    from vllm_omni.core.sched import omni_scheduler_mixin
-
-    monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
-
-    adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
-    scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
-    scheduler._process_pending_chunk_timeouts()
-
-    assert adapter.calls == []
-    assert scheduler.finish_calls == []
-
-
-class _FakeSendFailureAdapter:
-    """Stands in for the adapter's drained send-failure map (R1.2)."""
-
-    def __init__(self, failures):
-        self._failures = dict(failures)
-        self.collect_calls = 0
-
-    def collect_failed_send_request_ids(self):
-        self.collect_calls += 1
-        failures, self._failures = self._failures, {}
-        return failures
-
-
-class _FakeSendFailureScheduler(OmniSchedulerMixin):
-    def __init__(self, requests, adapter):
-        self.requests = requests
-        self.chunk_transfer_adapter = adapter
-
-
-def test_log_failed_chunk_sends_reports_a_live_request(caplog):
-    adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
-    scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
-
-    with caplog.at_level(logging.ERROR):
-        scheduler._log_failed_chunk_sends()
-
-    assert "req-live" in caplog.text
-    assert "No space left on device" in caplog.text
-
-
-def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
-    """``collect_*`` drains the map, so a failure for a request that already left
-    this scheduler is the one place a drop could still go unrecorded."""
-    adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
-    scheduler = _FakeSendFailureScheduler({}, adapter)
-
-    with caplog.at_level(logging.WARNING):
-        scheduler._log_failed_chunk_sends()
-
-    assert adapter.collect_calls == 1
-    assert "req-gone" in caplog.text
-    assert "No space left on device" in caplog.text
-    assert "finished or aborted" in caplog.text
-
-
-def test_log_failed_chunk_sends_noop_without_adapter():
-    scheduler = _FakeSendFailureScheduler({}, None)
-    scheduler._log_failed_chunk_sends()
-
-
-# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
-# these reloads the module under a different environment. One knob now gates the
-# deadline on both transports, which is why a silent 0 matters more than it did
-# when only the full-payload path used it.
-
-
-@pytest.fixture
-def reload_with(monkeypatch):
-    """Reload the mixin module under a given env value, then put it back.
-
-    ``importlib.reload`` re-executes into the live module dict, so whatever the
-    last test parsed is what every later test in the session sees -- ``-k
-    zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
-    silently disable both deadlines. The rejection cases are worse: the
-    ValueError fires *after* the rebind, so ``nan`` lingers in the module.
-    ``monkeypatch`` restores the environment but not the module, so reload once
-    more on the way out.
-    """
-    import vllm_omni.core.sched.omni_scheduler_mixin as mod
-
-    def _reload(raw):
-        if raw is None:
-            monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
-        else:
-            monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
-        return importlib.reload(mod)
-
-    yield _reload
-
-    monkeypatch.undo()
-    importlib.reload(mod)
-
-
-@pytest.mark.parametrize(
-    "raw,expected",
-    [
-        (None, 600.0),
-        ("30", 30.0),
-        ("0.5", 0.5),
-        ("0", 0.0),  # explicit opt-out is still honoured
-    ],
-)
-def test_timeout_accepts_valid_values(reload_with, raw, expected):
-    mod = reload_with(raw)
-    assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
-
-
-@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
-def test_negative_timeout_is_rejected(reload_with, raw):
-    """`-1` is a common "no limit" idiom; substituting 600 would give an operator
-    who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
-    with pytest.raises(ValueError, match="negative"):
-        reload_with(raw)
-
-
-@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
-def test_non_finite_timeout_is_rejected(reload_with, raw):
-    """nan compares false against every deadline check and inf is unreachable;
-    both disable the net while looking like a number."""
-    with pytest.raises(ValueError, match="not finite"):
-        reload_with(raw)
-
-
-def test_zero_warns_and_names_both_transports(reload_with, caplog):
-    with caplog.at_level(logging.WARNING):
-        mod = reload_with("0")
-    assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
-    text = caplog.text
-    assert "DISABLED" in text
-    assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
-
-
-def test_unparsable_timeout_falls_back(reload_with):
-    mod = reload_with("not-a-number")
-    assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
-
-
-def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
-    """The reason ``reload_with`` reloads on the way out, done here by hand.
-
-    ``nan`` is bound before the ValueError fires, so the reload that rejects it
-    still leaves it in the module -- and nan compares false against every
-    deadline check, which is exactly the silent-disable this item exists to
-    stop. Restoring the environment is not enough; the module has to be reloaded
-    too.
-    """
-    import vllm_omni.core.sched.omni_scheduler_mixin as mod
-
-    ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
-
-    monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
-    with pytest.raises(ValueError, match="not finite"):
-        importlib.reload(mod)
-    assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
-
-    monkeypatch.undo()
-    importlib.reload(mod)
-    assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
-
-
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/dfx/perf/tests/test_bagel_vllm_omni.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
@@ -61,7 +61,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
@@ -113,7 +113,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
@@ -164,7 +164,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",

tests/dfx/perf/tests/test_higgs_audio_v3.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "tts"
     ],
     "server_params": {

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-device baseline with enforce-eager (no torch.compile)",
@@ -64,7 +64,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-device with torch.compile (default, no enforce-eager)",
@@ -117,7 +117,7 @@
           "num_cards": 2
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "CFG-parallel=2 with enforce-eager",
@@ -172,7 +172,7 @@
           "num_cards": 2
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "CFG-parallel=2 with torch.compile",
@@ -226,7 +226,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "CacheDiT with enforce-eager",

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-device baseline (two input images)",
@@ -76,7 +76,7 @@
           "num_cards": 4
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -125,7 +125,7 @@
           "num_cards": 4
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Ulysses SP=2 + CFG=2 + CacheDiT",

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-device baseline",
@@ -73,7 +73,7 @@
           "num_cards": 4
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -121,7 +121,7 @@
           "num_cards": 4
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Ulysses SP=2 + CFG=2 + CacheDiT",

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "diffusion"
     ],
     "description": "Single-device baseline",

tests/dfx/perf/tests/test_voxcpm2.json

@@ -10,7 +10,7 @@
           "num_cards": 1
         }
       },
-      "local_model",
+      "full_model",
       "tts"
     ],
     "server_params": {

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -12,6 +12,7 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -335,7 +336,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                checkpoint_key="weight",
+                storage_key="weight",
                 file_path="unused",
             )
         },
@@ -394,6 +395,192 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            allgather_degree=1,
+            ulysses_mode="strict",
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_lock_timeout=1.0,
+        model="unused",
+        revision=None,
+        model_config={},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnquantizedLinearMethod()
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    def load_weights(_model):
+        events.append("load")
+        with torch.no_grad():
+            model.transformer.weight.fill_(1)
+
+    def process_weights(_model, _device):
+        events.append("process")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    def calibrate(_model):
+        events.append("calibrate")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    original_eval = model.eval
+
+    def eval_model():
+        events.append("eval")
+        return original_eval()
+
+    def build_cache(_pipeline, **kwargs):
+        events.append("cache")
+        assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
+        assert kwargs["tensor_parallel_rank"] == 0
+        assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
+        assert kwargs["quantization_config"] is None
+        return HostWeightPlanResult(runtime_plan)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._process_weights_after_loading = process_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = calibrate  # type: ignore[method-assign]
+    model.eval = eval_model  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "calibrate", "eval", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+
+    def build_cache(_pipeline, **kwargs):
+        assert kwargs["quantization_config"] == "module_quant_method"
+        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
+
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    original_weight = model.transformer.weight.detach().clone()
+    assert loader.load_model(load_device="cpu") is model
+    assert torch.equal(model.transformer.weight, original_weight)
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/model_loader/test_runtime_weight_cache.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fcntl
+import multiprocessing as mp
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+    build_runtime_weight_cache_plan,
+    default_runtime_cache_root,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _LoaderMarker:
+    pass
+
+
+class _Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(3, 2)
+        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
+        self.register_buffer("temporary", torch.ones(1), persistent=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
+            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
+
+
+class _Pipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer()
+
+
+def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
+    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
+
+    assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
+
+
+def _build(pipeline: nn.Module, cache_root: Path, **overrides):
+    kwargs = {
+        "dit_modules": (("transformer", pipeline.transformer),),
+        "loader_type": _LoaderMarker,
+        "cache_root": cache_root,
+        "lock_timeout_seconds": 2.0,
+        "max_shard_bytes": 32,
+        "model_identity": "test/model",
+        "revision": "revision",
+        "runtime_dtype": torch.float32,
+        "load_format": "default",
+        "loader_inputs": {"model_config": {"hidden_size": 3}},
+        "tensor_parallel_size": 1,
+        "tensor_parallel_rank": 0,
+        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
+        "use_hsdp": False,
+        "enable_expert_parallel": False,
+        "quantization_config": None,
+        "cfg_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+    }
+    kwargs.update(overrides)
+    return build_runtime_weight_cache_plan(pipeline, **kwargs)
+
+
+def _multiprocess_builder(cache_root: str, queue) -> None:
+    result = _build(_Pipeline(), Path(cache_root))
+    queue.put(
+        (
+            result.fallback_code,
+            result.plan.runtime_layout_key if result.plan is not None else None,
+            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
+        )
+    )
+
+
+def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    second = _build(_Pipeline(), tmp_path)
+
+    assert first.plan is not None
+    assert second.plan is not None
+    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
+    assert first.plan.post_load_complete
+    assert {binding.file_path for binding in first.plan.bindings.values()} == {
+        binding.file_path for binding in second.plan.bindings.values()
+    }
+    assert set(first.plan.bindings) == {
+        "transformer.persistent",
+        "transformer.proj.bias",
+        "transformer.proj.weight",
+    }
+
+
+def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
+    pipeline = _Pipeline()
+    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
+    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
+    sp2 = _build(
+        pipeline,
+        tmp_path,
+        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
+    )
+
+    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
+    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
+
+
+def test_final_runtime_content_is_authoritative(tmp_path):
+    first_pipeline = _Pipeline()
+    second_pipeline = _Pipeline()
+    with torch.no_grad():
+        second_pipeline.transformer.proj.weight[0, 0] += 1
+
+    first = _build(first_pipeline, tmp_path)
+    second = _build(second_pipeline, tmp_path)
+
+    assert first.plan is not None and second.plan is not None
+    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
+
+
+def test_expert_parallel_ownership_fails_closed(tmp_path):
+    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
+
+    assert result.plan is None
+    assert result.fallback_code == "unsupported_expert_parallel"
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    assert first.plan is not None
+    shard = Path(next(iter(first.plan.bindings.values())).file_path)
+    old_inode = shard.stat().st_ino
+    with shard.open("r+b") as handle:
+        handle.seek(-1, 2)
+        byte = handle.read(1)
+        handle.seek(-1, 2)
+        handle.write(bytes([byte[0] ^ 0xFF]))
+
+    second = _build(_Pipeline(), tmp_path)
+
+    assert second.plan is not None
+    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
+    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
+
+
+def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
+
+    with lock_path.open("a+b") as lock_handle:
+        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
+
+    assert result.plan is None
+    assert result.fallback_code == "lock_timeout"
+
+
+def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
+    stale_temp.mkdir(parents=True)
+    (stale_temp / "partial").write_bytes(b"partial")
+
+    result = _build(_Pipeline(), tmp_path)
+
+    assert result.plan is not None
+    assert not stale_temp.exists()
+
+
+@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
+def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
+    pipeline = _Pipeline()
+    if case == "alias":
+        pipeline.transformer.tied = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    elif case == "external_alias":
+        pipeline.encoder = nn.Module()
+        pipeline.encoder.shared = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    else:
+        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
+        expected_code = "unsupported_tensor"
+
+    result = _build(pipeline, tmp_path)
+
+    assert result.plan is None
+    assert result.fallback_code == expected_code
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
+    context = mp.get_context("spawn")
+    queue = context.Queue()
+    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
+    for process in processes:
+        process.start()
+    results = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert all(code is None for code, _, _ in results)
+    assert len({key for _, key, _ in results}) == 1
+    assert len({tuple(paths) for _, _, paths in results}) == 1
+    assert len(list((tmp_path / "v1").iterdir())) == 1

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,205 @@
+# 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 vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -473,6 +473,31 @@ class TestDistributedLayerwiseOffloadHook:
         assert staging[0][torch.float32].numel() == 12
         assert staging[1][torch.float32].numel() == 12
 
+    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        next_block.weight.data.copy_(expected)
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+        hook.registered_mmap = True
+
+        def fail_staging(_slot):
+            raise AssertionError("registered mmap must bypass host staging")
+
+        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
+        hook.prefetch_layer(slot=0, non_blocking=True)
+
+        assert torch.equal(next_block.weight, expected)
+
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -563,6 +588,24 @@ class TestPinnedResidentLayerGroup:
         group.offload()
         assert all(block.weight.numel() == 0 for block in blocks)
 
+    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
+        block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        block.weight.data.copy_(expected)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        group.registered_mmap = True
+        group._cpu_staging_buffers.clear()
+
+        group.load()
+
+        assert torch.equal(block.weight, expected)
+
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
     ):
@@ -692,6 +735,95 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
+    def test_runtime_cache_maps_final_weights_without_rerunning_post_load(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        class Transformer(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Linear(2, 2, bias=False)
+                self.post_load_calls = 0
+
+            def post_load_weights(self):
+                self.post_load_calls += 1
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        old_parameter = pipeline.transformer.proj.weight
+        old_parameter.runtime_marker = "preserved"
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.proj.weight": TensorBinding(
+                    storage_key="transformer.proj.weight",
+                    file_path=str(weight_file),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.post_load_calls == 0
+        assert pipeline.transformer.proj.weight is old_parameter
+        assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
+        assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
+        backend._release_mmap_handles()
+
+    def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        pipeline = nn.Module()
+        pipeline.transformer = nn.Linear(2, 2, bias=False)
+        old_parameter = pipeline.transformer.weight
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.weight": TensorBinding(
+                    storage_key="transformer.weight",
+                    file_path=str(tmp_path / "missing.safetensors"),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        with pytest.raises(Exception, match="missing.safetensors"):
+            backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.weight is old_parameter
+
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -715,7 +847,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    checkpoint_key=name,
+                    storage_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -786,6 +918,99 @@ class TestMmapWeightLoading:
         assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
         backend.disable()
 
+    def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file(weights, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        cleanup_calls: list[None] = []
+        backend._cleanup_after_loading = lambda: cleanup_calls.append(None)  # type: ignore[method-assign]
+
+        backend.enable(pipeline)
+
+        assert backend._using_rank_local_mmap
+        # One reclamation immediately after remap, plus the existing final
+        # post-hook cleanup.
+        assert len(cleanup_calls) == 2
+        backend.disable()
+
+    def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        events: list[str] = []
+
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                events.append("unregister")
+                return []
+
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
+
+        def create(sources, *, max_bytes):
+            create_calls.append((sources, max_bytes))
+            return Registration()
+
+        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.5,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._runtime_cache_mapped_sources = sources
+
+        assert backend._try_register_runtime_cache_mmap()
+        assert create_calls == [(sources, int(1.5 * 1024**3))]
+
+        backend.enabled = True
+        backend._using_rank_local_mmap = True
+        backend._mmap_file_cache = {}
+        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
+        backend.disable()
+
+        assert events == ["unregister", "mmap"]
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -1091,7 +1316,7 @@ class TestMmapValidation:
         assert result.fallback_reason is None
         assert result.plan is not None
         assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            checkpoint_key="blocks.0.weight",
+            storage_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1381,7 +1606,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            checkpoint_key="proj_in.weight",
+            storage_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_worker_wrapper_base.py

@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
     CustomPipelineWorkerExtension,
     DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
+        self,
+        mocker: MockerFixture,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        worker = object.__new__(DiffusionWorker)
+        offload_backend = mocker.Mock()
+        kv_transfer_manager = mocker.Mock()
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_transfer_manager,
+        )
+        events: list[str] = []
+        offload_backend.disable.side_effect = lambda: events.append("offloader")
+        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
+        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
+
+        worker.shutdown()
+
+        offload_backend.disable.assert_called_once_with()
+        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
+        assert events == ["offloader", "kv_transfer", "distributed"]
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -1728,112 +1728,3 @@ def test_purge_is_noop_on_empty_deques(build_adapter):
     adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
     assert running_queue == []
     assert waiting_queue == []
-
-
-# --------------------------------------------------------------------------- #
-#  Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
-#
-#  Before this, the async-chunk path had no deadline of any kind, so a dropped
-#  terminal chunk or an upstream stage that died mid-stream parked the request
-#  in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
-# --------------------------------------------------------------------------- #
-
-
-def _park_in_chunk_wait(adapter, request, *, waiting=True):
-    """Drive a request through one process_pending_chunks round into the wait."""
-    queue = DummyWaitingQueue([request]) if waiting else [request]
-    if waiting:
-        adapter.process_pending_chunks(queue, [])
-    else:
-        adapter.process_pending_chunks(DummyWaitingQueue(), queue)
-    return queue
-
-
-def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    request = _req("r1", RequestStatus.WAITING)
-
-    _park_in_chunk_wait(adapter, request)
-
-    assert request.status == RequestStatus.WAITING_FOR_CHUNK
-    assert "r1" in adapter._waiting_since
-
-
-def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
-
-    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
-    assert "r1" in adapter._waiting_since
-
-
-def test_dropped_terminal_chunk_expires_the_request(build_adapter):
-    """The #3833 scenario: upstream stops sending and never marks the stream done."""
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
-
-    adapter._waiting_since["r1"] -= 601.0
-
-    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
-    # Cleared, so a second sweep does not re-report the same request.
-    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
-
-
-def test_arriving_chunk_resets_the_clock(build_adapter):
-    """A slow but healthy stream must never expire: the deadline measures stall
-    time between chunks, not the lifetime of the stream."""
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    request = _req("r1", RequestStatus.WAITING)
-    queue = _park_in_chunk_wait(adapter, request)
-    adapter.restore_queues(queue, [], scheduler_requests={"r1": request})
-
-    # Age the wait to just under the deadline, then deliver a chunk.
-    adapter._waiting_since["r1"] -= 599.0
-    adapter._finished_load_reqs.add("r1")
-    adapter.process_pending_chunks(queue, [])
-
-    assert "r1" not in adapter._waiting_since
-    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
-
-
-def test_a_disabled_deadline_never_expires(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
-    adapter._waiting_since["r1"] -= 10_000.0
-
-    assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
-    assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()
-
-
-def test_finished_request_leaves_no_stale_timestamp(build_adapter):
-    """Otherwise a completed request would be reported as timed out later."""
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    request = _req("r1", RequestStatus.WAITING)
-    _park_in_chunk_wait(adapter, request)
-
-    adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})
-
-    assert "r1" not in adapter._waiting_since
-    assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()
-
-
-def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
-    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
-
-    adapter.cleanup_receiver("r1")
-
-    assert "r1" not in adapter._waiting_since
-
-
-def test_expiry_is_per_request(build_adapter):
-    """One stalled stream must not take down its healthy neighbours."""
-    adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
-    stalled = _req("stalled", RequestStatus.WAITING)
-    healthy = _req("healthy", RequestStatus.WAITING)
-    adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])
-
-    adapter._waiting_since["stalled"] -= 601.0
-
-    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
-    assert "healthy" in adapter._waiting_since

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -463,7 +463,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
     }
 
 
-@pytest.mark.slow
+@pytest.mark.full_model
 @pytest.mark.benchmark
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -591,7 +591,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
     assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
 
 
-@pytest.mark.slow
+@pytest.mark.full_model
 @pytest.mark.benchmark
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"}, num_cards=1)

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -544,7 +544,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
 
 
 @pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
-@pytest.mark.slow
+@pytest.mark.full_model
 @pytest.mark.benchmark
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"}, num_cards=1)

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -17,7 +17,7 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniServer
 
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
@@ -253,7 +253,6 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
     )
 
 
-@pytest.mark.slow
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_single_matches_diffusers(
@@ -279,7 +278,6 @@ def test_qwen_image_edit_single_matches_diffusers(
     )
 
 
-@pytest.mark.slow
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:

tests/e2e/accuracy/test_qwen_image_layered.py

@@ -17,7 +17,7 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer
 
-pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 MODEL_ID = "Qwen/Qwen-Image-Layered"

tests/e2e/accuracy/test_video_streaming_output_similarity.py

@@ -11,7 +11,7 @@ from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_vi
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "BestWishYsh/Helios-Distilled"
 PROMPT = "A serene lakeside sunrise with mist over the water."

tests/e2e/offline_inference/test_bagel_expansion.py

@@ -138,7 +138,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
 
 
 pytestmark = [
-    pytest.mark.slow,
+    pytest.mark.full_model,
     pytest.mark.diffusion,
     pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
 ]

tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py

@@ -109,7 +109,7 @@ def _diffusion_sampling_params(
 
 
 pytestmark = [
-    pytest.mark.slow,
+    pytest.mark.full_model,
     pytest.mark.diffusion,
 ]
 

tests/e2e/offline_inference/test_moss_tts_realtime.py

@@ -96,7 +96,7 @@ def _get_test_config() -> str:
 
 pytestmark = [
     pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
-    pytest.mark.slow,
+    pytest.mark.full_model,
     pytest.mark.tts,
     pytest.mark.parametrize(
         "omni_runner",
@@ -271,6 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.advanced_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
     """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
@@ -282,6 +283,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
     assert not torch.all(audio == 0), "Audio is silence"
 
 
+@pytest.mark.advanced_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
     """MossTTSRealtime: Chinese input produces non-empty audio."""

tests/e2e/online_serving/test_bagel_expansion.py

@@ -22,7 +22,7 @@ from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
 NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"

tests/e2e/online_serving/test_bagel_multi_replicas.py

@@ -1,8 +1,6 @@
-"""Weekly e2e coverage for BAGEL diffusion multi-replica serving.
+"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving.
 
 This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs.
-It is covered by the weekly Diffusion · H100 · Multi-GPU lane
-(``.buildkite/cuda/test-weekly.yml``, ``mirror_hardwares: h100_4``).
 """
 
 import os
@@ -19,7 +17,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml")
 ROUTE_STRESS_REQUESTS = 6
 
-pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 test_params = [
     OmniServerParams(

tests/e2e/online_serving/test_boogu_image.py

@@ -16,7 +16,7 @@ parity vs. the Diffusers baseline is covered separately (see the support plan, s
 From ``tests/``::
 
     pytest -s -v e2e/online_serving/test_boogu_image.py \
-        -m "slow and diffusion and H100"
+        -m "full_model and diffusion and H100" --run-level=full_model
 """
 
 import os
@@ -28,7 +28,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "Boogu/Boogu-Image-0.1-Base"
 T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"

tests/e2e/online_serving/test_boogu_image_edit.py

@@ -19,7 +19,7 @@ reference image is supported.
 From ``tests/``::
 
     pytest -s -v e2e/online_serving/test_boogu_image_edit.py \
-        -m "slow and diffusion and H100"
+        -m "full_model and diffusion and H100" --run-level=full_model
 """
 
 import base64
@@ -35,7 +35,7 @@ from tests.helpers.runtime import (
     dummy_messages_from_mix_data,
 )
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Change the style to a colored pencil drawing."
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})

tests/e2e/online_serving/test_boogu_image_edit_expansion.py

@@ -25,7 +25,7 @@ model and are intentionally omitted.
 From ``tests/``::
 
     pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \
-        -m "slow and diffusion and H100"
+        -m "full_model and diffusion and H100" --run-level=full_model
 """
 
 import os
@@ -38,7 +38,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "Boogu/Boogu-Image-0.1-Edit"
 EDIT_PROMPT = "Change the style to a colored pencil drawing."

tests/e2e/online_serving/test_boogu_image_expansion.py

@@ -30,7 +30,7 @@ omitted (add rows here when those features land).
 From ``tests/``::
 
     pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \
-        -m "slow and diffusion and H100"
+        -m "full_model and diffusion and H100" --run-level=full_model
 """
 
 import os
@@ -42,7 +42,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "Boogu/Boogu-Image-0.1-Base"
 T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"

tests/e2e/online_serving/test_glm_image_expansion.py

@@ -30,7 +30,7 @@ from tests.helpers.runtime import (
 )
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image")
 

tests/e2e/online_serving/test_lingbot_video_moe.py

@@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH
     )
 
 
-@pytest.mark.slow
+@pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
 def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:

tests/e2e/online_serving/test_ltx.py

@@ -10,7 +10,7 @@ from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
 BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")

tests/e2e/online_serving/test_ltx25.py

@@ -20,7 +20,7 @@ MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL)
 MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "")
 PROMPT = "A red fox walks through a snowy forest while the camera remains fixed."
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
 
 

tests/e2e/online_serving/test_ming_flash_omni_expansion.py

@@ -20,7 +20,7 @@ from tests.helpers.stage_config import get_deploy_config_path
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-pytestmark = [pytest.mark.omni, pytest.mark.slow]
+pytestmark = [pytest.mark.omni, pytest.mark.full_model]
 
 _SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
     reason="Requires 4x H100 GPUs; skipped in CI for now.",

tests/e2e/online_serving/test_qwen_image_edit.py

@@ -2,19 +2,20 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Online serving tests for ``Qwen/Qwen-Image-Edit-2511`` (image-to-image via chat completions).
+Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
 
-- ``test_single_image_to_image_001``: one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``slow`` — mixed input resolutions; ``extra_body`` uses
-  per-output ``width``/``height`` lists; the test client sends one scalar-size request
-  per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
+- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
+- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``slow`` — mixed input
+  resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
+  scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
 
-``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row.
+``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row per model.
 
 From ``tests/``::
 
-    pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "slow and diffusion" --run-level=full_model
+    pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "core_model and diffusion" --run-level=core_model
+    pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "advanced_model and diffusion" --run-level=advanced_model
 """
 
 import pytest
@@ -28,7 +29,6 @@ from tests.helpers.runtime import (
     dummy_messages_from_mix_data,
 )
 
-MODEL = "Qwen/Qwen-Image-Edit-2511"
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
@@ -39,7 +39,7 @@ NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 
 
-def _get_diffusion_feature_cases(model: str = MODEL):
+def _get_diffusion_feature_cases(model: str):
     """Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
     return [
         pytest.param(
@@ -56,11 +56,11 @@ def _get_diffusion_feature_cases(model: str = MODEL):
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases(),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
     indirect=True,
 )
 def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
+    """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -86,7 +86,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases(),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -117,7 +117,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases(),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -1,6 +1,8 @@
 """
 Comprehensive tests of diffusion features that are available in online serving mode
-and are supported by Qwen-Image-Edit-2511 (single image input and two image inputs).
+and are supported by the following models:
+- Qwen-Image-Edit: single image input
+- Qwen-Image-Edit-2511: single image input and two image inputs
 """
 
 import pytest
@@ -9,10 +11,10 @@ from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
-MODEL = "Qwen/Qwen-Image-Edit-2511"
-SINGLE_EDIT_PROMPT = "Restyle this image into a Vincent van Gogh style impressionist painting."
+EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -23,7 +25,9 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 
-def _get_diffusion_feature_cases(model: str = MODEL):
+# This test file targets two models, so I write a helper function.
+# If a similar test only involves one model, one can just define a global list variable.
+def _get_diffusion_feature_cases(model: str):
     return [
         pytest.param(
             OmniServerParams(
@@ -119,7 +123,35 @@ def _get_diffusion_feature_cases(model: str = MODEL):
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases(),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
+    indirect=True,
+)
+def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+
+    # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -134,7 +166,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
-    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT)
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
 
     request_config = {
         "model": omni_server.model,
@@ -154,7 +186,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases(),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -18,7 +18,7 @@ from tests.helpers.mark import hardware_marks
 from tests.helpers.media import decode_b64_image, generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "Qwen/Qwen-Image-Layered"
 EDIT_PROMPT = "Decompose this image into layers."

tests/e2e/online_serving/test_voxcpm2_tts_expansion.py

@@ -16,7 +16,7 @@ from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
-pytestmark = [pytest.mark.slow, pytest.mark.tts]
+pytestmark = [pytest.mark.full_model, pytest.mark.tts]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 

tests/e2e/online_serving/test_wan_2_1_vace_expansion.py

@@ -26,7 +26,7 @@ import pytest
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
 PROMPT = "A cat walking slowly across a sunlit garden path"

tests/engine/test_orchestrator_error_handling.py

@@ -411,51 +411,6 @@ async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orch
             orchestrator_fixture.thread.join(timeout=5)
 
 
-@pytest.mark.asyncio
-async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_request(
-    orchestrator_factory,
-) -> None:
-    """R1.3 of #4855: an async-chunk request whose stage-0 prompt carries no
-    token ids cannot prewarm anything, so it must fail with a client error
-    instead of stalling.
-
-    Non-fatal is the point: ``fatal=True`` reads as "the engine died" and the
-    entrypoints act on it (``OmniLLM.generate`` raises out of its batch loop,
-    the serving layer calls ``terminate_if_errored``). One client's embeds-only
-    prompt must not reach other requests.
-    """
-    stage0 = FakeStageClient(stage_type="llm", final_output=False)
-    stage1 = FakeStageClient(stage_type="llm", final_output=True)
-    orchestrator_fixture = orchestrator_factory([stage0, stage1], async_chunk=True)
-
-    try:
-        await _enqueue_add_request(
-            orchestrator_fixture,
-            request_id="req-embeds",
-            # No prompt_token_ids: what an embeds-only prompt looks like here.
-            prompt=SimpleNamespace(request_id="req-embeds"),
-            original_prompt={"prompt": "embeds only"},
-            sampling_params_list=[_sampling_params(), _sampling_params()],
-            final_stage_id=1,
-        )
-
-        error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-embeds")
-        assert error_msg.fatal is False
-        assert error_msg.status_code == 400
-        assert error_msg.error_type == "BadRequestError"
-        assert error_msg.stage_id == 0
-        assert "prompt_token_ids" in error_msg.error
-
-        # Nothing was prewarmed downstream, and the request is gone.
-        assert stage1.add_request_calls == []
-        await _wait_for(lambda: "req-embeds" not in orchestrator_fixture.orchestrator.request_states)
-        assert orchestrator_fixture.thread.is_alive()
-    finally:
-        if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-            orchestrator_fixture.thread.join(timeout=5)
-
-
 # ───────── Direct unit tests for the fault-isolation helpers ─────────
 
 

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -256,13 +256,12 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
         duplex_identity=SimpleNamespace(),
     )
 
-    prewarmed = await orchestrator._prewarm_async_chunk_stages(
+    await orchestrator._prewarm_async_chunk_stages(
         "req-prewarm",
         SimpleNamespace(prompt_token_ids=[1, 2], resumable=True),
         req_state,
     )
 
-    assert prewarmed is True
     assert stage1.submitted == []
     assert len(stage2.submitted) == 1
     assert 1 not in req_state.stage_submit_ts
@@ -278,7 +277,6 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
 @pytest.mark.asyncio
 async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
     port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
-    prewarm.return_value = True
 
     result = await port.submit(submission)
 
@@ -287,30 +285,6 @@ async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
     prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
 
 
-@pytest.mark.asyncio
-async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None:
-    """A failed prewarm has already aborted the request and popped its state.
-
-    Returning a success result here would hand the control plane a replica for a
-    request that no longer exists, and the trailing bookkeeping would re-register
-    a running counter the cleanup just released -- a leak that never decrements.
-    """
-    port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
-    counter = MagicMock()
-    port._running_counter = counter
-    prewarm.return_value = False
-    request_state = request_states["req-duplex"]
-
-    with pytest.raises(RuntimeError, match="prewarm failed"):
-        await port.submit(submission)
-
-    prewarm.assert_awaited_once()
-    assert request_state.duplex_stage_fences == {}
-    assert request_state.stage_submit_ts == {}
-    assert request_state.running_counter_registered is False
-    counter.increment.assert_not_called()
-
-
 @pytest.mark.asyncio
 async def test_async_route_forwards_to_outgoing_only_stage() -> None:
     orchestrator = object.__new__(Orchestrator)

tests/entrypoints/openai_api/test_audex_serving_guards.py

@@ -26,13 +26,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
-    from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter
-    from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
-
     serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
     serving._tts_model_type = model_type
-    adapter_cls = resolve_adapter(model_type)
-    serving._adapter = adapter_cls(SpeechServingContext(server=serving)) if adapter_cls is not None else None
     serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
     serving._discard_ref_audio_artifact_warmup = lambda request_id: None
     return serving
@@ -159,7 +154,7 @@ class TestAudexCfgValidation:
         assert _adapter().validate(_speech_request(cfg_scale=None)) is None
         serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": None})
-        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
         assert out.extra_args == {}
 
     @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
@@ -182,21 +177,21 @@ class TestAudexCfgInjection:
     def test_no_extra_args_is_untouched(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args=None)
-        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
         assert out.extra_args is None
         assert params.extra_args is None
 
     def test_scale_one_leaves_params_byte_identical(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
-        out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
         assert out.extra_args == {}
 
     def test_guided_request_gains_pair_contract(self):
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
         cond = self._cond_prompt("Some transcription text here.")
-        out = serving._adapter._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
+        out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
 
         extra = out.extra_args
         assert extra["cfg_scale"] == 1.5
@@ -212,14 +207,14 @@ class TestAudexCfgInjection:
         serving = self._serving_with_fake_tokenizer()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
         with pytest.raises(ValueError, match="adapter-built text prompt"):
-            serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
+            serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
 
     def test_injection_never_mutates_input_params(self):
         """The input may be the engine's SHARED default params (review P2):
         pair state must land only on the returned clone."""
         serving = self._serving_with_fake_tokenizer()
         shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
-        out = serving._adapter._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
+        out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
         assert out is not shared
         assert shared.extra_args == {"cfg_scale": 1.5}
         assert shared.temperature == 0.1
@@ -257,11 +252,11 @@ def test_guided_injection_sets_guided_temperature():
     serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
     cond = TestAudexCfgInjection()._cond_prompt()
     params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
-    out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
+    out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
     assert out.temperature == 0.05
 
     unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
-    out = serving._adapter._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
+    out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
     assert out.temperature == 0.1
 
 
@@ -345,7 +340,7 @@ class TestAudexTTAInjection:
     def test_default_injection_applies_rvq_and_cfg3(self):
         serving = self._serving()
         params = SimpleNamespace(extra_args=None)
-        out = serving._adapter._inject_audex_tta_args("req-9", self._prompt(), params)
+        out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
 
         extra = out.extra_args
         rvq = extra["tta_rvq"]
@@ -360,15 +355,15 @@ class TestAudexTTAInjection:
     def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
         serving = self._serving()
         params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
-        out = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), params)
+        out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
         assert "tta_rvq" in out.extra_args
         assert "cfg_scale" not in out.extra_args
         assert "cfg_role" not in out.extra_args
 
     def test_rvq_contract_cached_per_process(self):
         serving = self._serving()
-        out_a = serving._adapter._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
-        out_b = serving._adapter._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
+        out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
+        out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
         assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
 
     def test_injection_never_mutates_input_params(self):
@@ -377,8 +372,8 @@ class TestAudexTTAInjection:
         per-request pair state must land only on the returned clone."""
         serving = self._serving()
         shared = SimpleNamespace(extra_args=None)
-        out_1 = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), shared)
-        out_2 = serving._adapter._inject_audex_tta_args("req-2", self._prompt(), shared)
+        out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
+        out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
         assert shared.extra_args is None
         assert out_1 is not shared and out_2 is not shared
         assert out_1.extra_args["cfg_pair_id"] == "req-1"

tests/entrypoints/openai_api/test_serving_speech.py

@@ -52,7 +52,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
     KEY_SPEAKER_WAVEFORM,
     KEY_SPEAKER_WAVEFORM_LENGTHS,
     SPEAKER_EMBEDDING_DIM,
-    TEXT_EOS_TOKEN_ID,
 )
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -3619,21 +3618,7 @@ class TestCosyVoice3Serving:
         assert error is not None
         assert "max_new_tokens" in error
 
-    @pytest.mark.parametrize(
-        ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
-        [
-            (None, 10, 2000),
-            (5, 5, 5),
-        ],
-    )
-    def test_prepare_speech_generation_cosyvoice3(
-        self,
-        cosyvoice3_server,
-        mocker: MockerFixture,
-        max_new_tokens: int | None,
-        expected_min_tokens: int,
-        expected_max_tokens: int,
-    ):
+    def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture):
         cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock(
             return_value={
                 "prompt": "Hello",
@@ -3641,23 +3626,15 @@ class TestCosyVoice3Serving:
                 "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
             }
         )
-        cosyvoice3_server.model_config.hf_config = SimpleNamespace(
-            min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True
-        )
-        cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock()
-        mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10])
+        cosyvoice3_server._apply_cosyvoice3_dynamic_tokens = mocker.MagicMock(side_effect=lambda spl, req: spl)
 
         request = OpenAICreateSpeechRequest(
             input="Hello",
             ref_audio="data:audio/wav;base64,abc",
             ref_text="Reference text",
-            max_new_tokens=max_new_tokens,
         )
         request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request))
-        sampling_params = cosyvoice3_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
 
-        assert sampling_params.min_tokens == expected_min_tokens
-        assert sampling_params.max_tokens == expected_max_tokens
         assert request_id.startswith("speech-")
         assert generator == "generator"
         assert tts_params == {}
@@ -3741,90 +3718,6 @@ class TestGLMTTSServing:
         assert text_token_len == 3
         load_tokenizer.assert_called_once()
 
-    def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture):
-        glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock(
-            return_value={
-                "prompt": "Hello",
-                "multi_modal_data": {"audio": (np.zeros(24000), 24000)},
-                "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
-            }
-        )
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="Reference text",
-        )
-        glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200)
-        glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10)
-        request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request))
-        sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
-
-        assert sampling_params.min_tokens == 10
-        assert sampling_params.max_tokens == 2000
-        assert request_id.startswith("speech-")
-        assert generator == "generator"
-        assert tts_params == {}
-        glm_tts_server._build_glm_tts_prompt.assert_awaited_once()
-
-
-@pytest.fixture
-def ming_tts_server(mocker: MockerFixture):
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
-
-    mock_engine_client = mocker.MagicMock()
-    mock_engine_client.errored = False
-    mock_engine_client.model_config = mocker.MagicMock(
-        model="inclusionAI/Ming-omni-tts-0.5B",
-        hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
-    )
-    mock_engine_client.default_sampling_params_list = [
-        SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
-    ]
-    mock_engine_client.tts_batch_max_items = 32
-    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
-    mock_engine_client.stage_configs = [
-        SimpleNamespace(
-            engine_args=SimpleNamespace(
-                model_stage="llm", model_arch="MingTTSForConditionalGeneration", worker_type="ar"
-            ),
-            tts_args={},
-        )
-    ]
-
-    mock_models = mocker.MagicMock()
-    mock_models.is_base_model.return_value = True
-    return OmniOpenAIServingSpeech(
-        engine_client=mock_engine_client,
-        models=mock_models,
-        request_logger=mocker.MagicMock(),
-    )
-
-
-class TestMingTTSServing:
-    def test_prepare_speech_generation_ming(self, ming_tts_server, mocker: MockerFixture):
-        ming_tts_server._adapter.build = mocker.AsyncMock(
-            return_value=PreparedRequest(
-                prompt={"prompt": "hello"},
-                tts_params={},
-                model_type="ming_tts",
-            )
-        )
-
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="Reference text",
-            max_new_tokens=7,
-        )
-        asyncio.run(ming_tts_server._prepare_speech_generation(request))
-
-        assert ming_tts_server._tts_model_type == "ming_tts"
-        ming_tts_server._adapter.build.assert_awaited_once()
-        sampling_params = ming_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
-        assert sampling_params.stop_token_ids == [TEXT_EOS_TOKEN_ID]
-        assert sampling_params.max_tokens == 8
-
 
 @pytest.fixture
 def ming_flash_omni_tts_server(mocker: MockerFixture):
@@ -4008,9 +3901,6 @@ class TestTTSAsyncOffloading:
                     model_type=adapter_model_type,
                 )
 
-            def apply_sampling_overrides(self, sampling_params_list, request, prompt, request_id=None):
-                return sampling_params_list
-
         voxtral_server._tts_model_type = legacy_tts_model_type
         mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
         log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")

tests/entrypoints/openai_api/test_tts_detection.py

@@ -14,6 +14,7 @@ import itertools
 import pytest
 
 from vllm_omni.entrypoints.openai.tts_adapters import (
+    LEGACY_TTS_DETECTORS,
     TTS_ADAPTER_REGISTRY,
     all_tts_stage_keys,
     detect_tts_model_type,
@@ -281,17 +282,27 @@ def test_entry_stage_archs_is_ming_only():
     assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
 
 
-def test_every_detected_type_is_adapter_backed():
-    """Detection may only name a model that has a registered adapter."""
+def test_every_detected_type_is_adapter_backed_or_declared_legacy():
+    """Detection may only name a model that has an adapter or an explicit
+    :data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string."""
+    legacy_names = {d.name for d in LEGACY_TTS_DETECTORS}
     for stage, arch in itertools.product(_STAGES, _ARCHS):  # full cross product on purpose
         detected = detect_tts_model_type(stage, arch)
         if detected is None:
             continue
-        assert detected in TTS_ADAPTER_REGISTRY, (
+        assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, (
             f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}"
         )
 
 
+def test_legacy_detectors_have_no_adapter():
+    """A legacy entry that gained an adapter must be deleted from the list."""
+    for detector in LEGACY_TTS_DETECTORS:
+        assert resolve_adapter(detector.name) is None, (
+            f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS"
+        )
+
+
 def test_no_ambiguous_detectors_at_equal_priority():
     """Two detectors that can match the same deployment must be ordered.
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,6 +328,13 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
+            "--dlo-enable-runtime-cache",
+            "--dlo-runtime-cache-dir",
+            "/var/cache/vllm-omni-dlo",
+            "--dlo-runtime-cache-lock-timeout",
+            "30",
+            "--dlo-runtime-cache-pin-limit-gib",
+            "80",
             "--dlo-resident-layers",
             "20",
         ]
@@ -339,9 +346,17 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
+    assert args.dlo_enable_runtime_cache is True
+    assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo"
+    assert args.dlo_runtime_cache_lock_timeout == 30
+    assert args.dlo_runtime_cache_pin_limit_gib == 80
     assert args.dlo_resident_layers == 20
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
+    assert engine_args["dlo_enable_runtime_cache"] is True
+    assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo"
+    assert engine_args["dlo_runtime_cache_lock_timeout"] == 30
+    assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80
     assert engine_args["dlo_resident_layers"] == 20
 
 

tests/examples/offline_inference/test_x_to_video_audio.py

@@ -1,144 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import io
-from pathlib import Path
-from types import SimpleNamespace
-
-import av
-import numpy as np
-import pytest
-from PIL import Image
-
-from examples.offline_inference.x_to_video_audio import x_to_video_audio
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_builds_canonical_image_audio_to_video_prompt() -> None:
-    image = Image.new("RGB", (16, 8), "red")
-    audio = (np.arange(8, dtype=np.float32), 24000)
-
-    prompt = x_to_video_audio.build_x_to_video_audio_prompt(
-        "A person speaks to the camera.",
-        {"image": [image], "audio": [audio]},
-    )
-
-    assert prompt == {
-        "prompt": "A person speaks to the camera.",
-        "modalities": ["video"],
-        "multi_modal_data": {"image": [image], "audio": [audio]},
-    }
-
-
-def test_canonical_prompt_omits_empty_media() -> None:
-    assert x_to_video_audio.build_x_to_video_audio_prompt("A landscape.") == {
-        "prompt": "A landscape.",
-        "modalities": ["video"],
-    }
-
-
-def test_official_prompt_cleanup_stays_in_prompt_file_handling() -> None:
-    prompt = (
-        "[SPEAKER_TIMESTAMPS_START]metadata[SPEAKER_TIMESTAMPS_END]\n\n"
-        "A person waves.\n[AUDIO_DESCRIPTION_START]noise[AUDIO_DESCRIPTION_END]"
-    )
-
-    assert x_to_video_audio._clean_official_prompt(prompt) == "A person waves."
-
-
-def test_media_loader_preserves_complete_audio(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-) -> None:
-    image_path = tmp_path / "reference.png"
-    Image.new("RGBA", (4, 3), "blue").save(image_path)
-    complete_waveform = np.arange(80000, dtype=np.float32)
-    requested_rates: list[int | None] = []
-
-    def fake_load_audio(path: str, *, sr: int | None) -> tuple[np.ndarray, int]:
-        assert path == "reference.wav"
-        requested_rates.append(sr)
-        return complete_waveform, 16000
-
-    monkeypatch.setattr(x_to_video_audio, "load_audio", fake_load_audio)
-
-    images, audios = x_to_video_audio.load_image_and_audio(
-        [str(image_path)],
-        ["reference.wav"],
-        audio_sample_rate=16000,
-    )
-
-    assert images[0].mode == "RGB"
-    assert requested_rates == [16000]
-    assert audios[0][0] is complete_waveform
-    assert audios[0][0].shape == (80000,)
-    assert audios[0][1] == 16000
-
-
-@pytest.mark.parametrize(
-    "video",
-    [
-        np.zeros((3, 2, 4, 4), dtype=np.float32),
-        np.zeros((2, 4, 4, 3), dtype=np.uint8),
-    ],
-    ids=["channel_first_float", "frame_first_uint8"],
-)
-def test_output_contract_keeps_model_video_layout_opaque(video: np.ndarray) -> None:
-    audio = np.zeros(16, dtype=np.float32)
-    result = SimpleNamespace(
-        images=[video],
-        multimodal_output={"audio": audio, "fps": 25, "audio_sample_rate": 44100},
-    )
-
-    output = x_to_video_audio.extract_x_to_video_audio_output([result])
-
-    assert output.video is video
-    assert output.audio is audio
-    assert output.fps == 25.0
-    assert output.audio_sample_rate == 44100
-
-
-def test_output_contract_accepts_video_from_multimodal_output() -> None:
-    video = object()
-    result = SimpleNamespace(images=[], multimodal_output={"video": video, "fps": 12})
-
-    output = x_to_video_audio.extract_x_to_video_audio_output(result)
-
-    assert output == x_to_video_audio.XToVideoAudioOutput(
-        video=video,
-        audio=None,
-        fps=12.0,
-        audio_sample_rate=None,
-    )
-
-
-def test_output_contract_requires_model_neutral_metadata() -> None:
-    result = SimpleNamespace(images=[object()], multimodal_output={"audio": np.zeros(4)})
-
-    with pytest.raises(RuntimeError, match="must declare a valid 'fps'"):
-        x_to_video_audio.extract_x_to_video_audio_output(result)
-
-    with pytest.raises(RuntimeError, match="must declare a valid 'audio_sample_rate'"):
-        x_to_video_audio.extract_x_to_video_audio_output(result, fps=24)
-
-
-def test_output_encoder_normalizes_declared_range_and_channel_first_layout() -> None:
-    video = np.zeros((3, 2, 4, 6), dtype=np.float32)
-    output = x_to_video_audio.XToVideoAudioOutput(
-        video=video,
-        audio=None,
-        fps=8,
-        audio_sample_rate=None,
-        output_tensor_range="negative_one_to_one",
-    )
-
-    payload = x_to_video_audio.encode_x_to_video_audio_output(output)
-
-    with av.open(io.BytesIO(payload)) as container:
-        decoded = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
-    assert len(decoded) == 2
-    assert decoded[0].shape == (4, 6, 3)
-    np.testing.assert_allclose(decoded, 128, atol=2)

tests/model_extras/test_model_extras.py

@@ -15,10 +15,8 @@ from vllm_omni.model_extras import (
     build_image_to_video_prompt,
     build_text_to_image_prompt,
     build_x_to_text_prompt,
-    build_x_to_video_audio_prompt,
     get_extra_body_params,
     get_extra_output_params,
-    get_input_audio_sample_rate,
     get_output_tensor_range,
     get_x_to_text_model_family,
     should_init_extra_args_for_non_diffusion_stages,
@@ -143,50 +141,6 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
     assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_dreamid_x_to_video_audio_adapter() -> None:
-    image = Image.new("RGB", (16, 16))
-    waveform = object()
-    canonical_prompt = {
-        "prompt": "A person waves.",
-        "modalities": ["video"],
-        "multi_modal_data": {"image": [image], "audio": [(waveform, 16000)]},
-    }
-
-    result = build_x_to_video_audio_prompt(
-        "DreamIDOmniPipeline",
-        canonical_prompt,
-        {
-            "video_negative_prompt": "blur",
-            "audio_negative_prompt": "echo",
-        },
-    )
-
-    assert result == {
-        "prompt": "A person waves.",
-        "modalities": ["video"],
-        "video_negative_prompt": "blur",
-        "audio_negative_prompt": "echo",
-        "multi_modal_data": {"image": [image], "audio": [waveform]},
-    }
-    assert get_extra_body_params("DreamIDOmniPipeline") == frozenset({"solver_name", "shift"})
-    assert get_input_audio_sample_rate("DreamIDOmniPipeline") == 16000
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_generic_x_to_video_audio_adapter_is_model_neutral() -> None:
-    canonical_prompt = {
-        "prompt": "A person waves.",
-        "modalities": ["video"],
-        "multi_modal_data": {"audio": [(object(), 48000)]},
-    }
-
-    assert build_x_to_video_audio_prompt("FuturePipeline", canonical_prompt, {}) is canonical_prompt
-    assert get_input_audio_sample_rate("FuturePipeline") is None
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 @pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])

tests/tools/test_check_tts_adapter.py

@@ -17,7 +17,9 @@ import pytest
 sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit"))
 
 from check_tts_adapter import (  # noqa: E402
+    MAX_LEGACY_DETECTORS,
     MAX_MODEL_TYPE_BRANCHES,
+    _legacy_detectors,
     _model_type_branches,
     main,
 )
@@ -58,6 +60,12 @@ def test_dict_dispatch_is_a_known_gap():
     assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0
 
 
+def test_legacy_detector_counting():
+    src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))'
+    found = _legacy_detectors(ast.parse(src))
+    assert [name for _, name in found] == ["a", "b"]
+
+
 def test_gate_passes_on_current_tree():
     """The recorded budgets must match reality, or the ratchet is decorative."""
     assert main([]) == 0
@@ -66,4 +74,6 @@ def test_gate_passes_on_current_tree():
 def test_budgets_are_the_real_counts():
     root = Path(__file__).resolve().parents[2]
     serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py"
+    adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py"
     assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES
+    assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS

tests/worker/test_gpu_ar_model_runner.py

@@ -2,7 +2,6 @@ import asyncio
 import copy
 from collections.abc import Sequence
 from types import SimpleNamespace
-from typing import Any
 
 import numpy as np
 import pytest
@@ -201,15 +200,6 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
         async def build(self, request, sampling_params_list, has_inline_ref_audio):
             return PreparedRequest(prompt={"prompt": request.input}, tts_params={}, model_type="higgs_audio_v3")
 
-        def apply_sampling_overrides(
-            self,
-            sampling_params_list: list,
-            request: "OpenAICreateSpeechRequest",
-            prompt: dict[str, Any] | None = None,
-            request_id: str | None = None,
-        ) -> list:
-            return sampling_params_list
-
     engine = SimpleNamespace(
         errored=False,
         default_sampling_params_list=[SamplingParams(temperature=1.0, top_p=0.95, top_k=50)],

tools/nightly/buildkite_testcase_statistics.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 """
-Parse pytest commands from Buildkite .buildkite/cuda/test-{ready,merge,nightly,weekly}.yml;
-collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report.
+Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml,
+and test-weekly.yml; collect test cases (including parametrized) via pytest --collect-only -q
+and produce an HTML report.
 
 Leaf steps may live under ``group:`` blocks (nested ``steps``); those are flattened with the
 group title carried into the report.
@@ -29,7 +30,7 @@ import yaml
 
 # Repo root (parent of the directory containing this script)
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent
-BUILDKITE_DIR = REPO_ROOT / ".buildkite" / "cuda"
+BUILDKITE_DIR = REPO_ROOT / ".buildkite"
 PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml", "test-weekly.yml"]
 
 
@@ -255,34 +256,13 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l
         rel_dir = Path(target.rstrip("/"))
         dir_path = (repo_root / rel_dir).resolve()
         if dir_path.exists() and dir_path.is_dir():
-            return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 300
+            return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 120
 
     if not extra:
         raise RuntimeError(f"Failed to parse pytest sidecar args from line: {raw_line!r}")
     if not (repo_root / "tests").exists():
         raise FileNotFoundError("tests/ directory not found under repo root")
-    # Broad ``tests/`` collection loads many conftests; allow more time on busy CI agents.
-    return ["tests/"], "tests/", 300
-
-
-def _format_pytest_collect_failure(
-    path_args: list[str],
-    extra: list[str],
-    result: subprocess.CompletedProcess[str],
-) -> str:
-    """pytest often puts collection errors on stdout; include both streams + exit code."""
-    stdout = (result.stdout or "").strip()
-    stderr = (result.stderr or "").strip()
-    parts = [
-        f"pytest --collect-only failed for {path_args} with args {extra} (exit={result.returncode})",
-    ]
-    if stderr:
-        parts.append(f"stderr:\n{stderr}")
-    if stdout:
-        parts.append(f"stdout:\n{stdout}")
-    if not stderr and not stdout:
-        parts.append("(no stdout/stderr captured)")
-    return "\n".join(parts)
+    return ["tests/"], "tests/", 120
 
 
 def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]:
@@ -293,8 +273,6 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
     flags ``--ignore``, ``--test-config-file``, ``-m``, ``-k``, ``--run-level`` are taken from the
     full ``raw_line`` and appended to the same invocation so collection matches CI (perf scripts
     parametrize from ``--test-config-file`` at import time).
-
-    Pytest exit code 5 (no tests collected) is treated as an empty result, not a hard failure.
     """
     extra = _pytest_collect_sidecar_args(raw_line)
     path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line)
@@ -308,15 +286,8 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
         text=True,
         timeout=timeout_quiet,
     )
-    # 5 == no tests collected under the given -m/-k filters.
-    if result.returncode == 5:
-        print(
-            f"pytest --collect-only: no tests matched for {path_args} with args {extra}",
-            file=sys.stderr,
-        )
-        return []
     if result.returncode != 0:
-        raise RuntimeError(_format_pytest_collect_failure(path_args, extra, result))
+        raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}")
     print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}")
     return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "")
 
@@ -991,7 +962,7 @@ def main() -> None:
         "--buildkite-dir",
         type=Path,
         default=BUILDKITE_DIR,
-        help="Path to Buildkite CUDA pipeline directory (default: .buildkite/cuda)",
+        help="Path to .buildkite directory",
     )
     default_out = REPO_ROOT / "buildkite_testcase_statistics.html"
     parser.add_argument(
@@ -1052,12 +1023,7 @@ def main() -> None:
             names = []
             skipped_ids: set[str] = set()
             for target, raw in targets:
-                try:
-                    collected = collect_test_names(target, REPO_ROOT, raw)
-                except (RuntimeError, FileNotFoundError, subprocess.TimeoutExpired) as exc:
-                    # Keep generating the report for other steps; surface the root cause.
-                    print(f"WARNING: collect failed for {module_key} target={target!r}: {exc}", file=sys.stderr)
-                    collected = []
+                collected = collect_test_names(target, REPO_ROOT, raw)
                 skipped_ids.update(get_skip_status(target, raw, REPO_ROOT))
                 names.extend(collected)
             count = len(names) - len(skipped_ids)

tools/pre_commit/check_tts_adapter.py

@@ -6,8 +6,10 @@ Per-model behaviour is being moved out of ``serving_speech.py`` and into
 nothing stopped a new model from adding one more ``self._tts_model_type == ...``
 branch faster than the migration removed them.
 
-This gate makes the number of ``self._tts_model_type`` comparisons in
-``serving_speech.py`` monotonically non-increasing.
+This gate makes the two debt counters monotonically non-increasing:
+
+1. ``self._tts_model_type`` comparisons in ``serving_speech.py``
+2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter)
 
 Lowering a budget after removing branches is expected and encouraged — the check
 insists on it, so the recorded numbers cannot silently drift upward. Raising one
@@ -26,9 +28,11 @@ from pathlib import Path
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
+ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py")
 
 # Budgets. These may only go DOWN. See the module docstring.
-MAX_MODEL_TYPE_BRANCHES = 20
+MAX_MODEL_TYPE_BRANCHES = 27
+MAX_LEGACY_DETECTORS = 0
 
 _GUIDANCE = """
 Move the per-model behaviour into its adapter under
@@ -105,6 +109,20 @@ def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]:
     return found
 
 
+def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
+    """Every ``LegacyDetector(...)`` construction."""
+    found: list[tuple[int, str]] = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector":
+            name = "?"
+            for kw in node.keywords:
+                if kw.arg == "name" and isinstance(kw.value, ast.Constant):
+                    name = str(kw.value.value)
+            found.append((node.lineno, name))
+    found.sort()
+    return found
+
+
 def _check(
     label: str,
     path: Path,
@@ -133,18 +151,21 @@ def _check(
 
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description=__doc__)
-    # pre-commit passes the staged filenames; we always audit the same file,
+    # pre-commit passes the staged filenames; we always audit the same two files,
     # so they are accepted and ignored.
     parser.add_argument("filenames", nargs="*")
     parser.parse_args(argv)
 
     serving_path = REPO_ROOT / SERVING_SPEECH
-    if not serving_path.is_file():
-        print(f"check_tts_adapter: expected file not found: {serving_path}", file=sys.stderr)
-        return 1
+    adapters_path = REPO_ROOT / ADAPTERS_INIT
+    for path in (serving_path, adapters_path):
+        if not path.is_file():
+            print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
+            return 1
 
     notices: list[str] = []
     branches = _model_type_branches(ast.parse(serving_path.read_text()))
+    detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
 
     errors = _check(
         "`_tts_model_type` comparisons",
@@ -155,6 +176,15 @@ def main(argv: list[str] | None = None) -> int:
         "MAX_MODEL_TYPE_BRANCHES",
         notices,
     )
+    errors += _check(
+        "unmigrated model types in LEGACY_TTS_DETECTORS",
+        ADAPTERS_INIT,
+        len(detectors),
+        MAX_LEGACY_DETECTORS,
+        [f"line {lineno}: {name}" for lineno, name in detectors],
+        "MAX_LEGACY_DETECTORS",
+        notices,
+    )
 
     for notice in notices:
         print(f"check_tts_adapter: {notice}", file=sys.stderr)

vllm_omni/config/omni_config.py

@@ -576,6 +576,10 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0)
+    dlo_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
     dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -638,6 +642,13 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
+
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()
         elif isinstance(self.tf_model_config, Mapping):

vllm_omni/config/stage_config.py

@@ -382,6 +382,10 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
+    dlo_enable_runtime_cache: bool | None = None
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float | None = None
+    dlo_runtime_cache_pin_limit_gib: float | None = None
     dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import math
 import os
 import time
 from collections.abc import Iterable
@@ -43,56 +42,25 @@ logger = init_logger(__name__)
 
 _STATS_INTERVAL_S = 1.0
 
-# Upper bound on how long a request may wait for stage input before the
-# scheduler force-fails it.  Defends against stuck consumer-side requests when
-# the producer drops a payload, send fails, or recv never arrives.  Override
-# per-deployment via VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set exactly 0 to disable
-# the safety net. Negative and non-finite values are rejected at startup.
+# Upper bound on how long a request may sit in full-payload-input wait
+# (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
+# before the scheduler force-fails it.  Defends against stuck consumer-side
+# requests when the producer drops a full-payload, send fails, or recv
+# never arrives.  Override per-deployment via
+# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net.
 #
-# Scope: both transports.  The full-payload path measures time parked in
-# WAITING_FOR_INPUT (``OmniSchedulingCoordinator._waiting_since``); the
-# async-chunk path measures time stalled in WAITING_FOR_CHUNK
-# (``OmniChunkTransferAdapter._waiting_since``, reset on each chunk arrival).
-# One knob, because the question it answers -- "how long may a request wait for
-# stage input?" -- does not depend on which transport carries it.
-_INPUT_WAIT_TIMEOUT_DEFAULT = 600.0
+# Scope: this constant only covers the full-payload coordinator path
+# (``input_coordinator``).  The async-chunk path uses
+# ``chunk_transfer_adapter`` and is not affected by this constant.
 _INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")
 try:
     DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
 except ValueError:
     logger.warning(
-        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to %s seconds.",
+        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.",
         _INPUT_WAIT_TIMEOUT_RAW,
-        _INPUT_WAIT_TIMEOUT_DEFAULT,
-    )
-    DEFAULT_INPUT_WAIT_TIMEOUT_S = _INPUT_WAIT_TIMEOUT_DEFAULT
-
-if not math.isfinite(DEFAULT_INPUT_WAIT_TIMEOUT_S):
-    # nan compares false against everything, so every deadline check silently
-    # passes; inf makes the deadline unreachable. Both disable the net while
-    # looking like a number, which is the failure this item exists to stop.
-    raise ValueError(
-        f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is not finite. "
-        "Use a positive number of seconds, or exactly 0 to disable the deadline."
-    )
-if DEFAULT_INPUT_WAIT_TIMEOUT_S < 0:
-    # Fail at startup rather than substituting a default. `-1` is a common idiom
-    # for "no limit", so silently turning it into 600 would give an operator who
-    # meant "never time out" mysterious failures ten minutes in.
-    raise ValueError(
-        f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is negative. "
-        "Use a positive number of seconds, or exactly 0 to disable the deadline."
-    )
-elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0:
-    # Zero stays a supported opt-out, but it is no longer silent: it disables
-    # the deadline on BOTH transports, so a request whose stage input never
-    # arrives waits forever. Say which nets just went away.
-    logger.warning(
-        "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=0: stage-input deadlines are DISABLED for "
-        "both the full-payload path (WAITING_FOR_INPUT) and the async-chunk path "
-        "(WAITING_FOR_CHUNK). A request whose input never arrives will wait "
-        "indefinitely instead of failing."
     )
+    DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0
 
 
 class OmniSchedulerMixin:
@@ -216,8 +184,6 @@ class OmniSchedulerMixin:
                 self.running,
                 scheduler_requests=self.requests,
             )
-            self._process_pending_chunk_timeouts()
-            self._log_failed_chunk_sends()
 
     def _restore_omni_wait_queues(self) -> None:
         """Restore requests temporarily parked by Omni input gates."""
@@ -245,9 +211,9 @@ class OmniSchedulerMixin:
         Disabled when ``DEFAULT_INPUT_WAIT_TIMEOUT_S`` is <= 0.
 
         Scope: only covers ``input_coordinator`` (full-payload path).
-        Async-chunk requests park in ``chunk_transfer_adapter`` instead and
-        are handled by ``_process_pending_chunk_timeouts`` below, which shares
-        this timeout.
+        Async-chunk requests park in ``chunk_transfer_adapter`` instead
+        and are not handled here -- if a similar safety net is needed
+        for the chunk path, it belongs in the chunk adapter.
         """
         if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
             return
@@ -268,84 +234,6 @@ class OmniSchedulerMixin:
         )
         self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
 
-    def _log_failed_chunk_sends(self) -> None:
-        """Surface chunks the sender gave up on (R1.2 of #4855).
-
-        ``save_loop`` pops a task with ``popleft()`` and never re-queues it, and
-        ``connector.put`` can report failure without raising, so a dropped chunk
-        is a give-up rather than a retry. Until this change the only trace was a
-        warning that printed ``None`` for the request id, because the task dict
-        has no ``request_id`` key.
-
-        This only reports. Failing the request from here does not work: the
-        producer is not the ``final_output`` stage, so an abort synthesized by
-        ``finish_requests`` is not turned into a client-visible termination by
-        ``Orchestrator._route_output`` (see the ``final_output`` guards there).
-        Making a producer-side drop client-visible needs an orchestrator-side
-        path like ``_fail_request_dead_stage``; the consumer's R1.1 deadline is
-        what actually ends the request today.
-        """
-        adapter = getattr(self, "chunk_transfer_adapter", None)
-        if adapter is None:
-            return
-        failures = adapter.collect_failed_send_request_ids()
-        for request_id, reason in failures.items():
-            if request_id in self.requests:
-                logger.error(
-                    "[OmniScheduler] req=%s: chunk send gave up (%s); the consumer will "
-                    "fail on the stage-input deadline",
-                    request_id,
-                    reason,
-                )
-            else:
-                # Drained but no longer scheduled here: the request finished or
-                # was aborted between the failed send and this sweep. Still say
-                # so -- collect_* empties the map, so staying quiet would make
-                # this the one drop path in a change whose point is visibility.
-                logger.warning(
-                    "[OmniScheduler] req=%s: chunk send gave up (%s) after the request "
-                    "left this scheduler (finished or aborted)",
-                    request_id,
-                    reason,
-                )
-
-    def _process_pending_chunk_timeouts(self) -> None:
-        """Force-fail requests stalled in ``WAITING_FOR_CHUNK`` too long.
-
-        The async-chunk counterpart of ``_process_pending_input_timeouts``,
-        called right after ``chunk_transfer_adapter.process_pending_chunks``.
-        Until now the chunk path had no safety net at all: the scheduler-side
-        comment on the full-payload net said a similar guard "belongs in the
-        chunk adapter", while the worker mixin assumed the full-payload net
-        already covered it.  Only the first was true, so a dropped terminal
-        chunk or a producer that exhausted its send retries parked the request
-        forever (vllm-project/vllm-omni#3833).  Since ``async_chunk: true`` is
-        the default for the flagship TTS pipelines, that was the default
-        streaming path.
-
-        Shares ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` with the full-payload net:
-        one knob for "how long may a request wait for stage input", regardless
-        of which transport carries it.  Disabled when the value is <= 0.
-        """
-        if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
-            return
-        adapter = getattr(self, "chunk_transfer_adapter", None)
-        if adapter is None or not getattr(adapter, "receives_chunks", False):
-            return
-        timed_out_ids = adapter.collect_timed_out_request_ids(timeout_s=DEFAULT_INPUT_WAIT_TIMEOUT_S)
-        if not timed_out_ids:
-            return
-        present_ids = {req_id for req_id in timed_out_ids if req_id in self.requests}
-        if not present_ids:
-            return
-        logger.warning(
-            "Marking %d request(s) as FINISHED_ERROR after stalling > %.0fs waiting for a chunk: %s",
-            len(present_ids),
-            DEFAULT_INPUT_WAIT_TIMEOUT_S,
-            sorted(present_ids),
-        )
-        self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
-
     def _capture_omni_connector_output(self, model_runner_output: Any) -> None:
         """Stash the model runner's omni_connector_output for next schedule().
 

vllm_omni/diffusion/data.py

@@ -708,6 +708,18 @@ class OmniDiffusionConfig:
     # This avoids AllGather synchronization, while host memory follows the
     # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
+    # Build or join a node-local mmap cache after ordinary loading. This is
+    # used when direct checkpoint mmap is unavailable, or when direct CUDA
+    # registration requires a transform-complete final layout. The first
+    # version is no-AllGather only.
+    dlo_enable_runtime_cache: bool = False
+    # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights.
+    dlo_runtime_cache_dir: str | None = None
+    # Maximum wait for the per-layout POSIX writer lock.
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    # Maximum file-backed memory a worker may register with CUDA. Zero keeps
+    # the bounded mmap-to-pinned staging path.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
 
@@ -928,6 +940,16 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0:
+            raise ValueError("dlo_runtime_cache_lock_timeout must be positive")
+        if not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative")
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,6 +14,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
@@ -440,10 +441,16 @@ class DiffusersPipelineLoader:
                 _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
                 _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
                 _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
+                _prefer_registered_runtime_cache = (
+                    _dist_offload
+                    and not _use_ag
+                    and bool(getattr(self.od_config, "dlo_enable_runtime_cache", False))
+                    and float(getattr(self.od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) > 0
+                )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload:
+                if _dist_offload and not _prefer_registered_runtime_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -455,6 +462,16 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
+                elif _prefer_registered_runtime_cache:
+                    # Direct checkpoint bindings can carry deferred transforms
+                    # (for example grouped-QKV reordering), so registering the
+                    # raw checkpoint mapping would not make every recurrent H2D
+                    # source directly copyable. Materialize once, apply every
+                    # post-load mutation, then publish/register the final layout.
+                    logger.info(
+                        "DLO runtime-cache registration requested; using the ordinary loader "
+                        "once to publish a transform-complete final mmap layout"
+                    )
 
                 _skip_load = self.host_weight_plan is not None
 
@@ -491,7 +508,8 @@ class DiffusersPipelineLoader:
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
+                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
+                            plan_result.fallback_code or "unknown",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -516,7 +534,107 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        return model.eval()
+        model.eval()
+        self._maybe_build_runtime_weight_cache(model, load_format)
+        return model
+
+    def _maybe_build_runtime_weight_cache(self, model: nn.Module, load_format: str) -> None:
+        """Publish final ordinary-loader DiT tensors for no-AllGather DLO."""
+        if self.host_weight_plan is not None:
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            return
+        if not getattr(self.od_config, "dlo_enable_runtime_cache", False):
+            return
+        if getattr(self.od_config, "dlo_use_allgather", True):
+            logger.info("DLO runtime cache is no-AllGather only; leaving AllGather storage unchanged")
+            return
+
+        parallel_config = self.parallel_config
+        tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))
+        tp_rank = 0
+        if tp_size > 1:
+            try:
+                from vllm.distributed import get_tensor_model_parallel_rank
+
+                tp_rank = int(get_tensor_model_parallel_rank())
+            except Exception as exc:
+                logger.warning(
+                    "DLO runtime cache unavailable [parallel_rank_unavailable]: cannot resolve TP rank for TP=%d: %s",
+                    tp_size,
+                    exc,
+                )
+                return
+
+        modules = ModuleDiscovery.discover(model)
+        from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+            DEFAULT_LOCK_TIMEOUT_SECONDS,
+            DEFAULT_SHARD_SIZE_BYTES,
+            build_runtime_weight_cache_plan,
+        )
+
+        loader_inputs = {
+            "model_class_name": getattr(self.od_config, "model_class_name", None),
+            "model_config": getattr(self.od_config, "model_config", None),
+            "diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None),
+            "custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None),
+            "lora_path": getattr(self.od_config, "lora_path", None),
+            "lora_scale": getattr(self.od_config, "lora_scale", None),
+            "lora_backend": getattr(self.od_config, "lora_backend", None),
+        }
+        sp_guard = {
+            "sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)),
+            "ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)),
+            "ring_degree": int(getattr(parallel_config, "ring_degree", 1)),
+            "allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)),
+            "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
+        }
+        model_identity = getattr(self.od_config, "model", None)
+        runtime_quantization = self.quant_config
+        if runtime_quantization is None and any(
+            (quant_method := getattr(module, "quant_method", None)) is not None
+            and not isinstance(quant_method, UnquantizedLinearMethod)
+            for module in model.modules()
+        ):
+            runtime_quantization = "module_quant_method"
+        result = build_runtime_weight_cache_plan(
+            model,
+            dit_modules=tuple(zip(modules.dit_names, modules.dits)),
+            loader_type=type(self),
+            cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None),
+            lock_timeout_seconds=float(
+                getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS)
+            ),
+            max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES,
+            model_identity=str(model_identity) if model_identity is not None else None,
+            revision=getattr(self.od_config, "revision", None),
+            runtime_dtype=getattr(self.od_config, "dtype", None),
+            load_format=load_format,
+            loader_inputs=loader_inputs,
+            tensor_parallel_size=tp_size,
+            tensor_parallel_rank=tp_rank,
+            sequence_parallel_guard=sp_guard,
+            use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
+            enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
+            quantization_config=runtime_quantization,
+            cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
+            pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
+        )
+        if result.plan is not None:
+            self.host_weight_plan = result.plan
+            logger.info(
+                "DLO runtime cache selected (layout=%s, TP=%d/%d, SP=%d)",
+                result.plan.runtime_layout_key[:12] if result.plan.runtime_layout_key else "unknown",
+                tp_rank,
+                tp_size,
+                sp_guard["sequence_parallel_size"],
+            )
+        else:
+            logger.warning(
+                "DLO runtime cache unavailable [%s]; retaining ordinary host tensors: %s",
+                result.fallback_code or "unknown",
+                result.fallback_reason or "unspecified reason",
+            )
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -28,7 +28,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 class TensorBinding:
     """One runtime tensor backed by one safetensors entry."""
 
-    checkpoint_key: str
+    storage_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -40,6 +40,8 @@ class HostWeightPlan:
     backing_kind: str
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
+    runtime_layout_key: str | None = None
+    post_load_complete: bool = False
 
 
 @dataclass(frozen=True)
@@ -48,6 +50,7 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
+    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -258,11 +261,11 @@ def _validate_source_metadata(
         with safe_open(file_path, framework="pt", device="cpu") as handle:
             available = set(handle.keys())
             for runtime_name, target, binding in entries:
-                if binding.checkpoint_key not in available:
+                if binding.storage_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.checkpoint_key)
+                source = handle.get_slice(binding.storage_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -287,11 +290,19 @@ def build_checkpoint_mmap_plan(
 ) -> HostWeightPlanResult:
     """Build a complete direct-checkpoint plan or return a fallback reason."""
     if tensor_parallel_size != 1:
-        return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            f"TP={tensor_parallel_size} requires the ordinary loader",
+            "unsupported_tp",
+        )
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
     if online_quantization:
-        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+            "unsupported_online_quantization",
+        )
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -328,14 +339,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                checkpoint_key=checkpoint_key,
+                storage_key=checkpoint_key,
                 file_path=file_path,
                 transform=policy.transform if policy is not None else None,
             )
 
         _validate_source_metadata(required, bindings)
     except (OSError, ValueError, _PlanIncompatibleError) as exc:
-        return HostWeightPlanResult(None, str(exc))
+        return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible")
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -0,0 +1,700 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Node-local, content-addressed runtime-weight cache for diffusion DLO."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import errno
+import fcntl
+import hashlib
+import inspect
+import json
+import math
+import os
+import shutil
+import time
+import uuid
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    HostWeightPlanResult,
+    TensorBinding,
+)
+
+logger = init_logger(__name__)
+
+RUNTIME_CACHE_SCHEMA_VERSION = 1
+DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights"
+DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0
+DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
+_HASH_CHUNK_BYTES = 64 * 1024**2
+_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
+
+_SUPPORTED_DTYPES = {
+    torch.bool,
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+    torch.float64,
+}
+for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
+    if (_dtype := getattr(torch, _dtype_name, None)) is not None:
+        _SUPPORTED_DTYPES.add(_dtype)
+
+
+@dataclass(frozen=True)
+class _RuntimeTensor:
+    name: str
+    tensor: torch.Tensor
+    kind: str
+
+
+class _RuntimeCacheError(RuntimeError):
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+
+
+def default_runtime_cache_root() -> str:
+    """Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
+    return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT)
+
+
+def _persistent_buffer(module: nn.Module, local_name: str) -> bool:
+    parent_path, _, leaf_name = local_name.rpartition(".")
+    owner = module.get_submodule(parent_path)
+    return leaf_name not in owner._non_persistent_buffers_set
+
+
+def _named_parameters_with_duplicates(module: nn.Module):
+    try:
+        return module.named_parameters(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_parameters()
+
+
+def _named_buffers_with_duplicates(module: nn.Module):
+    try:
+        return module.named_buffers(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_buffers()
+
+
+def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
+    parent_path, _, leaf_name = runtime_name.rpartition(".")
+    parent = pipeline.get_submodule(parent_path)
+    tensor = parent._parameters.get(leaf_name)
+    if tensor is None:
+        tensor = parent._buffers.get(leaf_name)
+    return tensor
+
+
+def _collect_runtime_tensors(
+    pipeline: nn.Module,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> list[_RuntimeTensor]:
+    records: dict[str, _RuntimeTensor] = {}
+    for dit_name, dit_module in dit_modules:
+        for local_name, tensor in _named_parameters_with_duplicates(dit_module):
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "parameter")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+        for local_name, tensor in _named_buffers_with_duplicates(dit_module):
+            if not _persistent_buffer(dit_module, local_name):
+                continue
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "buffer")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+
+    if not records:
+        raise _RuntimeCacheError(
+            "ambiguous_ownership",
+            "no DiT parameters or persistent buffers were discovered",
+        )
+
+    storage_owners: dict[tuple[int, int], str] = {}
+    for record in records.values():
+        tensor = record.tensor
+        try:
+            pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
+        except AttributeError as exc:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} is not owned by the pipeline",
+            ) from exc
+        if pipeline_tensor is not tensor:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} does not resolve to the discovered object",
+            )
+        if tensor.device.type != "cpu" or tensor.is_meta:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
+            )
+        if tensor.layout != torch.strided:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported layout {tensor.layout}",
+            )
+        if not tensor.is_contiguous():
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
+            )
+        if tensor.dtype not in _SUPPORTED_DTYPES:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported dtype {tensor.dtype}",
+            )
+        if hasattr(tensor, "to_local"):
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a distributed tensor",
+            )
+
+        tensor_nbytes = tensor.numel() * tensor.element_size()
+        storage = tensor.untyped_storage()
+        if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a view into a larger storage",
+            )
+        if tensor_nbytes:
+            storage_id = (storage.data_ptr(), storage.nbytes())
+            owner = storage_owners.setdefault(storage_id, record.name)
+            if owner != record.name:
+                raise _RuntimeCacheError(
+                    "unsupported_alias",
+                    f"{record.name!r} shares storage with {owner!r}",
+                )
+
+    # The ownership boundary is the complete pipeline, not only the discovered
+    # DiT traversal. Reject a DiT storage that is also registered by an encoder,
+    # VAE, resident component, or non-persistent buffer.
+    for pipeline_name, tensor in chain(
+        _named_parameters_with_duplicates(pipeline),
+        _named_buffers_with_duplicates(pipeline),
+    ):
+        if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
+            continue
+        storage = tensor.untyped_storage()
+        owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
+        if owner is not None and owner != pipeline_name:
+            raise _RuntimeCacheError(
+                "unsupported_alias",
+                f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
+            )
+
+    return [records[name] for name in sorted(records)]
+
+
+def _json_normalize(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, torch.dtype):
+        return str(value)
+    if isinstance(value, os.PathLike):
+        return os.fspath(value)
+    if dataclasses.is_dataclass(value):
+        return _json_normalize(dataclasses.asdict(value))
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        try:
+            dumped = model_dump(mode="json")
+        except (TypeError, ValueError):
+            dumped = model_dump()
+        return _json_normalize(dumped)
+    if isinstance(value, Mapping):
+        return {str(key): _json_normalize(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))}
+    if isinstance(value, (set, frozenset)):
+        return sorted((_json_normalize(item) for item in value), key=lambda item: _canonical_json(item))
+    if isinstance(value, (list, tuple)):
+        return [_json_normalize(item) for item in value]
+    if inspect.isclass(value):
+        return f"{value.__module__}.{value.__qualname__}"
+    return {
+        "type": f"{type(value).__module__}.{type(value).__qualname__}",
+        "value": str(value),
+    }
+
+
+def _canonical_json(value: Any) -> bytes:
+    return json.dumps(
+        _json_normalize(value),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+    ).encode()
+
+
+def _implementation_fingerprint(
+    loader_type: type, pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]]
+) -> str:
+    objects: list[Any] = [loader_type, type(pipeline)]
+    load_weights = getattr(type(pipeline), "load_weights", None)
+    if load_weights is not None:
+        objects.append(load_weights)
+    for _, dit_module in dit_modules:
+        objects.append(type(dit_module))
+        post_load = getattr(type(dit_module), "post_load_weights", None)
+        if post_load is not None:
+            objects.append(post_load)
+
+    digest = hashlib.sha256()
+    identities: set[str] = set()
+    for obj in objects:
+        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
+        if identity in identities:
+            continue
+        identities.add(identity)
+        digest.update(identity.encode())
+        try:
+            digest.update(inspect.getsource(obj).encode())
+        except (OSError, TypeError):
+            pass
+    return digest.hexdigest()
+
+
+def _tensor_metadata(record: _RuntimeTensor) -> dict[str, Any]:
+    tensor = record.tensor
+    return {
+        "kind": record.kind,
+        "shape": list(tensor.shape),
+        "dtype": str(tensor.dtype),
+        "nbytes": tensor.numel() * tensor.element_size(),
+        "layout": "contiguous",
+    }
+
+
+def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None:
+    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
+    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
+        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
+        digest.update(memoryview(chunk.numpy()))
+
+
+def _runtime_content_digest(records: Sequence[_RuntimeTensor]) -> str:
+    digest = hashlib.sha256()
+    for record in records:
+        digest.update(_canonical_json({"name": record.name, **_tensor_metadata(record)}))
+        _update_hash_with_tensor(digest, record.tensor)
+    return digest.hexdigest()
+
+
+def _file_digest(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@contextlib.contextmanager
+def _exclusive_lock(path: Path, timeout_seconds: float):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    deadline = time.monotonic() + timeout_seconds
+    with path.open("a+b") as handle:
+        while True:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                break
+            except OSError as exc:
+                if exc.errno not in (errno.EACCES, errno.EAGAIN):
+                    raise _RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise _RuntimeCacheError(
+                        "lock_timeout",
+                        f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}",
+                    ) from exc
+                time.sleep(min(0.2, remaining))
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def _fsync_file(path: Path) -> None:
+    with path.open("rb") as handle:
+        os.fsync(handle.fileno())
+
+
+def _fsync_directory(path: Path) -> None:
+    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _split_shards(records: Sequence[_RuntimeTensor], max_shard_bytes: int) -> list[list[_RuntimeTensor]]:
+    shards: list[list[_RuntimeTensor]] = []
+    current: list[_RuntimeTensor] = []
+    current_bytes = 0
+    for record in records:
+        nbytes = record.tensor.numel() * record.tensor.element_size()
+        if current and current_bytes + nbytes > max_shard_bytes:
+            shards.append(current)
+            current = []
+            current_bytes = 0
+        current.append(record)
+        current_bytes += nbytes
+    if current:
+        shards.append(current)
+    return shards
+
+
+def _publish_entry(
+    *,
+    cache_root: Path,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    records: Sequence[_RuntimeTensor],
+    max_shard_bytes: int,
+) -> None:
+    total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
+    free_bytes = shutil.disk_usage(cache_root).free
+    required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
+    if free_bytes < required_bytes:
+        raise _RuntimeCacheError(
+            "insufficient_disk",
+            f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
+        )
+
+    tmp_root = cache_root / ".tmp"
+    tmp_root.mkdir(parents=True, exist_ok=True)
+    temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
+    temp_dir.mkdir()
+    try:
+        shards = _split_shards(records, max_shard_bytes)
+        tensor_manifest: dict[str, dict[str, Any]] = {}
+        file_manifest: dict[str, dict[str, Any]] = {}
+        shard_count = len(shards)
+        for index, shard in enumerate(shards, start=1):
+            filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
+            path = temp_dir / filename
+            save_file(
+                {record.name: record.tensor.detach() for record in shard},
+                str(path),
+                metadata={"format": "pt", "dlo_runtime_cache_schema": str(RUNTIME_CACHE_SCHEMA_VERSION)},
+            )
+            _fsync_file(path)
+            file_manifest[filename] = {
+                "size": path.stat().st_size,
+                "sha256": _file_digest(path),
+            }
+            for record in shard:
+                tensor_manifest[record.name] = {
+                    **_tensor_metadata(record),
+                    "file": filename,
+                    "storage_key": record.name,
+                }
+
+        manifest = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "cache_key": cache_key,
+            "layout_identity": layout_identity,
+            "runtime_content_sha256": content_digest,
+            "total_tensor_bytes": total_nbytes,
+            "created_at_unix_seconds": time.time(),
+            "writer_pid": os.getpid(),
+            "files": file_manifest,
+            "tensors": tensor_manifest,
+        }
+        manifest_path = temp_dir / "manifest.json"
+        with manifest_path.open("w", encoding="utf-8") as handle:
+            json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
+            handle.flush()
+            os.fsync(handle.fileno())
+        _fsync_directory(temp_dir)
+
+        invalid_dir: Path | None = None
+        if entry_dir.exists():
+            invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
+            os.replace(entry_dir, invalid_dir)
+        os.replace(temp_dir, entry_dir)
+        _fsync_directory(entry_dir.parent)
+        if invalid_dir is not None:
+            shutil.rmtree(invalid_dir, ignore_errors=True)
+    except Exception:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        raise
+
+
+def _validate_entry(
+    *,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    expected_records: Sequence[_RuntimeTensor],
+) -> HostWeightPlan:
+    manifest_path = entry_dir / "manifest.json"
+    try:
+        with manifest_path.open(encoding="utf-8") as handle:
+            manifest = json.load(handle)
+    except (OSError, ValueError) as exc:
+        raise _RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
+
+    if manifest.get("schema_version") != RUNTIME_CACHE_SCHEMA_VERSION:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch")
+    if manifest.get("cache_key") != cache_key:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache key mismatch")
+    if manifest.get("layout_identity") != layout_identity:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch")
+    if manifest.get("runtime_content_sha256") != content_digest:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch")
+
+    expected_metadata = {record.name: _tensor_metadata(record) for record in expected_records}
+    tensor_manifest = manifest.get("tensors")
+    file_manifest = manifest.get("files")
+    if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections")
+    if set(tensor_manifest) != set(expected_metadata):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT")
+
+    for filename, metadata in file_manifest.items():
+        if Path(filename).name != filename or not isinstance(metadata, dict):
+            raise _RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}")
+        path = entry_dir / filename
+        try:
+            stat = path.stat()
+        except OSError as exc:
+            raise _RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
+        if stat.st_size != metadata.get("size") or _file_digest(path) != metadata.get("sha256"):
+            raise _RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
+
+    bindings: dict[str, TensorBinding] = {}
+    mapped_records: list[_RuntimeTensor] = []
+    with contextlib.ExitStack() as stack:
+        handles = {
+            filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
+            for filename in file_manifest
+        }
+        expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest}
+        for record in expected_records:
+            stored = tensor_manifest.get(record.name)
+            if not isinstance(stored, dict):
+                raise _RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
+            filename = stored.get("file")
+            storage_key = stored.get("storage_key")
+            if filename not in handles or not isinstance(storage_key, str):
+                raise _RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
+            actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]}
+            if actual_metadata != expected_metadata[record.name]:
+                raise _RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
+            expected_keys_by_file[filename].add(storage_key)
+            try:
+                tensor = handles[filename].get_tensor(storage_key)
+            except Exception as exc:
+                raise _RuntimeCacheError(
+                    "entry_rejected",
+                    f"cannot map {record.name!r} from {filename}: {exc}",
+                ) from exc
+            if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype:
+                raise _RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
+            mapped_records.append(_RuntimeTensor(record.name, tensor, record.kind))
+            bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename))
+
+        for filename, expected_keys in expected_keys_by_file.items():
+            if set(handles[filename].keys()) != expected_keys:
+                raise _RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
+
+        if _runtime_content_digest(mapped_records) != content_digest:
+            raise _RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
+
+    return HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings=bindings,
+        runtime_layout_key=cache_key,
+        post_load_complete=True,
+    )
+
+
+def _remove_stale_temps(cache_root: Path, cache_key: str) -> None:
+    temp_root = cache_root / ".tmp"
+    if not temp_root.is_dir():
+        return
+    for path in temp_root.glob(f"{cache_key}.*"):
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+
+
+def build_runtime_weight_cache_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    loader_type: type,
+    cache_root: str | os.PathLike[str] | None,
+    lock_timeout_seconds: float,
+    max_shard_bytes: int,
+    model_identity: str | None,
+    revision: str | None,
+    runtime_dtype: Any,
+    load_format: str,
+    loader_inputs: Mapping[str, Any],
+    tensor_parallel_size: int,
+    tensor_parallel_rank: int,
+    sequence_parallel_guard: Mapping[str, Any],
+    use_hsdp: bool,
+    enable_expert_parallel: bool,
+    quantization_config: Any,
+    cfg_parallel_size: int,
+    pipeline_parallel_size: int,
+) -> HostWeightPlanResult:
+    """Publish or join one immutable cache entry for final CPU DiT tensors."""
+    if load_format == "dummy":
+        return HostWeightPlanResult(
+            None,
+            "dummy/random weights do not have a reusable runtime identity",
+            "unsupported_load_format",
+        )
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp")
+    if enable_expert_parallel:
+        return HostWeightPlanResult(
+            None,
+            "expert-parallel weight ownership is outside runtime-cache v1",
+            "unsupported_expert_parallel",
+        )
+    if quantization_config is not None:
+        return HostWeightPlanResult(
+            None,
+            "quantized runtime layouts are not supported by the first runtime-cache version",
+            "unsupported_quantization",
+        )
+    if cfg_parallel_size != 1:
+        return HostWeightPlanResult(None, "CFG-parallel component ownership is not proven", "unsupported_cfg_parallel")
+    if pipeline_parallel_size != 1:
+        return HostWeightPlanResult(None, "pipeline-parallel ownership is outside runtime-cache v1", "unsupported_pp")
+    if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0:
+        return HostWeightPlanResult(None, "runtime-cache lock timeout must be positive", "invalid_configuration")
+    if max_shard_bytes <= 0:
+        return HostWeightPlanResult(None, "runtime-cache shard size must be positive", "invalid_configuration")
+
+    try:
+        records = _collect_runtime_tensors(pipeline, dit_modules)
+        content_digest = _runtime_content_digest(records)
+        layout_identity = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "model": model_identity,
+            "revision": revision,
+            "runtime_dtype": str(runtime_dtype),
+            "load_format": load_format,
+            "loader_inputs": _json_normalize(loader_inputs),
+            "loader_implementation_sha256": _implementation_fingerprint(loader_type, pipeline, dit_modules),
+            "components": [name for name, _ in dit_modules],
+            "tensor_parallel_size": tensor_parallel_size,
+            "tensor_parallel_rank": tensor_parallel_rank,
+            "sequence_parallel_guard": _json_normalize(sequence_parallel_guard),
+            "runtime_content_sha256": content_digest,
+        }
+        cache_key = hashlib.sha256(_canonical_json(layout_identity)).hexdigest()
+        root = Path(cache_root or default_runtime_cache_root()).expanduser().resolve()
+        entries_root = root / f"v{RUNTIME_CACHE_SCHEMA_VERSION}"
+        entry_dir = entries_root / cache_key
+        entries_root.mkdir(parents=True, exist_ok=True)
+
+        if entry_dir.is_dir():
+            try:
+                plan = _validate_entry(
+                    entry_dir=entry_dir,
+                    cache_key=cache_key,
+                    layout_identity=layout_identity,
+                    content_digest=content_digest,
+                    expected_records=records,
+                )
+                logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                return HostWeightPlanResult(plan)
+            except _RuntimeCacheError as exc:
+                logger.warning("DLO runtime cache entry %s will be rebuilt: %s", cache_key[:12], exc)
+
+        lock_path = root / ".locks" / f"{cache_key}.lock"
+        with _exclusive_lock(lock_path, lock_timeout_seconds):
+            _remove_stale_temps(root, cache_key)
+            if entry_dir.is_dir():
+                try:
+                    plan = _validate_entry(
+                        entry_dir=entry_dir,
+                        cache_key=cache_key,
+                        layout_identity=layout_identity,
+                        content_digest=content_digest,
+                        expected_records=records,
+                    )
+                    logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                    return HostWeightPlanResult(plan)
+                except _RuntimeCacheError:
+                    pass
+
+            _publish_entry(
+                cache_root=root,
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                records=records,
+                max_shard_bytes=max_shard_bytes,
+            )
+            plan = _validate_entry(
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                expected_records=records,
+            )
+            logger.info(
+                "Published DLO runtime cache entry %s with %d tensors at %s",
+                cache_key[:12],
+                len(records),
+                entry_dir,
+            )
+            return HostWeightPlanResult(plan)
+    except _RuntimeCacheError as exc:
+        return HostWeightPlanResult(None, str(exc), exc.code)
+    except Exception as exc:
+        return HostWeightPlanResult(None, f"runtime-cache operation failed: {exc}", "cache_operation_failed")
+
+
+__all__ = [
+    "DEFAULT_LOCK_TIMEOUT_SECONDS",
+    "DEFAULT_RUNTIME_CACHE_ROOT",
+    "DEFAULT_SHARD_SIZE_BYTES",
+    "RUNTIME_CACHE_SCHEMA_VERSION",
+    "build_runtime_weight_cache_plan",
+    "default_runtime_cache_root",
+]

vllm_omni/diffusion/offloader/__init__.py

@@ -71,8 +71,9 @@ def get_offload_backend(
     Args:
         od_config: OmniDiffusionConfig with offload settings
         device: Target device (auto-detected if None)
-        host_weight_plan: Exact loader-produced backing plan, if ordinary
-            weight materialization was skipped.
+        host_weight_plan: Exact loader-produced backing plan. Checkpoint plans
+            may have skipped ordinary materialization; runtime-cache plans
+            describe already-valid ordinary tensors that can be remapped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -85,10 +86,14 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
+    plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete
     if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
-        raise RuntimeError(
-            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
-        )
+        if plan_requires_backend:
+            raise RuntimeError(
+                "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+            )
+        logger.warning("Ignoring optional runtime-cache plan because distributed layerwise offload is not selected")
+        host_weight_plan = None
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -96,7 +101,7 @@ def get_offload_backend(
 
     # Validate platform (CUDA required for now)
     if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
-        if host_weight_plan is not None:
+        if plan_requires_backend:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -112,7 +117,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if host_weight_plan is not None:
+            if plan_requires_backend:
                 raise RuntimeError(
                     "The loader skipped ordinary weight materialization for DLO, "
                     "but the target device could not be resolved"

vllm_omni/diffusion/offloader/base.py

@@ -31,6 +31,9 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Per-worker budget for registering shared runtime-cache mappings with CUDA.
+    # Zero retains bounded host staging.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -101,6 +104,7 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_runtime_cache_pin_limit_gib = float(getattr(od_config, "dlo_runtime_cache_pin_limit_gib", 0.0))
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -136,6 +140,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_runtime_cache_pin_limit_gib=dlo_runtime_cache_pin_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded CUDA registration for existing read-only host mappings."""
+
+from __future__ import annotations
+
+import mmap
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+
+
+class CudaHostRegistrationError(RuntimeError):
+    """Raised when an existing host range cannot be registered safely."""
+
+
+class CudaHostRegistrationBudgetError(CudaHostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class CudaHostRegistrationCleanupError(CudaHostRegistrationError):
+    """Raised when partial registration could not be rolled back safely."""
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _tensor_regions(
+    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+) -> tuple[_AddressRange, ...]:
+    """Resolve storage spans without merging unrelated file mappings."""
+    regions: list[_AddressRange] = []
+    for mapping_name, tensors in sources_by_mapping.items():
+        ranges: list[tuple[int, int]] = []
+        for tensor in tensors:
+            if tensor.device.type != "cpu":
+                raise CudaHostRegistrationError(
+                    f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
+                )
+            storage = tensor.untyped_storage()
+            ranges.append((storage.data_ptr(), storage.nbytes()))
+        regions.extend(_coalesce_ranges(ranges))
+    return tuple(regions)
+
+
+def _error_message(runtime: Any, error: Any) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for already-existing file-backed CPU tensors."""
+
+    def __init__(
+        self,
+        runtime: Any,
+        regions: tuple[_AddressRange, ...],
+    ) -> None:
+        self._runtime = runtime
+        self._regions = regions
+        self._closed = False
+
+    @classmethod
+    def create(
+        cls,
+        sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+        *,
+        max_bytes: int,
+    ) -> CudaHostRegistration:
+        if max_bytes <= 0:
+            raise CudaHostRegistrationBudgetError("CUDA host-registration budget is disabled")
+        regions = _tensor_regions(sources_by_mapping)
+        total_bytes = sum(region.size for region in regions)
+        if total_bytes > max_bytes:
+            raise CudaHostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not regions:
+            raise CudaHostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise CudaHostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise CudaHostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in regions:
+                error = runtime.cudaHostRegister(
+                    region.start,
+                    region.size,
+                    _CUDA_HOST_REGISTER_READ_ONLY,
+                )
+                if int(error) != 0:
+                    raise CudaHostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+
+            unpinned = [
+                mapping_name
+                for mapping_name, tensors in sources_by_mapping.items()
+                if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
+            ]
+            if unpinned:
+                raise CudaHostRegistrationError(
+                    f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
+                )
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}"
+                        )
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+            if rollback_errors:
+                raise CudaHostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
+                ) from exc
+            if isinstance(exc, CudaHostRegistrationError):
+                raise
+            raise CudaHostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, regions)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> list[str]:
+        """Unregister every range, returning errors after best-effort cleanup."""
+        if self._closed:
+            return []
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        self._closed = not self._regions
+        return errors

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -18,6 +18,7 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import time
 from itertools import chain
 from typing import Any
 
@@ -34,6 +35,11 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
+from .cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -93,6 +99,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -444,6 +451,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -456,18 +483,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -481,6 +497,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             staged[dtype] = destination
         return staged
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
         """Pack this block's mmap views into one bounded host staging slot."""
         previous_copy = self.cpu_staging_events[slot]
@@ -514,17 +556,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-            with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
-                    gw = gpu_weights[dtype]
-                    async_copy = non_blocking and cpu_shard.is_pinned()
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                evt.record(self.copy_stream)
-            if self.rank_local_mmap:
-                # The CPU slot may be overwritten only after this H2D copy has
-                # finished.  The shared event protects reuse by another hook.
-                self.cpu_staging_events[slot] = evt
+            if self.rank_local_mmap and self.registered_mmap:
+                with current_omni_platform.stream(self.copy_stream):
+                    self._copy_mmap_sources_to_device(
+                        self.cpu_sources,
+                        self.metadata,
+                        gpu_weights,
+                        non_blocking=non_blocking,
+                    )
+                    evt.record(self.copy_stream)
+            else:
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -734,6 +786,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -804,6 +857,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -901,8 +962,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
+        self._runtime_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
+        self._cuda_host_registration: CudaHostRegistration | None = None
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -1023,7 +1087,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1103,6 +1167,154 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    def _load_weights_from_runtime_cache(
+        self,
+        pipeline: nn.Module,
+        modules,
+        plan: HostWeightPlan,
+    ) -> None:
+        """Replace final ordinary-loader tensors with validated cache views."""
+        from safetensors import safe_open
+
+        if not plan.post_load_complete:
+            raise RuntimeError("A runtime-cache plan must describe fully post-processed weights")
+        if any(binding.transform is not None for binding in plan.bindings.values()):
+            raise RuntimeError("Runtime-cache bindings cannot carry deferred checkpoint transforms")
+
+        required_names: set[str] = set()
+        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
+            required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters())
+            for name, _ in dit_module.named_buffers():
+                parent_path, _, leaf_name = name.rpartition(".")
+                owner = dit_module.get_submodule(parent_path)
+                if leaf_name not in owner._non_persistent_buffers_set:
+                    required_names.add(f"{dit_name}.{name}")
+        if required_names != set(plan.bindings):
+            missing = sorted(required_names - set(plan.bindings))
+            unexpected = sorted(set(plan.bindings) - required_names)
+            raise RuntimeError(
+                "Runtime-cache plan coverage changed after loader validation: "
+                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
+            )
+
+        file_cache: dict[str, Any] = {}
+        replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
+        mapped_sources: dict[str, list[torch.Tensor]] = {}
+        try:
+            # Validate and map every tensor before mutating the pipeline. A
+            # missing/corrupt shard therefore leaves the ordinary tensors intact.
+            for runtime_name, binding in plan.bindings.items():
+                parent_path, _, leaf_name = runtime_name.rpartition(".")
+                try:
+                    parent = pipeline.get_submodule(parent_path)
+                except AttributeError as exc:
+                    raise RuntimeError(f"Runtime-cache target module {parent_path!r} no longer exists") from exc
+
+                target = parent._parameters.get(leaf_name)
+                if target is None:
+                    target = parent._buffers.get(leaf_name)
+                if target is None:
+                    raise RuntimeError(f"Runtime-cache target tensor {runtime_name!r} no longer exists")
+                if target.device.type != "cpu" or target.is_meta:
+                    raise RuntimeError(
+                        f"Runtime-cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}"
+                    )
+
+                if binding.file_path not in file_cache:
+                    file_cache[binding.file_path] = safe_open(
+                        binding.file_path,
+                        framework="pt",
+                        device="cpu",
+                    )
+                tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
+                if tensor.shape != target.shape or tensor.dtype != target.dtype:
+                    raise RuntimeError(
+                        f"Runtime-cache metadata changed for {runtime_name!r}: "
+                        f"cache={tuple(tensor.shape)}/{tensor.dtype}, "
+                        f"runtime={tuple(target.shape)}/{target.dtype}"
+                    )
+                replacements.append((target, tensor))
+                mapped_sources.setdefault(binding.file_path, []).append(tensor)
+        except Exception:
+            file_cache.clear()
+            raise
+
+        # Preserve Parameter objects and their loader/runtime attributes; only
+        # their storage changes. This also keeps any references held by model
+        # code valid while dropping the private allocator-backed storage.
+        for target, tensor in replacements:
+            set_tensor_storage(target, tensor)
+
+        self._mmap_transforms_by_tensor_id.clear()
+        self._mmap_file_cache = file_cache
+        self._runtime_cache_mapped_sources = mapped_sources
+        logger.info(
+            "Remapped %d final DiT tensors from runtime cache %s",
+            len(replacements),
+            plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown",
+        )
+
+    def _try_register_runtime_cache_mmap(self) -> bool:
+        """Register complete cache mappings when the operator budget permits."""
+        pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib
+        if pin_limit_gib <= 0:
+            return False
+        if self.device.type != "cuda":
+            logger.warning(
+                "DLO runtime-cache registration is CUDA-only; using bounded host staging on %s",
+                self.device,
+            )
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging")
+            return False
+        if not self._runtime_cache_mapped_sources:
+            logger.warning("DLO runtime-cache registration found no mapped sources; using bounded host staging")
+            return False
+
+        max_bytes = int(pin_limit_gib * 1024**3)
+        started = time.perf_counter()
+        try:
+            registration = CudaHostRegistration.create(
+                self._runtime_cache_mapped_sources,
+                max_bytes=max_bytes,
+            )
+        except CudaHostRegistrationCleanupError:
+            # Falling back could unmap a range that CUDA still owns. Abort the
+            # worker so process/context teardown provides the final cleanup.
+            logger.exception("DLO runtime-cache registration rollback failed")
+            raise
+        except CudaHostRegistrationError as exc:
+            logger.warning(
+                "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging",
+                exc,
+            )
+            return False
+
+        self._cuda_host_registration = registration
+        logger.info(
+            "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _release_registered_mmap(self) -> bool:
+        registration = self._cuda_host_registration
+        if registration is None:
+            return True
+        errors = registration.close()
+        if errors:
+            logger.error(
+                "DLO runtime-cache host unregistration failed; retaining mmap handles for retry/process exit: %s",
+                errors[:3],
+            )
+            return False
+        self._cuda_host_registration = None
+        logger.info("Unregistered DLO runtime-cache host mappings")
+        return True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 
@@ -1443,17 +1655,31 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = self.host_weight_plan is not None
         self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
         if self._using_mmap:
-            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
+            if self.host_weight_plan.backing_kind == "checkpoint_mmap":
+                self._load_weights_via_mmap(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+            elif self.host_weight_plan.backing_kind == "runtime_cache":
+                if self.config.dlo_use_allgather:
+                    raise ValueError("Runtime-cache backing is no-AllGather only in the first implementation")
+                self._using_rank_local_mmap = True
+                self._load_weights_from_runtime_cache(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+                # Assignments above dropped the model's private tensor
+                # references. Reclaim those allocations before hook setup so
+                # steady-state PSS reflects the shared mmap backing.
+                self._cleanup_after_loading()
+            else:
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
-            self._load_weights_via_mmap(
-                pipeline,
-                modules,
-                self.host_weight_plan,
-            )
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: checkpoint pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
+                    "transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1625,6 +1851,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
+            if self.enabled and self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+                self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+                self._runtime_cache_mapped_sources.clear()
+                assert self._resident_layer_group is not None
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+                if self._using_registered_mmap:
+                    self._resident_layer_group._cpu_staging_buffers.clear()
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
             return
@@ -1640,9 +1873,18 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
+
+        if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+            self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+            self._runtime_cache_mapped_sources.clear()
+            for hook in all_hooks:
+                hook.registered_mmap = self._using_registered_mmap
+            if self._resident_layer_group is not None:
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1724,6 +1966,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _release_mmap_handles(self) -> None:
         """Release safetensors mmap file handles."""
         self._mmap_transforms_by_tensor_id.clear()
+        self._runtime_cache_mapped_sources.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
@@ -1748,6 +1991,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if self._using_rank_local_mmap:
             current_omni_platform.synchronize()
+        registration_released = self._release_registered_mmap()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -1763,9 +2007,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        self._release_mmap_handles()
+        if registration_released:
+            self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -805,11 +805,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-            if mgr is not None:
-                mgr.shutdown_prefetch()
-        destroy_distributed_env()
+        try:
+            if self.model_runner is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                try:
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/distributed/omni_connectors/transfer_adapter/base.py

@@ -35,8 +35,6 @@ class OmniTransferAdapterBase:
         self.stop_event = threading.Event()
         self._recv_cond = threading.Condition()
         self._save_cond = threading.Condition()
-        self._send_failures: dict[str, str] = {}
-        self._send_failure_lock = threading.Lock()
 
         self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True)
         self.recv_thread.start()
@@ -84,24 +82,6 @@ class OmniTransferAdapterBase:
                 elif not any_success and not self.stop_event.is_set():
                     self._recv_cond.wait(timeout=0.001)
 
-    def record_send_failure(self, request_id: str | None, reason: str) -> None:
-        """Note that a chunk for *request_id* will never be delivered.
-
-        Called from the sender thread; drained by the scheduler thread, hence
-        the lock. ``_waiting_since`` needs none because only the scheduler
-        touches it.
-        """
-        if not request_id:
-            return
-        with self._send_failure_lock:
-            self._send_failures.setdefault(request_id, reason)
-
-    def collect_failed_send_request_ids(self) -> dict[str, str]:
-        """Drain and return the requests whose send gave up."""
-        with self._send_failure_lock:
-            failures, self._send_failures = self._send_failures, {}
-        return failures
-
     def save_loop(self):
         """Loop to send outgoing data."""
         while not self.stop_event.is_set():
@@ -110,18 +90,7 @@ class OmniTransferAdapterBase:
                 try:
                     self._send_single_request(task)
                 except Exception as e:
-                    # R1.2 of #4855. The task was popleft()-ed and is not re-queued,
-                    # so this is a give-up, not a retry: the consumer will wait for a
-                    # chunk that is never coming. Record it so the scheduler can fail
-                    # the request now instead of leaving it to the wait deadline.
-                    # The task dict carries the Request object, not a bare id --
-                    # `task.get("request_id")` is always None (it was that way in the
-                    # original warning too, which logged "None" for every failure).
-                    # Use the scheduler-side id: that is what `self.requests` is keyed
-                    # by, and what `finish_requests` expects.
-                    failed = getattr(task.get("request"), "request_id", None)
-                    logger.error("Send gave up for %s: %s", failed, e)
-                    self.record_send_failure(failed, f"{type(e).__name__}: {e}")
+                    logger.warning(f"Error saving data for {task.get('request_id')}: {e}")
 
             with self._save_cond:
                 if not self._pending_save_reqs and not self.stop_event.is_set():

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import importlib
-import time
 from collections import defaultdict, deque
 from collections.abc import Callable, Mapping
 from typing import Any
@@ -95,12 +94,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self._held_non_active: deque[Any] = deque()
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
         self._pending_streaming_prefills: dict[str, dict] = {}
-        # Monotonic timestamp of when each request last began waiting for a
-        # chunk, refreshed every time one arrives.  Read by
-        # collect_timed_out_request_ids() so a stream that stops advancing
-        # becomes a client-visible error instead of parking forever.  Mirrors
-        # OmniSchedulingCoordinator._waiting_since on the full-payload path.
-        self._waiting_since: dict[str, float] = {}
 
     @staticmethod
     def _is_truthy_scalar(value: Any) -> bool:
@@ -390,24 +383,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             # has been sent successfully. This avoids cleanup->save races.
             if is_payload_finished:
                 self.cleanup(request.request_id, external_req_id)
-        else:
-            # R1.2 of #4855. connector.put returning False is a silent drop: no
-            # exception, no retry, and the caller's `if success:` block simply
-            # does not run. /dev/shm exhaustion in SharedMemoryConnector.put is
-            # one way to get here. Record it so the request fails now rather
-            # than parking in WAITING_FOR_CHUNK until the deadline.
-            logger.error(
-                "Chunk send failed for %s (stage %s -> %s); giving up on this chunk",
-                external_req_id,
-                stage_id,
-                next_stage_id,
-            )
-            # Key on the scheduler-side id. `external_req_id` is the user-facing id
-            # (InputProcessor renames request_id to an internal UUID and keeps the
-            # original in external_req_id, see async_omni_engine.py), while
-            # `self.requests` -- and therefore `finish_requests` -- is keyed by the
-            # internal one.
-            self.record_send_failure(request.request_id, "connector.put reported failure")
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
@@ -463,7 +438,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         self._cancelled_load_reqs.add(request_id)
         self._finished_load_reqs.discard(request_id)
-        self._waiting_since.pop(request_id, None)
 
     @staticmethod
     def _discard_from_chunk_deque(deque_list: deque[Any], request_id: str) -> None:
@@ -605,42 +579,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self._active_streams[request_id] = request
         return True
 
-    def collect_timed_out_request_ids(self, timeout_s: float) -> set[str]:
-        """Return IDs whose chunk wait has exceeded *timeout_s*.
-
-        The async-chunk path had no deadline of any kind: a request parks in
-        ``WAITING_FOR_CHUNK`` and the receiver re-queues it on every failed
-        poll (``transfer_adapter/base.py``) with no attempt counter, so a
-        producer that crashed, gave up after its send retries, or simply never
-        emitted a terminal chunk left the request waiting forever
-        (vllm-project/vllm-omni#3833).  The full-payload path has had a net
-        since ``OmniSchedulingCoordinator.collect_timed_out_request_ids``; this
-        is its async-chunk counterpart, and both are driven by the same
-        ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
-
-        The clock measures *stall* time, not stream lifetime: it starts when a
-        request begins waiting for a chunk and resets each time one arrives, so
-        a long but healthy stream is never failed.
-
-        Clears ``_waiting_since`` for the expired IDs.  The caller marks them
-        ``FINISHED_ERROR`` via the scheduler's ``finish_requests``, which routes
-        back through this adapter's own ``finish_requests`` and releases the
-        rest of the per-request state.
-        """
-        if timeout_s <= 0 or not self._waiting_since:
-            return set()
-        now = time.monotonic()
-        timed_out_ids = {req_id for req_id, started in self._waiting_since.items() if now - started > timeout_s}
-        for req_id in timed_out_ids:
-            self._waiting_since.pop(req_id, None)
-            logger.warning(
-                "[Stage-%s] Request %s timed out waiting for a chunk (stalled > %.0fs)",
-                self.connector.stage_id,
-                req_id,
-                timeout_s,
-            )
-        return timed_out_ids
-
     @property
     def num_running_waiting_for_chunk(self) -> int:
         """Count running requests temporarily removed while awaiting a chunk."""
@@ -686,15 +624,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 # Requests that waiting for chunk
                 self.load_async(request)
                 request.status = RequestStatus.WAITING_FOR_CHUNK
-                self._waiting_since.setdefault(request.request_id, time.monotonic())
             else:
                 if request.request_id in finished_load_reqs:
                     request.status = target_status
                     finished_load_reqs.remove(request.request_id)
                     self.requests_with_ready_chunks.add(request.request_id)
-                    # A chunk landed: restart the clock for the next one, so the
-                    # deadline measures stall time rather than total stream time.
-                    self._waiting_since.pop(request.request_id, None)
                     continue
             queue.remove(request)
             self.requests_origin_status[request.request_id] = target_status
@@ -833,15 +767,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 # Requests that waiting for chunk
                 self.load_async(request)
                 request.status = RequestStatus.WAITING_FOR_CHUNK
-                self._waiting_since.setdefault(request.request_id, time.monotonic())
             else:
                 if request.request_id in finished_load_reqs:
                     request.status = target_status
                     finished_load_reqs.remove(request.request_id)
                     self.requests_with_ready_chunks.add(request.request_id)
-                    # A chunk landed: restart the clock for the next one, so the
-                    # deadline measures stall time rather than total stream time.
-                    self._waiting_since.pop(request.request_id, None)
                     continue
             queue.remove(request)
             self.requests_origin_status[request.request_id] = target_status
@@ -896,6 +826,5 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             self.upstream_exhausted_requests.discard(req_id)
             self._finished_load_reqs.discard(req_id)
             self._cancelled_load_reqs.add(req_id)
-            self._waiting_since.pop(req_id, None)
 
         return []

vllm_omni/engine/arg_utils.py

@@ -543,6 +543,10 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1000,6 +1000,10 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
+            "dlo_enable_runtime_cache": kwargs.get("dlo_enable_runtime_cache", False),
+            "dlo_runtime_cache_dir": kwargs.get("dlo_runtime_cache_dir", None),
+            "dlo_runtime_cache_lock_timeout": kwargs.get("dlo_runtime_cache_lock_timeout", 600.0),
+            "dlo_runtime_cache_pin_limit_gib": kwargs.get("dlo_runtime_cache_pin_limit_gib", 0.0),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (

vllm_omni/engine/orchestrator.py

@@ -15,7 +15,6 @@ import asyncio
 import time as _time
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from http import HTTPStatus
 from typing import TYPE_CHECKING, Any
 
 import janus
@@ -52,7 +51,6 @@ from vllm_omni.engine.messages import (
 from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError
-from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.outputs import OmniRequestOutput
@@ -223,7 +221,7 @@ class _OrchestratorDuplexStagePort:
         async_chunk: bool,
         prewarm_async_chunk_stages: Callable[
             [str, Any, OrchestratorRequestState],
-            Awaitable[bool],
+            Awaitable[None],
         ],
     ) -> None:
         self._stage_pools = stage_pools
@@ -309,23 +307,11 @@ class _OrchestratorDuplexStagePort:
         else:
             replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None)
             if self._async_chunk and context.stage_id == 0:
-                prewarmed = await self._prewarm_async_chunk_stages(
+                await self._prewarm_async_chunk_stages(
                     context.request_id,
                     request,
                     request_state,
                 )
-                if not prewarmed:
-                    # The prewarm already failed the request, aborted it, and
-                    # popped its state. Falling through would write fences and
-                    # a submit timestamp onto that orphaned object and
-                    # re-increment the running counter the cleanup just
-                    # released, then hand the control plane a success result
-                    # for a request that no longer exists. Raise instead:
-                    # ``DuplexControlPlane.handle_append`` turns this into an
-                    # error result for the append.
-                    raise RuntimeError(
-                        f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted"
-                    )
         request_state.duplex_stage_fences[context.stage_id] = context.fence
         request_state.stage_submit_ts[context.stage_id] = _time.time()
         if not request_state.running_counter_registered and self._running_counter is not None:
@@ -2231,54 +2217,18 @@ class Orchestrator:
         request_id: str,
         stage0_request: Any,
         req_state: OrchestratorRequestState,
-    ) -> bool:
-        """Pre-submit downstream stages for async-chunk mode.
-
-        Returns False when the request was failed and cleaned up in here, so a
-        caller still holding ``req_state`` stops instead of recording state on
-        an object the cleanup already popped from ``request_states``.
-        """
+    ) -> None:
+        """Pre-submit downstream stages for async-chunk mode."""
         if req_state.final_stage_id <= 0:
-            return True
+            return
 
         prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None)
         if prompt_token_ids is None:
-            # R1.3 of #4855. Skipping the prewarm leaves every downstream stage
-            # unsubmitted, so the request produces no output and no error -- it
-            # just stops. An embeds-only prompt in async-chunk mode is a config
-            # mistake, and the client should hear about it now rather than wait
-            # for the stage-input deadline to notice nothing ever arrived.
-            #
-            # Non-fatal, with a 4xx: `fatal=True` means "the engine is dead", and
-            # the entrypoints act on it -- `OmniBase._handle_output_message`
-            # raises `OmniEngineDeadError` out of `OmniLLM.generate`'s batch loop,
-            # failing every other request in that batch, and the serving layer
-            # calls `terminate_if_errored`. One client's bad prompt must not do
-            # either. This is the same shape `_handle_stage_error` uses for a
-            # request-scoped client error.
-            logger.error(
-                "[Orchestrator] req=%s: async_chunk prewarm needs stage0 prompt_token_ids "
-                "and none were provided; failing the request",
+            logger.warning(
+                "[Orchestrator] async_chunk prewarm skipped for req=%s: stage0 prompt_token_ids missing",
                 request_id,
             )
-            await self.output_async_queue.put(
-                ErrorMessage(
-                    error=(
-                        "async_chunk requires prompt_token_ids on the stage-0 request; "
-                        "an embeds-only prompt cannot prewarm downstream stages"
-                    ),
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    error_type=DEFAULT_CLIENT_ERROR_TYPE,
-                    request_id=request_id,
-                    stage_id=0,
-                )
-            )
-            await self._cleanup_request_ids(
-                [request_id, *self._cfg_tracker.cleanup_parent(request_id)],
-                abort=True,
-                close_duplex_sessions=True,
-            )
-            return False
+            return
 
         for next_stage_id in range(1, req_state.final_stage_id + 1):
             next_pool = self.stage_pools[next_stage_id]
@@ -2353,7 +2303,7 @@ class Orchestrator:
             # Attribute the failure to the stage that failed, not stage 0; the
             # request is already cleaned up, so stop prewarming.
             if not submitted:
-                return False
+                return
 
             # The guard swallows submit_initial's return value; the pool binding
             # it recorded carries the same replica.
@@ -2379,8 +2329,6 @@ class Orchestrator:
                 tx_ms=_tx_ms,
             )
 
-        return True
-
     def _build_kv_sender_info(
         self,
         sender_stage_ids: list[int],

vllm_omni/entrypoints/cli/serve.py

@@ -722,6 +722,39 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-enable-runtime-cache",
+            action="store_true",
+            help=(
+                "After ordinary loading, publish or join a node-local mmap cache "
+                "for final DiT weights. This first version requires "
+                "--dlo-no-use-allgather."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-dir",
+            type=str,
+            default=None,
+            help=(
+                "Shared local-disk directory for DLO runtime weights. Defaults "
+                "to ~/.cache/vllm-omni/dlo-runtime-weights."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-lock-timeout",
+            type=float,
+            default=600.0,
+            help="Seconds to wait for another process publishing the same runtime layout.",
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-pin-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Maximum GiB of read-only runtime-cache mappings each CUDA worker "
+                "may register for direct H2D. Zero keeps bounded pinned staging."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-resident-layers",
             type=int,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -92,7 +92,19 @@ logger = init_logger(__name__)
 # reference (issue #4644). Matches the offline example/test and upstream demo.
 _COSYVOICE3_PROMPT_DELIMITER = "<|endofprompt|>"
 _COSYVOICE3_PROMPT_PREFIX = f"You are a helpful assistant.{_COSYVOICE3_PROMPT_DELIMITER}"
-
+_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
+    "fish_tts",
+    "qwen3_tts",
+    "voxtral_tts",
+    "cosyvoice3",
+    "voxcpm2",
+    "higgs_audio_v2",
+    "higgs_audio_v3",
+    "indextts2",
+    "indextts2_5",
+    "audex",
+    "audex_tta",
+}
 # Audex contract: zero-codec / invalid generations arrive as empty terminal
 # payloads and must fail the request, never serialize as a successful empty
 # WAV. Covers both the TTS ("audex") and TTA ("audex_tta") pipelines.
@@ -432,7 +444,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model
         # path used for dynamic-token sizing. Without this, every request
         # re-ran snapshot_download + reloaded the tokenizer (~100 ms on the
-        # TTFP critical path) in the CosyVoice3 adapter's sampling override.
+        # TTFP critical path) in _apply_cosyvoice3_dynamic_tokens.
         self._cosyvoice3_tokenizer = None
 
         # Determine TTS model type or None
@@ -490,8 +502,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         # Resolve the per-model serving adapter (RFC #4327), keyed on the
         # detected model-type. Every dedicated TTS model has an adapter; the
-        # adapter owns request validation, prompt/param building, and sampling
-        # overrides; the model-type label remains available for compatibility.
+        # adapter owns request validation + prompt/param building. Sampling
+        # overrides and the model-type label remain in the orchestrator tail
+        # (keyed on ``_tts_model_type``) during this incremental migration.
         self._adapter = None
         if self._tts_stage is not None:
             adapter_cls = resolve_adapter(self._tts_model_type)
@@ -2812,6 +2825,154 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             self._audex_tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, folder))
         return self._audex_tokenizer
 
+    def _inject_audex_tta_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
+        """Complete the Audex TTA contract on a CLONE of the stage-0 params.
+
+        Always attaches the RVQ phase-mask contract (required for token
+        validity) and the CFG pair contract at the official default scale
+        3.0 unless the caller passed ``cfg_scale`` (1.0 disables guidance).
+
+        Returns the cloned params; the input is never mutated. On the
+        no-extra_params online path the input is an element of the engine's
+        SHARED ``default_sampling_params_list`` — writing per-request pair
+        state (``cfg_pair_id``/``cfg_null_prompt``) into it would race
+        concurrent requests and leak stale CFG metadata into later ones.
+        """
+        import copy
+
+        from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt
+        from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
+        if not isinstance(cond_prompt, str) or not cond_prompt:
+            raise ValueError("Audex TTA requires the adapter-built caption prompt")
+
+        stage0_params = copy.deepcopy(stage0_params)
+        tokenizer = self._get_audex_tokenizer()
+        if self._audex_tta_rvq is None:
+            phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+            self._audex_tta_rvq = {
+                "phase_token_ids": phase_token_ids,
+                "start_tid": start_tid,
+                "end_tid": end_tid,
+                # Official generation cap (decode truncates at 500 frames).
+                "codec_cap": 4000,
+                # The TTA prompt ends with <audiogen_start>.
+                "start_in_prompt": True,
+            }
+
+        extra_args = getattr(stage0_params, "extra_args", None)
+        if extra_args is None:
+            extra_args = {}
+            stage0_params.extra_args = extra_args
+        extra_args["tta_rvq"] = self._audex_tta_rvq
+
+        cfg_scale = extra_args.get("cfg_scale")
+        cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale)
+        if cfg_scale <= 1.0:
+            extra_args.pop("cfg_scale", None)
+            return stage0_params
+        extra_args.update(
+            {
+                "cfg_scale": cfg_scale,
+                "cfg_role": "cond",
+                "cfg_pair_id": request_id,
+                "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer),
+            }
+        )
+        return stage0_params
+
+    def _inject_audex_cfg_pair_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
+        """Complete the Audex CFG contract on a CLONE of the stage-0 params.
+
+        The adapter validated ``cfg_scale``; here (where the final request id
+        exists) the cond role, pair id, and the length-matched null prompt are
+        attached. ``cfg_scale`` absent/1.0 returns a clone value-identical to
+        the non-CFG flow (the key is dropped rather than forwarded).
+
+        Returns the cloned params; the input is never mutated (it may be an
+        element of the engine's shared ``default_sampling_params_list`` —
+        see ``_inject_audex_tta_args``).
+        """
+        import copy
+
+        stage0_params = copy.deepcopy(stage0_params)
+        extra_args = getattr(stage0_params, "extra_args", None)
+        if not extra_args or extra_args.get("cfg_scale") is None:
+            if extra_args:
+                extra_args.pop("cfg_scale", None)
+            return stage0_params
+        cfg_scale = float(extra_args["cfg_scale"])
+        if cfg_scale <= 1.0:
+            extra_args.pop("cfg_scale", None)
+            return stage0_params
+
+        from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
+        if not isinstance(cond_prompt, str) or not cond_prompt:
+            raise ValueError("Audex CFG requires the adapter-built text prompt")
+        null_prompt = build_null_prompt(cond_prompt, self._get_audex_tokenizer())
+        extra_args.update(
+            {
+                "cfg_scale": cfg_scale,
+                "cfg_role": "cond",
+                "cfg_pair_id": request_id,
+                "cfg_null_prompt": null_prompt,
+            }
+        )
+        # Guided decoding sharpens the distribution; the unguided default
+        # temperature (0.1) adds excess sampling noise under CFG. Measured
+        # on the en-24 gate corpus at cfg 1.5: temp 0.1 -> CER 7.31%,
+        # temp 0.05 -> CER 6.87% (unguided baseline 7.24%).
+        stage0_params.temperature = 0.05
+        return stage0_params
+
+    def _apply_cosyvoice3_dynamic_tokens(
+        self,
+        sampling_params_list: list,
+        request: OpenAICreateSpeechRequest,
+    ) -> list:
+        """Set min/max tokens from tokenized text length (ratios target tokens, not chars)."""
+        import copy
+
+        from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
+        from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
+
+        sampling_params_list = copy.deepcopy(sampling_params_list)
+        hf_cfg = self.model_config.hf_config
+        # Build the Qwen tokenizer once per process (resolving the model dir via
+        # snapshot_download at most once) and reuse it across requests.
+        tokenizer = self._cosyvoice3_tokenizer
+        if tokenizer is None:
+            model_path = self.engine_client.model_config.model
+            if not os.path.isdir(model_path):
+                from huggingface_hub import snapshot_download
+
+                model_path = snapshot_download(model_path)
+            tokenizer = get_qwen_tokenizer(
+                token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path),
+                skip_special_tokens=hf_cfg.skip_special_tokens,
+                version=hf_cfg.version,
+            )
+            self._cosyvoice3_tokenizer = tokenizer
+        _, text_token_len = extract_text_token(
+            request.input,
+            tokenizer,
+            hf_cfg.allowed_special,
+        )
+        min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
+        max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
+        sampling_params_list[0].min_tokens = max(1, int(text_token_len * min_ratio))
+        sampling_params_list[0].max_tokens = min(2048, int(text_token_len * max_ratio))
+        logger.info(
+            "CosyVoice3 dynamic tokens: text_tokens=%d, min_tokens=%d, max_tokens=%d",
+            text_token_len,
+            sampling_params_list[0].min_tokens,
+            sampling_params_list[0].max_tokens,
+        )
+        return sampling_params_list
+
     # ---- GLM-TTS helpers ----
 
     async def _build_glm_tts_prompt(
@@ -3011,8 +3172,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         # Build prompt + tts_params via the per-model adapter (RFC #4327). Every
         # dedicated TTS model resolves to an adapter that owns its validation,
-        # uploaded-speaker handling, prompt/param building, and sampling
-        # overrides. The model-type label remains available for compatibility.
+        # uploaded-speaker handling, and prompt/param building. Sampling
+        # overrides and the model-type label remain in the orchestrator tail
+        # below (keyed on ``_tts_model_type``) during this incremental migration.
         # Non-TTS deployments (no adapter) fall through to the rejection below.
         # Capture inline-ref-audio status BEFORE validate(): several adapters
         # apply uploaded speakers inside validate(), which sets request.ref_audio
@@ -3063,6 +3225,63 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type,
         )
 
+        # CosyVoice3: set dynamic min/max tokens based on text length.
+        # The official model requires min_token_text_ratio to prevent early
+        # EOS and max_token_text_ratio to cap generation length.
+        if self._tts_model_type == "cosyvoice3" and sampling_params_list:
+            sampling_params_list = self._apply_cosyvoice3_dynamic_tokens(sampling_params_list, request)
+
+        # GLM-TTS: set dynamic min/max tokens based on text length.
+        if self._tts_model_type == "glm_tts" and sampling_params_list:
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None
+            text_len_value = None
+            if isinstance(glm_metadata, dict):
+                text_len_value = glm_metadata.get("glm_tts_text_token_len")
+                if isinstance(text_len_value, list) and text_len_value:
+                    text_len_value = text_len_value[0]
+            text_token_len = (
+                int(text_len_value)
+                if text_len_value is not None
+                else self._estimate_glm_tts_text_token_len(request.input)
+            )
+            hf_cfg = self.model_config.hf_config
+            min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
+            max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
+            stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None)
+            stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None)
+            cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None]
+            hard_cap = min(cap_candidates) if cap_candidates else None
+
+            min_tokens = max(1, int(text_token_len * min_ratio))
+            if stage_min_tokens is not None:
+                min_tokens = max(min_tokens, int(stage_min_tokens))
+            if hard_cap is not None:
+                min_tokens = min(min_tokens, hard_cap)
+
+            max_tokens = max(min_tokens, int(text_token_len * max_ratio))
+            if hard_cap is not None:
+                max_tokens = min(max_tokens, hard_cap)
+            sampling_params_list[0].min_tokens = min_tokens
+            sampling_params_list[0].max_tokens = max_tokens
+            seed = getattr(request, "seed", None)
+            if seed is not None:
+                sampling_params_list[0].seed = seed
+            logger.info(
+                "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, "
+                "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d",
+                text_token_len,
+                min_ratio,
+                max_ratio,
+                stage_min_tokens,
+                stage_max_tokens,
+                request.max_new_tokens,
+                min_tokens,
+                max_tokens,
+            )
+
         # Apply model-specific extra parameters
         if request.extra_params is not None and sampling_params_list:
             if not isinstance(request.extra_params, dict):
@@ -3081,10 +3300,52 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             sampling_params_list[0].extra_args.update(request.extra_params)
             logger.info("Applied extra_params: %s", request.extra_params)
 
-        # Apply adapter-owned sampling overrides, including request-level token
-        # limits and model-specific dynamic token or stop-token configuration.
-        if sampling_params_list and (adapter := self._get_tts_adapter()) is not None:
-            sampling_params_list = adapter.apply_sampling_overrides(sampling_params_list, request, prompt, request_id)
+        # Audex CFG: turn an adapter-validated cfg_scale into the engine-side
+        # pair contract. This must run here (not in the adapter) because the
+        # pair id is the final request_id, which the adapter never sees.
+        # The injectors return a CLONE: without extra_params above,
+        # sampling_params_list still holds the engine's shared default params,
+        # which must never carry per-request CFG state.
+        if self._tts_model_type == "audex" and sampling_params_list:
+            sampling_params_list[0] = self._inject_audex_cfg_pair_args(request_id, prompt, sampling_params_list[0])
+        elif self._tts_model_type == "audex_tta" and sampling_params_list:
+            sampling_params_list[0] = self._inject_audex_tta_args(request_id, prompt, sampling_params_list[0])
+
+        # Some TTS model defaults come from deploy YAML. Their AR
+        # generation length is controlled by SamplingParams.max_tokens, so only
+        # override it when the caller explicitly requests max_new_tokens.
+        if (
+            self._tts_model_type in _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES
+            and request.max_new_tokens is not None
+            and sampling_params_list
+        ):
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            sampling_params_list[0].max_tokens = request.max_new_tokens
+            if self._tts_model_type == "cosyvoice3":
+                sampling_params_list[0].min_tokens = min(
+                    getattr(sampling_params_list[0], "min_tokens", 0),
+                    request.max_new_tokens,
+                )
+        elif self._tts_model_type == "ming_tts" and sampling_params_list:
+            import copy
+
+            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+                MOE_TEXT_EOS_TOKEN_ID,
+                TEXT_EOS_TOKEN_ID,
+            )
+
+            hf_config = self.engine_client.model_config.hf_config
+            is_moe = getattr(hf_config, "model_type", "") == "bailingmm"
+            stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            sampling_params_list[0].stop_token_ids = [int(stop_token_id)]
+            if request.max_new_tokens is not None:
+                # Ming emits TEXT_EOS after the latent decode budget is exhausted, so
+                # Stage-0 needs one extra token beyond ming_max_decode_steps.
+                sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
 
         if request.seed is not None and sampling_params_list:
             import copy

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -9,6 +9,8 @@ key alone cannot distinguish.
 That metadata is the single source of truth. :func:`detect_tts_model_type` and
 :func:`all_tts_stage_keys` derive the whole stage -> model-type mapping from it,
 so a new TTS model is one adapter file and needs no edit to ``serving_speech.py``.
+Models still awaiting an adapter are listed explicitly in
+:data:`LEGACY_TTS_DETECTORS`.
 """
 
 from vllm.logger import init_logger
@@ -16,6 +18,7 @@ from vllm.logger import init_logger
 from vllm_omni.entrypoints.openai.tts_adapters.base import (
     ARTTSAdapter,
     DiffusionTTSAdapter,
+    LegacyDetector,
     OutputPolicy,
     PreparedRequest,
     SpeechServingContext,
@@ -26,6 +29,10 @@ logger = init_logger(__name__)
 
 TTS_ADAPTER_REGISTRY: dict[str, type[TTSModelAdapter]] = {}
 
+#: Model types the serving layer detects that have no adapter yet. See
+#: :class:`LegacyDetector`. This list must only ever shrink.
+LEGACY_TTS_DETECTORS: tuple[LegacyDetector, ...] = ()
+
 
 def register_tts_adapter(cls: type[TTSModelAdapter]) -> type[TTSModelAdapter]:
     """Class decorator: index ``cls`` under its ``name`` (model-type)."""
@@ -50,13 +57,16 @@ def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None:
     return TTS_ADAPTER_REGISTRY.get(model_type)
 
 
-def iter_tts_detectors() -> list[type[TTSModelAdapter]]:
+def iter_tts_detectors() -> list[type[TTSModelAdapter] | LegacyDetector]:
     """Every detector, in resolution order.
 
     Sorted by ``detect_priority`` then ``name``, so the order is total and does
     not depend on adapter import order.
     """
-    detectors = list(TTS_ADAPTER_REGISTRY.values())
+    detectors: list[type[TTSModelAdapter] | LegacyDetector] = [
+        *TTS_ADAPTER_REGISTRY.values(),
+        *LEGACY_TTS_DETECTORS,
+    ]
     detectors.sort(key=lambda d: (d.detect_priority, d.name))
     return detectors
 
@@ -122,10 +132,12 @@ from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
 __all__ = [
     "ARTTSAdapter",
     "DiffusionTTSAdapter",
+    "LegacyDetector",
     "OutputPolicy",
     "PreparedRequest",
     "SpeechServingContext",
     "TTSModelAdapter",
+    "LEGACY_TTS_DETECTORS",
     "TTS_ADAPTER_REGISTRY",
     "all_tts_model_types",
     "all_tts_stage_keys",

vllm_omni/entrypoints/openai/tts_adapters/audex.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Audex (Nemotron-Labs-Audex-2B) TTS serving adapter."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
 
 if TYPE_CHECKING:
@@ -17,8 +17,8 @@ if TYPE_CHECKING:
 AUDEX_CFG_SCALE_MIN = 1.0
 AUDEX_CFG_SCALE_MAX = 10.0
 
-# Internal CFG pair-plumbing keys; injected by this adapter once it receives
-# the final request id, never accepted from callers.
+# Internal CFG pair-plumbing keys; injected by the serving layer once the
+# request id exists, never accepted from callers.
 _AUDEX_INTERNAL_CFG_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt")
 
 
@@ -74,65 +74,3 @@ class AudexAdapter(ARTTSAdapter):
     ) -> PreparedRequest:
         prompt = {"prompt": build_cond_prompt(request.input)}
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        if request_id is None:
-            raise ValueError("Audex CFG requires a request id")
-        stage0_params = sampling_params_list[0]
-        sampling_params_list[0] = self._inject_audex_cfg_pair_args(
-            request_id=request_id, prompt=prompt, stage0_params=stage0_params
-        )
-        return apply_max_new_tokens(sampling_params_list, request)
-
-    def _inject_audex_cfg_pair_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
-        """Complete the Audex CFG contract on a CLONE of the stage-0 params.
-
-        The adapter validated ``cfg_scale``; here (where the final request id
-        exists) the cond role, pair id, and the length-matched null prompt are
-        attached. ``cfg_scale`` absent/1.0 returns a clone value-identical to
-        the non-CFG flow (the key is dropped rather than forwarded).
-
-        Returns the cloned params; the input is never mutated (it may be an
-        element of the engine's shared ``default_sampling_params_list`` —
-        see ``_inject_audex_tta_args``).
-        """
-        server = self.ctx.server
-        import copy
-
-        stage0_params = copy.deepcopy(stage0_params)
-        extra_args = getattr(stage0_params, "extra_args", None)
-        if not extra_args or extra_args.get("cfg_scale") is None:
-            if extra_args:
-                extra_args.pop("cfg_scale", None)
-            return stage0_params
-        cfg_scale = float(extra_args["cfg_scale"])
-        if cfg_scale <= 1.0:
-            extra_args.pop("cfg_scale", None)
-            return stage0_params
-
-        from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
-
-        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
-        if not isinstance(cond_prompt, str) or not cond_prompt:
-            raise ValueError("Audex CFG requires the adapter-built text prompt")
-        null_prompt = build_null_prompt(cond_prompt, server._get_audex_tokenizer())
-        extra_args.update(
-            {
-                "cfg_scale": cfg_scale,
-                "cfg_role": "cond",
-                "cfg_pair_id": request_id,
-                "cfg_null_prompt": null_prompt,
-            }
-        )
-        # Guided decoding sharpens the distribution; the unguided default
-        # temperature (0.1) adds excess sampling noise under CFG. Measured
-        # on the en-24 gate corpus at cfg 1.5: temp 0.1 -> CER 7.31%,
-        # temp 0.05 -> CER 6.87% (unguided baseline 7.24%).
-        stage0_params.temperature = 0.05
-        return stage0_params

vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py

@@ -2,16 +2,16 @@
 """Audex (Nemotron-Labs-Audex-2B) text-to-audio serving adapter.
 
 Caption → general audio through the ``audex_tta`` pipeline.
-The adapter builds the TTA prompt and, after receiving the final request id,
-completes the sampling parameters with the RVQ phase-mask contract
+The adapter builds the TTA prompt; the serving layer completes the request
+after the final request id exists: the RVQ phase-mask contract
 (``extra_args["tta_rvq"]``) and the CFG pair contract (official default
 scale 3.0 — guidance is effectively mandatory for TTA quality).
 """
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt
 
 if TYPE_CHECKING:
@@ -21,7 +21,8 @@ AUDEX_TTA_DEFAULT_CFG_SCALE = 3.0
 AUDEX_TTA_CFG_SCALE_MIN = 1.0
 AUDEX_TTA_CFG_SCALE_MAX = 10.0
 
-# Internal plumbing keys; injected by this adapter, never accepted from callers.
+# Internal plumbing keys; injected by the serving layer, never accepted
+# from callers.
 _INTERNAL_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt", "tta_rvq")
 
 
@@ -66,76 +67,3 @@ class AudexTTAAdapter(ARTTSAdapter):
     ) -> PreparedRequest:
         prompt = {"prompt": build_tta_cond_prompt(request.input)}
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex_tta")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        if request_id is None:
-            raise ValueError("Audex TTA requires a request id")
-        stage0_params = sampling_params_list[0]
-        sampling_params_list[0] = self._inject_audex_tta_args(
-            request_id=request_id, prompt=prompt, stage0_params=stage0_params
-        )
-        return apply_max_new_tokens(sampling_params_list, request)
-
-    def _inject_audex_tta_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any:
-        """Complete the Audex TTA contract on a CLONE of the stage-0 params.
-
-        Always attaches the RVQ phase-mask contract (required for token
-        validity) and the CFG pair contract at the official default scale
-        3.0 unless the caller passed ``cfg_scale`` (1.0 disables guidance).
-
-        Returns the cloned params; the input is never mutated. On the
-        no-extra_params online path the input is an element of the engine's
-        SHARED ``default_sampling_params_list`` — writing per-request pair
-        state (``cfg_pair_id``/``cfg_null_prompt``) into it would race
-        concurrent requests and leak stale CFG metadata into later ones.
-        """
-        server = self.ctx.server
-        import copy
-
-        from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt
-        from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
-
-        cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt
-        if not isinstance(cond_prompt, str) or not cond_prompt:
-            raise ValueError("Audex TTA requires the adapter-built caption prompt")
-
-        stage0_params = copy.deepcopy(stage0_params)
-        tokenizer = server._get_audex_tokenizer()
-        if server._audex_tta_rvq is None:
-            phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
-            server._audex_tta_rvq = {
-                "phase_token_ids": phase_token_ids,
-                "start_tid": start_tid,
-                "end_tid": end_tid,
-                # Official generation cap (decode truncates at 500 frames).
-                "codec_cap": 4000,
-                # The TTA prompt ends with <audiogen_start>.
-                "start_in_prompt": True,
-            }
-
-        extra_args = getattr(stage0_params, "extra_args", None)
-        if extra_args is None:
-            extra_args = {}
-            stage0_params.extra_args = extra_args
-        extra_args["tta_rvq"] = server._audex_tta_rvq
-
-        cfg_scale = extra_args.get("cfg_scale")
-        cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale)
-        if cfg_scale <= 1.0:
-            extra_args.pop("cfg_scale", None)
-            return stage0_params
-        extra_args.update(
-            {
-                "cfg_scale": cfg_scale,
-                "cfg_role": "cond",
-                "cfg_pair_id": request_id,
-                "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer),
-            }
-        )
-        return stage0_params

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -7,7 +7,9 @@ model's request normalization, validation, prompt/param building, sampling
 overrides, and output policy, so adding a model means writing one adapter file
 instead of editing the shared serving module in ~10 scattered places.
 
-See the RFC for the full design (issue #4327).
+See the RFC for the full design (issue #4327). This is the foundation landed in
+the first migration PR; Qwen3-TTS is the first model routed through it while the
+remaining models stay on the legacy path until individually migrated.
 """
 
 from abc import ABC, abstractmethod
@@ -37,21 +39,6 @@ def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: di
     return _conditioning_cache_salt_fn(request, tts_params)
 
 
-def apply_max_new_tokens(
-    sampling_params_list: list,
-    request: "OpenAICreateSpeechRequest",
-) -> list:
-    """Apply a request-level ``max_new_tokens`` limit."""
-    if request.max_new_tokens is None:
-        return sampling_params_list
-
-    import copy
-
-    sampling_params_list = copy.deepcopy(sampling_params_list)
-    sampling_params_list[0].max_tokens = request.max_new_tokens
-    return sampling_params_list
-
-
 @dataclass
 class OutputPolicy:
     """How the orchestrator aggregates engine output for a model.
@@ -197,8 +184,6 @@ class TTSModelAdapter(ABC):
         self,
         sampling_params_list: list,
         request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
     ) -> list:
         """Apply model-specific sampling mutations.
 
@@ -243,10 +228,37 @@ class DiffusionTTSAdapter(TTSModelAdapter):
         return frozenset(params) if params is not None else frozenset()
 
 
+@dataclass(frozen=True)
+class LegacyDetector:
+    """A model type the serving layer detects but has no adapter for yet.
+
+    Detection has to keep naming these models so ``serving_speech.py`` can route
+    them down its legacy path, but they own none of the adapter contract. Each
+    entry is a migration debt: delete it when the model gets an adapter. The
+    pre-commit gate (``tools/pre_commit/check_tts_adapter.py``) fails if the list
+    grows.
+
+    Exposes the subset of the :class:`TTSModelAdapter` class surface that
+    detection reads, so both kinds live in one sorted sequence.
+    """
+
+    name: str
+    stage_keys: frozenset[str] = frozenset()
+    model_archs: frozenset[str] = frozenset()
+    arch_identifies_entry_stage: bool = False
+    detect_priority: int = 100
+
+    def matches(self, model_stage: str | None, model_arch: str | None) -> bool:
+        if model_arch is not None and model_arch in self.model_archs:
+            return True
+        return model_stage is not None and model_stage in self.stage_keys
+
+
 # Re-exported here to avoid import cycles at call sites.
 __all__ = [
     "ARTTSAdapter",
     "DiffusionTTSAdapter",
+    "LegacyDetector",
     "OutputPolicy",
     "PreparedRequest",
     "SpeechServingContext",

vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py

@@ -1,10 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """CosyVoice3 TTS serving adapter."""
 
-import os
-from typing import TYPE_CHECKING, Any
-
-from vllm.logger import init_logger
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
 from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
@@ -12,8 +9,6 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, Prepare
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
 
-logger = init_logger(__name__)
-
 
 @register_tts_adapter
 class CosyVoice3Adapter(ARTTSAdapter):
@@ -52,63 +47,6 @@ class CosyVoice3Adapter(ARTTSAdapter):
         self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
     ) -> PreparedRequest:
         prompt = await self.ctx.server._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
+        # NOTE: CosyVoice3 dynamic-token sampling stays in the orchestrator tail
+        # (keyed on _tts_model_type) during this incremental migration.
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="cosyvoice3")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        """Set dynamic min/max tokens based on tokenized text length.
-
-        The official model uses ``min_token_text_ratio`` to prevent early EOS
-        and ``max_token_text_ratio`` to cap the generated token count.
-        """
-        import copy
-
-        from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
-        from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-
-        server = self.ctx.server
-        sampling_params_list = copy.deepcopy(sampling_params_list)
-        hf_cfg = server.model_config.hf_config
-        # Build the Qwen tokenizer once per process (resolving the model dir via
-        # snapshot_download at most once) and reuse it across requests.
-        tokenizer = server._cosyvoice3_tokenizer
-        if tokenizer is None:
-            model_path = server.engine_client.model_config.model
-            if not os.path.isdir(model_path):
-                from huggingface_hub import snapshot_download
-
-                model_path = snapshot_download(model_path)
-            tokenizer = get_qwen_tokenizer(
-                token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path),
-                skip_special_tokens=hf_cfg.skip_special_tokens,
-                version=hf_cfg.version,
-            )
-            server._cosyvoice3_tokenizer = tokenizer
-        _, text_token_len = extract_text_token(
-            request.input,
-            tokenizer,
-            hf_cfg.allowed_special,
-        )
-        min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
-        max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
-        sampling_params_list[0].min_tokens = max(1, int(text_token_len * min_ratio))
-        sampling_params_list[0].max_tokens = min(2048, int(text_token_len * max_ratio))
-
-        if request.max_new_tokens is not None:
-            sampling_params_list[0].max_tokens = request.max_new_tokens
-            sampling_params_list[0].min_tokens = min(
-                getattr(sampling_params_list[0], "min_tokens", 0),
-                request.max_new_tokens,
-            )
-        logger.info(
-            "CosyVoice3 dynamic tokens: text_tokens=%d, min_tokens=%d, max_tokens=%d",
-            text_token_len,
-            sampling_params_list[0].min_tokens,
-            sampling_params_list[0].max_tokens,
-        )
-        return sampling_params_list

vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py

@@ -1,15 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Fish Speech serving adapter (retires the legacy ``_is_fish_speech`` flag)."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import (
-    ARTTSAdapter,
-    PreparedRequest,
-    apply_max_new_tokens,
-    conditioning_cache_salt,
-)
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -58,12 +53,3 @@ class FishSpeechAdapter(ARTTSAdapter):
         tts_params = {}
         prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="fish_speech")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py

@@ -1,9 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """GLM-TTS serving adapter."""
 
-from typing import TYPE_CHECKING, Any
-
-from vllm.logger import init_logger
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
 from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
@@ -11,8 +9,6 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, Prepare
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
 
-logger = init_logger(__name__)
-
 
 @register_tts_adapter
 class GlmTTSAdapter(ARTTSAdapter):
@@ -47,63 +43,6 @@ class GlmTTSAdapter(ARTTSAdapter):
         self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
     ) -> PreparedRequest:
         prompt = await self.ctx.server._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
+        # GLM-TTS dynamic-token sampling stays in the orchestrator tail
+        # (keyed on _tts_model_type) during this incremental migration.
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="glm_tts")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        # GLM-TTS: set dynamic min/max tokens based on text length.
-        import copy
-
-        server = self.ctx.server
-        sampling_params_list = copy.deepcopy(sampling_params_list)
-        glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None
-        text_len_value = None
-        if isinstance(glm_metadata, dict):
-            text_len_value = glm_metadata.get("glm_tts_text_token_len")
-            if isinstance(text_len_value, list) and text_len_value:
-                text_len_value = text_len_value[0]
-        text_token_len = (
-            int(text_len_value)
-            if text_len_value is not None
-            else server._estimate_glm_tts_text_token_len(request.input)
-        )
-        hf_cfg = server.model_config.hf_config
-        min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
-        max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
-        stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None)
-        stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None)
-        cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None]
-        hard_cap = min(cap_candidates) if cap_candidates else None
-
-        min_tokens = max(1, int(text_token_len * min_ratio))
-        if stage_min_tokens is not None:
-            min_tokens = max(min_tokens, int(stage_min_tokens))
-        if hard_cap is not None:
-            min_tokens = min(min_tokens, hard_cap)
-
-        max_tokens = max(min_tokens, int(text_token_len * max_ratio))
-        if hard_cap is not None:
-            max_tokens = min(max_tokens, hard_cap)
-        sampling_params_list[0].min_tokens = min_tokens
-        sampling_params_list[0].max_tokens = max_tokens
-        seed = getattr(request, "seed", None)
-        if seed is not None:
-            sampling_params_list[0].seed = seed
-        logger.info(
-            "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, "
-            "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d",
-            text_token_len,
-            min_ratio,
-            max_ratio,
-            stage_min_tokens,
-            stage_max_tokens,
-            request.max_new_tokens,
-            min_tokens,
-            max_tokens,
-        )
-        return sampling_params_list

vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Higgs-Audio v2 serving adapter."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -99,12 +99,3 @@ class HiggsAudioV2Adapter(ARTTSAdapter):
                 additional["voice_name"] = voice_lower
                 additional["voice_created_at"] = server._voice_created_at(voice_lower)
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v2")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Higgs-Audio v3 serving adapter."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -36,12 +36,3 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
     ) -> PreparedRequest:
         prompt = await self.ctx.server._build_higgs_audio_v3_params(request)
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v3")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/tts_adapters/indextts2.py

@@ -14,7 +14,7 @@ from vllm.inputs import tokens_input
 from vllm.utils import random_uuid
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
     INDEXTTS25_MAX_DURATION_FACTOR,
     INDEXTTS25_MIN_DURATION_FACTOR,
@@ -241,15 +241,6 @@ class IndexTTS2Adapter(ARTTSAdapter):
         prompt["cache_salt"] = indextts2_conditioning_cache_salt(request, tts_params)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
 
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: OpenAICreateSpeechRequest,
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)
-
     async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
         server = self.ctx.server
         params: dict[str, Any] = {"text": [request.input]}

vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Ming-TTS (dense) serving adapter."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
 
@@ -73,36 +73,9 @@ class MingTTSAdapter(ARTTSAdapter):
             voice_created_at=uploaded_audio_created_at,
         )
         tts_params = prompt.get("additional_information", {})
+        # Ming stop-token / max_tokens sampling stays in the orchestrator tail.
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="ming_tts")
 
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        import copy
-
-        server = self.ctx.server
-
-        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
-            MOE_TEXT_EOS_TOKEN_ID,
-            TEXT_EOS_TOKEN_ID,
-        )
-
-        hf_config = server.engine_client.model_config.hf_config
-        is_moe = getattr(hf_config, "model_type", "") == "bailingmm"
-        stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID
-
-        sampling_params_list = copy.deepcopy(sampling_params_list)
-        sampling_params_list[0].stop_token_ids = [int(stop_token_id)]
-        if request.max_new_tokens is not None:
-            # Ming emits TEXT_EOS after the latent decode budget is exhausted, so
-            # Stage-0 needs one extra token beyond ming_max_decode_steps.
-            sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
-        return sampling_params_list
-
     def _validate_ming_tts_single_speaker_request(self, request: "OpenAICreateSpeechRequest") -> str | None:
         server = self.ctx.server
         if request.ref_audio is not None:

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -2,12 +2,12 @@
 """Qwen3-TTS serving adapter."""
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -163,12 +163,3 @@ class Qwen3TTSAdapter(ARTTSAdapter):
             model_type=tts_params.get("task_type", ["unknown"])[0],
             warmup_artifact_key=warmup_key,
         )
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """VoxCPM2 serving adapter (AR base-LM + diffusion side-computation)."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -70,12 +70,3 @@ class VoxCPM2Adapter(ARTTSAdapter):
                 additional["voice_name"] = voice_lower
                 additional["voice_created_at"] = server._voice_created_at(voice_lower)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="voxcpm2")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/tts_adapters/voxtral.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Voxtral TTS serving adapter."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -48,12 +48,3 @@ class VoxtralTTSAdapter(ARTTSAdapter):
     ) -> PreparedRequest:
         prompt = await self.ctx.server._build_voxtral_prompt_async(request)
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts")
-
-    def apply_sampling_overrides(
-        self,
-        sampling_params_list: list,
-        request: "OpenAICreateSpeechRequest",
-        prompt: dict[str, Any] | None = None,
-        request_id: str | None = None,
-    ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -509,9 +509,9 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
     return frames_u8
 
 
-def encode_video_bytes(
+def _encode_video_bytes(
     video: Any,
-    fps: int | float,
+    fps: int,
     audio: Any | None = None,
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
@@ -530,10 +530,6 @@ def encode_video_bytes(
     )
 
 
-# Backward-compatible private alias for existing internal callers.
-_encode_video_bytes = encode_video_bytes
-
-
 class FragmentedMP4VideoEncoder:
     """Normalize video chunks and append them to one fragmented MP4 stream."""
 

vllm_omni/model_extras/__init__.py

@@ -6,10 +6,8 @@ from vllm_omni.model_extras.registry import (
     build_image_to_video_prompt,
     build_text_to_image_prompt,
     build_x_to_text_prompt,
-    build_x_to_video_audio_prompt,
     get_extra_body_params,
     get_extra_output_params,
-    get_input_audio_sample_rate,
     get_model_class_name,
     get_output_tensor_range,
     get_transformer_config_subfolder,
@@ -23,10 +21,8 @@ __all__ = [
     "build_image_to_video_prompt",
     "build_text_to_image_prompt",
     "build_x_to_text_prompt",
-    "build_x_to_video_audio_prompt",
     "get_extra_body_params",
     "get_extra_output_params",
-    "get_input_audio_sample_rate",
     "get_model_class_name",
     "get_output_tensor_range",
     "get_transformer_config_subfolder",

vllm_omni/model_extras/dreamid_omni.py

@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from typing import Any
-
-DREAMID_OMNI_EXTRA_BODY_PARAMS = frozenset({"solver_name", "shift"})
-DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE = 16000
-
-
-def build_x_to_video_audio_prompt(
-    prompt: dict[str, Any],
-    request_options: Mapping[str, Any],
-) -> dict[str, Any]:
-    """Translate the canonical X-to-video+audio envelope for DreamID-Omni."""
-    result = dict(prompt)
-
-    for key in ("video_negative_prompt", "audio_negative_prompt"):
-        if request_options.get(key) is not None:
-            result[key] = request_options[key]
-
-    media_inputs = prompt.get("multi_modal_data")
-    if isinstance(media_inputs, Mapping) and "audio" in media_inputs:
-        adapted_media = dict(media_inputs)
-        audio_inputs = media_inputs["audio"]
-        if not isinstance(audio_inputs, list):
-            audio_inputs = [audio_inputs]
-        adapted_media["audio"] = [_extract_waveform(item) for item in audio_inputs]
-        result["multi_modal_data"] = adapted_media
-    return result
-
-
-def _extract_waveform(audio: Any) -> Any:
-    if not isinstance(audio, tuple) or len(audio) != 2:
-        raise TypeError("Canonical audio inputs must be (waveform, sample_rate) tuples.")
-    waveform, sample_rate = audio
-    if sample_rate != DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE:
-        raise ValueError(
-            "DreamID-Omni audio inputs must be resampled to "
-            f"{DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE} Hz; got {sample_rate}."
-        )
-    return waveform

vllm_omni/model_extras/registry.py

@@ -28,13 +28,6 @@ from vllm_omni.model_extras.cosmos3 import (
     COSMOS3_EXTRA_BODY_PARAMS,
     COSMOS3_EXTRA_OUTPUT_PARAMS,
 )
-from vllm_omni.model_extras.dreamid_omni import (
-    DREAMID_OMNI_EXTRA_BODY_PARAMS,
-    DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE,
-)
-from vllm_omni.model_extras.dreamid_omni import (
-    build_x_to_video_audio_prompt as build_dreamid_x_to_video_audio_prompt,
-)
 from vllm_omni.model_extras.helios import (
     HELIOS_EXTRA_BODY_PARAMS,
     HELIOS_EXTRA_OUTPUT_PARAMS,
@@ -105,7 +98,6 @@ ImageToVideoPromptBuilder = Callable[
     dict[str, Any],
 ]
 XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]]
-XToVideoAudioPromptBuilder = Callable[[dict[str, Any], Mapping[str, Any]], dict[str, Any]]
 OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"]
 
 
@@ -211,11 +203,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
     },
-    "DreamIDOmniPipeline": {
-        "extra_body_params": DREAMID_OMNI_EXTRA_BODY_PARAMS,
-        "input_audio_sample_rate": DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE,
-        "x_to_video_audio_prompt_builder": build_dreamid_x_to_video_audio_prompt,
-    },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
         "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -308,12 +295,6 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
     return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()
 
 
-def get_input_audio_sample_rate(model_class_name: str | None) -> int | None:
-    """Return a pipeline's required input rate, or preserve the source rate."""
-    spec = _get_spec(model_class_name)
-    return spec.get("input_audio_sample_rate") if spec is not None else None
-
-
 def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
     """Return the declared range for floating-point tensor outputs.
 
@@ -420,14 +401,3 @@ def build_image_to_video_prompt(
         width=width,
         num_frames=num_frames,
     )
-
-
-def build_x_to_video_audio_prompt(
-    model_class_name: str | None,
-    prompt: dict[str, Any],
-    request_options: Mapping[str, Any],
-) -> dict[str, Any]:
-    """Apply a model adapter to an example-owned video+audio envelope."""
-    spec = _get_spec(model_class_name)
-    builder: XToVideoAudioPromptBuilder | None = spec.get("x_to_video_audio_prompt_builder") if spec else None
-    return builder(prompt, request_options) if builder is not None else prompt

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -406,13 +406,11 @@ class OmniConnectorModelRunnerMixin:
         # thread by one execute_model() cycle, especially when the request was
         # added after the current scheduler_output snapshot.
         #
-        # Orphaned pending recv entries (e.g. from upstream stage crash) are
-        # handled by collect_timed_out_request_ids() -- on
-        # OmniSchedulingCoordinator for full-payload requests, and on
-        # OmniChunkTransferAdapter for async-chunk ones -- which detect
-        # wait-time violations.  The scheduler then removes the request from
-        # its queues, sets FINISHED_ERROR, and calls _free_request() which
-        # ultimately triggers cleanup_finished_request() here.
+        # Orphaned pending recv entries (e.g. from upstream stage crash)
+        # are handled by OmniSchedulingCoordinator.collect_timed_out_request_ids()
+        # which detects wait-time violations.  The scheduler then removes the
+        # request from its queues, sets FINISHED_ERROR, and calls _free_request()
+        # which ultimately triggers cleanup_finished_request() here.
         for attr_name in (
             "_request_ids_mapping",
             "_get_req_chunk",