[Docs] Clarify TP1 DLO mmap support scope

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/.nav.yml

@@ -116,7 +116,6 @@ nav:
         - Async Diffusion Output: design/feature/async_diffusion_output.md
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
         - Automatic Prefix Caching: design/feature/prefix_caching.md
-        - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
       - Communication:
         - OmniConnector Implementations:
           - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md

docs/design/ar_diffusion_pipeline_capability.md

@@ -4,9 +4,6 @@
 > `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
 > direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
 > but it is not yet part of the stable module-design hierarchy.
->
-> The end-to-end request, identity, failure, and memory contracts are documented
-> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
 
 `ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
 autoregressive attention KV across requests. The engine only selects
@@ -54,13 +51,14 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
 those paths requires a batch-aware state-binding contract rather than silently
 bypassing `bind_ar_diffusion_state()`.
 
-The implementation may retain multiple sessions per model instance. It selects
-the largest feasible capacity up to the pipeline's `session_capacity`, based on
-self-attention pages, cross-attention allocations, scratch, and declared
-model-owned state. One session is admitted whenever it fits actual free device
-memory; `gpu_memory_fraction` controls additional residency. Requests beyond
-that capacity evict the least-recently-used session through the same lifecycle
-hook used by explicit close and failure cleanup.
+The current implementation retains one session per model instance. A request
+for a different `session_id` closes the resident session through the same LRU
+lifecycle used by explicit close and failure cleanup before allocating the new
+one. Self-attention and cross-attention memory are therefore sized for one
+resident session even if the pipeline declares a larger `session_capacity`.
+This prioritizes rollout throughput and prevents block-pool exhaustion before a
+count-based eviction can run. Multi-session residency can be added later with a
+memory-backed capacity calculation or allocation-pressure eviction.
 
 Scratch storage is also capability-driven. `frames_per_block` reserves space for
 an uncommitted current video block, while `max_scratch_tokens_per_branch`

docs/design/feature/realtime_ar_diffusion.md

@@ -1,213 +0,0 @@
-# Realtime AR-Diffusion sessions
-
-> **Status:** experimental. This document describes the internal execution
-> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
-> define a public HTTP or WebSocket schema.
-
-## Scope
-
-Realtime AR-Diffusion serves world models that generate one autoregressive
-latent block per request while retaining model state across requests. The
-generic runtime owns identity, event ordering, session lifecycle, routing, and
-KV allocation. A model adapter owns model-specific controls and conditioning.
-
-LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
-same runner capability with two logical CFG branches.
-
-## Identities
-
-Four identifiers have distinct lifetimes:
-
-| Identifier | Lifetime | Owner | Contract |
-| --- | --- | --- | --- |
-| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
-| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
-| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
-| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
-
-`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
-`request_id` remains inside the immutable AR-Diffusion snapshot and the
-standard output metadata, while caller-visible `OmniRequestOutput.request_id`
-continues to use AsyncOmni's existing external-ID contract. Protocol identity
-therefore does not require a private entrypoint override.
-
-## Request lifecycle
-
-1. `ARDiffusionSession.accept_event()` validates and queues an immutable
-   `ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
-   the same event.
-2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
-   `_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
-   of the current prompt, reduced controls, ordered event IDs, and chunk index.
-3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
-   parameters, stores the typed tick in
-   `sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
-   prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
-4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
-   `OmniDiffusionRequest` through the orchestrator.
-5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
-   `_get_or_create_session()` obtains runner-owned paged KV, and
-   `bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
-   of `execute_model()`.
-6. The model pipeline validates the tick and generates exactly one latent
-   block. LingBot interprets camera controls in
-   `LingBotCameraControlReducer` and
-   `LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
-   camera or action schemas.
-7. The pipeline returns one `DiffusionOutput` with the generated payload and an
-   `ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
-   `AsyncOmni` exposes it as
-   `OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
-8. The consumer parses the standard envelope. The session commits its reducer
-   state, prompt, controls, event IDs, and next chunk index only when returned
-   metadata exactly equals the submitted tick snapshot.
-
-The metadata envelope is:
-
-```json
-{
-  "metadata": {
-    "ar_diffusion": {
-      "session_id": "world-7",
-      "request_id": "ar-world-7-3-...",
-      "chunk_index": 3,
-      "applied_event_ids": [10, 11]
-    }
-  }
-}
-```
-
-## Transaction and failure semantics
-
-Accepted events are not acknowledged as applied until the model output and
-metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
-and the session snapshot form one logical commit. A model, metadata, or reducer
-failure leaves the chunk index and queued events uncommitted and transitions the
-session to `FAILED`.
-
-After a tick failure, worker state is closed because a forward may have
-partially modified or already released KV. The same chunk is not retried in
-place. A caller must explicitly reset the session, starting again at chunk
-zero, or close it.
-
-Explicit close follows this state machine:
-
-```text
-ACTIVE or FAILED
-       |
-       v
-    CLOSING ---- worker cleanup succeeds ----> CLOSED
-       |
-       +------- worker cleanup fails --------> CLEANUP_FAILED
-                                                   |
-                                                   +-- close retry --> CLOSING
-```
-
-`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
-creation of the same `session_id`, and rejects reset. Only an idempotent close
-retry that succeeds removes the session. Transport disconnect uses the same
-path.
-
-## KV ownership and capacity
-
-`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
-declares:
-
-- TP-local self-attention geometry, AR block size, sink, and recent window;
-- logical KV branches and their worker-local slots;
-- fixed-length cross-attention caches;
-- scratch tokens required by an uncommitted forward;
-- persistent model-owned CUDA bytes per resident session; and
-- the requested resident-session capacity.
-
-The runner computes:
-
-```text
-required(capacity) =
-    scratch
-  + managed_self_attention_pages(capacity)
-  + capacity * cross_attention_bytes_per_session
-  + capacity * model_owned_state_bytes_per_session
-```
-
-`gpu_memory_fraction` is a soft expansion budget:
-
-```text
-configured_budget = available_device_bytes * gpu_memory_fraction
-effective_budget = max(configured_budget, required(1))
-```
-
-If `required(1)` exceeds actual available device memory, initialization fails.
-Otherwise the runner always admits one viable session and uses the configured
-fraction to select additional resident sessions up to the model-declared cap.
-All model-owned reservations are deducted before the paged self-attention pool
-is allocated.
-
-LingBot reports its persistent full-horizon image-condition tensor. DreamZero
-reports the measured 603 MiB per-session Wan VAE causal-convolution state.
-
-## Concurrency and routing limits
-
-The current runner executes one request at a time:
-
-- `max_num_seqs` must be one;
-- request batching and diffusion step execution are rejected;
-- one AR block is generated per engine request; and
-- an AR-Diffusion stage must have exactly one replica.
-
-The single-replica restriction is intentional. Multi-replica support requires
-session-affine routing so every tick, reset, and close reaches the worker that
-owns the session. Random or round-robin routing is not correct.
-
-Multiple resident sessions may share one runner. Capacity selection accounts
-for all persistent pools, and the runner performs LRU eviction through the same
-model lifecycle hook used by explicit close.
-
-## Model boundary
-
-The generic runtime transports controls as opaque
-`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
-and reduce those values. This keeps target/velocity, key state, pose,
-intrinsics, and other model-specific conventions out of the session protocol.
-
-For LingBot World v2:
-
-- `lingbot_world/actions.py` parses key-state and trajectory controls;
-- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
-- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
-- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
-  produces one block plus the standard metadata envelope.
-
-## Non-goals
-
-This contract does not currently provide:
-
-- a public HTTP or WebSocket request schema (the structured-interaction
-  frontend is tracked separately in
-  [#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
-- camera/action semantics shared by every world model;
-- session migration or replication across workers;
-- retry of an ambiguous partially executed chunk;
-- cross-stage KV transfer; or
-- stateful streaming VAE decode.
-
-Those features can be layered on top without adding LingBot-specific identity
-or lifecycle fields to the generic runtime.
-
-## Required regression coverage
-
-CPU contract tests cover:
-
-- separation of tick correlation IDs from internal engine routing IDs;
-- event ordering, metadata equality, and reducer fail-closed commit;
-- close/disconnect failure tombstones and cleanup retry;
-- capacity one under the default fraction for shipped LingBot geometry;
-- capacity expansion under a tuned fraction;
-- DreamZero model-owned state accounting;
-- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
-- LingBot model registration, imports, camera/action reduction, and one-block
-  output metadata.
-
-GPU validation must additionally exercise real model loading, action input,
-prompt switching, rolling-window replacement, reset, close, finite output
-latents, and exact metadata/request identity.

docs/design/index.md

@@ -22,7 +22,6 @@ implementation contract; it is not, by itself, a general support claim.
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
 - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
-- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
 
 ### Communication
 

docs/models/supported_models.md

@@ -37,7 +37,6 @@ th {
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -118,10 +117,6 @@ support metadata until their recipe is audited.
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
 
-See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
-for offline and experimental realtime usage, supported controls, validation,
-and current limitations.
-
 <sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
 gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
 Other AMD SKUs (e.g. MI325X) are not listed until their

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/aura_omni/end2end.py

@@ -24,6 +24,7 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.media.audio import load_audio
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.stage_input_processors.aura_omni import (
     DEFAULT_QWEN3_TTS_REF_TEXT,
@@ -227,6 +228,7 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
     parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/diffusion/lingbot_world_v2.py

@@ -1,288 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Generate a LingBot-World v2 video from an image and camera trajectory.
-
-The official checkpoint is licensed separately under CC BY-NC-SA and is
-restricted to non-commercial use. This vLLM-Omni integration remains
-Apache-2.0 licensed.
-
-Example:
-    python examples/offline_inference/diffusion/lingbot_world_v2.py \
-        --prompt "The camera moves slowly forward through the scene." \
-        --image /path/to/first_frame.png \
-        --action-dir /path/to/actions/forward \
-        --output lingbot_world_v2.mp4
-
-``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
-example configures its parent as the trusted action root and sends the
-contained directory through ``extra_args["action_path"]``.
-"""
-
-from __future__ import annotations
-
-import argparse
-import math
-from collections.abc import Sequence
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
-_NUM_INFERENCE_STEPS = 4
-_MAX_SEQUENCE_LENGTH = 512
-_MAX_PIXEL_AREA = 480 * 832
-_MAX_RAW_FRAMES = 117
-_MAX_SOURCE_CAMERA_FRAMES = 4096
-_TEMPORAL_COMPRESSION = 4
-_LATENT_FRAMES_PER_BLOCK = 3
-_SPATIAL_ALIGNMENT = 16
-
-
-@dataclass(frozen=True)
-class LingBotPaths:
-    image: Path
-    action_dir: Path
-    action_root: Path
-    action_relative: Path
-    camera_frames: int
-    output: Path
-
-
-def _validate_camera_arrays(action_dir: Path) -> int:
-    """Validate the bounded camera contract before starting the engine."""
-
-    try:
-        poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
-        intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
-    except (OSError, TypeError, ValueError):
-        raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
-    if poses.ndim != 3 or poses.shape[1:] != (4, 4):
-        raise ValueError("poses.npy must have shape [frames, 4, 4].")
-    if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
-        raise ValueError("intrinsics.npy must have shape [frames, 4].")
-    if poses.shape[0] != intrinsics.shape[0]:
-        raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
-    if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
-        raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
-    if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
-        raise ValueError("Camera arrays must contain real numeric values.")
-    if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
-        raise ValueError("Camera arrays must contain only finite values.")
-    return int(poses.shape[0])
-
-
-def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
-    )
-    parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
-    parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
-    parser.add_argument("--image", required=True, help="Path to the first-frame image.")
-    parser.add_argument(
-        "--action-dir",
-        required=True,
-        help="Camera action directory containing poses.npy and intrinsics.npy.",
-    )
-    parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
-    parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=9,
-        help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
-    )
-    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument(
-        "--tensor-parallel-size",
-        type=int,
-        default=1,
-        help="Number of GPUs used for tensor parallelism inside the DiT.",
-    )
-    parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
-    parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
-    parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable compilation and force eager execution.",
-    )
-    return parser.parse_args(argv)
-
-
-def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
-    """Canonicalize local inputs and construct a safe trusted-root mapping."""
-
-    image = Path(args.image).expanduser().resolve()
-    if not image.is_file():
-        raise ValueError("--image must point to an existing file.")
-
-    action_dir = Path(args.action_dir).expanduser().resolve()
-    if not action_dir.is_dir():
-        raise ValueError("--action-dir must point to an existing directory.")
-    if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
-        raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
-    camera_frames = _validate_camera_arrays(action_dir)
-
-    action_root = action_dir.parent.resolve()
-    try:
-        action_relative = action_dir.relative_to(action_root)
-    except ValueError:
-        raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
-    if not action_relative.parts:
-        raise ValueError("--action-dir must name a directory contained by a trusted action root.")
-
-    output = Path(args.output).expanduser().resolve()
-    return LingBotPaths(
-        image=image,
-        action_dir=action_dir,
-        action_root=action_root,
-        action_relative=action_relative,
-        camera_frames=camera_frames,
-        output=output,
-    )
-
-
-def _positive_finite(value: float, flag: str) -> float:
-    result = float(value)
-    if not math.isfinite(result) or result <= 0:
-        raise ValueError(f"{flag} must be a positive finite number.")
-    return result
-
-
-def build_omni_kwargs(
-    args: argparse.Namespace,
-    paths: LingBotPaths,
-    *,
-    parallel_config: Any,
-) -> dict[str, Any]:
-    """Build Omni configuration without overriding checkpoint class discovery."""
-
-    if args.tensor_parallel_size <= 0:
-        raise ValueError("--tensor-parallel-size must be a positive integer.")
-    flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
-    model_path = Path(args.model).expanduser()
-    model = str(model_path.resolve()) if model_path.exists() else args.model
-    return {
-        "model": model,
-        "flow_shift": flow_shift,
-        "parallel_config": parallel_config,
-        "enforce_eager": args.enforce_eager,
-        "model_config": {"lingbot_action_root": str(paths.action_root)},
-    }
-
-
-def build_request(
-    args: argparse.Namespace,
-    paths: LingBotPaths,
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    """Build the prompt and fixed v1 causal-DMD sampling contract."""
-
-    prompt_text = args.prompt.strip()
-    if not prompt_text:
-        raise ValueError("--prompt must contain non-whitespace text.")
-    if args.height <= 0 or args.width <= 0:
-        raise ValueError("--height and --width must be positive integers.")
-    if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
-        raise ValueError("--height and --width must each be divisible by 16.")
-    if args.height * args.width > _MAX_PIXEL_AREA:
-        raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
-    if args.num_frames <= 0:
-        raise ValueError("--num-frames must be a positive integer.")
-    if args.num_frames > _MAX_RAW_FRAMES:
-        raise ValueError("--num-frames must not exceed 117 raw frames.")
-    if args.num_frames > paths.camera_frames:
-        raise ValueError("Camera arrays must contain at least --num-frames frames.")
-    if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
-        raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
-    latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
-    if latent_frames % _LATENT_FRAMES_PER_BLOCK:
-        raise ValueError("--num-frames must map to whole three-frame latent blocks.")
-    if args.fps <= 0:
-        raise ValueError("--fps must be a positive integer.")
-    flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
-
-    prompt = {
-        "prompt": prompt_text,
-        "multi_modal_data": {"image": str(paths.image)},
-    }
-    sampling_kwargs = {
-        "height": args.height,
-        "width": args.width,
-        "num_frames": args.num_frames,
-        "num_inference_steps": _NUM_INFERENCE_STEPS,
-        "max_sequence_length": _MAX_SEQUENCE_LENGTH,
-        "seed": args.seed,
-        "fps": args.fps,
-        "extra_args": {
-            "action_path": paths.action_relative.as_posix(),
-            "flow_shift": flow_shift,
-        },
-    }
-    return prompt, sampling_kwargs
-
-
-def extract_video_array(outputs: Any) -> np.ndarray:
-    """Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
-
-    if not isinstance(outputs, list) or len(outputs) != 1:
-        raise ValueError("Expected exactly one Omni output for the single LingBot request.")
-    output = outputs[0]
-    request_output = getattr(output, "request_output", None)
-    if request_output is not None and getattr(request_output, "images", None):
-        output = request_output
-
-    images = getattr(output, "images", None)
-    if not images:
-        multimodal_output = getattr(output, "multimodal_output", None) or {}
-        images = multimodal_output.get("video")
-    if images is None:
-        raise ValueError("No video frames found in Omni output.")
-    if isinstance(images, list) and len(images) == 1:
-        images = images[0]
-
-    video: np.ndarray = np.asarray(images)
-    if video.ndim == 5:
-        if video.shape[0] != 1:
-            raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
-        video = video[0]
-    if video.ndim != 4 or video.shape[-1] not in (3, 4):
-        raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
-    return video
-
-
-def main(argv: Sequence[str] | None = None) -> Path:
-    args = parse_args(argv)
-    paths = resolve_cli_paths(args)
-    prompt, sampling_kwargs = build_request(args, paths)
-    paths.output.parent.mkdir(parents=True, exist_ok=True)
-
-    # Keep vLLM imports below pure CLI/path validation so ``--help`` and the
-    # helper tests work on machines without the native vLLM runtime.
-    from diffusers.utils import export_to_video
-
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
-    from vllm_omni.entrypoints.omni import Omni
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
-    omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
-    omni = Omni(**omni_kwargs)
-    try:
-        outputs = omni.generate(
-            prompt,
-            OmniDiffusionSamplingParams(**sampling_kwargs),
-            use_tqdm=False,
-        )
-    finally:
-        omni.close()
-
-    frames = extract_video_array(outputs)
-    export_to_video(frames, str(paths.output), fps=args.fps)
-    print(f"Saved {frames.shape[0]} frames to {paths.output}")
-    return paths.output
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/diffusion/lingbot_world_v2_realtime.py

@@ -1,239 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
-
-This example exercises the internal session API. It is not a public HTTP or
-WebSocket protocol. Each JSONL event advances the world by one three-latent
-frame AR block and writes the latent plus its identity metadata. The current
-scoped implementation supports at most ten ticks per generation epoch because
-its image condition is bounded to 117 pixel frames; reset or create a session
-to start a new world.
-"""
-
-from __future__ import annotations
-
-import argparse
-import asyncio
-import json
-import math
-import time
-from collections.abc import Sequence
-from pathlib import Path
-from typing import Any
-
-_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
-_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
-_FRAMES_PER_BLOCK = 3
-# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
-_MAX_REALTIME_TICKS = 10
-
-
-def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
-    """Build the event-side script consumed by LingBotCameraControlReducer."""
-    return {"mode": "script", "frames": frames}
-
-
-def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
-    parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
-    parser.add_argument("--image", required=True, help="Initial RGB image.")
-    parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
-    parser.add_argument(
-        "--events",
-        required=True,
-        help="JSONL file with one prompt/action event per AR block; at most 10 events.",
-    )
-    parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
-    parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
-    parser.add_argument("--height", type=int, default=480)
-    parser.add_argument("--width", type=int, default=832)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
-    parser.add_argument("--tensor-parallel-size", type=int, default=1)
-    parser.add_argument("--enforce-eager", action="store_true")
-    return parser.parse_args(argv)
-
-
-def _load_events(path: Path) -> list[dict[str, Any]]:
-    events: list[dict[str, Any]] = []
-    for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
-        line = raw_line.strip()
-        if not line:
-            continue
-        try:
-            value = json.loads(line)
-        except json.JSONDecodeError as exc:
-            raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
-        if not isinstance(value, dict):
-            raise ValueError(f"events line {line_number} must be a JSON object.")
-        event_id = value.get("event_id")
-        if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
-            raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
-        prompt = value.get("prompt")
-        if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
-            raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
-        frames = value.get("frames")
-        if frames is not None:
-            if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
-                raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
-            for frame in frames:
-                if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
-                    raise ValueError(f"events line {line_number} frames must contain only action strings.")
-        if prompt is None and frames is None:
-            raise ValueError(f"events line {line_number} must update prompt and/or frames.")
-        events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
-    if not events:
-        raise ValueError("events file must contain at least one event.")
-    if len(events) > _MAX_REALTIME_TICKS:
-        raise ValueError(
-            "events file must contain at most "
-            f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
-            "image-condition horizon is 117 pixel frames."
-        )
-    event_ids = [event["event_id"] for event in events]
-    if event_ids != sorted(set(event_ids)):
-        raise ValueError("event_id values must be unique and strictly increasing.")
-    return events
-
-
-def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
-    image = Path(args.image).expanduser().resolve()
-    events = Path(args.events).expanduser().resolve()
-    output_dir = Path(args.output_dir).expanduser().resolve()
-    if not image.is_file():
-        raise ValueError("--image must point to an existing file.")
-    if not events.is_file():
-        raise ValueError("--events must point to an existing JSONL file.")
-    if not args.prompt.strip():
-        raise ValueError("--prompt must contain non-whitespace text.")
-    if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
-        raise ValueError("--height and --width must be positive multiples of 16.")
-    if args.tensor_parallel_size <= 0:
-        raise ValueError("--tensor-parallel-size must be positive.")
-    if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
-        raise ValueError("--gpu-memory-fraction must be in (0, 1].")
-    return image, events, output_dir
-
-
-async def run(argv: Sequence[str] | None = None) -> Path:
-    args = parse_args(argv)
-    image, events_path, output_dir = _validate_args(args)
-    events = _load_events(events_path)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Imports stay below pure input validation so --help and helper tests do not
-    # require a CUDA-enabled vLLM installation.
-    import torch
-
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
-    from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
-    from vllm_omni.entrypoints.async_omni import AsyncOmni
-    from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
-    from vllm_omni.experimental.ar_diffusion.session import (
-        ARDiffusionSessionEvent,
-        ARDiffusionSessionManager,
-        ARDiffusionWorkerLifecycle,
-    )
-    from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    engine = AsyncOmni(
-        model=args.model,
-        engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
-        enforce_eager=args.enforce_eager,
-        parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
-        max_num_seqs=1,
-        model_config={
-            "ar_diffusion_height": args.height,
-            "ar_diffusion_width": args.width,
-            "ar_diffusion_kv_config": {
-                "gpu_memory_fraction": args.gpu_memory_fraction,
-                "warmup_cudagraph": True,
-            },
-        },
-    )
-    sampling = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_frames=9,
-        num_inference_steps=4,
-        max_sequence_length=512,
-        seed=args.seed,
-        output_type="latent",
-        extra_args={"flow_shift": 5.0},
-    )
-    consumer = ARDiffusionOmniTickConsumer(
-        engine,
-        prompt_provider=lambda tick: {
-            "prompt": tick.prompt,
-            "multi_modal_data": {"image": str(image)},
-        },
-        sampling_params_list=[sampling],
-        diffusion_stage_id=0,
-    )
-    manager = ARDiffusionSessionManager(
-        tick_consumer=consumer,
-        lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
-        max_pending_events=32,
-        control_reducer_factory=LingBotCameraControlReducer,
-    )
-    session = await manager.create_session(args.session_id)
-    measurements: list[dict[str, Any]] = []
-    try:
-        for chunk_index, event in enumerate(events):
-            controls = ()
-            if event["frames"] is not None:
-                controls = (
-                    ARDiffusionControlInput(
-                        track="camera",
-                        schema=_CAMERA_ACTION_SCHEMA,
-                        data=_camera_event_data(event["frames"]),
-                    ),
-                )
-            await session.accept_event(
-                ARDiffusionSessionEvent(
-                    event_id=event["event_id"],
-                    prompt=event["prompt"]
-                    if event["prompt"] is not None
-                    else (args.prompt if chunk_index == 0 else None),
-                    controls=controls,
-                )
-            )
-            started = time.perf_counter()
-            output = await session.next_chunk()
-            elapsed = time.perf_counter() - started
-            if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
-                raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
-            latent = output.images[0].detach().float().cpu()
-            metadata = consumer.chunk_metadata(output).to_dict()
-            torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
-            (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
-                json.dumps(metadata, indent=2, sort_keys=True) + "\n"
-            )
-            measurements.append(
-                {
-                    "chunk_index": chunk_index,
-                    "latency_seconds": elapsed,
-                    "shape": list(latent.shape),
-                    "finite": bool(torch.isfinite(latent).all()),
-                    "metadata": metadata,
-                }
-            )
-            print(json.dumps(measurements[-1], sort_keys=True), flush=True)
-    finally:
-        try:
-            await manager.close_session(args.session_id)
-        finally:
-            engine.shutdown()
-
-    summary_path = output_dir / "summary.json"
-    summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
-    return summary_path
-
-
-def main(argv: Sequence[str] | None = None) -> Path:
-    return asyncio.run(run(argv))
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/lance/end2end.py

@@ -82,6 +82,9 @@ def parse_args():
     # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
     # and ``create_default_diffusion`` materializes the stage config.
 
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(p)
     return p.parse_args()
 
 

examples/offline_inference/nemotron_voicechat/end2end.py

@@ -67,7 +67,7 @@ def main() -> None:
     parser.add_argument("--output-dir", default="results/nemotron_voicechat")
     parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
     parser.add_argument(
-        "--deploy-config",
+        "--stage-configs-path",
         default=None,
         help=(
             "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
@@ -134,8 +134,8 @@ def main() -> None:
     ]
 
     omni_kwargs: dict = {}
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
     omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
     outputs = omni.generate(inputs, sampling_params_list)
 

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -38,6 +38,17 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = "fishaudio/s2-pro"
+DEFAULT_DEPLOY_CONFIG = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "..",
+    "..",
+    "..",
+    "vllm_omni",
+    "model_executor",
+    "stage_configs",
+    "fish_speech_s2_pro.yaml",
+)
 
 
 def build_prompt(
@@ -141,8 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
+    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
     return omni_kwargs
 
 

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -32,10 +32,11 @@ from end2end import _estimate_prompt_len
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
 
 
-def _default_deploy_config() -> str:
+def _default_stage_config() -> str:
     repo_root = Path(__file__).resolve().parents[4]
     return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
 
@@ -108,7 +109,7 @@ def main(args: Any) -> None:
         "output_dir",
     ):
         omni_kwargs.pop(key, None)
-    omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
+    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
     omni_kwargs["log_stats"] = args.log_stats
     omni = Omni(**omni_kwargs)
 
@@ -155,7 +156,7 @@ def parse_args() -> Any:
         default=None,
         help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
     )
-    parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
     parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
     parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
     parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
@@ -163,6 +164,7 @@ def parse_args() -> Any:
     parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
     parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

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__":

recipes/BosonAI/Higgs-Audio-V3-TTS.md

@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
 **Online serving:**
 
 ```bash
-vllm serve bosonai/higgs-audio-v3-tts-4b \
+vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
     --host 0.0.0.0 --port 8095 \
     --trust-remote-code --omni
 ```

recipes/NVIDIA/Nemotron-Labs-Audex.md

@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
 `model_type` auto-resolves `audex_tts.yaml`):
 
 ```bash
-vllm serve nvidia/Nemotron-Labs-Audex-2B \
+vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
     --host 0.0.0.0 --port 8097 \
     --trust-remote-code --omni
 # other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
@@ -109,7 +109,7 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
 auto-resolves to the 2B-tuned config otherwise):
 
 ```bash
-vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
     --host 0.0.0.0 --port 8097 \
     --trust-remote-code --omni \
     --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml

recipes/README.md

@@ -60,7 +60,6 @@ recipes/
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
-| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
 | [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
 | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |

recipes/Robbyant/LingBot-World-2.0.md

@@ -1,122 +0,0 @@
-# LingBot-World 2.0
-
-> Offline and experimental realtime interactive world generation
-
-## Summary
-
-- Vendor: Robbyant
-- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
-- Task: image-conditioned interactive world generation
-- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
-- Hardware validated: NVIDIA H200 and B200
-- Maintainer: Community
-
-The checkpoint is separately licensed under CC BY-NC-SA and restricted to
-non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
-
-## Offline generation
-
-The offline path consumes one source image and an action directory containing
-`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
-`[frames, 4]`.
-
-```bash
-python examples/offline_inference/diffusion/lingbot_world_v2.py \
-  --prompt "The camera moves slowly forward through the scene." \
-  --image /path/to/first_frame.png \
-  --action-dir /path/to/actions/forward \
-  --num-frames 81 \
-  --output lingbot_world_v2.mp4
-```
-
-The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
-counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
-longer bounded source trajectory (the official example contains 269 frames);
-the runtime consumes only the prefix needed by the request.
-
-## Realtime in-process generation
-
-The realtime example keeps AR-Diffusion KV and model-owned state across
-requests. Each JSONL line describes the prompt and/or three latent-frame camera
-actions applied at the next chunk boundary:
-
-```json
-{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
-{"event_id":2,"frames":[["w"],["w"],["w"]]}
-{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
-```
-
-Run:
-
-```bash
-python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
-  --image /path/to/first_frame.png \
-  --events /path/to/events.jsonl \
-  --output-dir /tmp/lingbot-realtime \
-  --gpu-memory-fraction 0.6
-```
-
-The example writes one latent tensor and one metadata JSON file per chunk. It
-exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
--> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
-transport.
-
-This PR intentionally does not define a public realtime HTTP/WebSocket schema.
-An online serving client should be added together with the public transport
-API rather than exposing LingBot-specific event fields from the generic model
-runtime. The structured camera-interaction frontend and transport work is
-tracked separately in
-[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
-
-## Realtime identity and controls
-
-- `session_id` identifies the persistent world and its worker-owned state.
-- `event_id` identifies a prompt/control update and remains monotonic across reset.
-- `chunk_index` is contiguous from zero and restarts from zero after reset.
-- `request_id` correlates one chunk snapshot with its output metadata.
-- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
-
-The generic runtime transports controls as opaque snapshots. LingBot's adapter
-accepts:
-
-- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
-  `s`, `d`, `i`, `j`, `k`, and `l`;
-- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
-
-## Validation
-
-Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
-The exercised matrix includes:
-
-- TP=1 and TP=2 execution;
-- two interleaved resident sessions;
-- action input and prompt switching;
-- seven contiguous chunks crossing the sink plus recent rolling window;
-- direct versus paged replay;
-- CUDA Graph execution;
-- reset, close, failure cleanup, and exact metadata matching; and
-- nine-frame VAE decode.
-
-Official `generate.py` accuracy and performance numbers, together with
-generation artifacts, are recorded in the PR validation section for the exact
-tested commit.
-
-## Current limitations
-
-- Only the 14B causal-fast checkpoint is supported.
-- The realtime control plane is internal; there is no public server transport yet.
-- AR-Diffusion stages currently require one replica because session-affine
-  routing across replicas is not implemented.
-- One AR block is generated per request and `max_num_seqs` must be one.
-- Stateful streaming VAE decode is not implemented; the realtime example emits
-  latent chunks.
-- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
-  Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
-- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
-
-## References
-
-- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
-- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
-- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
-- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)

recipes/StabilityAI/Stable-Audio-Open.md

@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
 Start the online serving endpoint:
 
 ```bash
-vllm serve /path/to/stable-audio-open-1.0 \
+vllm-omni serve /path/to/stable-audio-open-1.0 \
   --host 0.0.0.0 \
   --port 8091 \
   --gpu-memory-utilization 0.9 \

recipes/TEMPLATE.md

@@ -39,9 +39,6 @@ platform section, document one or more tested hardware configurations.
 
 #### Command
 
-Serve commands use `vllm serve <model> --omni`. The `--omni` flag is what
-selects the Omni pipeline, so it is required and cannot be dropped.
-
 ```bash
 # Add the exact command(s) here
 ```

recipes/inclusionAI/Ming-omni-tts.md

@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:
 
 ```bash
-vllm serve inclusionAI/Ming-omni-tts-0.5B \
+vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
     --deploy-config vllm_omni/deploy/ming_tts.yaml \
     --omni \
     --port 8091 \

recipes/zai-org/GLM-TTS.md

@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
 Start the server from the repository root:
 
 ```bash
-vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
 ```
 
 Async chunking is enabled by default in the bundled deployment config. For
 the sync (non-streaming) path:
 
 ```bash
-vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
 ```
 
 Use a custom deploy config for advanced cases:
 
 ```bash
-vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
   --deploy-config /path/to/your_glm_tts_overrides.yaml
 ```
 

tests/config/test_config_factory.py

@@ -1102,13 +1102,6 @@ class TestResolveScheduler:
 
 
 class TestDeployConfigLoading:
-    def test_rejects_legacy_stage_args_schema(self, tmp_path):
-        deploy_path = tmp_path / "legacy.yaml"
-        deploy_path.write_text("stage_args:\n  - stage_id: 0\n", encoding="utf-8")
-
-        with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
-            load_deploy_config(deploy_path)
-
     def test_load_minicpmo_duplex_deploy_config(self):
         deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
 
@@ -1384,8 +1377,7 @@ stages:
         assert async_stages[1].custom_process_input_func is None
 
     def test_no_bundled_legacy_stage_config_yamls(self):
-        repo_root = Path(__file__).resolve().parents[2]
-        stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
+        stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
         assert not list(stage_config_dir.glob("*.yaml"))
 
     def test_merge_pipeline_deploy(self):

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/ar_diffusion/test_capability_runner.py

@@ -18,7 +18,6 @@ from vllm_omni.experimental.ar_diffusion.capability import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
-from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
 
 BLOCK = 16
 POS = "positive"
@@ -57,21 +56,6 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
     )
 
 
-def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
-    return ARDiffusionKVCacheSpec(
-        num_layers=1,
-        num_kv_heads=1,
-        head_size=1,
-        tokens_per_frame=1,
-        frames_per_block=3,
-        window_frames=3,
-        sink_frames=3,
-        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
-        session_capacity=capacity,
-        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
-    )
-
-
 class CapablePipeline:
     def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
         self.spec = spec
@@ -119,7 +103,6 @@ def make_runner(
     *,
     available_bytes: int = 1 << 28,
     step_execution: bool = False,
-    gpu_memory_fraction: float = 0.1,
 ) -> ARDiffusionModelRunner:
     runner = object.__new__(ARDiffusionModelRunner)
     runner.od_config = SimpleNamespace(
@@ -130,10 +113,7 @@ def make_runner(
     )
     runner.device = torch.device("cpu")
     runner.pipeline = pipeline
-    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
-        enable=True,
-        gpu_memory_fraction=gpu_memory_fraction,
-    )
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
     runner.kv_cache = None
     runner._ar_diffusion_capability = None
     runner._ar_diffusion_kv_cache_spec = None
@@ -152,16 +132,6 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
     return state
 
 
-def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
-    state = runner._get_or_create_session(session_id)
-    assert runner.kv_cache is not None
-    seq_len = BLOCK * runner.kv_cache.frames_per_block
-    ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
-    ctx.ensure_video_slots(torch.device("cpu"))
-    state.commit_paged_context(kv_branch)
-    return state
-
-
 def test_ar_runner_rejects_pipeline_without_capability():
     runner = object.__new__(ARDiffusionModelRunner)
     runner.od_config = SimpleNamespace(max_num_seqs=1)
@@ -170,43 +140,6 @@ def test_ar_runner_rejects_pipeline_without_capability():
         runner._preallocate_kv_cache(available_bytes=1 << 20)
 
 
-def test_runner_uses_typed_tick_as_authoritative_session_contract():
-    tick = ARDiffusionTickRequest(
-        session_id="world-7",
-        request_id="request-3",
-        chunk_index=3,
-        reset=True,
-    )
-    req = SimpleNamespace(
-        request_id="request-3",
-        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
-    )
-
-    session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
-
-    assert session_id == "world-7"
-    assert extra_args == tick.to_extra_args()
-    assert parsed == tick
-
-
-def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
-    tick = ARDiffusionTickRequest(
-        session_id="world-7",
-        request_id="client-request-3",
-        chunk_index=3,
-    )
-    req = SimpleNamespace(
-        request_id="engine-request-uuid",
-        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
-    )
-
-    session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
-
-    assert session_id == "world-7"
-    assert req.request_id == "engine-request-uuid"
-    assert parsed.request_id == "client-request-3"
-
-
 def test_lingbot_like_single_branch_session_reuse_reset_and_close():
     pipeline = CapablePipeline(lingbot_like_spec())
     runner = make_runner(pipeline)
@@ -237,31 +170,6 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
     assert pipeline.closes == ["s1"]
 
 
-def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
-    runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
-    kv = runner.kv_cache
-    assert kv is not None
-
-    session_a = commit_one_block(runner, "world-a", "main")
-    a0_blocks = kv.window_block_ids(session_a.adapter("main"))
-    session_b = commit_one_block(runner, "world-b", "main")
-    b0_blocks = kv.window_block_ids(session_b.adapter("main"))
-    session_a_again = commit_one_block(runner, "world-a", "main")
-    a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
-
-    assert session_a_again is session_a
-    assert session_a.adapter("main").request_id == "ar::world-a::main"
-    assert session_b.adapter("main").request_id == "ar::world-b::main"
-    assert session_a.adapter("main").completed_chunks == 6
-    assert session_b.adapter("main").completed_chunks == 3
-    assert len(a0_blocks) == 3
-    assert len(b0_blocks) == 3
-    assert len(a1_blocks) == 6
-    assert set(a0_blocks) <= set(a1_blocks)
-    assert set(a1_blocks).isdisjoint(b0_blocks)
-    assert tuple(runner._sessions) == ("world-b", "world-a")
-
-
 def test_lingbot_like_sink_survives_sliding_window_eviction():
     runner = make_runner(CapablePipeline(lingbot_like_spec()))
     kv = runner.kv_cache
@@ -290,67 +198,24 @@ def test_dreamzero_like_two_branches_are_independent():
 
 
 def test_lru_eviction_releases_blocks_and_notifies_pipeline():
-    pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
     runner = make_runner(pipeline)
     kv = runner.kv_cache
     assert kv is not None
-    assert runner._session_capacity == 2
-    assert kv.session_capacity == 2
+    assert runner._session_capacity == 1
+    assert kv.session_capacity == 1
     free_total = kv.manager.block_pool.get_num_free_blocks()
     old = commit_one_frame(runner, "old", "main")
     k = torch.randn(1, 8, 4, 64)
     old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
     assert kv.manager.block_pool.get_num_free_blocks() < free_total
 
-    runner._get_or_create_session("new")
-    assert tuple(runner._sessions) == ("old", "new")
-    runner._get_or_create_session("newest")
-
-    assert tuple(runner._sessions) == ("new", "newest")
-    assert pipeline.closes == ["old"]
-    assert "old" not in kv._cross_sessions
-    assert kv.manager.block_pool.get_num_free_blocks() == free_total
-
-
-def test_budget_reduced_capacity_drives_runner_lru():
-    pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
-    runner = make_runner(
-        pipeline,
-        available_bytes=128,
-        gpu_memory_fraction=1.0,
-    )
-    kv = runner.kv_cache
-    assert kv is not None
-    assert kv.requested_session_capacity == 2
-    assert kv.session_capacity == 1
-    assert runner._session_capacity == 1
-
-    runner._get_or_create_session("old")
     runner._get_or_create_session("new")
 
     assert tuple(runner._sessions) == ("new",)
     assert pipeline.closes == ["old"]
-
-
-def test_dreamzero_like_requested_capacity_is_capped_by_budget():
-    spec = dreamzero_like_spec(capacity=64)
-    # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
-    # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
-    # cross-attention=1 page/session, for 44 pages total.
-    page_bytes = 65_536
-    pipeline = CapablePipeline(spec)
-    runner = make_runner(
-        pipeline,
-        available_bytes=44 * page_bytes,
-        gpu_memory_fraction=1.0,
-    )
-    kv = runner.kv_cache
-    assert kv is not None
-    assert kv.requested_session_capacity == 64
-    assert kv.session_capacity == 2
-    assert runner._session_capacity == 2
-    assert kv.cross_attention_reserved_bytes == 2 * page_bytes
+    assert "old" not in kv._cross_sessions
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
 
 
 def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
@@ -367,10 +232,7 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
         raise RuntimeError("layer exploded")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
-    request = SimpleNamespace(
-        request_id="broken-request",
-        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
-    )
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
 
     with pytest.raises(RuntimeError, match="layer exploded"):
         runner.execute_model(request)
@@ -401,10 +263,7 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
     monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
     runner.device = torch.device("cuda")
-    request = SimpleNamespace(
-        request_id="broken-request",
-        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
-    )
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
 
     with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
         runner.execute_model(request)

tests/diffusion/ar_diffusion/test_consumer.py

@@ -1,231 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from collections.abc import AsyncIterator, Sequence
-from typing import Any
-
-import pytest
-
-from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
-from vllm_omni.experimental.ar_diffusion.session import (
-    ARDiffusionSession,
-    ARDiffusionSessionEvent,
-)
-from vllm_omni.experimental.ar_diffusion.tick_protocol import (
-    ARDiffusionChunkMetadata,
-    ARDiffusionControlInput,
-    ARDiffusionTickRequest,
-)
-from vllm_omni.inputs.data import (
-    OmniDiffusionSamplingParams,
-    OmniPromptType,
-    OmniSamplingParams,
-)
-from vllm_omni.outputs import OmniRequestOutput
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeStagePool:
-    def __init__(self, num_replicas: int = 1) -> None:
-        self.num_replicas = num_replicas
-
-
-class FakeEngine:
-    def __init__(self, num_replicas: int = 1) -> None:
-        self.stage_pools = [FakeStagePool(num_replicas)]
-
-
-class FakeGenerateClient:
-    def __init__(
-        self,
-        *,
-        include_metadata: bool = True,
-        applied_event_ids: list[int] | None = None,
-        output_error: str | None = None,
-        num_replicas: int = 1,
-    ) -> None:
-        self.engine = FakeEngine(num_replicas)
-        self.include_metadata = include_metadata
-        self.applied_event_ids = applied_event_ids
-        self.output_error = output_error
-        self.calls: list[dict[str, Any]] = []
-
-    async def generate(
-        self,
-        prompt: OmniPromptType,
-        *,
-        request_id: str,
-        sampling_params_list: Sequence[OmniSamplingParams],
-        output_modalities: list[str] | None = None,
-    ) -> AsyncIterator[OmniRequestOutput]:
-        self.calls.append(
-            {
-                "prompt": prompt,
-                "request_id": request_id,
-                "sampling_params_list": sampling_params_list,
-                "output_modalities": output_modalities,
-            }
-        )
-        params = sampling_params_list[0]
-        assert isinstance(params, OmniDiffusionSamplingParams)
-        tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
-        assert tick is not None
-        chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
-        if self.applied_event_ids is not None:
-            chunk_metadata["applied_event_ids"] = self.applied_event_ids
-        metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
-        yield OmniRequestOutput(
-            request_id=request_id,
-            stage_id=0,
-            final_output_type="video",
-            images=["chunk"],  # type: ignore[list-item]
-            _multimodal_output=metadata,
-            finished=True,
-            error=self.output_error,
-        )
-
-
-class FakeLifecycle:
-    async def reset_session(self, session_id: str) -> None:
-        pass
-
-    async def close_session(self, session_id: str) -> None:
-        pass
-
-
-def make_tick() -> ARDiffusionTickRequest:
-    return ARDiffusionTickRequest(
-        session_id="world-1",
-        request_id="world-1-chunk-3",
-        chunk_index=3,
-        applied_event_ids=(7, 8),
-        prompt="turn left",
-        controls=(
-            ARDiffusionControlInput(
-                track="camera",
-                schema="example.camera.v1",
-                data={"samples": [{"x": 1.0}]},
-            ),
-        ),
-    )
-
-
-def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
-    client = FakeGenerateClient(num_replicas=2)
-
-    with pytest.raises(
-        ValueError,
-        match=r"require exactly one replica.*num_replicas=2",
-    ):
-        ARDiffusionOmniTickConsumer(
-            client,
-            prompt_provider=lambda tick: tick.prompt or "",
-            sampling_params_list=[OmniDiffusionSamplingParams()],
-            diffusion_stage_id=0,
-        )
-
-    assert client.calls == []
-
-
-@pytest.mark.asyncio
-async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
-    client = FakeGenerateClient()
-    template = OmniDiffusionSamplingParams(
-        output_type="latent",
-        extra_args={"deployment_option": "kept"},
-    )
-    consumer = ARDiffusionOmniTickConsumer(
-        client,
-        prompt_provider=lambda tick: {
-            "prompt": tick.prompt,
-            "multi_modal_data": {"image": "initial-image"},
-        },
-        sampling_params_list=[template],
-        diffusion_stage_id=0,
-        output_modalities=["video"],
-    )
-    tick = make_tick()
-
-    output = await consumer.execute_tick(tick)
-
-    assert output.images == ["chunk"]
-    assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
-    assert len(client.calls) == 1
-    call = client.calls[0]
-    assert call["request_id"] == tick.request_id
-    assert call["output_modalities"] == ["video"]
-    assert call["prompt"] == {
-        "prompt": "turn left",
-        "multi_modal_data": {"image": "initial-image"},
-    }
-    submitted_params = call["sampling_params_list"][0]
-    assert submitted_params.extra_args["deployment_option"] == "kept"
-    assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
-    assert "ar_diffusion_tick" not in template.extra_args
-
-
-@pytest.mark.asyncio
-async def test_consumer_rejects_missing_chunk_metadata() -> None:
-    consumer = ARDiffusionOmniTickConsumer(
-        FakeGenerateClient(include_metadata=False),
-        prompt_provider=lambda tick: tick.prompt or "",
-        sampling_params_list=[OmniDiffusionSamplingParams()],
-        diffusion_stage_id=0,
-    )
-
-    with pytest.raises(ValueError, match="multimodal_output.metadata"):
-        await consumer.execute_tick(make_tick())
-
-
-@pytest.mark.asyncio
-async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
-    consumer = ARDiffusionOmniTickConsumer(
-        FakeGenerateClient(applied_event_ids=[8, 7]),
-        prompt_provider=lambda tick: tick.prompt or "",
-        sampling_params_list=[OmniDiffusionSamplingParams()],
-        diffusion_stage_id=0,
-    )
-
-    with pytest.raises(ValueError, match="unique and strictly increasing"):
-        await consumer.execute_tick(make_tick())
-
-
-@pytest.mark.asyncio
-async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
-    consumer = ARDiffusionOmniTickConsumer(
-        FakeGenerateClient(include_metadata=False, output_error="model failed"),
-        prompt_provider=lambda tick: tick.prompt or "",
-        sampling_params_list=[OmniDiffusionSamplingParams()],
-        diffusion_stage_id=0,
-    )
-
-    with pytest.raises(RuntimeError, match="model failed"):
-        await consumer.execute_tick(make_tick())
-
-
-@pytest.mark.asyncio
-async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
-    client = FakeGenerateClient()
-    consumer = ARDiffusionOmniTickConsumer(
-        client,
-        prompt_provider=lambda tick: tick.prompt or "",
-        sampling_params_list=[OmniDiffusionSamplingParams()],
-        diffusion_stage_id=0,
-    )
-    session = ARDiffusionSession(
-        "world-1",
-        tick_consumer=consumer,
-        lifecycle=FakeLifecycle(),
-        max_pending_events=4,
-        request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
-    )
-    await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
-
-    output = await session.next_chunk()
-
-    assert output.request_id == "world-1-chunk-0"
-    assert session.chunk_index == 1
-    assert session.pending_event_count == 0
-    assert client.calls[0]["request_id"] == "world-1-chunk-0"

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -30,8 +30,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 BLOCK = 16
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 
 def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
     return ChunkWindowSpec(
@@ -97,6 +95,7 @@ def test_sliding_replace_snaps_to_chunk_boundary():
 def test_sink_chunks_protected():
     # sink = 1 chunk (16 tokens) is never skipped.
     skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
     assert skip == 32
 
 
@@ -202,149 +201,12 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
         kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
         session_capacity=1,
         cross_attention_lengths={"text": L},
-    )
-    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
-    page_bytes = kv.spec.page_size_bytes * 2
-    scratch_bytes = kv.scratch_num_blocks * page_bytes
-    expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
-    assert kv.num_blocks == expected
-    assert kv.cross_attention_reserved_bytes == cross_bytes
-    assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
-    assert expected > 1 * (2 + 1) + 2  # above the local-slot minimum, so the cross deduction is what's tested
-
-
-def _make_tiny_capacity_kv(
-    *,
-    requested_capacity: int,
-    available_bytes: int,
-    gpu_memory_fraction: float = 1.0,
-    model_owned_state_bytes_per_session: int = 0,
-) -> ARDiffusionKVCache:
-    return ARDiffusionKVCache(
-        ARDiffusionKVConfig(
-            enable=True,
-            chunk_size=1,
-            window_chunks=3,
-            sink_chunks=3,
-            gpu_memory_fraction=gpu_memory_fraction,
-        ),
-        num_layers=1,
-        num_kv_heads=1,
-        head_size=1,
-        dtype=torch.float32,
-        block_size=1,
-        max_model_len=128,
-        available_bytes=available_bytes,
-        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
-        session_capacity=requested_capacity,
-        cross_attention_lengths={"text": 2},
-        frames_per_block=3,
-        model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
         device=torch.device("cpu"),
     )
-
-
-def test_capacity_two_retains_both_windows_and_allocates_next_block():
-    # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
-    # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
-    # reservations: 192 bytes exactly.
-    kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
-    assert kv.requested_session_capacity == 2
-    assert kv.session_capacity == 2
-    assert kv.managed_num_blocks == 17
-
-    adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
-    for adapter in adapters:
-        for _ in range(6):  # three sink + three window blocks
-            kv.allocate_token_slots(adapter, 1)
-            kv.commit_chunk(adapter)
-        assert len(kv.window_block_ids(adapter)) == 6
-
-    # Only one request is in flight, so its three-frame block is counted once.
-    kv.allocate_token_slots(adapters[0], 3)
-
-
-def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
-    kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
-
-    assert kv.requested_session_capacity == 64
-    assert kv.session_capacity == 2
-    assert kv.cross_attention_bytes_per_session == 16
-    assert kv.cross_attention_reserved_bytes == 32
-    assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
-
-
-def test_capacity_rejects_budget_that_cannot_fit_one_session():
-    # One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
-    # 14 * 8 + 16 = 128 bytes.
-    with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
-        _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
-
-
-def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
-    kv = _make_tiny_capacity_kv(
-        requested_capacity=2,
-        available_bytes=1280,
-        gpu_memory_fraction=0.05,
-    )
-
-    assert kv.configured_memory_budget_bytes == 64
-    assert kv.memory_budget_bytes == 128
-    assert kv.session_capacity == 1
-
-
-def test_model_owned_state_reduces_effective_session_capacity():
-    kv = _make_tiny_capacity_kv(
-        requested_capacity=2,
-        available_bytes=200,
-        model_owned_state_bytes_per_session=16,
-    )
-
-    assert kv.session_capacity == 1
-    assert kv.model_owned_state_reserved_bytes == 16
-    assert (
-        kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
-        <= kv.memory_budget_bytes
-    )
-
-
-def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
-    return ARDiffusionKVCache(
-        ARDiffusionKVConfig(
-            enable=True,
-            chunk_size=1560,
-            window_chunks=9,
-            sink_chunks=9,
-            gpu_memory_fraction=gpu_memory_fraction,
-        ),
-        num_layers=40,
-        num_kv_heads=40,
-        head_size=128,
-        dtype=torch.bfloat16,
-        block_size=1560,
-        max_model_len=1 << 20,
-        available_bytes=100 * (1 << 30),
-        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
-        session_capacity=2,
-        cross_attention_lengths={"text": 512},
-        frames_per_block=3,
-        model_owned_state_bytes_per_session=7_488_000,
-    )
-
-
-def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
-    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
-
-    assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
-    assert kv.memory_budget_bytes == 33_653_670_400
-    assert kv.session_capacity == 1
-
-
-def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
-    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
-
-    assert kv.memory_budget_bytes == 60 * (1 << 30)
-    assert kv.session_capacity == 2
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
 
 
 def _make_kv(
@@ -359,24 +221,15 @@ def _make_kv(
         if local_branches == 1
         else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
     )
-    page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
-    declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
-    scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
-    managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
     return ARDiffusionKVCache(
-        ARDiffusionKVConfig(
-            enable=True,
-            chunk_size=BLOCK,
-            window_chunks=window_chunks,
-            gpu_memory_fraction=1.0,
-        ),
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
         num_layers=1,
         num_kv_heads=4,
         head_size=64,
         dtype=torch.float32,
         block_size=BLOCK,
         max_model_len=4096,
-        available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
+        available_bytes=1 << 16,  # tiny -> the floor binds
         device=torch.device("cpu"),
         kv_branches=kv_branches,
         session_capacity=1,

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -54,15 +54,13 @@ def test_full_request_lifecycle_and_eviction():
         used = {int(s) // kv.block_size for s in slots}
         assert kv.null_block_id not in used
         assert used <= set(block_table)
-        # Allocation may transiently add one in-flight chunk.
+        # Resident window stays bounded by window + the in-flight chunk.
         assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
-        kv.commit_chunk(adapter)
-        # A successful commit immediately prunes physical ownership.
-        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
         free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
 
     # Pool memory plateaus once the window is full (eviction recycles blocks).
-    assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
 
     kv.end_request(adapter)
     assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -332,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
 
 def test_prepare_is_idempotent_and_layers_share_metadata():
     device = torch.device("cpu")
-    kv, st = make_state(num_layers=2)
+    _, st = make_state(num_layers=2)
     contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
     fctx = contexts[0].forward_ctx
     fctx.prepare(device=device, action_len=0, query_len=BLOCK)
@@ -349,10 +349,6 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
     assert i0.block_table is i1.block_table
     assert i0.seq_lens is i1.seq_lens
     assert i0.video_slots is i1.video_slots
-    assert i0.key_pool is kv._k_pools[0]
-    assert i0.value_pool is kv._v_pools[0]
-    assert i1.key_pool is kv._k_pools[1]
-    assert i1.value_pool is kv._v_pools[1]
 
 
 def test_layer_inputs_before_prepare_raises():

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

@@ -174,19 +174,6 @@ def test_window_reset_can_keep_named_cross_attention_cache():
     assert not st.is_cross_attention_populated(POS, "text")
 
 
-def test_clear_cross_attention_preserves_self_attention_history():
-    kv, st = make_state(num_layers=1, cross_attn_length=8)
-    _prepare_and_commit(st, POS, 2)
-    k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
-    st.populate_cross_attention(POS, "text", [(k, k)])
-    resident_before = kv.window_block_ids(st.adapter(POS))
-
-    st.clear_cross_attention()
-
-    assert kv.window_block_ids(st.adapter(POS)) == resident_before
-    assert not st.is_cross_attention_populated(POS, "text")
-
-
 def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
     _, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
     assert st.kv_cache.num_local_kv_branches == 1

tests/diffusion/ar_diffusion/test_session.py

@@ -1,505 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import asyncio
-from typing import Any
-
-import pytest
-
-from vllm_omni.experimental.ar_diffusion.session import (
-    ARDiffusionChunkMetadataError,
-    ARDiffusionEventAcceptanceStatus,
-    ARDiffusionEventQueueFullError,
-    ARDiffusionPreparedControls,
-    ARDiffusionSession,
-    ARDiffusionSessionError,
-    ARDiffusionSessionEvent,
-    ARDiffusionSessionManager,
-    ARDiffusionSessionStateError,
-    ARDiffusionSessionStatus,
-    ARDiffusionStaleEventError,
-    ARDiffusionWorkerLifecycle,
-)
-from vllm_omni.experimental.ar_diffusion.tick_protocol import (
-    ARDiffusionChunkMetadata,
-    ARDiffusionControlInput,
-    ARDiffusionTickRequest,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeTickConsumer:
-    def __init__(self) -> None:
-        self.ticks: list[ARDiffusionTickRequest] = []
-        self.failures_remaining = 0
-        self.started: asyncio.Event | None = None
-        self.release: asyncio.Event | None = None
-        self.metadata_override: ARDiffusionChunkMetadata | None = None
-
-    async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
-        self.ticks.append(tick)
-        if self.started is not None:
-            self.started.set()
-        if self.release is not None:
-            await self.release.wait()
-        if self.failures_remaining:
-            self.failures_remaining -= 1
-            raise RuntimeError("fake tick failed")
-        return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
-
-    def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
-        return output
-
-
-class FakeLifecycle:
-    def __init__(self) -> None:
-        self.resets: list[str] = []
-        self.closes: list[str] = []
-        self.close_failures_remaining = 0
-
-    async def reset_session(self, session_id: str) -> None:
-        self.resets.append(session_id)
-
-    async def close_session(self, session_id: str) -> None:
-        self.closes.append(session_id)
-        if self.close_failures_remaining:
-            self.close_failures_remaining -= 1
-            raise RuntimeError("fake close failed")
-
-
-class FailingCommitReducer:
-    def __init__(self) -> None:
-        self.commit_calls = 0
-        self.reset_calls = 0
-
-    def prepare(
-        self,
-        *,
-        current_controls: dict[str, ARDiffusionControlInput],
-        events: tuple[ARDiffusionSessionEvent, ...],
-        chunk_index: int,
-    ) -> ARDiffusionPreparedControls:
-        del current_controls, events, chunk_index
-        return ARDiffusionPreparedControls(controls=())
-
-    def commit(self, prepared: ARDiffusionPreparedControls) -> None:
-        del prepared
-        self.commit_calls += 1
-        raise RuntimeError("fake reducer commit failed")
-
-    def reset(self) -> None:
-        self.reset_calls += 1
-
-
-def make_session(
-    *,
-    max_pending_events: int = 8,
-    consumer: FakeTickConsumer | None = None,
-    lifecycle: FakeLifecycle | None = None,
-) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
-    resolved_consumer = consumer or FakeTickConsumer()
-    resolved_lifecycle = lifecycle or FakeLifecycle()
-    session = ARDiffusionSession(
-        "world-1",
-        tick_consumer=resolved_consumer,
-        lifecycle=resolved_lifecycle,
-        max_pending_events=max_pending_events,
-        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
-    )
-    return session, resolved_consumer, resolved_lifecycle
-
-
-@pytest.mark.asyncio
-async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
-    session, consumer, _ = make_session()
-    event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
-    event_one = ARDiffusionSessionEvent(
-        event_id=1,
-        controls=(
-            ARDiffusionControlInput(
-                track="camera",
-                schema="example.camera.v1",
-                data={"samples": [{"x": 1.0}]},
-            ),
-        ),
-    )
-
-    accepted_two = await session.accept_event(event_two)
-    accepted_one = await session.accept_event(event_one)
-    duplicate = await session.accept_event(event_two)
-
-    assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
-    assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
-    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
-    assert consumer.ticks == []
-
-    metadata = await session.next_chunk()
-
-    assert metadata.applied_event_ids == (1, 2)
-    assert consumer.ticks[0] == ARDiffusionTickRequest(
-        session_id="world-1",
-        request_id="world-1-request-0",
-        chunk_index=0,
-        applied_event_ids=(1, 2),
-        prompt="turn",
-        controls=(
-            ARDiffusionControlInput(
-                track="camera",
-                schema="example.camera.v1",
-                data={"samples": [{"x": 1.0}]},
-            ),
-        ),
-    )
-    assert session.pending_event_count == 0
-
-    applied_duplicate = await session.accept_event(event_two)
-    assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
-
-    next_metadata = await session.next_chunk()
-    assert next_metadata.chunk_index == 1
-    assert next_metadata.applied_event_ids == ()
-    assert consumer.ticks[1].prompt == "turn"
-    assert consumer.ticks[1].controls == consumer.ticks[0].controls
-
-
-@pytest.mark.asyncio
-async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
-    session, consumer, _ = make_session(max_pending_events=1)
-    payload = {"samples": [{"x": 1.0}]}
-    composite = ARDiffusionSessionEvent(
-        event_id=5,
-        prompt="new prompt",
-        controls=(
-            ARDiffusionControlInput(
-                track="camera",
-                schema="example.camera.v1",
-                data=payload,
-            ),
-            ARDiffusionControlInput(
-                track="action",
-                schema="example.action.v1",
-                data={"keys": ["left"]},
-            ),
-        ),
-    )
-    await session.accept_event(composite)
-    payload["samples"][0]["x"] = 99.0
-
-    duplicate = await session.accept_event(composite)
-    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
-    with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
-        await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
-
-    await session.next_chunk()
-    assert consumer.ticks[0].applied_event_ids == (5,)
-    assert consumer.ticks[0].prompt == "new prompt"
-    assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
-    assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
-    assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
-    assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
-
-
-def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
-    with pytest.raises(ValueError, match="at most once"):
-        ARDiffusionSessionEvent(
-            event_id=1,
-            controls=(
-                ARDiffusionControlInput(track="camera", schema="one", data={}),
-                ARDiffusionControlInput(track="camera", schema="two", data={}),
-            ),
-        )
-
-    with pytest.raises(ValueError, match="transport-safe"):
-        ARDiffusionControlInput(
-            track="camera",
-            schema="example.camera.v1",
-            data={"object": object()},
-        )
-
-
-@pytest.mark.asyncio
-async def test_failed_tick_closes_worker_and_requires_reset() -> None:
-    consumer = FakeTickConsumer()
-    consumer.failures_remaining = 1
-    lifecycle = FakeLifecycle()
-    session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
-    await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
-
-    with pytest.raises(RuntimeError, match="fake tick failed"):
-        await session.next_chunk()
-
-    assert session.status is ARDiffusionSessionStatus.FAILED
-    assert session.chunk_index == 0
-    assert session.pending_event_count == 1
-    assert lifecycle.closes == ["world-1"]
-    with pytest.raises(ARDiffusionSessionStateError, match="failed"):
-        await session.next_chunk()
-
-    await session.reset()
-    assert session.status is ARDiffusionSessionStatus.ACTIVE
-    assert session.pending_event_count == 0
-    assert lifecycle.resets == ["world-1"]
-    await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
-    metadata = await session.next_chunk()
-    assert metadata.chunk_index == 0
-    assert metadata.applied_event_ids == (4,)
-    assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
-
-
-@pytest.mark.asyncio
-async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
-    consumer = FakeTickConsumer()
-    consumer.started = asyncio.Event()
-    consumer.release = asyncio.Event()
-    session, _, _ = make_session(consumer=consumer)
-    await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
-
-    first_chunk = asyncio.create_task(session.next_chunk())
-    await consumer.started.wait()
-    accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
-    assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
-    assert consumer.ticks[0].prompt == "first"
-    assert consumer.ticks[0].applied_event_ids == (10,)
-
-    with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
-        await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
-
-    consumer.release.set()
-    first_metadata = await first_chunk
-    second_metadata = await session.next_chunk()
-    assert first_metadata.applied_event_ids == (10,)
-    assert second_metadata.applied_event_ids == (11,)
-    assert consumer.ticks[1].prompt == "second"
-
-
-@pytest.mark.asyncio
-async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
-    consumer = FakeTickConsumer()
-    consumer.metadata_override = ARDiffusionChunkMetadata(
-        session_id="another-world",
-        request_id="wrong",
-        chunk_index=0,
-    )
-    lifecycle = FakeLifecycle()
-    session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
-    await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
-
-    with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
-        await session.next_chunk()
-
-    assert session.chunk_index == 0
-    assert session.pending_event_count == 1
-    assert session.status is ARDiffusionSessionStatus.FAILED
-    assert lifecycle.closes == ["world-1"]
-
-
-@pytest.mark.asyncio
-async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
-    consumer = FakeTickConsumer()
-    lifecycle = FakeLifecycle()
-    reducer = FailingCommitReducer()
-    session = ARDiffusionSession(
-        "world-1",
-        tick_consumer=consumer,
-        lifecycle=lifecycle,
-        max_pending_events=4,
-        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
-        control_reducer=reducer,
-    )
-    await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
-
-    with pytest.raises(RuntimeError, match="fake reducer commit failed"):
-        await session.next_chunk()
-
-    assert reducer.commit_calls == 1
-    assert session.status is ARDiffusionSessionStatus.FAILED
-    assert session.chunk_index == 0
-    assert session.pending_event_count == 1
-    assert lifecycle.closes == ["world-1"]
-    with pytest.raises(ARDiffusionSessionStateError, match="failed"):
-        await session.next_chunk()
-
-    await session.reset()
-    assert reducer.reset_calls == 1
-    assert session.status is ARDiffusionSessionStatus.ACTIVE
-    assert session.chunk_index == 0
-    assert session.pending_event_count == 0
-
-
-@pytest.mark.asyncio
-async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
-    consumer = FakeTickConsumer()
-    lifecycle = FakeLifecycle()
-    manager = ARDiffusionSessionManager(
-        tick_consumer=consumer,
-        lifecycle=lifecycle,
-        max_pending_events=4,
-        request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
-    )
-    session = await manager.create_session("world-reset")
-    await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
-    await session.next_chunk()
-    await manager.reset_session("world-reset")
-
-    assert lifecycle.resets == ["world-reset"]
-    assert session.status is ARDiffusionSessionStatus.ACTIVE
-    assert session.chunk_index == 0
-    duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
-    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
-
-    reset_tick = await session.next_chunk()
-    assert reset_tick.chunk_index == 0
-    assert consumer.ticks[-1].prompt is None
-    assert consumer.ticks[-1].controls == ()
-
-    await manager.close_session("world-reset")
-    assert lifecycle.closes == ["world-reset"]
-    assert session.status is ARDiffusionSessionStatus.CLOSED
-    with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
-        await manager.get_session("world-reset")
-
-    disconnected = await manager.create_session("world-disconnect")
-    await manager.disconnect("world-disconnect")
-    assert disconnected.status is ARDiffusionSessionStatus.CLOSED
-    assert lifecycle.closes[-1] == "world-disconnect"
-
-
-@pytest.mark.asyncio
-async def test_manager_keeps_interleaved_sessions_independent() -> None:
-    consumer = FakeTickConsumer()
-    manager = ARDiffusionSessionManager(
-        tick_consumer=consumer,
-        lifecycle=FakeLifecycle(),
-        max_pending_events=4,
-        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
-    )
-    session_a = await manager.create_session("world-a")
-    session_b = await manager.create_session("world-b")
-
-    await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
-    await session_a.next_chunk()
-    await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
-    await session_b.next_chunk()
-    await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
-    await session_a.next_chunk()
-
-    assert [
-        (
-            tick.session_id,
-            tick.request_id,
-            tick.chunk_index,
-            tick.applied_event_ids,
-        )
-        for tick in consumer.ticks
-    ] == [
-        ("world-a", "world-a-request-0", 0, (10,)),
-        ("world-b", "world-b-request-0", 0, (1,)),
-        ("world-a", "world-a-request-1", 1, (11,)),
-    ]
-    assert session_a.chunk_index == 2
-    assert session_b.chunk_index == 1
-    with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
-        await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
-    assert (
-        await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
-    ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
-
-
-@pytest.mark.asyncio
-async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
-    consumer = FakeTickConsumer()
-    lifecycle = FakeLifecycle()
-    lifecycle.close_failures_remaining = 1
-    manager = ARDiffusionSessionManager(
-        tick_consumer=consumer,
-        lifecycle=lifecycle,
-        max_pending_events=4,
-    )
-    session = await manager.create_session("world-close-retry")
-
-    with pytest.raises(RuntimeError, match="fake close failed"):
-        await manager.close_session("world-close-retry")
-
-    assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
-    assert await manager.get_session("world-close-retry") is session
-    with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
-        await manager.create_session("world-close-retry")
-    with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
-        await manager.reset_session("world-close-retry")
-
-    await manager.close_session("world-close-retry")
-
-    assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
-    assert session.status is ARDiffusionSessionStatus.CLOSED
-    with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
-        await manager.get_session("world-close-retry")
-    recreated = await manager.create_session("world-close-retry")
-    assert recreated is not session
-
-
-@pytest.mark.asyncio
-async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
-    lifecycle = FakeLifecycle()
-    lifecycle.close_failures_remaining = 1
-    manager = ARDiffusionSessionManager(
-        tick_consumer=FakeTickConsumer(),
-        lifecycle=lifecycle,
-        max_pending_events=4,
-    )
-    session = await manager.create_session("world-disconnect-retry")
-
-    with pytest.raises(RuntimeError, match="fake close failed"):
-        await manager.disconnect("world-disconnect-retry")
-
-    assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
-    assert await manager.get_session("world-disconnect-retry") is session
-
-    await manager.close_session("world-disconnect-retry")
-    assert session.status is ARDiffusionSessionStatus.CLOSED
-
-
-class FakeRPCClient:
-    def __init__(self, result: list[Any]) -> None:
-        self.result = result
-        self.calls: list[dict[str, Any]] = []
-
-    async def collective_rpc(self, **kwargs: Any) -> list[Any]:
-        self.calls.append(kwargs)
-        return self.result
-
-
-@pytest.mark.asyncio
-async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
-    client = FakeRPCClient([[True]])
-    lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
-
-    await lifecycle.reset_session("world-1")
-    await lifecycle.close_session("world-1")
-
-    assert client.calls == [
-        {
-            "method": "reset_ar_diffusion_session",
-            "timeout": 3.0,
-            "args": ("world-1",),
-            "stage_ids": [2],
-        },
-        {
-            "method": "close_ar_diffusion_session",
-            "timeout": 3.0,
-            "args": ("world-1",),
-            "stage_ids": [2],
-        },
-    ]
-
-
-@pytest.mark.asyncio
-async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
-    lifecycle = ARDiffusionWorkerLifecycle(
-        FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
-        stage_ids=[1],
-    )
-
-    with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
-        await lifecycle.close_session("world-1")

tests/diffusion/ar_diffusion/test_tick_protocol.py

@@ -1,108 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-
-from vllm_omni.experimental.ar_diffusion.tick_protocol import (
-    ARDiffusionChunkMetadata,
-    ARDiffusionControlInput,
-    ARDiffusionTickRequest,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
-    tick = ARDiffusionTickRequest(
-        session_id="world-1",
-        request_id="chunk-request-7",
-        chunk_index=7,
-        applied_event_ids=(10, 11),
-        prompt="turn left",
-        controls=(
-            ARDiffusionControlInput(
-                track="camera",
-                schema="lingbot.camera_trajectory.v1",
-                data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
-            ),
-        ),
-    )
-
-    restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
-
-    assert restored == tick
-    assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
-        "session_id": "world-1",
-        "request_id": "chunk-request-7",
-        "chunk_index": 7,
-        "applied_event_ids": [10, 11],
-    }
-
-
-@pytest.mark.parametrize(
-    ("field", "value", "message"),
-    [
-        ("session_id", "", "session_id"),
-        ("request_id", "", "request_id"),
-        ("chunk_index", -1, "chunk_index"),
-        ("applied_event_ids", (2, 1), "strictly increasing"),
-        ("applied_event_ids", (1, 1), "strictly increasing"),
-    ],
-)
-def test_tick_rejects_invalid_identity(
-    field: str,
-    value: object,
-    message: str,
-) -> None:
-    values = {
-        "session_id": "world-1",
-        "request_id": "request-1",
-        "chunk_index": 0,
-        "applied_event_ids": (),
-    }
-    values[field] = value
-
-    with pytest.raises(ValueError, match=message):
-        ARDiffusionTickRequest(**values)
-
-
-def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
-    payload = {"samples": [{"keys": ["w"]}]}
-    event_ids = [1, 2]
-    controls = [
-        ARDiffusionControlInput(
-            track="camera",
-            schema="example.camera.v1",
-            data=payload,
-        )
-    ]
-    tick = ARDiffusionTickRequest(
-        session_id="world-1",
-        request_id="request-a",
-        chunk_index=0,
-        applied_event_ids=event_ids,  # type: ignore[arg-type]
-        controls=controls,  # type: ignore[arg-type]
-    )
-    serialized = tick.to_extra_args()
-
-    payload["samples"][0]["keys"].append("j")
-    event_ids.append(3)
-    controls.clear()
-    serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
-
-    assert tick.applied_event_ids == (1, 2)
-    assert isinstance(tick.controls, tuple)
-    assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
-        {
-            "track": "camera",
-            "schema": "example.camera.v1",
-            "data": {"samples": [{"keys": ["w"]}]},
-        }
-    ]
-    with pytest.raises(TypeError):
-        tick.controls[0].data["samples"] = ()  # type: ignore[index]
-    with pytest.raises(AttributeError):
-        tick.controls[0].data["samples"][0]["keys"].append("j")
-
-
-def test_legacy_request_has_no_typed_tick() -> None:
-    assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None

tests/diffusion/ar_diffusion/test_worker_lifecycle.py

@@ -1,53 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-
-from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeARDiffusionRunner:
-    def __init__(self) -> None:
-        self.resets: list[str] = []
-        self.closes: list[str] = []
-
-    def reset_session(self, session_id: str) -> None:
-        self.resets.append(session_id)
-
-    def close_session(self, session_id: str) -> None:
-        self.closes.append(session_id)
-
-
-class FakeRegularRunner:
-    pass
-
-
-def make_worker(runner: object) -> DiffusionWorker:
-    worker = object.__new__(DiffusionWorker)
-    worker.model_runner = runner
-    return worker
-
-
-def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
-    runner = FakeARDiffusionRunner()
-    worker = make_worker(runner)
-
-    assert worker.reset_ar_diffusion_session("world-1") is True
-    assert worker.close_ar_diffusion_session("world-1") is True
-    assert runner.resets == ["world-1"]
-    assert runner.closes == ["world-1"]
-
-
-def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
-    worker = make_worker(FakeRegularRunner())
-
-    assert worker.reset_ar_diffusion_session("world-1") is False
-    assert worker.close_ar_diffusion_session("world-1") is False
-
-
-def test_worker_rejects_empty_session_id() -> None:
-    worker = make_worker(FakeARDiffusionRunner())
-
-    with pytest.raises(ValueError, match="session_id"):
-        worker.close_ar_diffusion_session("")

tests/diffusion/models/dreamzero/test_pipeline_state.py

@@ -6,10 +6,7 @@ from types import SimpleNamespace
 
 import pytest
 
-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
-    DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
-    DreamZeroPipeline,
-)
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -35,37 +32,6 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
     assert session_b.call_count == 3
 
 
-def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
-    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
-    pipeline.transformer = SimpleNamespace(
-        frame_seqlen=16,
-        blocks=[
-            SimpleNamespace(
-                self_attn=SimpleNamespace(
-                    max_attention_size=96,
-                    tp_num_heads=4,
-                )
-            )
-        ],
-        text_len=8,
-        model_type="t2v",
-        num_layers=2,
-        num_heads=4,
-        dim=256,
-        num_frame_per_block=4,
-        num_action_per_block=2,
-        num_state_per_block=1,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
-        lambda: 1,
-    )
-
-    spec = pipeline.ar_diffusion_kv_cache_spec()
-
-    assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
-
-
 def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
     pipeline = _empty_pipeline()
 

tests/diffusion/models/lingbot_world/__init__.py

@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture

@@ -1,13 +0,0 @@
-{
-  "_class_name": "LingBotWorldCausalDMDPipeline",
-  "_diffusers_version": "0.35.0.dev0",
-  "boundary_ratio": 0.9,
-  "image_encoder": [null, null],
-  "image_processor": [null, null],
-  "scheduler": ["diffusers", "UniPCMultistepScheduler"],
-  "text_encoder": ["transformers", "UMT5EncoderModel"],
-  "tokenizer": ["transformers", "T5TokenizerFast"],
-  "transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
-  "transformer_2": [null, null],
-  "vae": ["diffusers", "AutoencoderKLWan"]
-}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture

@@ -1,42 +0,0 @@
-{
-  "sources": {
-    "config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
-    "shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
-    "shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
-  },
-  "dtype": "F32",
-  "shard_representatives": {
-    "diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
-    "diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
-    "diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
-  },
-  "parameter_shapes": {
-    "blocks.0.cam_injector_layer1.bias": [5120],
-    "blocks.4.cam_injector_layer1.bias": [5120],
-    "blocks.10.cam_injector_layer1.bias": [5120],
-    "blocks.14.cam_scale_layer.bias": [5120],
-    "blocks.20.cam_injector_layer1.bias": [5120],
-    "blocks.25.cam_injector_layer1.bias": [5120],
-    "blocks.30.cam_injector_layer1.bias": [5120],
-    "blocks.35.cam_injector_layer1.bias": [5120],
-    "patch_embedding.weight": [5120, 36, 1, 2, 2],
-    "patch_embedding_wancamctrl.weight": [5120, 1536],
-    "c2ws_hidden_states_layer1.weight": [5120, 5120],
-    "c2ws_hidden_states_layer2.weight": [5120, 5120],
-    "blocks.0.modulation": [1, 6, 5120],
-    "blocks.0.self_attn.q.weight": [5120, 5120],
-    "blocks.0.ffn.0.weight": [13824, 5120],
-    "blocks.0.ffn.2.weight": [5120, 13824],
-    "blocks.0.cam_injector_layer1.weight": [5120, 5120],
-    "blocks.0.cam_injector_layer2.weight": [5120, 5120],
-    "blocks.0.cam_scale_layer.weight": [5120, 5120],
-    "blocks.0.cam_shift_layer.weight": [5120, 5120],
-    "head.modulation": [1, 2, 5120],
-    "head.head.weight": [64, 5120]
-  }
-}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture

@@ -1,30 +0,0 @@
-{
-  "_class_name": "UniPCMultistepScheduler",
-  "_diffusers_version": "0.35.0.dev0",
-  "beta_end": 0.02,
-  "beta_schedule": "linear",
-  "beta_start": 0.0001,
-  "disable_corrector": [],
-  "dynamic_thresholding_ratio": 0.995,
-  "final_sigmas_type": "zero",
-  "flow_shift": 5.0,
-  "lower_order_final": true,
-  "num_train_timesteps": 1000,
-  "predict_x0": true,
-  "prediction_type": "flow_prediction",
-  "rescale_betas_zero_snr": false,
-  "sample_max_value": 1.0,
-  "solver_order": 2,
-  "solver_p": null,
-  "solver_type": "bh2",
-  "steps_offset": 0,
-  "thresholding": false,
-  "time_shift_type": "exponential",
-  "timestep_spacing": "linspace",
-  "trained_betas": null,
-  "use_beta_sigmas": false,
-  "use_dynamic_shifting": false,
-  "use_exponential_sigmas": false,
-  "use_flow_sigmas": true,
-  "use_karras_sigmas": false
-}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture

@@ -1,112 +0,0 @@
-{
-  "source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
-  "index_parameter_count": 1421,
-  "canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
-  "canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
-  "canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
-  "shards": {
-    "diffusion_pytorch_model-00001-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
-    },
-    "diffusion_pytorch_model-00002-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
-    },
-    "diffusion_pytorch_model-00003-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
-    },
-    "diffusion_pytorch_model-00004-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
-    },
-    "diffusion_pytorch_model-00005-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
-    },
-    "diffusion_pytorch_model-00006-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
-    },
-    "diffusion_pytorch_model-00007-of-00008.safetensors": {
-      "parameter_count": 180,
-      "parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
-    },
-    "diffusion_pytorch_model-00008-of-00008.safetensors": {
-      "parameter_count": 161,
-      "parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
-    }
-  },
-  "top_level_modules": [
-    "blocks",
-    "c2ws_hidden_states_layer1",
-    "c2ws_hidden_states_layer2",
-    "head",
-    "patch_embedding",
-    "patch_embedding_wancamctrl",
-    "text_embedding",
-    "time_embedding",
-    "time_projection"
-  ],
-  "non_block_parameters": [
-    "c2ws_hidden_states_layer1.bias",
-    "c2ws_hidden_states_layer1.weight",
-    "c2ws_hidden_states_layer2.bias",
-    "c2ws_hidden_states_layer2.weight",
-    "head.head.bias",
-    "head.head.weight",
-    "head.modulation",
-    "patch_embedding.bias",
-    "patch_embedding.weight",
-    "patch_embedding_wancamctrl.bias",
-    "patch_embedding_wancamctrl.weight",
-    "text_embedding.0.bias",
-    "text_embedding.0.weight",
-    "text_embedding.2.bias",
-    "text_embedding.2.weight",
-    "time_embedding.0.bias",
-    "time_embedding.0.weight",
-    "time_embedding.2.bias",
-    "time_embedding.2.weight",
-    "time_projection.1.bias",
-    "time_projection.1.weight"
-  ],
-  "block_parameter_suffixes": [
-    "cam_injector_layer1.bias",
-    "cam_injector_layer1.weight",
-    "cam_injector_layer2.bias",
-    "cam_injector_layer2.weight",
-    "cam_scale_layer.bias",
-    "cam_scale_layer.weight",
-    "cam_shift_layer.bias",
-    "cam_shift_layer.weight",
-    "cross_attn.k.bias",
-    "cross_attn.k.weight",
-    "cross_attn.norm_k.weight",
-    "cross_attn.norm_q.weight",
-    "cross_attn.o.bias",
-    "cross_attn.o.weight",
-    "cross_attn.q.bias",
-    "cross_attn.q.weight",
-    "cross_attn.v.bias",
-    "cross_attn.v.weight",
-    "ffn.0.bias",
-    "ffn.0.weight",
-    "ffn.2.bias",
-    "ffn.2.weight",
-    "modulation",
-    "norm3.bias",
-    "norm3.weight",
-    "self_attn.k.bias",
-    "self_attn.k.weight",
-    "self_attn.norm_k.weight",
-    "self_attn.norm_q.weight",
-    "self_attn.o.bias",
-    "self_attn.o.weight",
-    "self_attn.q.bias",
-    "self_attn.q.weight",
-    "self_attn.v.bias",
-    "self_attn.v.weight"
-  ]
-}

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -1,294 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import asyncio
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.lingbot_world.actions import (
-    LINGBOT_CAMERA_ACTION_SCHEMA,
-    LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
-    LingBotCameraControlReducer,
-    integrate_lingbot_camera_actions,
-    parse_lingbot_camera_action_frames,
-)
-from vllm_omni.experimental.ar_diffusion.session import (
-    ARDiffusionSession,
-    ARDiffusionSessionEvent,
-)
-from vllm_omni.experimental.ar_diffusion.tick_protocol import (
-    ARDiffusionChunkMetadata,
-    ARDiffusionControlInput,
-    ARDiffusionTickRequest,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _action_control(mode: str, **data) -> ARDiffusionControlInput:
-    return ARDiffusionControlInput(
-        track="camera",
-        schema=LINGBOT_CAMERA_ACTION_SCHEMA,
-        data={"mode": mode, **data},
-    )
-
-
-def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
-    assert len(prepared.controls) == 1
-    return prepared.controls[0].data["frames"]
-
-
-def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
-    reducer = LingBotCameraControlReducer()
-    pressed = reducer.prepare(
-        current_controls={},
-        events=(
-            ARDiffusionSessionEvent(
-                event_id=1,
-                controls=(
-                    _action_control(
-                        "state",
-                        transitions=[{"actions": ["W"], "client_ts_ms": 10}],
-                    ),
-                ),
-            ),
-        ),
-        chunk_index=0,
-    )
-    assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
-    assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
-    reducer.commit(pressed)
-
-    held = reducer.prepare(
-        current_controls={"camera": pressed.controls[0]},
-        events=(),
-        chunk_index=1,
-    )
-    assert _prepared_frames(held) == (("w",), ("w",), ("w",))
-    reducer.commit(held)
-
-    pulse = reducer.prepare(
-        current_controls={"camera": held.controls[0]},
-        events=(
-            ARDiffusionSessionEvent(
-                event_id=2,
-                controls=(
-                    _action_control(
-                        "state",
-                        transitions=[
-                            {"actions": ["j"], "client_ts_ms": 20},
-                            {"actions": [], "client_ts_ms": 21},
-                        ],
-                    ),
-                ),
-            ),
-        ),
-        chunk_index=2,
-    )
-    assert _prepared_frames(pulse) == (("j",), (), ())
-    reducer.commit(pulse)
-
-    released = reducer.prepare(
-        current_controls={"camera": pulse.controls[0]},
-        events=(),
-        chunk_index=3,
-    )
-    assert _prepared_frames(released) == ((), (), ())
-
-
-def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
-    reducer = LingBotCameraControlReducer()
-    event = ARDiffusionSessionEvent(
-        event_id=1,
-        controls=(
-            _action_control(
-                "script",
-                frames=[["w"], ["d"], ["l"], ["i"]],
-            ),
-        ),
-    )
-    first_attempt = reducer.prepare(
-        current_controls={},
-        events=(event,),
-        chunk_index=0,
-    )
-    assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
-
-    # An uncommitted failed attempt must not consume the reducer's script.
-    second_attempt = reducer.prepare(
-        current_controls={},
-        events=(event,),
-        chunk_index=0,
-    )
-    assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
-    reducer.commit(second_attempt)
-
-    tail = reducer.prepare(
-        current_controls={"camera": second_attempt.controls[0]},
-        events=(),
-        chunk_index=1,
-    )
-    assert _prepared_frames(tail) == (("i",), (), ())
-    reducer.commit(tail)
-    exhausted = reducer.prepare(
-        current_controls={"camera": tail.controls[0]},
-        events=(),
-        chunk_index=2,
-    )
-    assert _prepared_frames(exhausted) == ((), (), ())
-
-
-def test_trajectory_control_remains_an_expert_override() -> None:
-    reducer = LingBotCameraControlReducer()
-    trajectory = ARDiffusionControlInput(
-        track="camera",
-        schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
-        data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
-    )
-    prepared = reducer.prepare(
-        current_controls={},
-        events=(
-            ARDiffusionSessionEvent(
-                event_id=1,
-                controls=(trajectory,),
-            ),
-        ),
-        chunk_index=0,
-    )
-    assert prepared.controls == (trajectory,)
-
-
-def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
-    trajectory, pitch = integrate_lingbot_camera_actions(
-        [["w"], ["d"], ["l"]],
-        width=832,
-        height=480,
-    )
-
-    assert pitch == 0.0
-    torch.testing.assert_close(
-        trajectory.poses[0, :3, 3],
-        torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
-        atol=1e-6,
-        rtol=0,
-    )
-    torch.testing.assert_close(
-        trajectory.poses[1, :3, 3],
-        torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
-        atol=1e-6,
-        rtol=0,
-    )
-    expected_yaw = torch.tensor(
-        [
-            [0.9945219, 0.0, 0.10452846],
-            [0.0, 1.0, 0.0],
-            [-0.10452846, 0.0, 0.9945219],
-        ],
-        dtype=torch.float64,
-    )
-    torch.testing.assert_close(
-        trajectory.poses[2, :3, :3],
-        expected_yaw,
-        atol=1e-6,
-        rtol=0,
-    )
-    torch.testing.assert_close(
-        trajectory.intrinsics[0],
-        torch.tensor([500.0, 500.0, 416.0, 240.0]),
-    )
-
-    resized, _ = integrate_lingbot_camera_actions(
-        [[]],
-        width=416,
-        height=240,
-    )
-    scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
-    torch.testing.assert_close(
-        scaled,
-        torch.tensor([500.0, 500.0, 208.0, 120.0]),
-    )
-
-
-@pytest.mark.parametrize(
-    ("data", "message"),
-    [
-        ({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
-        ({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
-        ({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
-    ],
-)
-def test_model_tick_action_payload_validation(data, message: str) -> None:
-    with pytest.raises(ValueError, match=message):
-        parse_lingbot_camera_action_frames(data, expected_frames=3)
-
-
-class _TickConsumer:
-    def __init__(self) -> None:
-        self.ticks: list[ARDiffusionTickRequest] = []
-
-    async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
-        self.ticks.append(tick)
-        return ARDiffusionChunkMetadata.from_tick(tick)
-
-    def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
-        return output
-
-
-class _Lifecycle:
-    async def reset_session(self, session_id: str) -> None:
-        del session_id
-
-    async def close_session(self, session_id: str) -> None:
-        del session_id
-
-
-def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
-    async def exercise() -> None:
-        consumer = _TickConsumer()
-        reducer = LingBotCameraControlReducer()
-        session = ARDiffusionSession(
-            "world-actions",
-            tick_consumer=consumer,
-            lifecycle=_Lifecycle(),
-            max_pending_events=8,
-            request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
-            control_reducer=reducer,
-        )
-        await session.accept_event(
-            ARDiffusionSessionEvent(
-                event_id=1,
-                prompt="turn toward the window",
-                controls=(
-                    _action_control(
-                        "state",
-                        transitions=[{"actions": ["d"]}],
-                    ),
-                ),
-            )
-        )
-        await session.next_chunk()
-
-        tick = consumer.ticks[0]
-        assert tick.prompt == "turn toward the window"
-        assert tick.applied_event_ids == (1,)
-        assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
-
-        await session.reset()
-        await session.accept_event(
-            ARDiffusionSessionEvent(
-                event_id=2,
-                controls=(
-                    _action_control(
-                        "state",
-                        transitions=[{"actions": []}],
-                    ),
-                ),
-            )
-        )
-        await session.next_chunk()
-        assert consumer.ticks[1].chunk_index == 0
-        assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
-
-    asyncio.run(exercise())

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -1,543 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import sys
-from pathlib import Path
-from types import ModuleType
-
-import pytest
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
-_MODULE_NAME = "_lingbot_world_attention_under_test"
-_MISSING_MODULE = object()
-_STUBBED_MODULE_NAMES = (
-    "vllm",
-    "vllm.distributed",
-    "vllm.model_executor",
-    "vllm.model_executor.layers",
-    "vllm.model_executor.layers.conv",
-    "vllm.model_executor.layers.linear",
-    "vllm.model_executor.model_loader",
-    "vllm.model_executor.model_loader.weight_utils",
-    "vllm.model_executor.utils",
-    "vllm_omni",
-    "vllm_omni.diffusion",
-    "vllm_omni.diffusion.attention",
-    "vllm_omni.diffusion.attention.layer",
-    "vllm_omni.diffusion.layers",
-    "vllm_omni.diffusion.layers.norm",
-    "vllm_omni.diffusion.layers.rope",
-    "vllm_omni.experimental",
-    "vllm_omni.experimental.ar_diffusion",
-    "vllm_omni.experimental.ar_diffusion.kv_cache",
-    "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
-)
-
-
-@pytest.fixture(autouse=True)
-def _inference_context():
-    with torch.inference_mode():
-        yield
-
-
-def _install_vllm_stubs() -> None:
-    for name in _STUBBED_MODULE_NAMES:
-        sys.modules[name] = ModuleType(name)
-
-    distributed = sys.modules["vllm.distributed"]
-    distributed.get_tensor_model_parallel_rank = lambda: 0
-    distributed.get_tensor_model_parallel_world_size = lambda: 1
-    distributed.tensor_model_parallel_all_reduce = lambda value: value
-
-    def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
-        for name, value in attrs.items():
-            setattr(weight, name, value)
-
-    sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
-    sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
-        param.data.copy_(loaded_weight)
-    )
-
-    class _Linear(nn.Module):
-        def __init__(
-            self,
-            input_size: int,
-            output_size: int,
-            *,
-            bias: bool = True,
-            return_bias: bool = False,
-            **kwargs,
-        ) -> None:
-            super().__init__()
-            del kwargs
-            self.return_bias = return_bias
-            self.weight = nn.Parameter(torch.empty(output_size, input_size))
-            self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
-            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-            if self.bias is not None:
-                nn.init.zeros_(self.bias)
-            self.calls = 0
-
-        def forward(self, value: torch.Tensor):
-            self.calls += 1
-            output = F.linear(value, self.weight, self.bias)
-            return (output, self.bias) if self.return_bias else output
-
-    linear = sys.modules["vllm.model_executor.layers.linear"]
-    linear.ColumnParallelLinear = _Linear
-    linear.RowParallelLinear = _Linear
-
-    class _QKVLinear(nn.Module):
-        def __init__(
-            self,
-            hidden_size: int,
-            head_size: int,
-            total_num_heads: int,
-            total_num_kv_heads: int | None = None,
-            *,
-            bias: bool = True,
-            **kwargs,
-        ) -> None:
-            super().__init__()
-            del kwargs
-            self.num_heads = total_num_heads
-            self.num_kv_heads = total_num_kv_heads or total_num_heads
-            output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
-            self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
-            self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
-            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-            if self.bias is not None:
-                nn.init.zeros_(self.bias)
-            self.calls = 0
-
-            def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
-                shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
-                shard_size = param.shape[0] // 3
-                param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
-
-            self.weight.weight_loader = load_shard
-            if self.bias is not None:
-                self.bias.weight_loader = load_shard
-
-        def forward(self, value: torch.Tensor):
-            self.calls += 1
-            return F.linear(value, self.weight, self.bias), None
-
-    linear.QKVParallelLinear = _QKVLinear
-
-    class _Conv3dLayer(nn.Conv3d):
-        def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
-            super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
-
-    sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
-
-    class _LayerNorm(nn.LayerNorm):
-        def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
-            super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
-
-    sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
-
-    class _Attention(nn.Module):
-        def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
-            super().__init__()
-            del args
-            self.softmax_scale = softmax_scale
-            self.causal = kwargs["causal"]
-            self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
-
-        def forward(
-            self,
-            query: torch.Tensor,
-            key: torch.Tensor,
-            value: torch.Tensor,
-            attn_metadata=None,
-        ) -> torch.Tensor:
-            del attn_metadata
-            self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
-            scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
-            weights = scores.softmax(dim=-1)
-            return torch.einsum("bhqk,bkhd->bqhd", weights, value)
-
-    sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
-
-    class _RotaryEmbeddingWan(nn.Module):
-        def __init__(self, *args, **kwargs) -> None:
-            super().__init__()
-            del args, kwargs
-            self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
-
-        def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-            self.calls.append((cos.detach().clone(), sin.detach().clone()))
-            return value
-
-    sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
-
-    paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
-
-    class _PagedLayerContext:
-        pass
-
-    class _PagedLayerInputs:
-        pass
-
-    paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
-    paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
-    paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
-        AssertionError("CUDA paged attention was not expected in this CPU test")
-    )
-    paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
-        AssertionError("paged path was not expected in this direct-cache test")
-    )
-
-
-def _load_module():
-    assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
-    module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
-    previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
-    try:
-        _install_vllm_stubs()
-        spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
-        assert spec is not None and spec.loader is not None
-        module = importlib.util.module_from_spec(spec)
-        sys.modules[spec.name] = module
-        spec.loader.exec_module(module)
-        return module
-    finally:
-        for name, original in previous.items():
-            if original is _MISSING_MODULE:
-                sys.modules.pop(name, None)
-            else:
-                sys.modules[name] = original
-
-
-def test_load_module_restores_process_module_state() -> None:
-    module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
-    previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
-
-    _load_module()
-
-    for name, original in previous.items():
-        assert sys.modules.get(name, _MISSING_MODULE) is original
-
-
-def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
-    module = _load_module()
-
-    cache = module.allocate_lingbot_cache(
-        batch_size=2,
-        num_layers=3,
-        max_tokens=7,
-        num_local_heads=2,
-        head_dim=4,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    assert len(cache.self_attention) == 3
-    assert len(cache.cross_attention) == 3
-    assert cache.cross_attention == [None, None, None]
-    assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
-    assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
-    assert cache.self_attention[0].key.dtype == torch.float32
-    assert cache.self_attention[0].end == 0
-    assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
-
-
-def _allocate_single_layer(module, *, max_tokens: int):
-    return module.allocate_lingbot_cache(
-        batch_size=1,
-        num_layers=1,
-        max_tokens=max_tokens,
-        num_local_heads=1,
-        head_dim=2,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-
-def _set_identity_attention(attention: nn.Module) -> None:
-    qkv = getattr(attention, "qkv", None)
-    if qkv is not None:
-        projections = (qkv, attention.o)
-    else:
-        projections = (attention.q, attention.k, attention.v, attention.o)
-    for projection in projections:
-        with torch.no_grad():
-            identity = torch.eye(2)
-            projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
-            if projection.bias is not None:
-                projection.bias.zero_()
-    attention.norm_q = nn.Identity()
-    attention.norm_k = nn.Identity()
-
-
-def _tokens(*values: float) -> torch.Tensor:
-    return torch.tensor([[[value, 0.0] for value in values]])
-
-
-def _record_inputs(module: nn.Module) -> list[tuple]:
-    records: list[tuple] = []
-    module.register_forward_pre_hook(lambda _module, args: records.append(args))
-    return records
-
-
-def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
-    records: list[torch.Tensor] = []
-
-    def record(_module, _args, output) -> None:
-        records.append(output[0] if isinstance(output, tuple) else output)
-
-    module.register_forward_hook(record)
-    return records
-
-
-def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
-    return (
-        cache.key.detach().view(torch.uint8).clone(),
-        cache.value.detach().view(torch.uint8).clone(),
-        (cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
-    )
-
-
-def _assert_cache_unchanged(cache, snapshot) -> None:
-    key, value, metadata = snapshot
-    assert torch.equal(cache.key.detach().view(torch.uint8), key)
-    assert torch.equal(cache.value.detach().view(torch.uint8), value)
-    assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
-
-
-def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
-
-    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
-    attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
-
-    assert cache.end == 2
-    assert cache.absolute_end == 2
-    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
-
-    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
-
-    assert cache.end == 4
-    assert cache.absolute_end == 4
-    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
-
-
-def test_self_attention_uses_one_fused_qkv_projection() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
-
-    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
-
-    assert attention.qkv.calls == 1
-
-
-def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    paged_inputs = module.ARDiffusionPagedLayerInputs()
-    calls = []
-
-    def paged_write(inputs, query, key, value, k_act, v_act, scale):
-        calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
-        return query
-
-    module.paged_write_attn = paged_write
-    output = attention(
-        _tokens(1, 2),
-        cache=paged_inputs,
-        current_start=0,
-        sink_tokens=0,
-        update_cache=False,
-    )
-
-    assert output.shape == (1, 2, 2)
-    assert len(calls) == 1
-    assert calls[0][0] is paged_inputs
-    assert calls[0][4:6] == (None, None)
-    assert attention.attn.calls == []
-    assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
-
-
-def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
-    attention_inputs = _record_inputs(attention.attn)
-
-    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
-    first_keys = attention_inputs[-1][1].detach().clone()
-    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
-    second_keys = attention_inputs[-1][1].detach().clone()
-
-    assert attention.attn.causal is False
-    torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
-    torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
-
-
-def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
-    attention_inputs = _record_inputs(attention.attn)
-
-    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
-    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
-    attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
-
-    visible_keys = attention_inputs[-1][1].detach().clone()
-    assert cache.end == 4
-    assert cache.absolute_end == 6
-    torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
-    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
-
-
-def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
-    module = _load_module()
-    attention = module.LingBotCrossAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    key_outputs = _record_outputs(attention.k)
-    value_outputs = _record_outputs(attention.v)
-    attention_inputs = _record_inputs(attention.attn)
-
-    output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
-    second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
-
-    assert output.shape == second_output.shape == (1, 1, 2)
-    assert reused_cache is cache
-    assert len(key_outputs) == 1
-    assert len(value_outputs) == 1
-    assert cache.end == 2
-    torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
-    torch.testing.assert_close(attention_inputs[-1][1], cache.key)
-
-
-def test_self_attention_cache_is_isolated_between_requests() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    _set_identity_attention(attention)
-    first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
-    second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
-
-    attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
-
-    assert first.end == 2
-    assert second.end == 0
-    assert torch.count_nonzero(second.key) == 0
-    attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
-    torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
-    torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
-
-
-def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
-    module = _load_module()
-    attention = module.LingBotCrossAttention(dim=4, num_heads=2)
-
-    names = set(attention.state_dict())
-
-    assert {
-        "q.weight",
-        "k.weight",
-        "v.weight",
-        "o.weight",
-        "norm_q.weight",
-        "norm_k.weight",
-    } <= names
-    assert callable(attention.norm_q.weight.weight_loader)
-
-
-def test_tp_world_size_one_attention_output_shapes() -> None:
-    module = _load_module()
-    self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
-    cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
-    cache = module.allocate_lingbot_cache(
-        batch_size=2,
-        num_layers=1,
-        max_tokens=5,
-        num_local_heads=2,
-        head_dim=2,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-    hidden_states = torch.randn(2, 3, 4)
-    encoder_hidden_states = torch.randn(2, 5, 4)
-
-    self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
-    cross_output, cache.cross_attention[0] = cross_attention(
-        hidden_states,
-        encoder_hidden_states,
-        cache=cache.cross_attention[0],
-    )
-
-    assert self_output.shape == hidden_states.shape
-    assert cross_output.shape == hidden_states.shape
-    assert cache.cross_attention[0] is not None
-
-
-def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
-    module = _load_module()
-    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
-    cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
-    cos = torch.ones(2, 1)
-    sin = torch.zeros(2, 1)
-    rotary_inputs = _record_inputs(attention.rotary_embedding)
-
-    attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
-
-    assert len(rotary_inputs) == 2
-    torch.testing.assert_close(rotary_inputs[0][1], cos)
-    torch.testing.assert_close(rotary_inputs[1][2], sin)
-
-
-def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
-    module = _load_module()
-    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
-
-    names = set(block.state_dict())
-    expected = {
-        *(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
-        *(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
-    }
-
-    assert expected <= names
-
-
-def test_cross_attention_caches_are_isolated_between_requests() -> None:
-    module = _load_module()
-    attention = module.LingBotCrossAttention(dim=2, num_heads=1)
-
-    _, first = attention(_tokens(1), _tokens(2, 3), cache=None)
-    _, second = attention(_tokens(1), _tokens(8, 9), cache=None)
-
-    assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
-    assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
-    second_key = second.key.clone()
-    first.key.zero_()
-    torch.testing.assert_close(second.key, second_key)
-
-
-def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
-    module = _load_module()
-    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
-    monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
-    norm = module._LingBotRMSNorm(2, eps=1e-6)
-
-    norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
-
-    torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -1,525 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
-_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
-assert _SPEC is not None and _SPEC.loader is not None
-_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
-sys.modules[_SPEC.name] = _CAMERA_MODULE
-_SPEC.loader.exec_module(_CAMERA_MODULE)
-
-CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
-build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
-interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
-load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
-resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _identity_poses(num_frames: int) -> np.ndarray:
-    return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
-
-
-def _write_trajectory(
-    directory: Path,
-    poses: np.ndarray,
-    intrinsics: np.ndarray,
-) -> None:
-    np.save(directory / "poses.npy", poses)
-    np.save(directory / "intrinsics.npy", intrinsics)
-
-
-def _trusted(directory: Path):
-    return resolve_trusted_action_directory(".", directory)
-
-
-def _trajectory(
-    poses: torch.Tensor,
-    intrinsics: torch.Tensor | None = None,
-) -> CameraTrajectory:
-    if intrinsics is None:
-        intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
-    return CameraTrajectory(poses=poses, intrinsics=intrinsics)
-
-
-@pytest.mark.parametrize("present_file", [None, "poses.npy"])
-def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
-    if present_file == "poses.npy":
-        np.save(tmp_path / present_file, _identity_poses(1))
-
-    missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
-    with pytest.raises(FileNotFoundError, match=missing_file):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-@pytest.mark.parametrize(
-    ("poses", "intrinsics", "message"),
-    [
-        (np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
-        (_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
-        (_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
-        (_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
-    ],
-)
-def test_load_camera_trajectory_rejects_invalid_shapes(
-    tmp_path: Path,
-    poses: np.ndarray,
-    intrinsics: np.ndarray,
-    message: str,
-) -> None:
-    _write_trajectory(tmp_path, poses, intrinsics)
-
-    with pytest.raises(ValueError, match=message):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
-def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
-    poses = _identity_poses(1)
-    intrinsics = np.ones((1, 4), dtype=np.float64)
-    if array_name == "poses":
-        poses[0, 0, 0] = np.nan
-    else:
-        intrinsics[0, 0] = np.inf
-    _write_trajectory(tmp_path, poses, intrinsics)
-
-    with pytest.raises(ValueError, match="finite"):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
-    first_c2w = np.array(
-        [
-            [0.0, -1.0, 0.0, 2.0],
-            [1.0, 0.0, 0.0, 3.0],
-            [0.0, 0.0, 1.0, 4.0],
-            [0.0, 0.0, 0.0, 1.0],
-        ]
-    )
-    second_c2w = np.array(
-        [
-            [0.0, 0.0, 1.0, 5.0],
-            [0.0, 1.0, 0.0, 4.0],
-            [-1.0, 0.0, 0.0, 6.0],
-            [0.0, 0.0, 0.0, 1.0],
-        ]
-    )
-    poses = np.stack((first_c2w, second_c2w))
-    intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
-    _write_trajectory(tmp_path, poses, intrinsics)
-
-    trajectory = load_camera_trajectory(_trusted(tmp_path))
-
-    assert trajectory.poses.dtype == torch.float32
-    assert trajectory.intrinsics.dtype == torch.float32
-    torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
-    torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
-
-
-def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
-    poses = _identity_poses(1).astype(object)
-    intrinsics = np.ones((1, 4), dtype=np.float32)
-    _write_trajectory(tmp_path, poses, intrinsics)
-
-    with pytest.raises(ValueError, match="poses.npy.*real numeric"):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
-    _write_trajectory(
-        tmp_path,
-        _identity_poses(269),
-        np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
-    )
-
-    trajectory = load_camera_trajectory(_trusted(tmp_path))
-
-    assert trajectory.poses.shape == (117, 4, 4)
-    assert trajectory.intrinsics.shape == (117, 4)
-    np.testing.assert_array_equal(
-        trajectory.intrinsics.numpy(),
-        np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
-    )
-
-
-def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
-    _write_trajectory(
-        tmp_path,
-        _identity_poses(4097),
-        np.ones((4097, 4), dtype=np.float32),
-    )
-
-    with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
-    _write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
-    poses_path = tmp_path / "poses.npy"
-    poses_path.write_bytes(poses_path.read_bytes()[:-1])
-
-    with pytest.raises(ValueError, match="poses.npy byte size"):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
-    header_size = 20_000
-    (tmp_path / "poses.npy").write_bytes(
-        np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
-    )
-    np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
-
-    with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
-        load_camera_trajectory(_trusted(tmp_path))
-
-
-def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
-    trusted_root = tmp_path / "trusted"
-    action_directory = trusted_root / "action"
-    outside = tmp_path / "outside"
-    action_directory.mkdir(parents=True)
-    outside.mkdir()
-    _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
-    _write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
-    resolved = resolve_trusted_action_directory("action", trusted_root)
-    action_directory.rename(trusted_root / "original-action")
-    action_directory.symlink_to(outside, target_is_directory=True)
-
-    with pytest.raises(ValueError, match="securely open.*action directory"):
-        load_camera_trajectory(resolved)
-
-
-def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
-    trusted_root = tmp_path / "trusted"
-    action_directory = trusted_root / "action"
-    outside = tmp_path / "outside"
-    action_directory.mkdir(parents=True)
-    outside.mkdir()
-    _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
-    np.save(outside / "poses.npy", _identity_poses(1))
-    resolved = resolve_trusted_action_directory("action", trusted_root)
-    (action_directory / "poses.npy").unlink()
-    (action_directory / "poses.npy").symlink_to(outside / "poses.npy")
-
-    with pytest.raises(ValueError, match="poses.npy.*opened securely"):
-        load_camera_trajectory(resolved)
-
-
-@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
-def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
-    import os
-
-    os.mkfifo(tmp_path / "poses.npy")
-    np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
-    script = """
-import importlib.util
-import sys
-
-spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
-module = importlib.util.module_from_spec(spec)
-sys.modules[spec.name] = module
-spec.loader.exec_module(module)
-action = module.resolve_trusted_action_directory(".", sys.argv[2])
-try:
-    module.load_camera_trajectory(action)
-except ValueError as error:
-    print(error)
-else:
-    raise AssertionError("FIFO camera input was accepted")
-"""
-
-    completed = subprocess.run(
-        [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
-        capture_output=True,
-        text=True,
-        timeout=3,
-        check=False,
-    )
-
-    assert completed.returncode == 0, completed.stderr
-    assert "poses.npy must be a regular NPY file" in completed.stdout
-
-
-def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
-    original_poses = _identity_poses(1)
-    replacement_poses = original_poses.copy()
-    replacement_poses[0, 0, 3] = 123.0
-    _write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
-    poses_path = tmp_path / "poses.npy"
-    original_load = _CAMERA_MODULE.np.load
-    mutation_count = 0
-
-    def mutate_source_then_decode(source, *args, **kwargs):
-        nonlocal mutation_count
-        if mutation_count == 0:
-            mutation_count += 1
-            np.save(poses_path, replacement_poses)
-        return original_load(source, *args, **kwargs)
-
-    monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
-
-    trajectory = load_camera_trajectory(_trusted(tmp_path))
-
-    assert mutation_count == 1
-    torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
-
-
-def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
-    poses = _identity_poses(1)
-    intrinsics = np.ones((1, 4), dtype=np.float32)
-    _write_trajectory(tmp_path, poses, intrinsics)
-    np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
-
-    trajectory = load_camera_trajectory(_trusted(tmp_path))
-
-    torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
-
-
-def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
-    poses = torch.eye(4).repeat(2, 1, 1)
-    poses[1, :3, :3] = torch.tensor(
-        [
-            [0.0, -1.0, 0.0],
-            [1.0, 0.0, 0.0],
-            [0.0, 0.0, 1.0],
-        ]
-    )
-    poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
-    intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
-    trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
-
-    result = interpolate_camera_trajectory(trajectory, 3)
-
-    torch.testing.assert_close(result.poses[0], poses[0])
-    torch.testing.assert_close(result.poses[-1], poses[-1])
-    torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
-    expected_mid_rotation = torch.tensor(
-        [
-            [math.sqrt(0.5), -math.sqrt(0.5), 0.0],
-            [math.sqrt(0.5), math.sqrt(0.5), 0.0],
-            [0.0, 0.0, 1.0],
-        ]
-    )
-    torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
-    torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
-
-
-def test_build_plucker_embedding_has_forward_center_ray() -> None:
-    trajectory = _trajectory(torch.eye(4).unsqueeze(0))
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=3,
-        width=3,
-        target_height=3,
-        target_width=3,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
-    torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
-
-
-def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
-    trajectory = _trajectory(
-        torch.eye(4).unsqueeze(0),
-        torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
-    )
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=4,
-        width=8,
-        target_height=4,
-        target_width=8,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
-    expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
-    torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
-
-
-def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
-    trajectory = _trajectory(
-        torch.eye(4).unsqueeze(0),
-        torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
-    )
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=6,
-        width=10,
-        target_height=2,
-        target_width=5,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    # Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
-    expected_direction = torch.tensor([-0.8, -0.5, 1.0])
-    expected_direction /= torch.linalg.vector_norm(expected_direction)
-    torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
-
-
-def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
-    poses = torch.eye(4).repeat(3, 1, 1)
-    poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
-    poses[1, :3, :3] = torch.tensor(
-        [
-            [0.0, -1.0, 0.0],
-            [1.0, 0.0, 0.0],
-            [0.0, 0.0, 1.0],
-        ]
-    )
-    poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
-    poses[2, :3, :3] = torch.tensor(
-        [
-            [-1.0, 0.0, 0.0],
-            [0.0, -1.0, 0.0],
-            [0.0, 0.0, 1.0],
-        ]
-    )
-    poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
-    trajectory = _trajectory(
-        poses,
-        torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
-    )
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=1,
-        width=2,
-        target_height=1,
-        target_width=2,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
-    torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
-    torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
-    torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
-    torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
-    torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
-
-
-def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
-    poses = torch.eye(4).repeat(3, 1, 1)
-    poses[:, :3, 3] = torch.tensor(
-        [
-            [10.0, -2.0, 5.0],
-            [13.0, 2.0, 5.0],
-            [19.0, 10.0, 5.0],
-        ]
-    )
-
-    result = build_plucker_embedding(
-        _trajectory(poses),
-        height=1,
-        width=1,
-        target_height=1,
-        target_width=1,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
-    torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
-
-
-def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
-    poses = torch.eye(4).repeat(2, 1, 1)
-    poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
-    poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
-    raw = _trajectory(poses)
-
-    interpolated = interpolate_camera_trajectory(raw, 3)
-    result = build_plucker_embedding(
-        interpolated,
-        height=1,
-        width=1,
-        target_height=1,
-        target_width=1,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    expected_delta = torch.tensor([0.6, 0.8, 0.0])
-    torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
-    torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
-    torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
-
-
-def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
-    poses = torch.eye(4).repeat(2, 1, 1)
-    poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
-    trajectory = _trajectory(
-        poses,
-        torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
-    )
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=1,
-        width=2,
-        target_height=1,
-        target_width=2,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    expected_direction = torch.tensor([0.25, 0.5, 1.0])
-    expected_direction /= torch.linalg.vector_norm(expected_direction)
-    torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
-    torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
-
-
-def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
-    trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
-
-    result = build_plucker_embedding(
-        trajectory,
-        height=12,
-        width=20,
-        target_height=3,
-        target_width=5,
-        device=torch.device("cpu"),
-        dtype=torch.float64,
-    )
-
-    assert result.shape == (2, 6, 3, 5)
-    assert result.dtype == torch.float64
-    assert result.device == torch.device("cpu")
-
-
-def test_build_plucker_embedding_is_deterministic() -> None:
-    trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
-    kwargs = {
-        "height": 12,
-        "width": 20,
-        "target_height": 3,
-        "target_width": 5,
-        "device": torch.device("cpu"),
-        "dtype": torch.float32,
-    }
-
-    first = build_plucker_embedding(trajectory, **kwargs)
-    second = build_plucker_embedding(trajectory, **kwargs)
-
-    torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py

@@ -1,685 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import hashlib
-import inspect
-import json
-import math
-from pathlib import Path
-
-import pytest
-import torch
-
-from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
-_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
-
-
-@pytest.fixture(autouse=True)
-def _inference_context():
-    with torch.inference_mode():
-        yield
-
-
-def _tiny_model(
-    module,
-    *,
-    num_layers: int = 2,
-    num_frames_per_block: int = 1,
-    sliding_window_num_frames: int = 3,
-):
-    return module.CausalLingBotWorldTransformer3DModel(
-        patch_size=(1, 2, 2),
-        num_attention_heads=2,
-        attention_head_dim=2,
-        in_channels=36,
-        out_channels=2,
-        text_dim=6,
-        freq_dim=4,
-        ffn_dim=8,
-        num_layers=num_layers,
-        cross_attn_norm=True,
-        eps=1e-6,
-        rope_max_seq_len=16,
-        sink_size=1,
-        num_frames_per_block=num_frames_per_block,
-        sliding_window_num_frames=sliding_window_num_frames,
-        local_attn_size=-1,
-    )
-
-
-def _cache(module, model, *, max_tokens: int | None = None):
-    if max_tokens is None:
-        window_frames = (
-            model.config.local_attn_size
-            if model.config.local_attn_size != -1
-            else model.config.sliding_window_num_frames
-        )
-        max_tokens = window_frames * 2 * 2
-    return module.allocate_lingbot_cache(
-        batch_size=1,
-        num_layers=len(model.blocks),
-        max_tokens=max_tokens,
-        num_local_heads=2,
-        head_dim=2,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-
-def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
-    specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
-    for name, parameter in model.named_parameters():
-        marker = ".self_attn.qkv."
-        if marker not in name:
-            specs[name] = (parameter.shape, parameter.dtype)
-            continue
-        shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
-        for projection in ("q", "k", "v"):
-            specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
-    return specs
-
-
-def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
-    checkpoint_weights: dict[str, torch.Tensor] = {}
-    expected_parameters: dict[str, torch.Tensor] = {}
-    next_value = 1
-    for name, parameter in model.named_parameters():
-        marker = ".self_attn.qkv."
-        if marker not in name:
-            value = torch.full_like(parameter, next_value)
-            checkpoint_weights[name] = value
-            expected_parameters[name] = value
-            next_value += 1
-            continue
-        shards = []
-        shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
-        for projection in ("q", "k", "v"):
-            value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
-            checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
-            shards.append(value)
-            next_value += 1
-        expected_parameters[name] = torch.cat(shards)
-    return checkpoint_weights, expected_parameters
-
-
-def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
-    return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
-
-
-def _assert_self_cache_unchanged(cache, snapshot) -> None:
-    for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
-        attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
-
-
-def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
-    torch.manual_seed(7)
-    module = attention_tests._load_module()
-    model = _tiny_model(module).eval()
-    cache = _cache(module, model)
-    encoder_hidden_states = torch.randn(1, 3, 6)
-    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
-    cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
-    text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
-    outputs = []
-
-    assert model.patch_embedding.in_channels == 36
-    assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
-
-    for start_frame in range(4):
-        hidden_states = torch.randn(1, 36, 1, 4, 4)
-        timestep = torch.tensor([float(start_frame + 1)])
-        snapshot = _self_cache_snapshot(cache)
-
-        transient_output = model(
-            hidden_states,
-            timestep,
-            encoder_hidden_states,
-            camera_hidden_states,
-            cache=cache,
-            start_frame=start_frame,
-            update_cache=False,
-        )
-
-        assert transient_output.shape == (1, 2, 1, 4, 4)
-        _assert_self_cache_unchanged(cache, snapshot)
-        assert all(layer_cache is not None for layer_cache in cache.cross_attention)
-
-        committed_output = model(
-            hidden_states,
-            timestep,
-            encoder_hidden_states,
-            camera_hidden_states,
-            cache=cache,
-            start_frame=start_frame,
-            update_cache=True,
-        )
-        torch.testing.assert_close(committed_output, transient_output)
-        outputs.append(committed_output)
-
-        expected_token_end = (start_frame + 1) * 2 * 2
-        assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
-
-    assert len(outputs) == 4
-    assert all(len(outputs) == 1 for outputs in cross_key_outputs)
-    assert len(text_projection_outputs) == 1
-
-    alternate_cache = _cache(module, model)
-    alternate_output = model(
-        torch.zeros(1, 36, 1, 4, 4),
-        torch.tensor([1.0]),
-        encoder_hidden_states,
-        torch.zeros_like(camera_hidden_states),
-        cache=alternate_cache,
-        start_frame=0,
-        update_cache=False,
-    )
-    camera_output = model(
-        torch.zeros(1, 36, 1, 4, 4),
-        torch.tensor([1.0]),
-        encoder_hidden_states,
-        torch.ones_like(camera_hidden_states),
-        cache=_cache(module, model),
-        start_frame=0,
-        update_cache=False,
-    )
-    assert not torch.equal(camera_output, alternate_output)
-
-
-def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=2).eval()
-    cache = _cache(module, model)
-    hidden_states = torch.randn(1, 36, 1, 4, 4)
-    timestep = torch.tensor([1.0])
-    encoder_hidden_states = torch.randn(1, 3, 6)
-    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
-
-    model(
-        hidden_states,
-        timestep,
-        encoder_hidden_states,
-        camera_hidden_states,
-        cache=cache,
-        start_frame=0,
-        update_cache=False,
-    )
-    cache.cross_attention[1] = None
-    text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
-    cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
-
-    model(
-        hidden_states,
-        timestep,
-        encoder_hidden_states,
-        camera_hidden_states,
-        cache=cache,
-        start_frame=0,
-        update_cache=False,
-    )
-
-    assert len(text_projection_outputs) == 1
-    assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
-    assert all(layer_cache is not None for layer_cache in cache.cross_attention)
-
-
-@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
-def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(
-        module,
-        num_layers=1,
-        num_frames_per_block=3,
-        sliding_window_num_frames=6,
-    )
-    cache = _cache(module, model)
-    snapshot = _self_cache_snapshot(cache)
-    patch_inputs = attention_tests._record_inputs(model.patch_embedding)
-
-    with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
-        model(
-            torch.randn(1, 36, frames, 4, 4),
-            torch.tensor([1.0]),
-            torch.randn(1, 3, 6),
-            torch.randn(1, 6 * 8 * 8, frames, 4, 4),
-            cache=cache,
-            start_frame=0,
-            update_cache=True,
-        )
-
-    assert patch_inputs == []
-    _assert_self_cache_unchanged(cache, snapshot)
-    assert cache.cross_attention == [None]
-
-
-def test_forward_accepts_exactly_one_configured_frame_block() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(
-        module,
-        num_layers=1,
-        num_frames_per_block=3,
-        sliding_window_num_frames=6,
-    )
-
-    output = model(
-        torch.randn(1, 36, 3, 4, 4),
-        torch.tensor([1.0]),
-        torch.randn(1, 3, 6),
-        torch.randn(1, 6 * 8 * 8, 3, 4, 4),
-        cache=_cache(module, model),
-        start_frame=0,
-        update_cache=False,
-    )
-
-    assert output.shape == (1, 2, 3, 4, 4)
-
-
-def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(
-        module,
-        num_layers=2,
-        num_frames_per_block=3,
-        sliding_window_num_frames=6,
-    )
-
-    cache = model.allocate_cache(
-        batch_size=2,
-        latent_height=8,
-        latent_width=12,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-
-    assert len(cache.self_attention) == 2
-    assert len(cache.cross_attention) == 2
-    assert cache.cross_attention == [None, None]
-    assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
-    assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
-
-
-def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
-    hidden_states = torch.zeros(1, 36, 2, 2, 4)
-    hidden_states[0, 0] = torch.tensor(
-        [
-            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
-            [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
-        ]
-    )
-    with torch.no_grad():
-        model.patch_embedding.weight.zero_()
-        model.patch_embedding.bias.zero_()
-        model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
-
-    tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
-
-    torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
-
-
-def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
-    hidden_states = (
-        torch.stack(
-            (
-                torch.arange(0, 8),
-                torch.arange(10, 18),
-                torch.arange(20, 28),
-                torch.arange(30, 38),
-            )
-        )
-        .unsqueeze(0)
-        .float()
-    )
-    expected = torch.tensor(
-        [
-            [
-                [[0, 2, 10, 12], [4, 6, 14, 16]],
-                [[20, 22, 30, 32], [24, 26, 34, 36]],
-            ],
-            [
-                [[1, 3, 11, 13], [5, 7, 15, 17]],
-                [[21, 23, 31, 33], [25, 27, 35, 37]],
-            ],
-        ],
-        dtype=torch.float32,
-    ).unsqueeze(0)
-
-    output = model._unpatchify(
-        hidden_states,
-        batch_size=1,
-        frames=2,
-        height=1,
-        width=2,
-    )
-
-    torch.testing.assert_close(output, expected)
-
-
-def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
-    model.head.norm = torch.nn.Identity()
-    with torch.no_grad():
-        model.head.modulation.zero_()
-        model.head.head.weight.zero_()
-        model.head.head.bias.zero_()
-        model.head.head.weight[:, 0] = 1
-    hidden_states = torch.ones(1, 4, model.dim)
-    timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
-    expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
-
-    output = model.head(hidden_states, timestep_embedding)
-
-    torch.testing.assert_close(output, expected)
-
-
-def test_constructor_defaults_match_official_checkpoint_config() -> None:
-    module = attention_tests._load_module()
-    parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
-    expected = {
-        "in_channels": 36,
-        "out_channels": 16,
-        "num_layers": 40,
-        "num_attention_heads": 40,
-        "attention_head_dim": 128,
-        "patch_size": (1, 2, 2),
-        "text_dim": 4096,
-        "ffn_dim": 13824,
-        "freq_dim": 256,
-        "rope_max_seq_len": 1024,
-        "eps": 1e-6,
-        "cross_attn_norm": True,
-        "sink_size": 9,
-        "num_frames_per_block": 3,
-        "sliding_window_num_frames": 18,
-        "local_attn_size": -1,
-    }
-
-    assert {name: parameters[name].default for name in expected} == expected
-
-
-def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module)
-
-    assert model.dtype == next(model.parameters()).dtype
-
-
-def test_transformer_declares_regional_compile_block() -> None:
-    module = attention_tests._load_module()
-
-    assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
-
-
-def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
-    module = attention_tests._load_module()
-    reduced_values: list[torch.Tensor] = []
-
-    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
-
-    def all_reduce(value: torch.Tensor) -> torch.Tensor:
-        reduced_values.append(value.detach().clone())
-        return value * 2
-
-    monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
-    norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
-    value = torch.tensor([[3.0, 4.0]])
-
-    output = norm(value)
-
-    torch.testing.assert_close(output, value / math.sqrt(12.5))
-    assert len(reduced_values) == 1
-    torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
-
-
-def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
-    module = attention_tests._load_module()
-    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
-
-    with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
-        module.LingBotSelfAttention(dim=4, num_heads=2)
-
-
-def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
-    module = attention_tests._load_module()
-
-    with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
-        module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
-
-
-@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
-def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
-    module = attention_tests._load_module()
-
-    with pytest.raises(ValueError, match=field):
-        module.CausalLingBotWorldTransformer3DModel(**{field: 4})
-
-
-def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
-    module = attention_tests._load_module()
-
-    with pytest.raises(RuntimeError, match="quant_config.*not supported"):
-        module.CausalLingBotWorldTransformer3DModel(quant_config=object())
-
-
-def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
-    module = attention_tests._load_module()
-
-    class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
-        def __init__(self, **kwargs) -> None:
-            self.received_kwargs = kwargs
-
-    config = {
-        "_class_name": "CausalLingBotWorldTransformer3DModel",
-        "_diffusers_version": "0.35.0.dev0",
-        "patch_size": [1, 2, 2],
-        "num_attention_heads": 40,
-        "attention_head_dim": 128,
-        "in_channels": 36,
-        "out_channels": 16,
-        "text_dim": 4096,
-        "freq_dim": 256,
-        "ffn_dim": 13824,
-        "num_layers": 40,
-        "cross_attn_norm": True,
-        "eps": 1e-6,
-        "image_dim": None,
-        "added_kv_proj_dim": None,
-        "rope_max_seq_len": 1024,
-        "pos_embed_seq_len": None,
-        "qk_norm": "rms_norm_across_heads",
-        "sink_size": 9,
-        "num_frames_per_block": 3,
-        "sliding_window_num_frames": 18,
-        "local_attn_size": -1,
-    }
-
-    model = ConfigProbe.from_config(config, prefix="transformer")
-
-    assert model.received_kwargs["patch_size"] == (1, 2, 2)
-    assert model.received_kwargs["num_layers"] == 40
-    assert model.received_kwargs["prefix"] == "transformer"
-
-
-@pytest.mark.parametrize(
-    ("field", "value"),
-    [
-        ("in_channels", 35),
-        ("num_layers", 39),
-        ("num_attention_heads", 32),
-        ("attention_head_dim", 64),
-        ("ffn_dim", 10240),
-        ("num_frames_per_block", 4),
-        ("sliding_window_num_frames", 24),
-        ("sink_size", 0),
-    ],
-)
-def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
-    module = attention_tests._load_module()
-    config = {
-        "_class_name": "CausalLingBotWorldTransformer3DModel",
-        "_diffusers_version": "0.35.0.dev0",
-        "patch_size": [1, 2, 2],
-        "num_attention_heads": 40,
-        "attention_head_dim": 128,
-        "in_channels": 36,
-        "out_channels": 16,
-        "text_dim": 4096,
-        "freq_dim": 256,
-        "ffn_dim": 13824,
-        "num_layers": 40,
-        "cross_attn_norm": True,
-        "eps": 1e-6,
-        "image_dim": None,
-        "added_kv_proj_dim": None,
-        "rope_max_seq_len": 1024,
-        "pos_embed_seq_len": None,
-        "qk_norm": "rms_norm_across_heads",
-        "sink_size": 9,
-        "num_frames_per_block": 3,
-        "sliding_window_num_frames": 18,
-        "local_attn_size": -1,
-    }
-    config[field] = value
-
-    with pytest.raises(ValueError, match=field):
-        module.CausalLingBotWorldTransformer3DModel.from_config(config)
-
-
-def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
-    module = attention_tests._load_module()
-    config = {
-        "_class_name": "CausalLingBotWorldTransformer3DModel",
-        "architectures": ["CausalLingBotWorldTransformer3DModel"],
-        "future_diffusers_metadata": {"version": 1},
-    }
-
-    with torch.device("meta"):
-        model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
-
-    assert model.config.in_channels == 36
-    assert model.config.num_layers == 40
-
-
-def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=1)
-    checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
-    loader_calls: list[str] = []
-    qkv_weight = model.blocks[0].self_attn.qkv.weight
-
-    def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
-        loader_calls.append(shard_id)
-        shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
-        shard_size = param.shape[0] // 3
-        param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
-
-    qkv_weight.weight_loader = record_loader
-    loaded = model.load_weights(iter(checkpoint_weights.items()))
-
-    assert set(checkpoint_weights) <= loaded
-    assert set(dict(model.named_parameters())) <= loaded
-    assert loader_calls == ["q", "k", "v"]
-    for name, param in model.named_parameters():
-        torch.testing.assert_close(param, expected_parameters[name])
-
-    first_name = next(iter(checkpoint_weights))
-    partial = _tiny_model(module, num_layers=1)
-    partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
-    expected_partial = {first_name}
-    for projection in ("q", "k", "v"):
-        marker = f".self_attn.{projection}."
-        if marker in first_name:
-            expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
-            break
-    assert partial_loaded == expected_partial
-
-    with pytest.raises(KeyError, match="unexpected_model.weight"):
-        model.load_weights([("unexpected_model.weight", torch.ones(1))])
-
-
-def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
-    module = attention_tests._load_module()
-    model = _tiny_model(module, num_layers=1)
-    first_name, first_param = next(iter(model.named_parameters()))
-    loaded_weight = torch.full_like(first_param, 17)
-
-    def weights():
-        yield first_name, loaded_weight
-        torch.testing.assert_close(first_param, loaded_weight)
-
-    assert model.load_weights(weights()) == {first_name}
-
-
-def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
-    utils = pytest.importorskip("vllm.model_executor.models.utils")
-    module = attention_tests._load_module()
-    transformer = _tiny_model(module, num_layers=1)
-    checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
-
-    class Pipeline(torch.nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.transformer = transformer
-
-    pipeline = Pipeline()
-    loaded = utils.AutoWeightsLoader(pipeline).load_weights(
-        (f"transformer.{name}", value) for name, value in checkpoint_weights.items()
-    )
-
-    assert set(dict(pipeline.named_parameters())) <= loaded
-
-
-def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
-    module = attention_tests._load_module()
-    fixture = json.loads(_FIXTURE_PATH.read_text())
-    model = _tiny_model(module, num_layers=1)
-    parameter_names = set(_checkpoint_parameter_specs(model))
-
-    assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
-    assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
-    assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
-        fixture["block_parameter_suffixes"]
-    )
-    expanded_names = set(fixture["non_block_parameters"])
-    expanded_names.update(
-        f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
-    )
-    canonical_names = "\n".join(sorted(expanded_names))
-
-    assert len(expanded_names) == fixture["index_parameter_count"] == 1421
-    assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
-    assert fixture["canonical_name_to_shard_sha256"] == (
-        "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
-    )
-    assert list(fixture["shards"]) == [
-        f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
-    ]
-    assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
-    assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
-
-
-def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
-    module = attention_tests._load_module()
-    fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
-
-    with torch.device("meta"):
-        model = module.CausalLingBotWorldTransformer3DModel()
-    parameter_specs = _checkpoint_parameter_specs(model)
-    expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
-
-    assert len(fixture["shard_representatives"]) == 8
-    assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
-
-    for name, expected_shape in fixture["parameter_shapes"].items():
-        shape, dtype = parameter_specs[name]
-        assert tuple(shape) == tuple(expected_shape), name
-        assert dtype == expected_dtype

tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py

@@ -1,39 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import ast
-from pathlib import Path
-
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_ROOT = Path(__file__).parents[4]
-_PRODUCTION_FILES = (
-    _ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
-    _ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
-    _ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
-)
-
-
-def test_all_lingbot_production_functions_have_complete_annotations() -> None:
-    missing: list[str] = []
-    for path in _PRODUCTION_FILES:
-        tree = ast.parse(path.read_text())
-        for node in ast.walk(tree):
-            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                continue
-            arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
-            for argument in arguments:
-                if argument.arg not in {"self", "cls"} and argument.annotation is None:
-                    missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
-            if node.args.vararg is not None and node.args.vararg.annotation is None:
-                missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
-            if node.args.kwarg is not None and node.args.kwarg.annotation is None:
-                missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
-            if node.returns is None:
-                missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
-
-    assert missing == []

tests/diffusion/test_result_pump.py

@@ -418,77 +418,3 @@ class TestShutdownCleansUpFutures:
 
         assert len(executor._rpc_futures) == 0
         assert len(executor._output_futures) == 0
-
-
-class _RacyFuture(concurrent.futures.Future):
-    # done() always lies and reports False, to deterministically force the
-    # pump's check-then-act race window without needing real thread timing.
-    def done(self) -> bool:
-        return False
-
-
-def _racy_cancelled_future() -> concurrent.futures.Future:
-    fut = _RacyFuture()
-    fut.cancel()
-    assert fut.cancelled()
-    return fut
-
-
-class TestResultPumpCancelledFutureRace:
-    """Regression for #5793: a future cancelled concurrently with the pump's
-    resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must
-    be dropped, not raise InvalidStateError and kill the pump thread.
-    """
-
-    def test_compute_done_racing_cancel_does_not_crash_pump(self):
-        executor = _make_executor()
-        fut = _racy_cancelled_future()
-        with executor._futures_lock:
-            executor._rpc_futures["1"] = fut
-
-        msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc")
-        _feed_one_msg_to_pump(executor, msg)
-
-        assert fut.cancelled()
-
-    def test_output_ready_racing_cancel_does_not_crash_pump(self):
-        executor = _make_executor()
-        fut = _racy_cancelled_future()
-        with executor._futures_lock:
-            executor._output_futures["abc123"] = fut
-
-        msg = AsyncDiffusionOutput(
-            kind=AsyncOutputKind.OUTPUT_READY,
-            async_output_id="abc123",
-            output=DiffusionOutput(output="late"),
-        )
-        _feed_one_msg_to_pump(executor, msg)
-
-        assert fut.cancelled()
-
-    def test_batch_split_racing_cancel_does_not_crash_pump(self):
-        executor = _make_executor()
-        cancelled_fut = _racy_cancelled_future()
-        healthy_fut = concurrent.futures.Future()
-        with executor._futures_lock:
-            executor._output_futures["batch-1/r-aborted"] = cancelled_fut
-            executor._output_futures["batch-1/r-healthy"] = healthy_fut
-            executor._batch_split_map["batch-1"] = {
-                "batch-1/r-aborted": "r-aborted",
-                "batch-1/r-healthy": "r-healthy",
-            }
-
-        outputs = {
-            "r-aborted": DiffusionOutput(output="late"),
-            "r-healthy": DiffusionOutput(output="ok"),
-        }
-        msg = AsyncDiffusionOutput(
-            kind=AsyncOutputKind.OUTPUT_READY,
-            async_output_id="batch-1",
-            output=_FakeBatchOutput(outputs),
-        )
-        _feed_one_msg_to_pump(executor, msg)
-
-        assert cancelled_fut.cancelled()
-        assert healthy_fut.done()
-        assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]

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_hunyuanimage3_i2t_expansion.py

@@ -3,18 +3,19 @@
 """Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline."""
 
 from collections.abc import Generator
+from pathlib import Path
 
 import pytest
 import torch
 from PIL import Image
 
 from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import Omni
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt
 
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
+REPO_ROOT = Path(__file__).resolve().parents[3]
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml"
 
 # First 20 generated token IDs from the HF greedy reference on this input.
 # vllm-omni AR output matches this prefix bitwise; the two implementations
@@ -44,14 +45,14 @@ EXPECTED_PREFIX_TOKEN_IDS: list[int] = [
 # Decoded form, kept only for human-readable assertion messages.
 EXPECTED_PREFIX_TEXT = "The image is a solid, uniform green color with no variations, objects, or details present. It"
 
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cuda]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 @pytest.fixture(scope="module")
 def omni() -> Generator[Omni, None, None]:
     with OmniRunner(
         MODEL_NAME,
-        deploy_config=DEPLOY_CONFIG_PATH,
+        deploy_config=str(STAGE_CONFIG_PATH),
         trust_remote_code=True,
     ) as runner:
         yield runner.omni

tests/e2e/offline_inference/test_lingbot_world_v2.py

@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end smoke test for LingBot-World v2 generation."""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-MODEL = os.environ.get(
-    "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
-    "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
-)
-_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
-_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
-_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
-_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
-
-_OMNI_RUNNER_PARAM = (
-    MODEL,
-    None,
-    {
-        "flow_shift": 5.0,
-        "parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
-        "model_config": {
-            "lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
-        },
-    },
-)
-
-pytestmark = [
-    pytest.mark.slow,
-    pytest.mark.diffusion,
-    pytest.mark.skipif(
-        not _HAS_INPUT_ASSETS,
-        reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
-    ),
-    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
-]
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
-    """Load the real checkpoint and generate one causal block on one GPU."""
-
-    assert _IMAGE_PATH is not None
-    assert _ACTION_DIR is not None
-    response = omni_runner_handler.send_diffusion_request(
-        {
-            "model": omni_runner_handler.runner.model_name,
-            "prompt": "The camera moves slowly forward through the scene.",
-            "multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
-            "sampling_params": OmniDiffusionSamplingParams(
-                height=480,
-                width=832,
-                num_frames=9,
-                num_inference_steps=4,
-                max_sequence_length=512,
-                seed=42,
-                fps=16,
-                extra_args={
-                    "action_path": _ACTION_DIR.name,
-                    "flow_shift": 5.0,
-                },
-            ),
-        }
-    )
-
-    assert response.success
-    assert response.images is not None
-    assert len(response.images) == 1
-    video = np.asarray(response.images[0])
-    assert video.shape == (1, 9, 480, 832, 3)
-    assert np.isfinite(video).all()
-    assert float(np.ptp(video)) > 1e-6
-    assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6

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

@@ -492,24 +492,6 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
 
 
-@pytest.mark.parametrize(
-    ("legacy_arg", "value"),
-    [
-        ("stage_configs_path", "legacy.yaml"),
-        ("stage_configs", [{"stage_id": 0}]),
-    ],
-)
-def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value):
-    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-
-    with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"):
-        engine._resolve_stage_configs(
-            "dummy-model",
-            {legacy_arg: value},
-            trust_remote_code=False,
-        )
-
-
 def test_default_stage_config_includes_quantization_config():
     """Ensure structured quantization_config survives default diffusion-stage creation."""
     quantization_config = {

tests/entrypoints/test_utils.py

@@ -27,7 +27,7 @@ from vllm_omni.entrypoints.utils import (
     filter_dataclass_kwargs,
     filter_stages,
     load_and_resolve_stage_configs,
-    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
 
@@ -417,9 +417,8 @@ class TestTryResolveOmniModelType:
 
 class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
-        """Ensure that kwargs survive default stage creation."""
-        engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
-        kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
+        """Ensure that dtype survives default stage creation."""
+        kwargs = {"dtype": torch.float32}
         mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
         mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
 
@@ -432,7 +431,6 @@ class TestLoadAndResolveStageConfigs:
         assert config_path is None
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
-        assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
 
     def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
         deploy_path = tmp_path / "qwen3_multi.yaml"
@@ -520,20 +518,71 @@ class TestLoadAndResolveStageConfigs:
         assert filtered[0].final_output_type is None
 
 
-class TestLoadStageConfigsFromModel:
-    def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
+class TestLoadStageConfigsFromYaml:
+    """Regression tests for stage-config loading and merging."""
+
+    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "async_chunk": True,
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "runtime": {"device": 0},
+                        "engine_args": {
+                            "parallel_config": {"tensor_parallel_size": 4},
+                        },
+                    }
+                ],
+            }
+        )
         mocker.patch(
-            "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
-            return_value=(None, None),
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
         )
-        resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
 
-        result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={
+                "parallel_config": {
+                    "tensor_parallel_size": 1,
+                    "pipeline_parallel_size": 2,
+                },
+                "model": "base-model",
+            },
+        )
 
-        assert result == ([], None)
-        resolve_path.assert_not_called()
-        assert "No registered PipelineConfig resolved" in caplog.text
-        assert "deploy_config" in caplog.text
+        merged_engine_args = stages[0]["engine_args"]
+        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
+        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
+        assert merged_engine_args["model"] == "base-model"
+        assert merged_engine_args["async_chunk"] is True
+
+    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "engine_args": {
+                            "nested": {"override": 2},
+                        },
+                    }
+                ],
+            }
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
+        )
+
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={"nested": {"base": 1}},
+        )
+
+        assert stages[0]["engine_args"]["nested"]["base"] == 1
+        assert stages[0]["engine_args"]["nested"]["override"] == 2
 
 
 class TestCumulativeStreamingCoercion:

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/model_tests/diffusion/test_alignment.py

@@ -64,9 +64,6 @@ EXCLUDED_MODELS = [
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
     "LingBotVideoPipeline",
-    # Requires camera/action assets and persistent AR state; covered by its
-    # dedicated contract tests and real-checkpoint E2E smoke instead.
-    "LingBotWorldCausalDMDPipeline",
     "MagiHumanPipeline",
     "OmniVoicePipeline",
     "OmniVoice",

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)