[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup. Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Diffusion] Document runtime cache hardware tradeoffs
2026-08-17 11:50:32 +08:00 · 2026-08-17 00:43:06 +08:00 · 2026-08-16 23:26:25 +08:00 · 2026-08-16 23:15:48 +08:00 · 2026-08-16 11:22:28 +08:00 · 2026-08-16 11:02:38 +08:00
235 changed files with 27595 additions and 2070 deletions
@@ -76,6 +76,9 @@ steps:
          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
          - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -92,6 +95,9 @@ steps:
          - tests/e2e/online_serving/test_qwen3_tts_base.py
          - tests/e2e/offline_inference/test_qwen3_tts_base.py
          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
          - .buildkite/common/scripts/run_cov_split.sh
@@ -109,27 +115,15 @@ steps:
              -- -m 'advanced_model and cuda' --run-level 'advanced_model'
        mirror_hardwares: l4_1

-      - label: "TTS · Higgs-Audio-V3 Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_higgs_audio_v3.py
-          - vllm_omni/model_executor/models/higgs_audio_v3/
-          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
-          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_USE_DEEP_GEMM=0
-              export VLLM_MOE_USE_DEEP_GEMM=0
-              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model'
-            "
-        mirror_hardwares: l4_1
-
      - label: "Omni · Qwen3-Omni Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_qwen3_omni.py
          - tests/e2e/online_serving/test_qwen3_omni.py
          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
          - vllm_omni/deploy/qwen3_omni_moe.yaml
        timeout_in_minutes: 50
@@ -141,6 +135,11 @@ steps:
        source_file_dependencies:
          - tests/e2e/offline_inference/test_minicpmo_4_5.py
          - tests/e2e/online_serving/test_minicpmo_4_5.py
+          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/models/step_audio2/
+          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
+          - vllm_omni/model_executor/models/cosyvoice3/utils.py
+          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        timeout_in_minutes: 50
        commands:
@@ -162,70 +161,35 @@ steps:
          - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
+          - vllm_omni/deploy/minicpmo_4_5.yaml
          - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
          - vllm_omni/experimental/fullduplex/
          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/models/step_audio2/
+          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
+          - vllm_omni/model_executor/models/cosyvoice3/utils.py
          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/core/sched/
          - vllm_omni/outputs/multimodal_accumulation.py
          - vllm_omni/engine/stage_init_utils.py
          - vllm_omni/engine/async_omni_engine.py
+          - vllm_omni/engine/orchestrator.py
          - vllm_omni/entrypoints/cli/serve.py
+          - vllm_omni/entrypoints/openai/api_server.py
+          - vllm_omni/entrypoints/openai/duplex_capability.py
+          - vllm_omni/config/stage_config.py
        timeout_in_minutes: 50
        commands:
          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
        mirror_hardwares: h100_1

-      - label: "TTS · VoxCPM2 Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_voxcpm2_tts.py
-          - tests/e2e/online_serving/test_voxcpm2_tts.py
-          - vllm_omni/model_executor/models/voxcpm2/
-          - vllm_omni/deploy/voxcpm2.yaml
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        mirror_hardwares: l4_1
-
-      - label: "Diffusion · Qwen Image Edit Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_qwen_image_edit.py
-          - vllm_omni/diffusion/models/qwen_image/
-        timeout_in_minutes: 20
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-        mirror_hardwares: h100_1
-      - label: "Diffusion · Bagel Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_bagel.py
-          - tests/e2e/online_serving/test_bagel.py
-          - vllm_omni/model_executor/models/bagel/
-          - vllm_omni/diffusion/models/bagel/
-          - vllm_omni/model_executor/stage_input_processors/bagel.py
-          - vllm_omni/deploy/bagel.yaml
-          - .buildkite/common/scripts/run_cov_split.sh
-          - pyproject.toml
-        # 20 -> 40: the coverage split runs the model twice.
-        timeout_in_minutes: 40
-        commands:
-          - |
-            .buildkite/common/scripts/run_cov_split.sh \
-              --model-id bagel \
-              --offline tests/e2e/offline_inference/test_bagel.py \
-              --online tests/e2e/online_serving/test_bagel.py \
-              -- -m 'advanced_model and cuda' --run-level 'advanced_model'
-        mirror_hardwares: h100_1
-
      - label: "Diffusion · Qwen Image Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen_image.py
          - vllm_omni/diffusion/models/qwen_image/
+          - vllm_omni/diffusion/models/dmd2/
+          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
        commands:
          - |
            timeout 40m bash -c '
@@ -234,23 +198,12 @@ steps:
            '
        mirror_hardwares: h100_1

-      - label: "Diffusion · Qwen Image Layered Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_qwen_image_layered.py
-          - vllm_omni/diffusion/models/qwen_image/
-        commands:
-          - |
-            timeout 40m bash -c "
-              export VLLM_IMAGE_FETCH_TIMEOUT=60
-              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        mirror_hardwares: h100_1
-
      - label: "Diffusion · Wan22 Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_wan22_t2v.py
          - tests/e2e/online_serving/test_wan22_t2v.py
          - vllm_omni/diffusion/models/wan2_2/
+          - vllm_omni/diffusion/models/dmd2/
        commands:
          - |
            timeout 40m bash -c "
@@ -258,21 +211,3 @@ steps:
              pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
            "
        mirror_hardwares: h100_1
-
-      - label: "Diffusion · Z Image Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_z_image.py
-          - tests/e2e/online_serving/test_z_image.py
-          - vllm_omni/diffusion/models/z_image/
-        timeout_in_minutes: 30
-        commands:
-          - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
-        mirror_hardwares: l4_1
-      - label: "Diffusion · FLUX.2 Klein Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_flux2_klein.py
-          - vllm_omni/diffusion/models/flux2_klein/
-        timeout_in_minutes: 30
-        commands:
-          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
-        mirror_hardwares: l4_1
@@ -459,6 +459,12 @@ steps:
          - 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:
@@ -487,6 +493,13 @@ steps:
          - 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
@@ -169,6 +169,10 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_omni.py
          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
          - vllm_omni/deploy/qwen3_omni_moe.yaml
        commands:
@@ -181,6 +185,11 @@ steps:
      - label: "Omni · MiniCPM-o 4.5 Offline Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_minicpmo_4_5.py
+          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/models/step_audio2/
+          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
+          - vllm_omni/model_executor/models/cosyvoice3/utils.py
+          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        commands:
          - |
@@ -192,6 +201,11 @@ steps:
      - label: "Omni · MiniCPM-o 4.5 Online Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_minicpmo_4_5.py
+          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/models/step_audio2/
+          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
+          - vllm_omni/model_executor/models/cosyvoice3/utils.py
+          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        commands:
          - |
@@ -210,16 +224,24 @@ steps:
          - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
+          - vllm_omni/deploy/minicpmo_4_5.yaml
          - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
          - vllm_omni/experimental/fullduplex/
          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/models/step_audio2/
+          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
+          - vllm_omni/model_executor/models/cosyvoice3/utils.py
          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/core/sched/
          - vllm_omni/outputs/multimodal_accumulation.py
          - vllm_omni/engine/stage_init_utils.py
          - vllm_omni/engine/async_omni_engine.py
+          - vllm_omni/engine/orchestrator.py
          - vllm_omni/entrypoints/cli/serve.py
+          - vllm_omni/entrypoints/openai/api_server.py
+          - vllm_omni/entrypoints/openai/duplex_capability.py
+          - vllm_omni/config/stage_config.py
        commands:
          - |
            timeout 20m bash -c "
@@ -239,25 +261,13 @@ steps:
          - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
        mirror_hardwares: h100_1

-      - label: "TTS · VoxCPM2 Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_voxcpm2_tts.py
-          - vllm_omni/model_executor/models/voxcpm2/
-          - vllm_omni/deploy/voxcpm2.yaml
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: l4_1
-
      - label: "TTS · Qwen3-TTS CustomVoice Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -272,6 +282,9 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_tts_base.py
          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -284,52 +297,11 @@ steps:
            "
        mirror_hardwares: l4_1

-      - label: "TTS · Higgs-Audio-V3 Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_higgs_audio_v3.py
-          - vllm_omni/model_executor/models/higgs_audio_v3/
-          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
-          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_USE_DEEP_GEMM=0
-              export VLLM_MOE_USE_DEEP_GEMM=0
-              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: l4_1
-
-      - label: "Diffusion · Qwen Image Edit Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_qwen_image_edit.py
-          - vllm_omni/diffusion/models/qwen_image/
-        commands:
-          - |
-            timeout 20m bash -c "
-              pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: h100_1
-
-      - label: "Diffusion · Bagel Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_bagel.py
-          - vllm_omni/model_executor/models/bagel/
-          - vllm_omni/diffusion/models/bagel/
-          - vllm_omni/model_executor/stage_input_processors/bagel.py
-          - vllm_omni/deploy/bagel.yaml
-        commands:
-          - |
-            timeout 40m bash -c "
-              export VLLM_IMAGE_FETCH_TIMEOUT=60
-              pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: h100_1
-
      - label: "Diffusion · Wan22 Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_wan22_t2v.py
          - vllm_omni/diffusion/models/wan2_2/
+          - vllm_omni/diffusion/models/dmd2/
        commands:
          - |
            timeout 40m bash -c "
@@ -342,6 +314,9 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_cosmos3.py
          - vllm_omni/diffusion/models/cosmos3/
+          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
+          - vllm_omni/experimental/world_models/
+          - vllm_omni/model_extras/cosmos3.py
        commands:
          - |
            timeout 40m bash -c "
@@ -354,6 +329,8 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen_image.py
          - vllm_omni/diffusion/models/qwen_image/
+          - vllm_omni/diffusion/models/dmd2/
+          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
        commands:
          - |
            timeout 40m bash -c "
@@ -361,15 +338,3 @@ steps:
              pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
            "
        mirror_hardwares: h100_1
-
-      - label: "Diffusion · Qwen Image Layered Test"
-        source_file_dependencies:
-          - tests/e2e/online_serving/test_qwen_image_layered.py
-          - vllm_omni/diffusion/models/qwen_image/
-        commands:
-          - |
-            timeout 40m bash -c "
-              export VLLM_IMAGE_FETCH_TIMEOUT=60
-              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
-            "
-        mirror_hardwares: h100_1
@@ -108,6 +108,22 @@ steps:
            --run-level "full_model"
        mirror_hardwares: h100_1

+      # Own step: the image has no NanoCodec, so this file import-skips out of
+      # the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts]
+      # from downgrading transformers; see
+      # examples/offline_inference/text_to_speech/README.md.
+      - label: "TTS · Gepard-1.0 · L4"
+        timeout_in_minutes: 60
+        commands:
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
+          - uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
+          - >-
+            pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
+            -m "tts and cuda"
+            --run-level "full_model"
+        mirror_hardwares: l4_1
+
      - label: "Diffusion · L4"
        timeout_in_minutes: 60
        commands:
@@ -254,6 +254,7 @@ outputs/
 results/
 generated/
 output_*/
+/tests/e2e/accuracy/artifacts/

 # Configuration overrides
 configs/local.yaml
@@ -116,6 +116,7 @@ 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
@@ -329,11 +329,13 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for

 ### Per-model coverage

-Pilot (v1): two Merge-tier (L3) jobs in `.buildkite/cuda/test-merge.yml` upload
-per-model, per-entry-mode coverage as Buildkite artifacts — "Diffusion · Bagel
-Test" and "TTS · Qwen3-TTS Base Test". Both run on a single GPU so the pilot is
-cheap to reproduce. This is a pilot, not a rollout: no dashboard/visualization
-lives in this repo, and nothing is gated on coverage.
+Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
+per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
+Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
+diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
+another diffusion coverage job if a second pilot is needed again.) This is a
+pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
+is gated on coverage.

 **Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
 `<model_id>` is the model's directory name under
@@ -410,7 +412,8 @@ runs a coverage job needs the same.
 List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
 `source_file_dependencies` — both change what the job measures, so without them a
 change there is filtered out of normal PR builds and only surfaces in a later
-nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots.
+nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
+a synthetic job (it does not pin real merge labels).
 Editing only the surrounding CI YAML still does not schedule them, so a PR that
 touches just the wiring needs a full E2E run (or the commands run on a GPU host)
 to produce artifacts. When checking a new model's
@@ -4,6 +4,9 @@
 > `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
@@ -51,14 +54,13 @@ 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 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.
+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.

 Scratch storage is also capability-driven. `frames_per_block` reserves space for
 an uncommitted current video block, while `max_scratch_tokens_per_branch`
@@ -117,7 +117,22 @@ After (async D2H):

 ### Batch Split

-When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
+When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
+
+### Reliability, Lifecycle & Timeout Behavior
+
+1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
+   When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
+
+2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
+   The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
+
+3. **Queue Serialization & Atomic Future Resolution**:
+   - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
+   - `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
+
+4. **Storage-Aware IPC Packing**:
+   `pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.

 ## Model Coverage

@@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat
 |---|---|---|
 | **HunyuanImage-3.0** | Image | `False` |
 | **Qwen-Image** | Image | `True` |
+| **LTX-2.3** | Video / Audio | `False` |

 Other models with `step_execution=False` are also supported but not yet verified.

@@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion
 models. DLO keeps only a small number of DiT blocks on the accelerator and
 streams the remaining blocks from host memory. The distributed backend can
 either shard those host-side weights across an existing parallel group or keep
-the standard loader's rank-local weights and avoid an additional collective.
+complete rank-local block sources and avoid an additional collective.

 For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -14,10 +14,22 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec

 DLO is implemented for multi-device diffusion execution. The default
 AllGather path is the primary path for DP and SP deployments. The
-`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
-for standard-loader sharding, workstation bring-up, and systems where an
-additional DLO collective is undesirable, but it does not reduce host weight
-storage across ranks.
+`--dlo-no-use-allgather` path streams complete blocks independently and adds no
+DLO weight collective.
+
+Host storage is selected separately from the transfer protocol. The loader can
+produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
+otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
+runtime cache can normalize those final ordinary-loader DiT tensors into an
+immutable node-local mmap entry. Consequently, replicas can share either
+proven-compatible checkpoint pages or equivalent final runtime layouts. An
+opt-in CUDA registration budget can make the complete final mapping directly
+H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
+
+The Phase A shared-mmap support boundary is TP1. TP greater than one is an
+ordinary-loader compatibility path: DLO can consume the resulting TP-local
+tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
+can share each final TP coordinate among equivalent processes.

 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -41,6 +53,120 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own
 parameter-sharding lifecycle and is not allowed to be sharded a second time by
 DLO's AllGather path.

+### The loader owns host-weight planning
+
+Before it decides whether ordinary weight materialization can be skipped, the
+diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
+accepted only when preflight proves all of the following:
+
+- every required DiT parameter and persistent buffer has exactly one source;
+- runtime names, checkpoint keys, shapes, and dtypes match;
+- the runtime topology is TP1 without HSDP or online quantization; and
+- every custom loader operation is represented by a loader-owned checkpoint
+  adapter.
+
+The exact plan object is handed to DLO. The backend does not rescan checkpoint
+files, repeat the capability decision, or reconstruct names from its block
+topology. If preflight fails, the loader materializes weights normally and DLO
+consumes those runtime tensors.
+
+The plan owns only dedicated DiT component sources. If a pipeline also exposes
+ordinary sources for a text encoder or another non-DiT component, the loader
+still consumes those sources and includes their loaded names in its strict
+coverage check. Only the source prefixes covered by the plan skip ordinary
+materialization. A source that mixes DiT and non-DiT weights fails closed to
+the complete ordinary-loader path because it cannot be skipped safely as a
+unit.
+
+This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
+flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
+transforms. Model-specific direct-layout knowledge, when required, lives in a
+checkpoint adapter beside the ordinary loader.
+
+### Post-load runtime cache
+
+The Phase B cache preserves the same ownership boundary. DLO neither derives
+cache keys nor interprets model loaders. When enabled with no-AllGather, every
+rank first completes the ordinary loader, custom callbacks, post-load
+processing, strict validation, and calibration. The loader then:
+
+1. discovers final CPU DiT parameters and persistent buffers;
+2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
+   ambiguous layouts before changing the model;
+3. hashes tensor names, metadata, and final bytes;
+4. builds or joins one content-addressed entry under a strict POSIX lock; and
+5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
+   manifest, and remapped-tensor content validation succeeds.
+
+The identity excludes DP size/rank and SP rank. It includes TP world size/rank
+and conservative SP implementation/world-size guards. DP replicas and
+equivalent SP ranks can therefore share; different TP coordinates cannot.
+Final tensor content is authoritative, so an unmodeled loader/config change
+that changes bytes creates a different entry instead of risking silent
+sharing.
+
+Publication writes sharded safetensors and a JSON manifest into a same-
+filesystem temporary directory, flushes them, and atomically renames the
+directory. The manifest includes per-file SHA-256 digests and the final tensor
+content digest. Kernel `flock` ownership distinguishes a live writer from a
+stale lock file: process death releases the lock automatically. Stale temporary
+entries are removed only by the next lock owner. A bounded wait or any cache
+error retains the already-valid ordinary tensors.
+
+DLO maps this plan without moving the DiT to `meta` and without rerunning any
+post-load hook. It changes tensor storage while preserving `Parameter` objects,
+then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
+streaming hooks. The source mappings remain open through DLO disable. With the
+default zero registration budget they feed the existing bounded two-slot host
+staging path. Every streamed block is first packed from pageable mmap storage
+into one of those pinned slots and is then copied to the device. This removes
+model-sized private pinned storage, but adds a recurrent host-to-host copy and
+can amplify TP synchronization waits when several ranks compete for host
+bandwidth.
+
+This version is deliberately opt-in and steady-state-oriented. It requires
+roughly one local-disk model copy per runtime identity, has no eviction, and
+still performs ordinary loading plus full content-hash passes in every rank.
+Skip-load-on-hit is a separate follow-up after loader side effects can be
+modeled safely.
+
+### Registered runtime-cache transfer
+
+`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
+runtime-cache registration. It also changes loader selection: no-AllGather
+uses the ordinary loader and final runtime cache even when a TP1 direct-
+checkpoint plan is available. A checkpoint binding can carry a deferred
+adapter such as grouped-QKV reordering. Registering that raw source would not
+remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
+contains the transform-complete final tensors.
+
+After mapping and validating every final tensor, DLO groups source storages by
+backing file, page-aligns their address ranges, and coalesces overlapping or
+adjacent ranges within that mapping. The complete byte count is checked
+against the per-worker budget before any CUDA call. It then registers every
+range with the read-only CUDA host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. Failure rolls back any ranges already
+registered and preserves the existing two-slot staging path. Partial direct
+transfer is deliberately unsupported. A rollback failure aborts worker
+initialization so a still-registered range is retained until process/context
+teardown rather than being unsafely unmapped.
+
+On success, the hooks copy each immutable tensor view directly into its offset
+in the existing flattened rotating device buffer on the copy stream. Parameter
+repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
+implementation currently issues more H2D copy operations than block packing,
+but avoids the much larger recurrent CPU copy. A packed cache schema should be
+considered only if launch-count measurements justify its added format and
+publication complexity.
+
+Registration is per process and CUDA context; physical file-cache pages remain
+shared and are not copied into one private host allocation per worker. The
+operator must nevertheless budget page-locked memory and satisfy OS/CUDA
+registration limits for every worker. During shutdown, DLO first synchronizes
+outstanding device work, unregisters all ranges, and then closes the safetensors
+mappings. Worker teardown invokes this cleanup before destroying the
+distributed environment and CUDA context.
+
 ### AllGather path

 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -61,16 +187,46 @@ Buffers:    [current slot]       [prefetch slot]       [current slot]
 The backend uses two shared device buffers, so accelerator weight residency is
 bounded by the largest streamed blocks rather than the complete model.

+When direct checkpoint mmap is selected, the checkpoint mappings are only the
+source used to prepare each rank's persistent shard. They can be closed after
+shard preparation. Across the AllGather group, those private shards total
+approximately one runtime model copy.
+
+An effective DLO group size of one performs no collective, even when
+`dlo_use_allgather=True`; it follows the rank-local transfer path described
+below.
+
 When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
 requests must take the same execution path at every denoising step.

+Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
+independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
+plus shared mmap storage avoids coupling their request schedules and failure
+domains.
+
 ### Rank-local path without DLO AllGather

 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one. The regular model loader remains responsible for preparing each rank's
-weights, including TP-local tensors or HSDP-managed parameters. DLO then
-streams those rank-local tensors block by block using H2D copies only.
+one and streams complete blocks using H2D copies only. The host backing may be
+a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
+ordinary runtime tensors.
+
+For direct mmap, each process retains immutable safetensors views and uses two
+bounded pinned host staging slots. Processes on the same node that map the same
+files share physical checkpoint pages through the OS page cache. This removes
+the persistent private full-model copy per pure-DP process, but each process
+still packs and transfers every complete block. Sharing is node-local; each
+node has its own page cache.
+
+For a runtime-cache mapping, a positive registration budget replaces those
+host slots with direct copies from the shared final tensor views. A zero budget
+or failed registration retains the bounded staging behavior.
+
+When direct mmap preflight fails, the regular model loader remains responsible
+for preparing each rank's weights, including TP-local tensors or HSDP-managed
+parameters. Supported CPU layouts may then enter the opt-in runtime cache.
+Unsupported layouts keep one private runtime copy per process.

 This mode means:

@@ -80,21 +236,18 @@ This mode means:
  not shard weights across SP ranks.
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
  DLO's additional weight AllGather is disabled.
- Pure DP deployments keep a full host-side model copy per rank, subject to
-  shared OS page-cache behavior.
+- Pure DP deployments share one checkpoint-backed copy per node when direct
+  mmap is selected. The runtime cache can share equivalent ordinary-loader
+  outputs; its fallback keeps one private runtime copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.

-This path is intentionally not implemented by reusing the AllGather mmap
-loader. It relies on the standard loader so model-specific TP/HSDP transforms
-remain intact.
-
 ## Parallelism compatibility

 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
-| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
+| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |

 ### Combined dimensions
@@ -122,11 +275,13 @@ AllGather DP multi-concurrency requires:
 The no-AllGather path does not impose these DLO-specific synchronized-wave
 requirements.

-The mmap loader is used only by the supported DLO+AllGather path when the
-model has a compatible checkpoint layout. Online quantization is incompatible
-with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
-quantization. Models that do not meet the mmap requirements use the regular
-loader path.
+Direct checkpoint mmap can back either transfer path. It is currently limited
+to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
+ordinary loader. Online quantization remains incompatible with DLO AllGather;
+use `--dlo-no-use-allgather` or disable online quantization. Runtime-cache v1
+is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
+expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
+tied/shared storage.

 ## Validation coverage

@@ -134,15 +289,170 @@ Current source-level validation includes:

 - HSDP + DLO + AllGather rejection;
 - HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
+- loader preflight fallback for TP, HSDP, online quantization, unknown custom
+  loaders, missing keys, and shape/dtype mismatches;
+- exact loader-to-backend plan transfer and ordinary-loader fallback;
+- rank-local mmap source retention, bounded two-slot staging, and adapter
+  transforms without parameter-side flags;
+- runtime-cache content identity, sharded atomic publication, concurrent
+  writer election, corruption rebuild, lock timeout, and fail-closed layout
+  rejection;
+- loader ordering through final post-load mutation and DLO remapping without
+  rerunning post-load hooks;
+- all-or-nothing page-aligned registration, budget rejection, partial-failure
+  rollback, direct-copy staging bypass, and explicit shutdown cleanup;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
  tests.

-The highest-value missing coverage is end-to-end numerical comparison against
-ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
-on the target CUDA/NCCL or CANN/HCCL hardware.
+### B300 parallel-topology smoke matrix
+
+A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
+CUDNN attention backend, 256x256 output, two denoising steps, and
+`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
+VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
+or shared-mmap host-memory savings.
+
+| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
+|---|---:|---:|---:|---:|
+| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
+| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
+| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
+| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
+
+Within each topology, the AllGather and no-AllGather video and audio outputs
+were byte-identical. All four runs completed without an `ERROR` or traceback
+and released their device allocations. For DP4xTP1, no-AllGather direct mmap
+reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
+125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
+ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
+AllGather, because DP replicas did not share checkpoint-backed runtime
+weights. This is a functional and memory smoke test, not a production-quality
+performance or output-quality benchmark.
+
+### Host-memory measurement
+
+A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
+ordinary-loader fallback with direct mmap. Both runs used
+DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
+256x256 four-second request. The ordinary-loader workers were sampled after
+initialization. The mmap workers were sampled after one completed request, so
+the checkpoint working set had been faulted into the page cache; this is the
+more conservative point for mmap.
+
+The values below come from `/proc/<worker>/smaps_rollup` and include the whole
+worker, not only the DiT. The stable rank-to-rank difference comes from other
+pipeline components, so each worker should be compared with the same worker in
+the other storage mode.
+
+| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
+|---|---:|---:|---:|---:|---:|
+| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
+| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
+| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
+
+The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
+31.20 GiB `Pss_File`, which is the proportional charge expected when the same
+resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
+from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
+1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
+because it counts a shared physical page in every process that maps it; summed
+PSS is the appropriate node-memory comparison.
+
+### L20X runtime-cache smoke matrix
+
+Phase B was also exercised on four L20X GPUs in one NVLink domain with
+MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
+outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
+VAE TP1. This is a functional, memory, and profiler smoke test rather than a
+production-quality performance benchmark. Host values below are node-wide
+maxima during request waves from descendant-process `smaps_rollup`; they do
+not represent the private startup peak. HBM is the engine-reported request
+peak because sampled `nvidia-smi` data can miss short-lived allocations.
+
+One internal DP2xTP2 engine compared the three host-storage/transfer choices.
+Each measured wave contained two requests. NVLink TX is the physical per-link
+counter delta summed across the four GPUs; RX had the same value and is not
+added again here.
+
+| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
+|---|---:|---:|---:|---:|---:|---:|
+| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
+| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
+
+The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
+footprint while reducing node PSS by about 134 GiB relative to private loader
+weights. Cold and warm labels describe cache publication or reuse during
+initialization; both request phases ran after initialization. A warm hit still
+performed ordinary loading and full content hashing, so it was not a startup
+fast path. All measured no-AllGather private/cache outputs were byte-identical
+for matching seeds.
+
+The deployment shape targeted by the cache was tested separately as two
+independent DP1xTP2 engines on the same node. Both engines mapped the same two
+TP-coordinate entries, including the same file device/inode tuples; the cache
+occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
+ahead of profiling.
+
+| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
+|---|---:|---:|---:|---:|---:|
+| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
+| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
+
+The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
+`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
+69.3% in this host-contention smoke. Across one profiled request per engine,
+aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
+payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
+rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
+3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
+summed across ranks and can overlap; they identify the mechanism, not wall
+time. Matching private/cache outputs were byte-identical.
+
+The result confirms that zero-budget runtime-cache staging is a host-capacity
+option, not a latency-neutral replacement for private pinned weights. Broader
+SP, model, platform, and production-quality validation remains tracked in
+[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
+historical B300 rows above predate the runtime cache and must not be treated as
+Phase B memory validation.
+
+### L20X registered-runtime-cache smoke
+
+The read-only registered path was then tested with the same MiniMax-H3
+workload. These are functional and mechanism-validation measurements, not a
+production benchmark. The TP1 pair used one request per wave. Both modes moved
+exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
+and audio for both measured seeds.
+
+| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
+| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
+
+Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
+fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
+unchanged, identifying recurrent host transform/packing as the removed work.
+The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
+that launch increase was much smaller than the eliminated CPU cost.
+
+Two independent TP2 engines also registered the same two TP-coordinate cache
+entries across four workers. Compared with the zero-budget cache, mean wave
+latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
+the private-loader reference was 3.61/3.72 seconds. Registered request-time
+host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
+Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
+staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
+payload were effectively unchanged. Ordinary TP communication remained, but
+the extra wait caused by staging-driven rank skew largely disappeared.
+
+All workers explicitly unregistered before releasing their mappings, the
+system's acquired-minus-released pin counter returned to its pre-run value,
+and every GPU allocation was released. Cold runtime-cache publication remained
+expensive because it still performs ordinary loading and multiple full-content
+passes; that startup problem is outside the registered transfer scope.

 ## Recommendations

@@ -150,7 +460,13 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
  scaling path.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
  not the goal.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
-  the expectation of higher host-memory use and lower validation confidence.
+- Use **no-AllGather** when independent replica execution is required. TP1
+  direct-mmap deployments can share checkpoint pages per node. Enable the
+  runtime cache only when independent replicas have repeated final layouts and
+  benchmark the target CPU/memory topology, especially for matching
+  coordinates across multiple TP engines. On CUDA, provide an explicit
+  registration budget large enough for the complete final layout when
+  recurrent staging cost is unacceptable; leave it at zero when page-locking
+  that layout is not operationally acceptable.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
  HSDP + DLO no-AllGather path has broader end-to-end coverage.
@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
 ModelOpt checkpoint metadata is detected from its method, producer metadata, or
 `quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.

+Unless the input is a per-component dictionary, the resolved configuration is
+passed unchanged to every quantization-aware component constructed by the
+pipeline. It affects only layers supported by that component's quantization
+method; it does not rewrite arbitrary `torch.nn` modules. A per-component
+dictionary narrows this global scope through runtime layer-prefix routing.
+
 The factory uses Omni-specific overrides for methods whose diffusion behavior
 or platform dispatch is not provided by the vLLM registry. The current override
 set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
 Model code must keep precision-sensitive or unsupported modules unquantized.
 The shared `safe_quant_config()` helper is used by model integrations for
 normalization and modulation layers where a generic FP8 configuration would be
-unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
-are not implicitly quantized just because the main transformer is quantized;
-they require explicit component routing and method support.
+unsafe. Encoders exposed through the quantization factory receive a global
+configuration just like the main transformer, but only their supported vLLM
+quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
+schedulers, and tokenizers are not rewritten automatically. A component map can
+select a narrower scope, but it cannot add quantization support to a layer.

 Weight quantization and quantized KV cache are separate paths. The former is
 resolved through `quantization_config`; runtime attention/KV-cache quantization
@@ -0,0 +1,213 @@
+# 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.
@@ -22,6 +22,7 @@ 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

@@ -101,7 +101,7 @@ outputs = omni.generate(
 )

 # Access custom trajectory data
-output = outputs[0].request_output
+output = outputs[0]
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```
@@ -47,7 +47,7 @@ if __name__ == "__main__":
    omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].images
    images[0].save("coffee.png")
 ```

@@ -75,8 +75,7 @@ if __name__ == "__main__":
    ]
    omni_outputs = omni.generate(prompts)
    for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output
-        this_images = this_request_output.images
+        this_images = prompt_output.images
        for i_image, image in enumerate(this_images):
            image.save(f"p{i_prompt}-img{i_image}.jpg")
            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
@@ -37,13 +37,19 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
 | `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -98,6 +104,7 @@ th {
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
 | `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |

 The `Recipe` column links the corresponding deployment recipe. `Published` is
@@ -111,6 +118,10 @@ 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
@@ -3,7 +3,13 @@
 Distributed layerwise offloading (DLO) extends block streaming to multi-device
 deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
 the host weights and reconstructs each layer at runtime. Without AllGather,
-each rank streams its standard-loader rank-local weights independently.
+each rank streams a complete block independently. Compatible TP1 deployments
+can share checkpoint-backed host pages among processes on the same node;
+otherwise DLO streams the ordinary loader's rank-local tensors. The optional
+runtime-weight cache extends node-local sharing to final ordinary-loader
+layouts, including matching TP coordinates across independent engines. On
+CUDA, an explicit registration budget can make those shared mappings directly
+H2D-copyable and avoid recurrent mmap-to-pinned host packing.

 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -19,8 +25,10 @@ AllGather:       [AG N+1]           [AG N+2]
 Slots:           slot 0: Layer N    slot 1: Layer N+1
 ```

-AllGather communicates only request-independent weight shards, so data-
-parallel ranks may process different requests concurrently.
+AllGather communicates request-independent weight shards, but every member of
+the transfer group must request the same next block in the same collective
+order. A synchronized DP wave may contain different prompts only when those
+requests follow the same denoising and block-execution path.

 ## Usage

@@ -36,6 +44,15 @@ vllm serve /path/to/model --omni \
  --data-parallel-size 4 \
  --dlo-no-use-allgather

+# Independently scheduled replicas sharing final runtime weights on one node
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --tensor-parallel-size 2 \
+  --dlo-no-use-allgather \
+  --dlo-enable-runtime-cache \
+  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
+  --dlo-runtime-cache-pin-limit-gib 80
+
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
  --enable-distributed-layerwise-offload \
@@ -48,7 +65,10 @@ from vllm_omni import Omni
 omni = Omni(
    model="/path/to/model",
    enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
+    dlo_use_allgather=False,
+    dlo_enable_runtime_cache=True,
+    dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
+    dlo_runtime_cache_pin_limit_gib=80,
 )
 ```

@@ -59,21 +79,124 @@ omni = Omni(
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
 | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
-| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
+| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
+| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
+| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
+| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |

-## mmap weight loading
+## Host-weight loading

-The DLO plus AllGather path:
+The diffusion loader chooses host storage before DLO is enabled. It first
+attempts to build a complete, validated direct-checkpoint mmap plan. If names,
+coverage, shape, dtype, topology, or loader-callback compatibility cannot be
+proven, it runs the ordinary model loader instead. DLO consumes that result and
+does not make a second checkpoint-compatibility decision.
+
+The shared-mmap optimization in this phase is supported only with TP1. TP
+greater than one falls back before model mutation to ordinary TP-aware loading.
+DLO may still consume those TP-local tensors, but this is a compatibility path:
+it does not share checkpoint-backed runtime weights across DP replicas and
+provides no shared-mmap host-memory guarantee.
+
+When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
+gains a second storage opportunity. Every rank completes loader callbacks,
+post-load processing, validation, and calibration. It then hashes the final
+CPU DiT parameters and persistent buffers. One process per equivalent runtime
+layout publishes sharded safetensors plus a checksummed manifest; peers
+validate and mmap the same files. DLO replaces only tensor storage, preserves
+the final `Parameter` objects, and releases the private allocator-backed copies
+before installing its streaming hooks.
+
+A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
+runtime layout even when a direct TP1 checkpoint plan is available. Raw
+checkpoint mappings may require deferred loader transforms such as grouped-QKV
+reordering; registering them would still leave recurrent host work. Publishing
+the fully transformed layout once makes every streamed source directly
+copyable.
+
+The runtime layout identity excludes DP rank and SP rank. It includes TP world
+size/rank and conservative SP implementation/world-size guards. Therefore two
+independent `TP=2` engines normally create two entries: both TP0 workers map
+one entry and both TP1 workers map the other. Cache sharing does not create a
+process group or an inference-time collective.
+
+With a zero registration budget, cached tensors remain pageable mmap views.
+DLO packs each streamed block into one of two bounded pinned host slots before
+the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
+host-to-host copies; competing TP ranks or engines can turn DDR contention and
+staging skew into collective wait time.
+
+With a positive budget on CUDA, DLO instead attempts to register the complete
+page-aligned cache mapping as read-only pinned host memory. Registration is
+all-or-nothing: the total mapped range must fit the per-worker budget, and any
+CUDA registration or verification failure rolls back partial registrations and
+uses the bounded staging path. If CUDA cannot roll back a partial registration,
+worker initialization aborts instead of unmapping memory that CUDA may still
+own. On success, each immutable tensor view is copied directly into the
+existing rotating device buffers on the H2D stream, so no private host staging
+slots are allocated.
+
+CUDA registration is process-local, but it does not clone file contents. The
+same physical page-cache pages remain shared among equivalent workers; each
+worker does, however, hold its own registration reference and must satisfy the
+host's page-locking limits. Shutdown synchronizes outstanding GPU work,
+unregisters every range, and only then closes the mmap handles.
+
+The mmap plan skips only dedicated DiT weight sources. Other component sources,
+such as a text encoder loaded through the shared diffusion loader, continue to
+use their ordinary component loader. A checkpoint source that mixes DiT and
+non-DiT weights falls back completely rather than leaving an unplanned
+component uninitialized.
+
+With direct checkpoint mmap, the loader:

 1. saves non-persistent buffers such as RoPE frequencies;
 2. moves the normally created transformer to the meta device;
 3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
-4. calls model-specific `post_load_weights()` conversions; and
-5. restores the saved non-persistent buffers.
+4. applies any loader-owned bounded layout adapters while packing blocks;
+5. restores saved non-persistent buffers; and
+6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
+   hooks.

-This avoids one full checkpoint RSS copy per rank and does not require
-model-specific loading code.
+For AllGather with a group larger than one, each process copies only its
+persistent shard and then releases the source mapping. For no-AllGather, each
+process keeps the mapping open and packs complete blocks through two bounded
+pinned staging slots. Processes mapping the same files on one node share the
+immutable pages; no-AllGather still performs a complete-block H2D copy in each
+process.
+
+When the effective DLO group size is one, `dlo_use_allgather=True` does not
+perform a collective and uses the same rank-local transfer behavior.
+
+## Runtime-cache lifecycle
+
+Runtime-cache publication uses a per-layout POSIX file lock, a temporary
+directory on the same filesystem, file and final-tensor SHA-256 validation,
+`fsync`, and an atomic directory rename. A crashed writer releases its kernel
+lock; the next writer removes only that layout's stale temporary directories
+while holding the lock. Lock timeout, unsupported layout, insufficient disk,
+or validation failure keeps the already-valid ordinary tensors.
+
+The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
+vLLM-Omni may isolate that variable between stage replicas. All consumers that
+should share must receive the same `--dlo-runtime-cache-dir` on a local,
+disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
+host-memory PSS, and do not use a cross-node directory in this version.
+
+There is no automatic eviction. Budget roughly one final DiT copy per distinct
+TP coordinate, dtype/layout, and model version, plus temporary space while an
+entry is published. Operators may remove entries or the complete cache root
+only after all workers using its mmap files have stopped. The feature is
+opt-in for this reason.
+
+For direct H2D, set the registration limit above the complete page-aligned
+layout size reported in the worker log. Partial registration is intentionally
+unsupported because mixing direct and staged tensors would complicate buffer
+lifetime and make performance topology-dependent. Registration is CUDA-only;
+an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
+registration error is reported and falls back to bounded staging.

 ## Declarative topology

@@ -103,16 +226,37 @@ must enter each collective.

 ## Limitations

- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
-  Use `--dlo-no-use-allgather` or disable online quantization.
- Tensor parallel size greater than one is rejected in the mmap path because
-  it bypasses TP-aware loader callbacks. The no-AllGather path retains the
-  standard loader and remains experimental with TP.
+- Direct checkpoint mmap currently requires TP1. TP greater than one falls
+  back to the ordinary TP-aware loader. With the optional runtime cache,
+  equivalent processes at the same TP coordinate can then share that final
+  rank-local layout in no-AllGather mode. A positive registration budget also
+  chooses the final runtime cache at TP1 so deferred checkpoint transforms are
+  not repeated during inference.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
  AllGather has limited end-to-end validation.
+- Online quantization uses the ordinary loader with no-AllGather. It remains
+  incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also
+  rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
+  expert-parallel, CFG-parallel, and PP layouts.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
  `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.

+The runtime cache improves steady-state host PSS, not startup peak: every rank
+loads private weights before remapping, and validation adds full content-hash
+passes. Direct H2D removes recurrent staging but does not change that startup
+contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+
+Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
+smoke documented in the L20X runtime-cache matrix in the
+[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
+two independent TP2 engines reduced request-time host PSS by 36.6% with the
+cache, but recurrent mmap-to-pinned staging reduced combined throughput by
+69.3%. The registered path recovered near-private per-engine latency while
+retaining the shared-memory benefit; a matched TP1 run improved mean wave
+latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
+feature design for the full smoke-test caveats and metrics.
+
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.
@@ -161,6 +161,7 @@ The following tables show which models support each feature:
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -174,7 +175,7 @@ The following tables show which models support each feature:
 **Frame Interpolation Support**

 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni

 ### AudioGen

@@ -43,7 +43,7 @@ if __name__ == "__main__":
            ),
        ],
    )
-    outputs[0].request_output.images[0].save("output.png")
+    outputs[0].images[0].save("output.png")
 ```

 ## Image-to-Image (Image Editing)
@@ -72,7 +72,7 @@ if __name__ == "__main__":
            ),
        ],
    )
-    outputs[0].request_output.images[0].save("output.png")
+    outputs[0].images[0].save("output.png")
 ```

 ## Generation Parameters
@@ -55,7 +55,7 @@ if __name__ == "__main__":
    omni = Omni(model="Qwen/Qwen-Image")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].images
    images[0].save("coffee.png")
 ```

@@ -179,7 +179,7 @@ if __name__ == "__main__":
    ]
    outputs = omni.generate(prompts)
    for i, output in enumerate(outputs):
-        output.request_output.images[0].save(f"{i}.jpg")
+        output.images[0].save(f"{i}.jpg")
 ```

 !!! info
@@ -207,7 +207,7 @@ if __name__ == "__main__":
        }
    ])
    for i, output in enumerate(outputs):
-        output.request_output.images[0].save(f"{i}.jpg")
+        output.images[0].save(f"{i}.jpg")
 ```

 You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
 | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
-| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
+| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |

 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)

@@ -138,6 +138,39 @@ outputs = omni.generate(
 )
 ```

+### Global and per-component scope
+
+A plain, global quantization configuration is passed unchanged to every
+quantization-aware component constructed by a pipeline. This includes an
+eligible encoder when that encoder is implemented with vLLM quantizable
+layers; it does not rewrite arbitrary `torch.nn` modules. A structured
+component map is the only way to narrow that scope. Pipeline integrations that
+do not yet expose an encoder through the quantization factory remain DiT-only.
+
+For a pipeline that exposes both a transformer and a quantization-aware text
+encoder, the scope is:
+
+| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
+|---------------|-------------|--------------|-------------------------------------------------|
+| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
+| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
+| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
+
+Use `quantization_config` for component-selective Python configuration. For
+example, quantize only the text decoder:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization_config={"text_encoder": {"method": "fp8"}},
+)
+```
+
+Component keys are runtime prefixes exposed by the pipeline integration; common
+keys include `transformer` and `text_encoder`. Entries may be combined, and
+`ignored_layers` can keep named eligible layers in checkpoint precision. Check
+the model recipe for supported components and their runtime prefixes.
+
 CLI:

 ```bash
@@ -24,7 +24,6 @@ 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,
@@ -228,7 +227,6 @@ 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()


@@ -0,0 +1,288 @@
+# 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()
@@ -0,0 +1,239 @@
+# 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()
@@ -96,7 +96,7 @@ if __name__ == "__main__":
        ),
    )
    from diffusers.utils import export_to_video
-    frames = outputs[0].request_output.images
+    frames = outputs[0].images
    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
 ```

@@ -68,6 +68,7 @@ from vllm_omni.model_extras import (
 from vllm_omni.model_extras import (
    get_extra_body_params,
    get_model_class_name,
+    should_preserve_reference_image_size,
 )
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -454,7 +455,13 @@ def main():

    media_inputs: dict[str, Any] = {}
    if image is not None:
-        media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        preserve_image_size = should_preserve_reference_image_size(
+            model_class_name,
+            model=args.model,
+        )
+        media_inputs["image"] = (
+            image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        )
    if last_image is not None:
        media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
    if mask_image is not None:
@@ -82,9 +82,6 @@ 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()


@@ -335,7 +335,7 @@ def main(args):
    print(f"use_tts (tts_bos): {use_tts}")

    for stage_outputs in omni_generator:
-        output = stage_outputs.request_output
+        output = stage_outputs
        if stage_outputs.final_output_type == "text":
            request_id = output.request_id
            text_output = output.outputs[0].text
@@ -1,21 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
+"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
+
+The DLO modes accept any positive data-parallel size.  With AllGather enabled,
+the script submits one complete DP wave and verifies recovery after an invalid
+wave.  It also checks that every mode shuts down all worker processes cleanly.

 Examples:

    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
+    CUDA_VISIBLE_DEVICES=0,1,2,3 \
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
+
+    # Rank-local mmap storage, with no DLO collective or DP wave scheduling.
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    CUDA_VISIBLE_DEVICES=0,1 \
-    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2

    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    CUDA_VISIBLE_DEVICES=0,1 \
-    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode request
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
 """

 from __future__ import annotations
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
 )


+def positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed < 1:
+        raise argparse.ArgumentTypeError("must be a positive integer")
+    return parsed
+
+
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
    parser.add_argument(
        "--mode",
-        choices=("dlo-dp2", "request"),
+        choices=("dlo", "dlo-no-allgather", "request"),
        required=True,
-        help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
+        help="DLO (with or without AllGather) or ordinary non-DLO request mode",
+    )
+    parser.add_argument(
+        "--dp-size",
+        type=positive_int,
+        default=2,
+        help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
+    )
+    parser.add_argument(
+        "--tp-size",
+        type=positive_int,
+        default=2,
+        help="Tensor-parallel size used only by request mode (default: 2)",
    )
    parser.add_argument("--steps", type=int, default=2)
    parser.add_argument("--duration", type=float, default=5.0)
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:


 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
    common: dict[str, Any] = {
        "model": args.model,
        "trust_remote_code": True,
-        "num_gpus": 2,
+        "num_gpus": args.dp_size if is_dlo else args.tp_size,
        "usp": 1,
        "ring": 1,
        "vae_parallel_mode": "tile",
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
        "stage_init_timeout": args.init_timeout,
        "init_timeout": args.init_timeout,
    }
-    if args.mode == "dlo-dp2":
+    if is_dlo:
        common.update(
            tensor_parallel_size=1,
-            data_parallel_size=2,
+            data_parallel_size=args.dp_size,
            text_encoder_tp_size=1,
            vae_patch_parallel_size=1,
            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
+            dlo_use_allgather=args.mode == "dlo",
            dlo_resident_layers=0,
        )
    else:
        common.update(
-            tensor_parallel_size=2,
+            tensor_parallel_size=args.tp_size,
            data_parallel_size=1,
-            text_encoder_tp_size=2,
-            vae_patch_parallel_size=2,
+            text_encoder_tp_size=args.tp_size,
+            vae_patch_parallel_size=args.tp_size,
            enable_distributed_layerwise_offload=False,
        )
    return common
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
        "engine_kwargs": engine_kwargs(args),
    }
    try:
-        if args.mode == "dlo-dp2":
+        uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
+        if uses_dp_wave:
            started = time.perf_counter()
            asymmetric = await asyncio.gather(
                generate_one(
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
                    prompt="",
                    seed=1001,
                ),
-                generate_one(
-                    engine,
-                    args,
-                    request_id="invalid-peer",
-                    prompt=DEFAULT_PROMPTS[0],
-                    seed=1002,
+                *(
+                    generate_one(
+                        engine,
+                        args,
+                        request_id=f"invalid-peer-{index}",
+                        prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
+                        seed=1001 + index,
+                    )
+                    for index in range(1, args.dp_size)
                ),
                return_exceptions=True,
            )
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
            summary["asymmetric_errors"] = [
                f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
            ]
-            if len(summary["asymmetric_errors"]) != 2:
-                raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
+            if len(summary["asymmetric_errors"]) != args.dp_size:
+                raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")

        started = time.perf_counter()
-        request_count = 2 if args.mode == "dlo-dp2" else 1
+        request_count = args.dp_size if uses_dp_wave else 1
        outputs = await asyncio.gather(
            *(
                generate_one(
                    engine,
                    args,
                    request_id=f"recovery-{index}",
-                    prompt=DEFAULT_PROMPTS[index],
+                    prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
                    seed=2000 + index,
                )
                for index in range(request_count)
@@ -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(
-        "--stage-configs-path",
+        "--deploy-config",
        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.stage_configs_path:
-        omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
    omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
    outputs = omni.generate(inputs, sampling_params_list)

@@ -74,7 +74,7 @@ if __name__ == "__main__":
    omni = Omni(model="Qwen/Qwen-Image")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].images
    images[0].save("coffee.png")
 ```

@@ -265,7 +265,7 @@ if __name__ == "__main__":
    ]
    outputs = omni.generate(prompts)
    for i, output in enumerate(outputs):
-        output.request_output.images[0].save(f"{i}.jpg")
+        output.images[0].save(f"{i}.jpg")
 ```

 !!! info
@@ -301,7 +301,7 @@ if __name__ == "__main__":
        }
    ])
    for i, output in enumerate(outputs):
-        output.request_output.images[0].save(f"{i}.jpg")
+        output.images[0].save(f"{i}.jpg")
 ```

 You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
 |---|---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
+| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
 | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -395,6 +396,58 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \

 ---

+## Gepard-1.0
+
+Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
+
+### Prerequisites
+The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
+
+```bash
+uv pip install "nemo_toolkit[tts]==2.7.3" \
+    --overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
+```
+
+The overrides are what make that command work. `nemo_toolkit[tts]` declares
+`transformers~=4.57.0`, so a plain install downgrades the version
+`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
+which `import vllm` fails. Both packages are named here rather than left to resolve on
+their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
+that would make the second repair an accident of the first. `--constraint` cannot be
+used instead, since it narrows a version range rather than overriding a dependency a
+package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
+against transformers 5.x, so nothing is lost by holding it there.
+
+(The same clash is why this is documented here rather than declared in
+`pyproject.toml`: an extra resolves together with the base dependencies, so
+`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
+SVS documents `nemo_toolkit[asr]` the same way.)
+
+On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
+consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
+sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
+dies inside `profile_run` long after every import has succeeded, which reads as a model
+failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
+here.
+
+### Quick start (zero-shot, default voice)
+```bash
+python examples/offline_inference/text_to_speech/gepard/end2end.py \
+    --text "Hello, this is a Gepard demo."
+```
+
+### Voice cloning
+Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
+
+### Notes
+- Output: 22.05 kHz mono WAV.
+- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
+- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
+- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
+- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
+- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
+- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
+
 ## VoxCPM2

 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
@@ -38,17 +38,6 @@ 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(
@@ -152,7 +141,8 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )
-    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
    return omni_kwargs


@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
+
+Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
+frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
+waveform. Zero-shot: the default learned voice, no reference audio.
+
+Generation length and seed are stage settings rather than flags here. One
+output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
+frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
+
+Usage:
+  python end2end.py --text "Hello, this is Gepard speaking."
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+
+import vllm_omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+from vllm_omni import Omni  # noqa: E402
+
+MODEL = "nineninesix/gepard-1.0"
+SAMPLE_RATE = 22050
+# Absolute: the loader hands the path straight to OmegaConf, which resolves a
+# bare name against the current working directory, not the package.
+DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
+
+
+def build_request(text: str) -> dict:
+    """Build an Omni request payload for Gepard (zero-shot, no ref audio).
+
+    The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
+    a bare text prompt would have no speaker slots and no SOS.
+    """
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+    from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+
+    # From the checkpoint's own sidecar, which is what the worker builds its
+    # config from too. A layout that disagrees with the trained one does not
+    # fail, it just degrades the speech.
+    cfg = GepardConfig.from_checkpoint(MODEL)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+    prompt_token_ids = build_gepard_prompt_ids(
+        tokenizer(text, add_special_tokens=False)["input_ids"],
+        config=cfg,
+    )
+    return {
+        "prompt_token_ids": prompt_token_ids,
+        "additional_information": {"text": [text]},
+    }
+
+
+def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
+    sf.write(path, waveform.float().numpy(), sample_rate)
+    seconds = waveform.numel() / sample_rate
+    print(
+        f"  Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
+        f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
+    )
+
+
+def main(args) -> None:
+    omni = Omni(
+        model=MODEL,
+        deploy_config=args.deploy_config,
+        stage_init_timeout=args.stage_init_timeout,
+        init_timeout=args.init_timeout,
+    )
+
+    # No explicit SamplingParams: one supplied here would replace the stage
+    # defaults rather than merge over them, dropping the pipeline's stop token.
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Synthesizing: {args.text!r}")
+    inputs = build_request(text=args.text)
+
+    for i, stage_outputs in enumerate(omni.generate(inputs)):
+        for j, out in enumerate(stage_outputs.outputs):
+            mm = out.multimodal_output
+            if mm is None:
+                print(f"  [req {i}] No audio output.")
+                continue
+            # The consolidation path renames "model_outputs" to "audio".
+            audio = mm.get("audio")
+            if audio is None:
+                audio = mm.get("model_outputs")
+            if isinstance(audio, list):
+                audio = audio[0] if len(audio) else None
+            if audio is None:
+                print(f"  [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
+                continue
+            sr_tensor = mm.get("sr")
+            sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
+            save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
+
+    print("Done.")
+
+
+def parse_args():
+    parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
+    parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
+    parser.add_argument(
+        "--output-dir",
+        default=os.path.join(
+            os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
+            "gepard_output",
+        ),
+        help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        default=DEFAULT_DEPLOY_CONFIG,
+        help=(
+            "Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
+            "stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
+            "Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
+        ),
+    )
+    # The upstream defaults are too tight for a cold start.
+    parser.add_argument("--stage-init-timeout", type=int, default=900)
+    parser.add_argument("--init-timeout", type=int, default=1800)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())
@@ -32,11 +32,10 @@ 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_stage_config() -> str:
+def _default_deploy_config() -> str:
    repo_root = Path(__file__).resolve().parents[4]
    return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")

@@ -109,7 +108,7 @@ def main(args: Any) -> None:
        "output_dir",
    ):
        omni_kwargs.pop(key, None)
-    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+    omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
    omni_kwargs["log_stats"] = args.log_stats
    omni = Omni(**omni_kwargs)

@@ -156,7 +155,7 @@ def parse_args() -> Any:
        default=None,
        help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
    )
-    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--deploy-config", 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")
@@ -164,7 +163,6 @@ 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()


@@ -570,6 +570,8 @@ def main():
        num_frames=args.num_frames,
        extra_args=extra_args,
    )
+    if args.frame_rate is not None:
+        sampling_kwargs["frame_rate"] = args.frame_rate
    if args.guidance_scale_high is not None:
        sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high

@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
 **Online serving:**

 ```bash
-vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
+vllm serve bosonai/higgs-audio-v3-tts-4b \
    --host 0.0.0.0 --port 8095 \
    --trust-remote-code --omni
 ```
@@ -0,0 +1,175 @@
+# LTX-2.5
+
+> Text-to-video and first-frame image-to-video generation with synchronized audio
+
+## Overview
+
+vLLM-Omni supports the gated
+[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+checkpoint for T2V and first-frame I2V. Every output includes synchronized
+48 kHz stereo audio.
+
+The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
+repository is not directly loadable with `--model`. It supplies the official
+upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
+both model licenses and authenticate before first use.
+
+## Choose a pipeline
+
+| Pipeline | Mode | Default size | Schedule |
+|---|---|---:|---:|
+| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
+| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
+| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
+| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
+
+Both two-stage pipelines generate at half resolution, apply the official x2
+latent upsampler, and run a three-step refinement stage. Select the class with
+`--model-class-name`; no `--task-type` flag is required. Supplying one initial
+image selects I2V, while omitting it selects T2V.
+
+## Prerequisites
+
+```bash
+hf auth login
+export MODEL=Lightricks/LTX-2.5-Diffusers
+```
+
+Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
+`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
+`libx264` encoder.
+
+## Hardware
+
+| GPU | Status | Recommended scope |
+|---|---|---|
+| NVIDIA B300 | Verified | All four canonical pipelines |
+| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
+
+The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
+80 GB GPUs do not have enough safety margin for the canonical two-stage
+configuration; use a larger GPU or reduce the output configuration.
+
+### H100 80 GB FP8
+
+Use the distilled one-stage pipeline with FP8 and cuDNN attention:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+  --omni \
+  --model-class-name LTX2DistilledOneStagePipeline \
+  --quantization fp8 \
+  --diffusion-attention-backend CUDNN_ATTN \
+  --host 0.0.0.0 \
+  --port 8000
+```
+
+Use the T2V or I2V request below after the server is ready.
+
+## Offline inference
+
+Choose values from the pipeline table. For example, the distilled two-stage
+path uses:
+
+```bash
+export PIPELINE=LTX2DistilledTwoStagePipeline
+export WIDTH=1920
+export HEIGHT=1088
+export STEPS=8
+```
+
+### Text-to-video
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model "${MODEL}" \
+  --model-class-name "${PIPELINE}" \
+  --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
+  --width "${WIDTH}" \
+  --height "${HEIGHT}" \
+  --num-frames 121 \
+  --num-inference-steps "${STEPS}" \
+  --frame-rate 24 \
+  --fps 24 \
+  --seed 42 \
+  --output ltx25-t2v.mp4
+```
+
+### First-frame image-to-video
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model "${MODEL}" \
+  --model-class-name "${PIPELINE}" \
+  --image /absolute/path/to/first-frame.png \
+  --prompt "The red fox walks forward while the camera tracks alongside." \
+  --width "${WIDTH}" \
+  --height "${HEIGHT}" \
+  --num-frames 121 \
+  --num-inference-steps "${STEPS}" \
+  --frame-rate 24 \
+  --fps 24 \
+  --seed 42 \
+  --output ltx25-i2v.mp4
+```
+
+LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
+
+## Online serving
+
+Start one server for the selected pipeline:
+
+```bash
+vllm serve "${MODEL}" \
+  --omni \
+  --model-class-name "${PIPELINE}" \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --stage-init-timeout 900
+```
+
+T2V request:
+
+```bash
+curl -sS --fail-with-body \
+  -X POST http://127.0.0.1:8000/v1/videos/sync \
+  -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
+  -o ltx25-online-t2v.mp4
+```
+
+For I2V, add exactly one first frame:
+
+```bash
+export FIRST_FRAME=/absolute/path/to/first-frame.png
+
+curl -sS --fail-with-body \
+  -X POST http://127.0.0.1:8000/v1/videos/sync \
+  -F "input_reference=@${FIRST_FRAME};type=image/png" \
+  -F 'prompt=The red fox walks forward while the camera tracks alongside.' \
+  -o ltx25-online-i2v.mp4
+```
+
+Restart the server after changing `PIPELINE` because each class selects a
+different weight layout and execution topology.
+
+## Constraints
+
+- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
+  are optional and use the selected pipeline defaults when omitted. `seed` is
+  optional and random when omitted.
+- `num_frames` must be `8k+1` when overridden.
+- One-stage width and height must be divisible by 32; two-stage final
+  dimensions must be divisible by 64.
+- Full/SFT pipelines accept negative prompts. Distilled pipelines are
+  positive-only and reject negative prompts.
+- I2V accepts exactly one initial image per prompt.
+- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
+
+## References
+
+- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
+- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
+- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)
@@ -8,10 +8,12 @@
 |---|---|---|
 | `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
 | `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
-| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
 | `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
 | `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
-| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
+| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |

 Repositories in the table are download units. A full pipeline repository
 contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
@@ -26,9 +28,12 @@ image selects I2V. Both one-stage repositories declare this class, so
 the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.

 `LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
-then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
-merged distilled Transformer in both stages. Both entries support T2V and I2V;
-select their class explicitly.
+then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
+merged distilled Transformer without upsampling, while
+`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
+T2V and I2V; select their class explicitly. The deprecated
+`LTX2DistilledPipeline` name remains an alias for
+`LTX2DistilledTwoStagePipeline`.

 ## API Migration

@@ -52,8 +57,8 @@ The consolidation also removes these registry names without aliases:
 | `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
 | `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
 | `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
-| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
-| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
+| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
+| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |

 Passing any second positional argument now raises `TypeError`. These changes
 affect direct Python callers and explicit `--model-class-name` overrides;
@@ -112,10 +117,10 @@ Start a two-stage checkpoint with its explicit class:

 ```bash
 vllm serve rootonchair/LTX-2-19b-distilled --omni \
-  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
 # LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
 vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
-  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
 # Ordinary LTX-2.3 two-stage
 vllm serve diffusers/LTX-2.3-Diffusers --omni \
  --model-class-name LTX2TwoStagePipeline \
@@ -0,0 +1,195 @@
+# MiniMax-H3 on a single NPU 950PR
+
+This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
+NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
+1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
+[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
+
+## Capacity requirements
+
+| Resource | Requirement |
+| --- | ---: |
+| NPU | 1x NPU 950PR |
+| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
+| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
+| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
+| Container shared memory | 8 GiB, or a writable local filesystem for spill |
+
+The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
+sized for exactly one resident task partition at 1024x576. Do not expect a
+larger output shape or concurrency greater than one to fit.
+
+Because a single 950PR holds the whole partition in HBM, this route does not
+need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
+GPU recipes does not apply. The validated run completed inside a container
+limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
+actively harmful: it stages weights in non-reclaimable host memory and the
+container's OOM killer terminates the server with exit code -9.
+
+
+## Environment
+
+- Host architecture: x86_64
+-  driver: 25.7.rc1.6 (hal 7.35.23)
+-  firmware: 9.0.0.105.229
+- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
+- npu-smi: 25.7.rc1.6
+- Python: 3.12.13
+- PyTorch: 2.10.0+cpu
+- torch_npu: 2.10.0.post2
+- vLLM: 0.26.0
+- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
+
+Install vLLM-Omni from a checkout with MiniMax-H3 support:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e .
+```
+
+## Start a server
+
+```bash
+export MODEL=/path/to/MiniMax-H3/FL2VA
+export PORT=8000
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
+
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 1 \
+  --tensor-parallel-size 1 \
+  --usp 1 \
+  --ring 1 \
+  --text-encoder-tp-size 1 \
+  --vae-patch-parallel-size 1 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --init-timeout 14400 \
+  --stage-init-timeout 14400 \
+  --quantization int8
+```
+
+Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
+NPU implementation and fails at the first denoise step. Leave the backend
+unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
+installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
+
+H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
+
+## Validated evidence
+
+Measured on one NPU 950PR with the server command above plus
+`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
+the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
+requested 60 denoise steps and executed 59 denoise updates, so per-step latency
+is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
+124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
+
+Five requests were issued back to back with no idle time between them. The card
+is measurably faster on the first request than in steady state, so both regimes
+are reported.
+
+| Stage | First request, cold card | Thermal steady state |
+| --- | ---: | ---: |
+| End-to-end request | 462.13 s | 506.38 s |
+| Text encode | 0.55 s | 0.05 s |
+| Denoise (59 updates) | 429.29 s | 471.74 s |
+| Per denoise update | 7,276 ms | 7,996 ms |
+| VAE decode | 8.43 s | 8.14 s |
+| Worker-to-server handoff | 24 s | 27 s |
+| MP4 muxing | 0.57 s | 0.50 s |
+| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
+| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
+| Average NPU power | 566.9 W | 556.9 W |
+| Peak NPU temperature | 104 C | 104 C |
+
+Sample counts differ per row. The cold column is a single request. In the steady
+column, text encode, denoise and VAE decode are means over the four subsequent
+requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
+end-to-end, handoff and muxing are means over the two of those four that
+survived the handoff timeout described above.
+
+**Report the steady-state column.** The cold card completes denoise 9.9% faster
+while drawing only 1.8% more average power, and both regimes reach the same
+104 C ceiling, so the first request is buying a short window of higher clocks
+before the die saturates. Any benchmark that issues a single request against an
+idle card will overstate throughput by about 10%.
+
+**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
+the entire stage wall time by the 60 *requested* steps, so it absorbs text
+encoding, VAE decode and the 837 MiB handoff. All three successful requests
+satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
+per-update figure in the table is `diffuse / 59` from the pipeline profiler.
+
+Denoise accounts for 93% of the steady-state request and the handoff accounts
+for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
+to roughly 480 s.
+
+Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
+reserved high-water mark rather than live allocation: it reads 111,272 MiB after
+weight loading and before the first request, and does not fall back after
+requests complete. Re-measure for longer outputs, a different output shape, or
+concurrency greater than one.
+
+### Startup
+
+Weight loading took **31 min 50 s** on the validated host (first log line
+08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
+shared GlusterFS network mount. This is why the server command sets
+`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
+1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
+Reduce both timeouts when the partition is staged on local disk.
+
+## T2VA request example
+
+```bash
+export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
+
+curl -sS --max-time 1800 -X POST "${API_URL}" \
+  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
+  -F 'width=1024' \
+  -F 'height=576' \
+  -F 'aspect_ratio=16:9' \
+  -F 'fps=24' \
+  -F 'num_inference_steps=60' \
+  -F 'flow_shift=12' \
+  -F 'seed=1101' \
+  -F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
+  -o t2va.mp4
+```
+
+## Known limitations
+
+- INT8 online quantization is validated for T2VA on this card. Use BF16 for
+  FL2VA and Ref2VA, or re-measure before relying on it.
+- Single-card serving loads one task partition at a time. For Ref2VA, stop the
+  server and restart it against the `Ref2VA` directory.
+- Sustained load drives the die to 104 C and costs about 10% of denoise
+  throughput relative to a cold card. Warm the server with at least one full
+  request before measuring, and treat steady state as the reportable number.
+- Requests can fail during result handoff on hosts where the IPC spill path is
+  slow. See [§ Result handoff and shared
+  memory](#result-handoff-and-shared-memory).
+- Loading the checkpoint from a network mount takes over 30 minutes and forces
+  very large init timeouts. Stage the partition on local disk when possible.
+- The configuration measured here is 1024x576 at 60 steps, which differs from
+  the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
+  cost scales superlinearly with token count, so the numbers are not directly
+  comparable across recipes.
+- The image ships a `triton` package whose  backend is not built
+  (`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
+  at startup and disables Triton; the diffusion path does not need it and the
+  run completes normally.
+
+## Additional resources
+
+- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
+- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)
@@ -357,12 +357,19 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while

 ### Online FP8 quantization

-MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
-remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
-uses dynamic activation scaling during inference. By default, attention and
-MLP linears in the token refiner and main DiT blocks, the condition
-projection, and all AdaLN projections use FP8. Patch, timestep, and final
-projections remain FP32; the text encoder and VAEs are unchanged.
+MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
+text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
+weights at runtime and uses dynamic activation scaling during inference. By
+default, `--quantization fp8` quantizes eligible attention and MLP linears in
+the text decoder, token refiner, and main DiT blocks, as well as the DiT
+condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
+RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
+keep checkpoint precision.
+
+To select a component, use `--diffusion-quantization-config` with
+`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
+`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
+entries can be combined. The shorthand below enables both components.

 Add this option to an existing H3 server command:

@@ -370,6 +377,28 @@ Add this option to an existing H3 server command:
 --quantization fp8
 ```

+#### Single 96 GB GPU, no-offload capacity check
+
+Use the FL2VA-only partition for this capacity test. Loading the combined
+service would also load the Ref2VA DiT and would test a different memory
+budget. A no-offload capacity check should contain none of
+`--enable-cpu-offload`, `--enable-layerwise-offload`, or
+`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
+but does not offload model weights to the CPU.
+
+The run passes the capacity check when the server initializes, the request
+finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
+card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
+32 kHz stereo AAC audio. Report the measured headroom rather than assuming
+that every nominal 96 GB SKU exposes the same MiB total.
+
+As a capacity proxy only, the same five-second case on one B300 measured a
+92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
+encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
+6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
+they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
+reported device capacity differ across GPUs.
+
 Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
 resolves the `transformer` component before constructing the DiT, so names do
 not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -386,10 +415,10 @@ For example, keep the first main block's attention projections in BF16 with:
  '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
 ```

-The structured option replaces `--quantization fp8`. Online FP8 is currently
-incompatible with H3 layerwise offload because the offload path produces a
-weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
-parallelism and VAE tiling instead.
+The structured option replaces `--quantization fp8`. Online FP8 can be used
+with H3 layerwise offload and with distributed layerwise offload's full-weight
+per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
+supported for runtime-created FP8 weights.

 ## AMD ROCm (gfx942 / gfx950)

@@ -750,6 +779,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
 resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
 reduction.

+For direct human inspection, see the external
+[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
+which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
+[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
+also record per-task fidelity metrics and provenance without storing generated
+media in this repository.
+
 ## TeaCache acceleration

 TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -834,7 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
  H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
  pure Ulysses.
- Online FP8 is not compatible with layerwise offload.
+- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
+  use `--dlo-no-use-allgather` when combining online FP8 with DLO.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone
@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
 `model_type` auto-resolves `audex_tts.yaml`):

 ```bash
-vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
+vllm 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-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+vllm 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
@@ -55,10 +55,12 @@ recipes/
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
+| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
+| [`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 |
@@ -0,0 +1,122 @@
+# 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)
@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
 Start the online serving endpoint:

 ```bash
-vllm-omni serve /path/to/stable-audio-open-1.0 \
+vllm serve /path/to/stable-audio-open-1.0 \
  --host 0.0.0.0 \
  --port 8091 \
  --gpu-memory-utilization 0.9 \
@@ -39,6 +39,9 @@ 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
 ```
@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:

 ```bash
-vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+vllm serve inclusionAI/Ming-omni-tts-0.5B \
    --deploy-config vllm_omni/deploy/ming_tts.yaml \
    --omni \
    --port 8091 \
@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
 Start the server from the repository root:

 ```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+vllm 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-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
 ```

 Use a custom deploy config for advanced cases:

 ```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
  --deploy-config /path/to/your_glm_tts_overrides.yaml
 ```

@@ -1,15 +1,15 @@
 # Common dependencies for all platforms
-# transformers: kept exactly aligned with the upstream vLLM constraint.
-# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
-# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
-# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
+# transformers: required unified multimodal model classes are unavailable
+# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
+# regression. transformers 5.9 also changed
+# _adjust_missing_and_unexpected_keys from
 # list+list to set|set, breaking models that declare
 # _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
 # with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
 # That breakage is now handled in code instead of by pinning:
 # MossTTSNanoForGeneration wraps weight loading in
 # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 5.5.3
+transformers >= 5.10.1, < 5.15
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
@@ -39,8 +39,6 @@ def _render(changed_files: list[str]) -> str:


 def test_bootstrap_if_injected_by_step_key() -> None:
-    import yaml
-
    rendered = _render_bootstrap_pipeline(
        BOOTSTRAP_STEPS_TEMPLATE,
        decision=resolve_ci_decision([]),
@@ -142,12 +140,7 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
    ]


-COVERAGE_PILOT_LABELS = ("Diffusion · Bagel Test", "TTS · Qwen3-TTS Base Test")
-MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
-
-
-def _surviving_labels(changed_files: list[str]) -> set[str]:
-    doc = yaml.safe_load(MERGE_PIPELINE.read_text(encoding="utf-8"))
+def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
    rendered = _render_test_pipeline(doc, changed_files=changed_files)
    labels: set[str] = set()

@@ -163,6 +156,29 @@ def _surviving_labels(changed_files: list[str]) -> set[str]:
    return labels


+# Synthetic coverage-style job: shared inputs that change what the split measures.
+_COVERAGE_SHARED_INPUTS_DOC = {
+    "steps": [
+        {
+            "label": "Coverage Pilot",
+            "source_file_dependencies": [
+                "tests/e2e/online_serving/test_example.py",
+                ".buildkite/common/scripts/run_cov_split.sh",
+                "pyproject.toml",
+            ],
+            "commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
+        },
+        {
+            "label": "Unrelated Model Test",
+            "source_file_dependencies": [
+                "tests/e2e/online_serving/test_other.py",
+            ],
+            "commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
+        },
+    ],
+}
+
+
@pytest.mark.parametrize(
    "changed_file",
    [
@@ -170,18 +186,17 @@ def _surviving_labels(changed_files: list[str]) -> set[str]:
        "pyproject.toml",
    ],
 )
-def test_coverage_pilots_selected_by_their_shared_inputs(changed_file: str) -> None:
-    """The split helper and the coverage config both change what these jobs measure.
-
-    Without them in source_file_dependencies the jobs are filtered out of normal
-    PR builds, so a regression would only surface in a later nightly.
-    """
-    labels = _surviving_labels([changed_file])
-    for pilot in COVERAGE_PILOT_LABELS:
-        assert pilot in labels, f"{pilot} dropped when {changed_file} changed"
+def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
+    """Jobs that list coverage shared inputs must stay selected when those files change."""
+    labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
+    assert "Coverage Pilot" in labels
+    assert "Unrelated Model Test" not in labels


-def test_coverage_pilots_not_selected_by_unrelated_change() -> None:
-    labels = _surviving_labels(["vllm_omni/entrypoints/openai/serving_chat.py"])
-    for pilot in COVERAGE_PILOT_LABELS:
-        assert pilot not in labels
+def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
+    labels = _surviving_labels(
+        _COVERAGE_SHARED_INPUTS_DOC,
+        ["vllm_omni/entrypoints/openai/serving_chat.py"],
+    )
+    assert "Coverage Pilot" not in labels
+    assert "Unrelated Model Test" not in labels
@@ -14,7 +14,8 @@ import pytest
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig

 from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
-from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config import config_factory as config_factory_module
+from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
 from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -55,6 +56,7 @@ def clear_config_factory_caches():
    yield
    StageConfigFactory.get_hf_config.cache_clear()
    StageConfigFactory.try_infer_model_type.cache_clear()
+    _materialize_object_storage_configs.cache_clear()


 Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
@@ -1100,6 +1102,13 @@ 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"))

@@ -1375,7 +1384,8 @@ stages:
        assert async_stages[1].custom_process_input_func is None

    def test_no_bundled_legacy_stage_config_yamls(self):
-        stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
+        repo_root = Path(__file__).resolve().parents[2]
+        stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
        assert not list(stage_config_dir.glob("*.yaml"))

    def test_merge_pipeline_deploy(self):
@@ -2765,3 +2775,89 @@ class TestPipelineConfigResolvers:
            pass

        assert resolver(NotTheRightHfConfig()) is None
+
+
+class TestObjectStorageConfigResolution:
+    """Object-storage URIs must be materialized locally before any HF-style reads.
+
+    Regression coverage for review on the Run:AI PR: ``get_config`` and the
+    config.json/model_index.json fallbacks crashed with ``HFValidationError``
+    for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
+    a bucket named after another pipeline could hijack pipeline selection.
+    """
+
+    @pytest.fixture
+    def fake_object_storage(self, monkeypatch, tmp_path):
+        """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
+
+        Returns a recorder: calls to the returned callable write files into
+        the materialized directory, and ``recorder.pulls`` records every
+        ``pull_files`` invocation.
+        """
+        pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
+        materialized_files: list[tuple[str, str]] = []
+
+        class FakeObjectStorageModel:
+            def __init__(self, url: str):
+                self.dir = str(tmp_path)
+
+            def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
+                pulls.append((model_path, allow_pattern, ignore_pattern))
+                for name, content in materialized_files:
+                    (tmp_path / name).write_text(content)
+
+        monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
+
+        class Recorder:
+            def set_files(self, files: list[tuple[str, str]]) -> None:
+                materialized_files.clear()
+                materialized_files.extend(files)
+
+        recorder = Recorder()
+        recorder.pulls = pulls
+        return recorder
+
+    def test_passthrough_for_non_uri(self, fake_object_storage):
+        assert _materialize_object_storage_configs("org/model") == "org/model"
+        assert _materialize_object_storage_configs("/local/model") == "/local/model"
+        assert fake_object_storage.pulls == []
+
+    def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
+        uri = "s3://bucket/model"
+
+        first = _materialize_object_storage_configs(uri)
+        second = _materialize_object_storage_configs(uri)
+
+        assert first == second
+        assert len(fake_object_storage.pulls) == 1
+        pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
+        assert pulled_uri == uri
+        assert "*.json" in allow_pattern
+        assert ignore_pattern is None
+
+    def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
+        """A misleading bucket name must not beat the real config.json content."""
+        fake_object_storage.set_files(
+            [
+                (
+                    "config.json",
+                    '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
+                )
+            ]
+        )
+        uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
+
+        model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
+
+        assert model == "qwen3_omni_moe"
+        assert fake_object_storage.pulls  # materialization happened
+
+    def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
+        # Empty config.json (CosyVoice3 style) forces the name-match fallback.
+        fake_object_storage.set_files([("config.json", "{}")])
+
+        deceptive = "s3://qwen3-tts-models/plain-checkpoint"
+        assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
+
+        matching = "s3://any-bucket/my-cosyvoice3-model"
+        assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
@@ -424,10 +424,10 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              82.1886
+              82.6344
            ],
            "median_audio_rtf": [
-              0.1241
+              0.1277
            ]
          }
        }
@@ -453,13 +453,13 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              160.9168
+              170.3737
            ],
            "median_audio_rtf": [
-              0.1907
+              0.2076
            ],
            "audio_throughput": [
-              39.2332
+              37.0834
            ]
          }
        }
@@ -485,13 +485,13 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              100.1265
+              101.0293
            ],
            "median_audio_rtf": [
-              0.1437
+              0.1498
            ],
            "audio_throughput": [
-              11.7605
+              10.1186
            ]
          }
        }
@@ -18,6 +18,7 @@ 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"
@@ -56,6 +57,21 @@ 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
@@ -103,6 +119,7 @@ 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(
@@ -113,7 +130,10 @@ def make_runner(
    )
    runner.device = torch.device("cpu")
    runner.pipeline = pipeline
-    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
+        enable=True,
+        gpu_memory_fraction=gpu_memory_fraction,
+    )
    runner.kv_cache = None
    runner._ar_diffusion_capability = None
    runner._ar_diffusion_kv_cache_spec = None
@@ -132,6 +152,16 @@ 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)
@@ -140,6 +170,43 @@ 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)
@@ -170,6 +237,31 @@ 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
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():


 def test_lru_eviction_releases_blocks_and_notifies_pipeline():
-    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
    runner = make_runner(pipeline)
    kv = runner.kv_cache
    assert kv is not None
-    assert runner._session_capacity == 1
-    assert kv.session_capacity == 1
+    assert runner._session_capacity == 2
+    assert kv.session_capacity == 2
    free_total = kv.manager.block_pool.get_num_free_blocks()
    old = commit_one_frame(runner, "old", "main")
    k = torch.randn(1, 8, 4, 64)
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
    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",)
+    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
+
+
 def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
    pipeline = CapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline)
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
        raise RuntimeError("layer exploded")

    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
-    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+    request = SimpleNamespace(
+        request_id="broken-request",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+    )

    with pytest.raises(RuntimeError, match="layer exploded"):
        runner.execute_model(request)
@@ -263,7 +401,10 @@ 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(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+    request = SimpleNamespace(
+        request_id="broken-request",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+    )

    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
        runner.execute_model(request)
@@ -0,0 +1,231 @@
+# 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"
@@ -30,6 +30,8 @@ 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(
@@ -95,7 +97,6 @@ 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


@@ -201,12 +202,149 @@ 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},
-        device=torch.device("cpu"),
    )
    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)
+    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 expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
+    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


 def _make_kv(
@@ -221,15 +359,24 @@ 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),
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=BLOCK,
+            window_chunks=window_chunks,
+            gpu_memory_fraction=1.0,
+        ),
        num_layers=1,
        num_kv_heads=4,
        head_size=64,
        dtype=torch.float32,
        block_size=BLOCK,
        max_model_len=4096,
-        available_bytes=1 << 16,  # tiny -> the floor binds
+        available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
        device=torch.device("cpu"),
        kv_branches=kv_branches,
        session_capacity=1,
@@ -54,13 +54,15 @@ 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)
-        # Resident window stays bounded by window + the in-flight chunk.
+        # Allocation may transiently add one in-flight chunk.
        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
-        free_after.append(kv.manager.block_pool.get_num_free_blocks())
        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())

    # Pool memory plateaus once the window is full (eviction recycles blocks).
-    assert free_after[-1] == free_after[kv.spec.window_chunks]
+    assert free_after[-1] == free_after[kv.spec.window_chunks - 1]

    kv.end_request(adapter)
    assert kv.manager.block_pool.get_num_free_blocks() == free_total
@@ -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")
-    _, st = make_state(num_layers=2)
+    kv, 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,6 +349,10 @@ 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():
@@ -174,6 +174,19 @@ 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
@@ -0,0 +1,505 @@
+# 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")
@@ -0,0 +1,108 @@
+# 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
@@ -0,0 +1,53 @@
+# 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("")
@@ -5,6 +5,7 @@ from contextlib import nullcontext

 import pytest
 import torch
+from torch.nn.attention import SDPBackend

 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]


+@pytest.mark.parametrize(
+    ("query_sequence_length", "key_sequence_length", "expected_backend"),
+    [
+        (1, 1, SDPBackend.MATH),
+        (8, 1, SDPBackend.MATH),
+        (1, 8, SDPBackend.MATH),
+        (8, 8, SDPBackend.CUDNN_ATTENTION),
+    ],
+)
+def test_cudnn_routes_single_token_attention_to_math(
+    monkeypatch, query_sequence_length, key_sequence_length, expected_backend
+):
+    selected_backends = []
+
+    def fake_sdpa_kernel(backends):
+        selected_backends.append(backends)
+        return nullcontext()
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
+        fake_sdpa_kernel,
+    )
+    monkeypatch.setattr(
+        torch.nn.functional,
+        "scaled_dot_product_attention",
+        lambda query, _key, _value, **_kwargs: query,
+    )
+    impl = CuDNNAttentionImpl(
+        num_heads=2,
+        head_size=4,
+        softmax_scale=0.5,
+    )
+    query = torch.randn(1, query_sequence_length, 2, 4)
+    key = torch.randn(1, key_sequence_length, 2, 4)
+    value = torch.randn_like(key)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert output.shape == query.shape
+    assert selected_backends == [[expected_backend]]
+
+
 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
    observed = {}

@@ -36,7 +36,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

 SEPARATE_CFG_TRANSFORMERS = [
    DreamIdOmniModel,
-    LTX2VideoTransformer3DModel,
    HeliosTransformer3DModel,
    LongCatImageTransformer2DModel,
    Cosmos3VFMTransformer,
@@ -124,6 +123,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
    assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True


+def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
+    """LTX batches all guidance passes into one Transformer invocation."""
+    adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
+
+    assert adapter_config.has_separate_cfg is False
+
+
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
@@ -53,14 +53,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
 def _extract_single_image(outputs) -> Image.Image:
    first_output = outputs[0]
    assert first_output.final_output_type == "image"
-    if not hasattr(first_output, "request_output") or not first_output.request_output:
-        raise ValueError("No request_output found in OmniRequestOutput")
+    if not isinstance(first_output, OmniRequestOutput):
+        raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")

-    req_out = first_output.request_output
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
-
-    images = req_out.images
+    images = first_output.images
    if images is None or len(images) != 1:
        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
    return images[0]
@@ -72,8 +72,7 @@ def _to_float_array(data: Any) -> np.ndarray:


 def _extract_output_array(output: Any) -> np.ndarray:
-    req_out = output.request_output
-    payload = req_out.images[0]
+    payload = output.images[0]
    return _to_float_array(payload)


@@ -12,10 +12,16 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod

 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    HostWeightPlanResult,
+    TensorBinding,
+)
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model

@@ -110,6 +116,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
        loader.load_weights(model)


+def test_stream_online_quant_weights_offloads_layers_after_processing():
+    from vllm.model_executor.model_loader.reload.layerwise import (
+        get_layerwise_info,
+    )
+
+    events: list[str] = []
+
+    class _OnlineQuantMethod:
+        uses_meta_device = True
+
+    class _TrackedLayer(nn.Linear):
+        def __init__(self, name: str):
+            super().__init__(2, 2, bias=False)
+            self.name = name
+            self.quant_method = _OnlineQuantMethod()
+            get_layerwise_info(self).load_numel_total = self.weight.numel()
+
+        def to(self, *args, **kwargs):
+            events.append(self.name)
+            return super().to(*args, **kwargs)
+
+    class _StreamingModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.first = _TrackedLayer("first")
+            self.second = _TrackedLayer("second")
+
+    model = _StreamingModel()
+    weights = iter(
+        [
+            ("first.weight", torch.zeros((2, 2))),
+            ("second.weight", torch.zeros((2, 2))),
+        ]
+    )
+    streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
+
+    assert next(streamed)[0] == "first.weight"
+    get_layerwise_info(model.first).reset()
+    assert next(streamed)[0] == "second.weight"
+    assert events == ["first"]
+
+    get_layerwise_info(model.second).reset()
+    with pytest.raises(StopIteration):
+        next(streamed)
+    assert events == ["first", "second"]
+
+
+def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+    from vllm.model_executor.model_loader.reload import layerwise
+
+    class _TrackedLayer(nn.Linear):
+        def __init__(self):
+            super().__init__(2, 2, bias=False)
+            self.to_calls: list[object] = []
+
+        def to(self, *args, **kwargs):
+            self.to_calls.append(args[0] if args else kwargs.get("device"))
+            return super().to(*args, **kwargs)
+
+    model = nn.Module()
+    model.layer = _TrackedLayer()
+    quant_method = Mock(spec=QuantizeMethodBase)
+    quant_method.uses_meta_device = True
+    model.layer.quant_method = quant_method
+    model.layer._already_called_process_weights_after_loading = True
+    finalize = Mock()
+    monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
+
+    loader = _make_loader_with_weights([])
+    loader._process_weights_after_loading(model, torch.device("cuda"))
+
+    finalize.assert_called_once_with(model, model_config=None)
+    quant_method.process_weights_after_loading.assert_not_called()
+    assert model.layer.to_calls == []
+
+
 class _ConfigAwareModel(nn.Module):
    def __init__(self, *, od_config):
        super().__init__()
@@ -179,6 +264,323 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
    assert get_current_diffusion_config_or_none() is None


+def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    plan = HostWeightPlan(
+        backing_kind="checkpoint_mmap",
+        bindings={},
+    )
+    loaded_ordinary_weights = False
+
+    def load_weights(_model):
+        nonlocal loaded_ordinary_weights
+        loaded_ordinary_weights = True
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert not loaded_ordinary_weights
+    assert loader.take_host_weight_plan() is plan
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    class MixedSourceModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+            self.text_encoder = nn.Linear(2, 2, bias=False)
+            self.weights_sources = (
+                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
+                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
+            )
+            self.loaded_weight_names: list[str] = []
+
+        def load_weights(self, weights):
+            self.loaded_weight_names = [name for name, _ in weights]
+            return set(self.loaded_weight_names)
+
+    model = MixedSourceModel()
+    plan = HostWeightPlan(
+        backing_kind="checkpoint_mmap",
+        bindings={
+            "transformer.weight": TensorBinding(
+                storage_key="weight",
+                file_path="unused",
+            )
+        },
+        planned_source_prefixes=frozenset({"transformer."}),
+    )
+    requested_prefixes: list[str] = []
+
+    def get_weights(source, model=None):
+        del model
+        requested_prefixes.append(source.prefix)
+        yield source.prefix + "weight", torch.ones(2, 2)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader._get_weights_iterator = get_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert requested_prefixes == ["text_encoder."]
+    assert model.loaded_weight_names == ["text_encoder.weight"]
+    assert loader.take_host_weight_plan() is plan
+
+
+def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            allgather_degree=1,
+            ulysses_mode="strict",
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_lock_timeout=1.0,
+        model="unused",
+        revision=None,
+        model_config={},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnquantizedLinearMethod()
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    def load_weights(_model):
+        events.append("load")
+        with torch.no_grad():
+            model.transformer.weight.fill_(1)
+
+    def process_weights(_model, _device):
+        events.append("process")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    def calibrate(_model):
+        events.append("calibrate")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    original_eval = model.eval
+
+    def eval_model():
+        events.append("eval")
+        return original_eval()
+
+    def build_cache(_pipeline, **kwargs):
+        events.append("cache")
+        assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
+        assert kwargs["tensor_parallel_rank"] == 0
+        assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
+        assert kwargs["quantization_config"] is None
+        return HostWeightPlanResult(runtime_plan)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._process_weights_after_loading = process_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = calibrate  # type: ignore[method-assign]
+    model.eval = eval_model  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "calibrate", "eval", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+
+    def build_cache(_pipeline, **kwargs):
+        assert kwargs["quantization_config"] == "module_quant_method"
+        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
+
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    original_weight = model.transformer.weight.detach().clone()
+    assert loader.load_model(load_device="cpu") is model
+    assert torch.equal(model.transformer.weight, original_weight)
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+
+import huggingface_hub
+import pytest
+
+from vllm_omni.diffusion.model_loader import hub_prefetch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_prefetch_subfolders_propagates_revision(monkeypatch):
+    calls = []
+
+    def fake_snapshot_download(**kwargs):
+        calls.append(kwargs)
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
+
+    hub_prefetch.prefetch_subfolders(
+        "org/model",
+        ["transformer"],
+        revision="pinned-revision",
+    )
+
+    assert calls == [
+        {
+            "repo_id": "org/model",
+            "revision": "pinned-revision",
+            "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
+        }
+    ]
+
+
+def test_from_pretrained_retry_preserves_revision(monkeypatch):
+    factory_calls = []
+    prefetch_calls = []
+
+    def fake_factory(model, **kwargs):
+        factory_calls.append((model, kwargs))
+        if len(factory_calls) == 1:
+            raise OSError("partial cache")
+        return "loaded"
+
+    def fake_prefetch(model, subfolders, **kwargs):
+        prefetch_calls.append((model, tuple(subfolders), kwargs))
+
+    monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
+    monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
+
+    result = hub_prefetch.from_pretrained_with_prefetch(
+        fake_factory,
+        "org/model",
+        subfolder="transformer",
+        prefetch_list=["transformer", "vae"],
+        revision="pinned-revision",
+        max_attempts=2,
+    )
+
+    assert result == "loaded"
+    assert len(factory_calls) == 2
+    assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
+    assert prefetch_calls == [
+        (
+            "org/model",
+            ("transformer", "vae"),
+            {"local_files_only": False, "revision": "pinned-revision"},
+        )
+    ]
@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fcntl
+import multiprocessing as mp
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+    build_runtime_weight_cache_plan,
+    default_runtime_cache_root,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _LoaderMarker:
+    pass
+
+
+class _Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(3, 2)
+        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
+        self.register_buffer("temporary", torch.ones(1), persistent=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
+            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
+
+
+class _Pipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer()
+
+
+def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
+    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
+
+    assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
+
+
+def _build(pipeline: nn.Module, cache_root: Path, **overrides):
+    kwargs = {
+        "dit_modules": (("transformer", pipeline.transformer),),
+        "loader_type": _LoaderMarker,
+        "cache_root": cache_root,
+        "lock_timeout_seconds": 2.0,
+        "max_shard_bytes": 32,
+        "model_identity": "test/model",
+        "revision": "revision",
+        "runtime_dtype": torch.float32,
+        "load_format": "default",
+        "loader_inputs": {"model_config": {"hidden_size": 3}},
+        "tensor_parallel_size": 1,
+        "tensor_parallel_rank": 0,
+        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
+        "use_hsdp": False,
+        "enable_expert_parallel": False,
+        "quantization_config": None,
+        "cfg_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+    }
+    kwargs.update(overrides)
+    return build_runtime_weight_cache_plan(pipeline, **kwargs)
+
+
+def _multiprocess_builder(cache_root: str, queue) -> None:
+    result = _build(_Pipeline(), Path(cache_root))
+    queue.put(
+        (
+            result.fallback_code,
+            result.plan.runtime_layout_key if result.plan is not None else None,
+            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
+        )
+    )
+
+
+def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    second = _build(_Pipeline(), tmp_path)
+
+    assert first.plan is not None
+    assert second.plan is not None
+    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
+    assert first.plan.post_load_complete
+    assert {binding.file_path for binding in first.plan.bindings.values()} == {
+        binding.file_path for binding in second.plan.bindings.values()
+    }
+    assert set(first.plan.bindings) == {
+        "transformer.persistent",
+        "transformer.proj.bias",
+        "transformer.proj.weight",
+    }
+
+
+def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
+    pipeline = _Pipeline()
+    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
+    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
+    sp2 = _build(
+        pipeline,
+        tmp_path,
+        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
+    )
+
+    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
+    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
+
+
+def test_final_runtime_content_is_authoritative(tmp_path):
+    first_pipeline = _Pipeline()
+    second_pipeline = _Pipeline()
+    with torch.no_grad():
+        second_pipeline.transformer.proj.weight[0, 0] += 1
+
+    first = _build(first_pipeline, tmp_path)
+    second = _build(second_pipeline, tmp_path)
+
+    assert first.plan is not None and second.plan is not None
+    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
+
+
+def test_expert_parallel_ownership_fails_closed(tmp_path):
+    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
+
+    assert result.plan is None
+    assert result.fallback_code == "unsupported_expert_parallel"
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    assert first.plan is not None
+    shard = Path(next(iter(first.plan.bindings.values())).file_path)
+    old_inode = shard.stat().st_ino
+    with shard.open("r+b") as handle:
+        handle.seek(-1, 2)
+        byte = handle.read(1)
+        handle.seek(-1, 2)
+        handle.write(bytes([byte[0] ^ 0xFF]))
+
+    second = _build(_Pipeline(), tmp_path)
+
+    assert second.plan is not None
+    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
+    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
+
+
+def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
+
+    with lock_path.open("a+b") as lock_handle:
+        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
+
+    assert result.plan is None
+    assert result.fallback_code == "lock_timeout"
+
+
+def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
+    stale_temp.mkdir(parents=True)
+    (stale_temp / "partial").write_bytes(b"partial")
+
+    result = _build(_Pipeline(), tmp_path)
+
+    assert result.plan is not None
+    assert not stale_temp.exists()
+
+
+@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
+def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
+    pipeline = _Pipeline()
+    if case == "alias":
+        pipeline.transformer.tied = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    elif case == "external_alias":
+        pipeline.encoder = nn.Module()
+        pipeline.encoder.shared = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    else:
+        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
+        expected_code = "unsupported_tensor"
+
+    result = _build(pipeline, tmp_path)
+
+    assert result.plan is None
+    assert result.fallback_code == expected_code
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
+    context = mp.get_context("spawn")
+    queue = context.Queue()
+    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
+    for process in processes:
+        process.start()
+    results = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert all(code is None for code, _, _ in results)
+    assert len({key for _, key, _ in results}) == 1
+    assert len({tuple(paths) for _, _, paths in results}) == 1
+    assert len(list((tmp_path / "v1").iterdir())) == 1
@@ -6,7 +6,10 @@ from types import SimpleNamespace

 import pytest

-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+    DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
+    DreamZeroPipeline,
+)
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +35,37 @@ 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()

@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
    Flux2Pipeline,
    _resolve_component_quant_config,
 )
+from vllm_omni.quantization import ComponentQuantizationConfig

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


 def test_resolve_component_quant_config_routes_components():
    text_encoder_config = object()
-    component_config = Mock()
-    component_config.resolve.side_effect = {
-        "text_encoder": text_encoder_config,
-        "transformer": None,
-    }.get
+    component_config = ComponentQuantizationConfig(
+        component_configs={"text_encoder": text_encoder_config},
+        default_config=None,
+    )

    assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
    assert _resolve_component_quant_config(component_config, "transformer") is None
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,13 @@
+{
+  "_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"]
+}
@@ -0,0 +1,42 @@
+{
+  "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]
+  }
+}
@@ -0,0 +1,30 @@
+{
+  "_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
+}
@@ -0,0 +1,112 @@
+{
+  "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"
+  ]
+}
@@ -0,0 +1,294 @@
+# 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())
@@ -0,0 +1,543 @@
+# 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]))
@@ -0,0 +1,525 @@
+# 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)
@@ -0,0 +1,685 @@
+# 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
@@ -0,0 +1,39 @@
+# 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 == []
@@ -0,0 +1,881 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Focused LTX-2.5 pipeline correctness tests."""
+
+import json
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
+from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+    LTX25_DISTILLED_COMPONENT_PROFILE,
+    LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+    LTX25_FULL_COMPONENT_PROFILE,
+    LTX25_TWO_STAGE_COMPONENT_PROFILE,
+    detect_ltx_model_version,
+    resolve_ltx_checkpoint_kind,
+    resolve_ltx_component_profile,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
+    LTXPhaseResult,
+    _first_frame_keyframes_mask,
+    _official_ltx_sigmas,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+    LTX_GUIDANCE_EXECUTOR,
+    LTXGuidanceSpec,
+    LTXModalityGuidance,
+    velocity_from_x0,
+    x0_from_velocity,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_TWO_STAGE_RECIPE,
+    LTX2_ONE_STAGE_RECIPE,
+    LTX23_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    LTX25_FULL_RECIPE,
+    LTX25_TWO_STAGE_RECIPE,
+    LTX_DISTILLED_SIGMAS,
+    LTX_POSITIVE_ONLY_RECIPE,
+    resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
+    LTX2Pipeline,
+)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+    LTX2DistilledTwoStagePipeline,
+    LTX2TwoStagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
+def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
+    pipe = object.__new__(LTX2Pipeline)
+    pipe.model_version = model_version
+    pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+    ts = torch.tensor([750.0, 500.0])
+
+    kwargs = LTXRuntime._denoise_timestep_kwargs(
+        pipe,
+        ts,
+        SimpleNamespace(),
+        SimpleNamespace(),
+        video_token_count=4,
+        audio_token_count=3,
+    )
+
+    torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
+    torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
+    assert kwargs["sigma"] is ts
+    assert kwargs["audio_sigma"] is ts
+
+
+def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
+    state = LTXAVState(
+        video=torch.zeros(1, 2, 2),
+        audio=torch.zeros(1, 4, 2),
+    )
+    forward_ctx = SimpleNamespace(
+        original_audio_num_frames=3,
+        sampler="euler_ancestral",
+    )
+    denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
+    pipe = SimpleNamespace(
+        _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
+    )
+
+    def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
+        assert actual_forward_ctx.sampler == "euler_ancestral"
+        updated_audio = torch.ones_like(state.audio)
+        updated_audio[:, 3:] = 99.0
+        return torch.ones_like(state.video), updated_audio
+
+    monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
+
+    actual = LTXRuntime._denoise_step(
+        pipe,
+        0,
+        torch.tensor(1.0),
+        state,
+        forward_ctx,
+        denoise_ctx,
+    )
+
+    torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
+    torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
+
+
+def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
+    reference = torch.zeros(2, 12, 8)
+
+    mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
+
+    assert mask.shape == (2, 12, 1)
+    torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
+    torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
+
+
+def _make_ltx_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    pipe.vae_spatial_compression_ratio = 32
+    pipe.vae_temporal_compression_ratio = 8
+    return pipe
+
+
+def _resolve_request_inputs_for_test(
+    pipe,
+    req,
+    *,
+    guidance_scale=4.0,
+    negative_prompt=None,
+):
+    return pipe._resolve_request_inputs(
+        req,
+        prompt=None,
+        negative_prompt=negative_prompt,
+        height=None,
+        width=None,
+        num_frames=None,
+        frame_rate=None,
+        num_inference_steps=None,
+        guidance_scale=guidance_scale,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=None,
+    )
+
+
+def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
+    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
+    od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
+
+    with pytest.raises(ImportError) as exc_info:
+        ltx2_components.initialize_pipeline_components(pipe, od_config)
+
+    assert str(exc_info.value) == (
+        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
+    )
+
+
+def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
+    revision = "pinned-revision"
+    calls = {"components": []}
+    profile = replace(
+        LTX25_DISTILLED_COMPONENT_PROFILE,
+        text_encoder_cls=object,
+        video_vae_cls=object,
+        vocoder_cls=object,
+        vocoder_fallback_cls=None,
+        scheduler_use_dynamic_shifting=False,
+    )
+    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
+    od_config = SimpleNamespace(
+        model="org/converted-ltx25",
+        revision=revision,
+        dtype=torch.bfloat16,
+        quantization_config=None,
+    )
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
+
+    def fake_prefetch(model, subfolders, **kwargs):
+        calls["prefetch"] = (model, tuple(subfolders), kwargs)
+
+    def fake_tokenizer(*args, **kwargs):
+        calls["tokenizer"] = (args, kwargs)
+        return SimpleNamespace(model_max_length=1024)
+
+    def fake_component(component_cls, model, subfolder, **kwargs):
+        calls["components"].append((component_cls, model, subfolder, kwargs))
+        if subfolder == "vae":
+            return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
+        if subfolder == "audio_vae":
+            return SimpleNamespace(
+                mel_compression_ratio=4,
+                temporal_compression_ratio=4,
+                config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
+            )
+        return object()
+
+    def fake_transformer_config(model, subfolder, local_files_only, *, revision):
+        calls["transformer_config"] = (model, subfolder, local_files_only, revision)
+        return {"component": "transformer"}
+
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
+    monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+    monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+    monkeypatch.setattr(
+        ltx2_components,
+        "create_transformer_from_config",
+        lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
+    )
+
+    def fake_scheduler(*args, **kwargs):
+        calls["scheduler"] = (args, kwargs)
+        return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
+
+    monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
+
+    ltx2_components.initialize_pipeline_components(pipeline, od_config)
+
+    assert pipeline.weights_sources[0].revision == revision
+    assert calls["prefetch"][2]["revision"] == revision
+    assert calls["tokenizer"][1]["revision"] == revision
+    assert {call[2] for call in calls["components"]} == {
+        "text_encoder",
+        "connectors",
+        "vae",
+        "audio_vae",
+        "vocoder",
+        "latent_upsampler",
+    }
+    assert all(call[3]["revision"] == revision for call in calls["components"])
+    assert calls["transformer_config"] == (
+        od_config.model,
+        profile.transformer_subfolder,
+        False,
+        revision,
+    )
+    assert calls["scheduler"][1]["revision"] == revision
+
+
+def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
+    (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
+    (tmp_path / "transformer").mkdir()
+    (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
+
+    assert detect_ltx_model_version(str(tmp_path)) == "2.3"
+
+
+def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+
+    assert pipe.model_version == "2.5"
+    assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
+    assert pipe.preserve_sp_padded_audio_duration
+    assert pipe.reports_stage_durations
+
+
+def test_ltx25_full_recipe_matches_official_sft_defaults():
+    (phase,) = LTX25_FULL_RECIPE.phases
+
+    assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
+    assert LTX25_FULL_RECIPE.num_frames == 121
+    assert LTX25_FULL_RECIPE.frame_rate == 24.0
+    assert LTX25_FULL_RECIPE.num_inference_steps == 30
+    assert phase.guidance.video == LTXModalityGuidance(
+        cfg_scale=3.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.guidance.audio == LTXModalityGuidance(
+        cfg_scale=7.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.sigmas is None
+    assert phase.noise_scale == 1.0
+    assert phase.use_official_sigma_schedule
+    assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
+
+
+def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
+    scheduler = SimpleNamespace(
+        config={
+            "base_image_seq_len": 1024,
+            "max_image_seq_len": 4096,
+            "base_shift": 0.95,
+            "max_shift": 2.05,
+            "shift_terminal": 0.1,
+        }
+    )
+
+    sigmas = _official_ltx_sigmas(
+        scheduler,
+        steps=30,
+        device=torch.device("cpu"),
+    )
+
+    torch.testing.assert_close(
+        sigmas[[0, 1, -2, -1]],
+        torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2DistilledTwoStagePipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
+
+    assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+
+def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
+    for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
+        path = tmp_path / subfolder
+        path.mkdir()
+        (path / "config.json").write_text(json.dumps({"marker": marker}))
+
+    config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
+    config.enrich_config()
+
+    assert config.tf_model_config.get("marker") == "full"
+
+
+def test_ltx25_four_public_pipeline_semantics_are_disjoint():
+    assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
+    assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+    assert LTX2Pipeline.pipeline_kind == "one_stage"
+    assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
+    assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
+    assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
+
+
+@pytest.mark.parametrize(
+    ("pipeline_kind", "expected"),
+    [
+        ("one_stage", "regular"),
+        ("two_stage", "regular"),
+        ("distilled_one_stage", "distilled"),
+        ("distilled_two_stage", "distilled"),
+        ("dmd2", None),
+    ],
+)
+def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
+    assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
+
+
+def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
+    with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
+        resolve_ltx_checkpoint_kind("unknown")
+
+
+def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
+    stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert stage1.spatial_downscale == 2
+    assert stage1.sigmas == LTX_DISTILLED_SIGMAS
+    assert stage1.sampler == "euler_ancestral"
+    assert stage2.sampler == "euler"
+    assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+    assert stage2.input_transform == "spatial_upsample"
+
+    assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
+    assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
+
+
+def test_ltx25_distilled_one_stage_uses_native_low_resolution():
+    (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
+    assert phase.spatial_downscale == 1
+    assert phase.sigmas == LTX_DISTILLED_SIGMAS
+    assert phase.sampler == "euler_ancestral"
+    assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
+
+
+def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
+    stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
+    assert stage1.spatial_downscale == 2
+    assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
+    assert stage2.adapter_slot == "ltx_distilled"
+    assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX2_ONE_STAGE_RECIPE,
+        LTX23_ONE_STAGE_RECIPE,
+        LTX_POSITIVE_ONLY_RECIPE,
+    ],
+)
+def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
+    assert recipe.supports_cache_dit
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX25_FULL_RECIPE,
+        LTX25_DISTILLED_ONE_STAGE_RECIPE,
+        LTX2_DISTILLED_TWO_STAGE_RECIPE,
+        LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    ],
+)
+def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
+    assert not recipe.supports_cache_dit
+
+
+def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+    sigma = torch.tensor(0.5)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        sigma,
+        LTXModalityGuidance(),
+    )
+    expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
+
+    assert actual.dtype == torch.float32
+    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+
+
+def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
+    item_calls = []
+
+    class TrackingTensor(torch.Tensor):
+        @staticmethod
+        def __new__(cls):
+            return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
+
+        def item(self):
+            item_calls.append(True)
+            return super().item()
+
+    sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
+    x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
+
+    velocity_from_x0(sample, x0, TrackingTensor())
+
+    assert item_calls == [True]
+
+
+def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        torch.tensor(0.5),
+        LTXModalityGuidance(),
+        preserve_positive_velocity=True,
+    )
+
+    assert actual is velocity
+
+
+@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
+def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
+    expected_dtype = torch.bfloat16
+    captured = {}
+
+    def prepare_latents(**kwargs):
+        captured["video_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
+
+    def prepare_audio_latents(*_args, **kwargs):
+        captured["audio_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
+
+    pipeline = SimpleNamespace(
+        model_version=model_version,
+        transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
+        prepare_latents=prepare_latents,
+        prepare_audio_latents=prepare_audio_latents,
+        audio_sampling_rate=48_000,
+        audio_hop_length=160,
+        audio_vae_temporal_compression_ratio=4,
+        audio_vae_mel_compression_ratio=4,
+        audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
+        _resolve_audio_latent_length=lambda length, _latents: length,
+    )
+    request = SimpleNamespace(
+        num_videos_per_prompt=1,
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        audio_latents_normalized=False,
+    )
+    prompt = SimpleNamespace(
+        batch_size=1,
+        positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+        positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+    )
+
+    LTXRuntime._prepare_video_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+    LTXRuntime._prepare_audio_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+
+    assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
+
+
+def test_ltx_seeded_latents_match_official_packed_rng_layout():
+    pipeline = SimpleNamespace(
+        model_version="2.5",
+        vae_spatial_compression_ratio=2,
+        vae_temporal_compression_ratio=2,
+        transformer_spatial_patch_size=1,
+        transformer_temporal_patch_size=1,
+        audio_vae_mel_compression_ratio=2,
+        od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
+    )
+    device = torch.device("cpu")
+    actual_generator = torch.Generator(device=device).manual_seed(123)
+    expected_generator = torch.Generator(device=device).manual_seed(123)
+
+    actual_video = ltx2_latents.prepare_video_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        height=4,
+        width=6,
+        num_frames=3,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+    actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        audio_latent_length=3,
+        num_mel_bins=4,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+
+    expected_video = torch.randn((1, 12, 2), generator=expected_generator)
+    expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
+
+    torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
+    assert (original_length, padded_length) == (3, 3)
+
+
+def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
+    request_inputs = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="",
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        num_inference_steps=8,
+        guidance=LTXGuidanceSpec.positive_only(),
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.25,
+        decode_noise_scale=0.5,
+        output_type="np",
+        max_sequence_length=16,
+    )
+    prompt_context = object()
+    phase_calls = []
+
+    def resolve_request_inputs(req, **kwargs):
+        return request_inputs
+
+    def run_phase(req, inputs, *, prompt_context=None, **kwargs):
+        phase_recipe = kwargs["phase_recipe"]
+        assert kwargs["image"] is source_image
+        phase_calls.append((phase_recipe, inputs, prompt_context))
+        if len(phase_calls) == 1:
+            assert prompt_context is None
+            assert inputs.decode_timestep == 0.25
+            assert inputs.decode_noise_scale == 0.5
+            context = prompt_context_sentinel
+            assert phase_recipe.name == "generate_lowres"
+            assert (inputs.height, inputs.width) == (32, 32)
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_1_sigmas
+            assert kwargs["noise_scale"] == 1.0
+            video = torch.ones(1, 128, 1, 1, 1)
+            audio = torch.full((1, 8, 1, 2), 2.0)
+            audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
+        else:
+            assert prompt_context is prompt_context_sentinel
+            assert phase_recipe.name == "refine"
+            assert (inputs.height, inputs.width) == (64, 64)
+            assert inputs.latents.shape == (1, 128, 1, 2, 2)
+            torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
+            assert inputs.audio_latents_normalized
+            assert inputs.guidance_scale == 1.0
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_2_sigmas
+            assert kwargs["noise_scale"] == stage_2_sigmas[0]
+            assert inputs.decode_timestep == 0.0
+            assert inputs.decode_noise_scale is None
+            context = prompt_context
+            video = torch.full((1, 128, 1, 2, 2), 3.0)
+            audio = torch.full((1, 8, 1, 2), 4.0)
+            audio_for_next_phase = None
+        return LTXPhaseResult(
+            forward_context=SimpleNamespace(prompt_context=context),
+            video=video,
+            audio=audio,
+            audio_for_next_phase=audio_for_next_phase,
+        )
+
+    def decode_phase(phase):
+        return DiffusionOutput(output=(phase.video, phase.audio))
+
+    class FakeUpsampler(torch.nn.Module):
+        dtype = torch.float32
+
+        def forward(self, latents):
+            return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+
+    prompt_context_sentinel = prompt_context
+    stage_1_sigmas = [1.0, 0.5, 0.0]
+    stage_2_sigmas = [0.8, 0.2, 0.0]
+    pipeline = object.__new__(LTX2TwoStagePipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
+    pipeline.device = torch.device("cpu")
+    pipeline.vae_spatial_compression_ratio = 32
+    pipeline.vae_temporal_compression_ratio = 8
+    pipeline.latent_upsampler = FakeUpsampler()
+    object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
+    object.__setattr__(pipeline, "run_phase", run_phase)
+    object.__setattr__(pipeline, "decode_phase", decode_phase)
+
+    source_image = object()
+    output = pipeline.forward(
+        SimpleNamespace(sampling_params_list=[], prompts=[]),
+        image=source_image,
+        stage_1_sigmas=stage_1_sigmas,
+        stage_2_sigmas=stage_2_sigmas,
+    )
+
+    assert len(phase_calls) == 2
+    assert phase_calls[1][2] is prompt_context_sentinel
+    torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
+    torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
+
+
+class TestLTXRequestParsing:
+    def test_request_resolves_custom_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        sigmas = [1.0, 0.5, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
+                    request_id="ltx-custom-sigmas",
+                )
+            ]
+        )
+
+        assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
+
+    def test_request_resolves_per_stage_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        stage_1_sigmas = [1.0, 0.4, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
+                    request_id="ltx-custom-phase-sigmas",
+                )
+            ]
+        )
+
+        resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
+            req, [1.0, 0.0], [0.8, 0.0]
+        )
+
+        assert resolved == (stage_1_sigmas, [0.8, 0.0])
+
+    @pytest.mark.parametrize("value", [-1, 52, 1.5, True])
+    def test_request_rejects_invalid_image_crf(self, value):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
+                    request_id="ltx-invalid-image-crf",
+                )
+            ]
+        )
+
+        with pytest.raises(ValueError, match="image_crf"):
+            _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
+
+
+class TestLTXForwardStages:
+    def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
+        pipe = _make_ltx_request_pipe(LTX2Pipeline)
+
+        class FakeTokenizer:
+            bos_token_id = 2
+            eos_token_id = 1
+            eos_token = "<eos>"
+            pad_token_id = 0
+            pad_token = "<pad>"
+            padding_side = "right"
+
+            def __init__(self):
+                self.calls = []
+
+            def __call__(self, text, **kwargs):
+                self.calls.append((text, kwargs))
+                ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
+                return SimpleNamespace(input_ids=ids)
+
+            def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
+                assert padding == "max_length"
+                assert return_tensors == "pt"
+                assert return_attention_mask is True
+                rows = []
+                masks = []
+                for ids in encoded["input_ids"]:
+                    pad = max_length - len(ids)
+                    rows.append([self.pad_token_id] * pad + ids)
+                    masks.append([0] * pad + [1] * len(ids))
+                return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
+
+        class FakeTextEncoder(torch.nn.Module):
+            dtype = torch.float32
+
+            def __init__(self):
+                super().__init__()
+                self.seen_input_ids = None
+                self.seen_attention_mask = None
+
+            def forward(self, *, input_ids, attention_mask, output_hidden_states):
+                assert output_hidden_states is True
+                self.seen_input_ids = input_ids.clone()
+                self.seen_attention_mask = attention_mask.clone()
+                return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
+
+        tokenizer = FakeTokenizer()
+        text_encoder = FakeTextEncoder()
+        object.__setattr__(pipe, "tokenizer", tokenizer)
+        object.__setattr__(pipe, "text_encoder", text_encoder)
+
+        pipe._get_gemma_prompt_embeds(
+            [" long ", "short"],
+            max_sequence_length=4,
+            device=torch.device("cpu"),
+            dtype=torch.float32,
+        )
+
+        torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
+        torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
+        assert tokenizer.padding_side == "left"
+        assert [call[0] for call in tokenizer.calls] == ["long", "short"]
@@ -3,6 +3,7 @@

 """Unit tests for LTX image-to-video input and conditioning behavior."""

+import sys
 from types import SimpleNamespace

 import pytest
@@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls):


 class TestLTXImageToVideoForwardStages:
-    def test_i2v_pil_preprocessing_matches_official_center_crop(self):
+    def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images

        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages:

        torch.testing.assert_close(actual, expected)

+    def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+
+        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
+        pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
+        image = Image.fromarray(pixels.numpy())
+        monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
+
+        actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
+        expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
+
+        torch.testing.assert_close(actual, expected)
+
+    def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_crfs = []
+        captured_prepare_kwargs = []
+
+        def fake_preprocess(image, *, height, width, crf, device, dtype):
+            del image, device
+            captured_crfs.append(crf)
+            return torch.zeros(1, 3, height, width, dtype=dtype)
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=Image.new("RGB", (8, 4)),
+        )
+
+        assert captured_crfs == [18]
+        assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
+        assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
+
+    def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+
+        with pytest.raises(ValueError, match="requires PIL images"):
+            pipe._prepare_video_latents_stage(
+                SimpleNamespace(),
+                SimpleNamespace(),
+                device=torch.device("cpu"),
+                noise_scale=0.0,
+                image=torch.zeros(3, 4, 4),
+            )
+
+    def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_prepare_kwargs = []
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+            image_crf=0,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=torch.zeros(3, 4, 4),
+        )
+
+        assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
+
+    def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        assert _apply_image_conditioning_crf(image, 18) is image
+
+    def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
+            _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
+
    def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
        from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages:
        with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
            pipe._resolve_request_image(mixed_req, None, request_inputs)

-    def test_denoise_timestep_kwargs_masks_video_only(self):
+    def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
@@ -189,13 +320,18 @@ class TestLTXImageToVideoConditioning:
        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))

-    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
        torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
        pipe.vae_spatial_compression_ratio = 1
        pipe.vae_temporal_compression_ratio = 1
        pipe.transformer_spatial_patch_size = 1
@@ -237,14 +373,20 @@ class TestLTXImageToVideoConditioning:

        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]

-    def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
        torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
        pipe.vae_spatial_compression_ratio = 1
        pipe.vae_temporal_compression_ratio = 1
        pipe.transformer_spatial_patch_size = 1
@@ -266,6 +408,7 @@ class TestLTXImageToVideoConditioning:
            return torch.ones(shape, device=device, dtype=dtype)

        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)

        out, conditioning_mask = pipe.prepare_latents(
            image=torch.zeros(1, 3, 1, 1),
@@ -279,10 +422,54 @@ class TestLTXImageToVideoConditioning:
            device=torch.device("cpu"),
        )

-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]
        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
        torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))

+    def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        object.__setattr__(
+            pipe,
+            "_encode_i2v_image_latents",
+            lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
+        )
+        captured = {}
+
+        def fake_create_noised_state(latents, noise_scale, generator=None):
+            del generator
+            captured["noise_scale"] = noise_scale
+            return latents
+
+        monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
+        pipe.prepare_latents(
+            image=torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=0.909375,
+            dtype=torch.bfloat16,
+            device=torch.device("cpu"),
+            latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
+        )
+
+        assert captured["noise_scale"] == 0.909375
+
    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

@@ -385,6 +572,99 @@ class TestLTXImageToVideoConditioning:
        torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
        assert pipeline.step_indices == [0, 1]

+    def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+        device = torch.device("cpu")
+        sigmas = torch.tensor([1.0, 0.5, 0.0])
+        pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
+        mask = torch.tensor([[1.0, 0.0]])
+        adapter = LTXVideoAudioStepAdapter(
+            pipeline,
+            SimpleNamespace(sigmas=sigmas.clone()),
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            image_conditioned=True,
+            sampler="euler_ancestral",
+            generator=torch.Generator(device=device).manual_seed(42),
+            conditioning_mask=mask,
+        )
+        video = torch.tensor([[[3.0], [2.0]]])
+        audio = torch.tensor([[[4.0]]])
+        video_velocity = torch.tensor([[[7.0], [0.25]]])
+        audio_velocity = torch.tensor([[[0.5]]])
+
+        reference_generator = torch.Generator(device=device).manual_seed(10042)
+        video_noise = torch.randn(video.shape, generator=reference_generator)
+        audio_noise = torch.randn(audio.shape, generator=reference_generator)
+
+        def official_step(sample, velocity, noise):
+            sigma = sigmas[0]
+            sigma_next = sigmas[1]
+            denoised = sample.float() - velocity.float() * sigma
+            sigma_down = sigma_next * (sigma_next / sigma)
+            ratio = sigma_down / sigma
+            deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
+            alpha_next = 1.0 - sigma_next
+            alpha_down = 1.0 - sigma_down
+            coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+            return (alpha_next / alpha_down) * deterministic + noise * coeff
+
+        expected_video = official_step(video, video_velocity, video_noise)
+        expected_video[:, :1] = video[:, :1]
+        expected_audio = official_step(audio, audio_velocity, audio_noise)
+
+        ((actual_video, actual_audio),) = adapter.step(
+            (video_velocity, audio_velocity),
+            (sigmas[0], sigmas[0]),
+            (video, audio),
+        )
+        torch.testing.assert_close(actual_video, expected_video)
+        torch.testing.assert_close(actual_audio, expected_audio)
+
+        terminal_video_velocity = torch.full_like(actual_video, 0.125)
+        terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
+        ((terminal_video, terminal_audio),) = adapter.step(
+            (terminal_video_velocity, terminal_audio_velocity),
+            (sigmas[1], sigmas[1]),
+            (actual_video, actual_audio),
+        )
+        expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
+        expected_terminal_video[:, :1] = actual_video[:, :1]
+        expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
+        torch.testing.assert_close(terminal_video, expected_terminal_video)
+        torch.testing.assert_close(terminal_audio, expected_terminal_audio)
+
+    def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
+
+        sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
+        velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
+        sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
+        generator = torch.Generator(device="cpu").manual_seed(10042)
+        reference_generator = torch.Generator(device="cpu").manual_seed(10042)
+
+        sigma, sigma_next = sigmas[:2]
+        denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
+        sigma_down = sigma_next * (sigma_next / sigma)
+        sigma_down_ratio = sigma_down / sigma
+        expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
+        alpha_next = 1.0 - sigma_next
+        alpha_down = 1.0 - sigma_down
+        coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+        noise = torch.randn(
+            sample.shape,
+            generator=reference_generator,
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+        expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
+
+        actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
    def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -202,3 +202,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
            device=torch.device("cpu"),
            latents=latents,
        )
+
+
+def test_create_noised_state_matches_official_fp32_lerp():
+    latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
+        latents.dtype
+    )
+
+    actual = latent_ops.create_noised_state(
+        latents,
+        torch.tensor(0.15, dtype=torch.bfloat16),
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_create_conditioned_noised_state_matches_official_two_lerps():
+    latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
+    clean_latents = latents.clone()
+    clean_latents[:, 0] = 3.0
+    denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    noised = torch.lerp(latents.float(), noise.float(), 0.909375)
+    expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
+
+    actual = latent_ops.create_conditioned_noised_state(
+        latents,
+        clean_latents,
+        denoise_mask,
+        0.909375,
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_prepare_supplied_video_latents_uses_official_token_major_layout():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
+    actual_video = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        num_frames=9,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=supplied,
+    )
+
+    torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
+    assert actual_video.stride()[1:] == (1, actual_video.shape[1])
+
+
+def test_clear_audio_padding_keeps_padding_outside_sampler_state():
+    updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
+
+    actual = latent_ops.clear_audio_padding(updated, 2)
+
+    torch.testing.assert_close(actual[:, :2], updated[:, :2])
+    torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))
@@ -28,6 +28,15 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
 class TestCFGParallelHelpers:
    """Test LTX-2.3 CFG helper math without loading model weights."""

+    def test_single_guidance_pass_preserves_official_token_layout(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
+
+        tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
+        repeated = _repeat_batch(tokens, 1)
+
+        assert repeated is tokens
+        assert repeated.stride() == tokens.stride()
+
    def test_combine_cfg_noise_matches_x0_space_formula(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

@@ -20,6 +20,7 @@ from vllm_omni.diffusion.models.ltx2 import (
 from vllm_omni.diffusion.models.ltx2.ltx2_components import (
    LTX2_COMPONENT_PROFILE,
    LTX2_DISTILLED_COMPONENT_PROFILE,
+    LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
    LTX2_TWO_STAGE_COMPONENT_PROFILE,
    LTX23_COMPONENT_PROFILE,
    LTX23_DISTILLED_COMPONENT_PROFILE,
@@ -47,6 +48,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
 from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_ONE_STAGE_RECIPE,
    LTX2_DISTILLED_TWO_STAGE_RECIPE,
    LTX2_ONE_STAGE_RECIPE,
    LTX2_TWO_STAGE_RECIPE,
@@ -64,12 +66,14 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
    LTX2I2VDMD2Pipeline,
    LTX2Pipeline,
    LTX2T2VDMD2Pipeline,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
    LTX2DistilledPipeline,
+    LTX2DistilledTwoStagePipeline,
    LTX2TwoStagePipeline,
 )

@@ -125,15 +129,21 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():

    assert issubclass(LTX2Pipeline, LTXRuntime)
    assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
    assert issubclass(LTX2DistilledPipeline, LTXRuntime)
+    assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
    assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
    assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
    assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
    assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
+    assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
    assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
    for pipeline_cls, profile in (
        (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
        (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
+        (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
        (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
    ):
        assert pipeline_cls._dit_modules == list(profile.dit_modules)
@@ -254,7 +264,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa

    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)

-    with pytest.raises(ValueError, match="only one-stage LTX recipes"):
+    with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
        LTX2TwoStagePipeline(
            od_config=SimpleNamespace(
                model=str(tmp_path),
@@ -709,15 +719,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):

    assert len(installed) == 2
    assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
-    assert all(processor.has_learned_registers for processor in installed)
    assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
    assert all(kernel["role_category"] == "self" for kernel in kernels)
    assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
    assert all(kernel["disable_kv_quant"] for kernel in kernels)
+    assert all(processor.has_learned_registers for processor in installed)
+    assert not any(processor.preserve_learned_register_mask for processor in installed)


-@pytest.mark.parametrize("has_learned_registers", [False, True])
-def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
+def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
+    installed = []
+
+    class OmniAttention:
+        def __init__(self, **_kwargs):
+            pass
+
+    class Attention:
+        heads = 32
+        head_dim = 128
+        inner_kv_dim = 4096
+
+        def set_processor(self, processor):
+            installed.append(processor)
+
+    monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
+    connectors = SimpleNamespace(
+        video_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+        audio_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+    )
+
+    _install_connector_attention(connectors, preserve_learned_register_mask=True)
+
+    assert all(processor.has_learned_registers for processor in installed)
+    assert all(processor.preserve_learned_register_mask for processor in installed)
+
+
+def test_ltx_connector_attention_dispatches_through_omni_kernel():
    calls = []

    class Backend:
@@ -732,12 +775,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
            calls.append((query, key, value, metadata))
            return query

+    class UpcastingNorm(torch.nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
    attention = SimpleNamespace(
        heads=2,
        head_dim=4,
        inner_kv_dim=8,
-        norm_q=torch.nn.Identity(),
-        norm_k=torch.nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
        to_q=torch.nn.Identity(),
        to_k=torch.nn.Identity(),
        to_v=torch.nn.Identity(),
@@ -746,24 +793,29 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
        rope_type="split",
        omni_attention=OmniAttention(),
    )
-    hidden_states = torch.randn(2, 3, 8)
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
    additive_mask = torch.zeros(2, 1, 3, 3)
+    rotary = (
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+    )

-    output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
+    output = ltx2_components._LTXConnectorAttnProcessor()(
        attention,
        hidden_states,
        attention_mask=additive_mask,
+        query_rotary_emb=rotary,
    )

    query, key, value, metadata = calls[0]
    assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
-    if has_learned_registers:
-        assert metadata is None
-    else:
-        assert metadata.attn_mask.shape == (2, 3)
-        assert metadata.attn_mask.dtype == torch.bool
-        assert metadata.attn_mask.all()
-    torch.testing.assert_close(output, hidden_states)
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata.attn_mask.shape == (2, 3)
+    assert metadata.attn_mask.dtype == torch.bool
+    assert metadata.attn_mask.all()
+    expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
+    expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)


 def test_ltx2_distilled_pipeline_shares_recipe_runtime():
@@ -1707,6 +1759,7 @@ class TestLTXForwardStages:
        assert seen["request_inputs"].output_type == "latent"
        assert seen["kwargs"] == {
            "request_sigmas": request_sigmas,
+            "request_phase_sigmas": None,
            "image": None,
        }

@@ -1824,6 +1877,16 @@ class TestRegistryIntegration:
            "pipeline_ltx2_two_stage",
            "LTX2TwoStagePipeline",
        )
+        assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2",
+            "LTX2DistilledOneStagePipeline",
+        )
+        assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2_two_stage",
+            "LTX2DistilledTwoStagePipeline",
+        )
        assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
            "ltx2",
            "pipeline_ltx2_two_stage",
@@ -1848,6 +1911,8 @@ class TestRegistryIntegration:
        expected = [
            "LTX2Pipeline",
            "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
            "LTX2DistilledPipeline",
            "LTX2T2VDMD2Pipeline",
            "LTX2I2VDMD2Pipeline",
@@ -1861,6 +1926,8 @@ class TestRegistryIntegration:
        [
            "LTX2Pipeline",
            "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
            "LTX2DistilledPipeline",
            "LTX2T2VDMD2Pipeline",
            "LTX2I2VDMD2Pipeline",
@@ -7,18 +7,54 @@ import pytest
 import torch
 from torch import nn

-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
+from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
-    LTX2AudioVideoAttnProcessor,
+    LTX2FeedForward,
    LTX2VideoTransformer3DModel,
-    _LTX2ParallelAttention,
+    LTX2VideoTransformerBlock,
    _make_rms_norm,
+    apply_interleaved_rotary_emb,
+    apply_keyframes_absolute_embedding,
+    apply_split_rotary_emb,
 )

 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


+def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(0)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 8, dtype=x.dtype)
+    sin = torch.randn(2, 3, 8, dtype=x.dtype)
+    x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
+    rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    expected = x * cos + rotated * sin
+
+    actual = apply_interleaved_rotary_emb(x, (cos, sin))
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(1)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 4, dtype=x.dtype)
+    sin = torch.randn(2, 3, 4, dtype=x.dtype)
+    split_x = x.reshape(2, 3, 2, 4)
+    first_x, second_x = split_x.chunk(2, dim=-2)
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+    expected = split_x * cos_u
+    expected[..., :1, :].addcmul_(-sin_u, second_x)
+    expected[..., 1:, :].addcmul_(sin_u, first_x)
+    expected = expected.reshape_as(x)
+
+    actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)

@@ -36,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
    assert "weight" in norm.state_dict()


-def test_ltx_transformer_has_separate_cfg_cache_dit_config():
+def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
+    calls = []
+
+    class UpcastingNorm(nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
+    class RecordingAttention:
+        def __call__(self, query, key, value, metadata):
+            calls.append((query, key, value, metadata))
+            return query
+
+    attention = SimpleNamespace(
+        heads=2,
+        head_dim=4,
+        to_qkv=None,
+        to_q=nn.Identity(),
+        to_k=nn.Identity(),
+        to_v=nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
+        to_gate_logits=None,
+        attn=RecordingAttention(),
+        to_out=(nn.Identity(), nn.Identity()),
+        rope_type="split",
+    )
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    rotary = (
+        torch.ones(2, 2, 3, 2, dtype=torch.float32),
+        torch.zeros(2, 2, 3, 2, dtype=torch.float32),
+    )
+
+    processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
+    processor._slice_rope_for_tp = lambda rope, _: rope
+    output = processor(attention, hidden_states, query_rotary_emb=rotary)
+
+    query, key, value, metadata = calls[0]
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata is None
+    torch.testing.assert_close(output, hidden_states)
+
+
+def test_ltx_transformer_has_fused_cfg_cache_dit_config():
    adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")

    assert isinstance(adapter_config, CacheDiTAdapterConfig)
-    assert adapter_config.has_separate_cfg
+    assert not adapter_config.has_separate_cfg


 def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -76,8 +154,9 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):


 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
-    monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
-    processor = LTX2AudioVideoAttnProcessor()
+    processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
+    monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
+    processor = processor_cls()
    hidden_states = torch.zeros(1, 3, 4)
    encoder_hidden_states = torch.zeros(1, 4, 4)
    key_padding_mask = torch.tensor([[True, True, True, False]])
@@ -113,7 +192,7 @@ def _fake_ltx_parallel_attention():

        return forward

-    layer = object.__new__(_LTX2ParallelAttention)
+    layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
    nn.Module.__init__(layer)
    layer.attn_backend = Backend
    layer.attention = SimpleNamespace(forward=run("native"))
@@ -123,33 +202,281 @@ def _fake_ltx_parallel_attention():


@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
-def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
+def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
    layer, calls = _fake_ltx_parallel_attention()
    query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
    key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
    key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+    metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])

    layer._run_local_attention(query, key, value, metadata)

    assert calls == ["sdpa"]


-def test_ltx_masked_self_attention_uses_sdpa():
-    layer, calls = _fake_ltx_parallel_attention()
-    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
-    key_padding_mask = torch.tensor([[True, True, True, False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
-
-    layer._run_local_attention(query, key, value, metadata)
-
-    assert calls == ["sdpa"]
-
-
-def test_ltx_unmasked_cross_attention_keeps_native_backend():
+def test_ltx_unmasked_attention_keeps_native_backend():
    layer, calls = _fake_ltx_parallel_attention()
    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)

    layer._run_local_attention(query, key, value, None)

    assert calls == ["native"]
+
+
+def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
+    class FakeParallelLinear(nn.Linear):
+        def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
+            super().__init__(dim_in, dim_out, bias=bias)
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
+
+    without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
+    with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
+
+    assert without_bias.net[0].proj.bias is None
+    assert without_bias.net[2].bias is None
+    assert with_bias.net[0].proj.bias is not None
+    assert with_bias.net[2].bias is not None
+
+
+def _build_tiny_ltx_transformer(monkeypatch, **overrides):
+    class FakeBlock(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.init_kwargs = kwargs
+
+        def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
+            return hidden_states, audio_hidden_states
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
+    kwargs = {
+        "in_channels": 4,
+        "out_channels": 4,
+        "num_attention_heads": 2,
+        "attention_head_dim": 4,
+        "cross_attention_dim": 8,
+        "vae_scale_factors": (1, 1, 1),
+        "pos_embed_max_pos": 4,
+        "base_height": 8,
+        "base_width": 8,
+        "audio_in_channels": 4,
+        "audio_out_channels": 4,
+        "audio_num_attention_heads": 2,
+        "audio_attention_head_dim": 4,
+        "audio_cross_attention_dim": 8,
+        "audio_scale_factor": 1,
+        "audio_pos_embed_max_pos": 4,
+        "audio_sampling_rate": 16,
+        "audio_hop_length": 4,
+        "num_layers": 1,
+        "caption_channels": 8,
+        "use_prompt_embeddings": False,
+        "cross_attn_mod": True,
+        "audio_cross_attn_mod": True,
+    }
+    kwargs.update(overrides)
+    return LTX2VideoTransformer3DModel(**kwargs)
+
+
+def test_ltx25_keyframe_absolute_embedding_matches_official():
+    hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
+    keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
+    embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
+
+    actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
+    expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
+
+
+def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
+    legacy_model = _build_tiny_ltx_transformer(monkeypatch)
+
+    assert legacy_model.config.ff_bias is True
+    assert legacy_model.config.audio_ff_bias is True
+    assert legacy_model.config.use_prompt_adaln_single is True
+    assert legacy_model.config.use_keyframes_abs_pos_embedding is False
+    assert hasattr(legacy_model, "prompt_adaln")
+    assert hasattr(legacy_model, "audio_prompt_adaln")
+    assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
+    assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
+    assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    ltx25_model = _build_tiny_ltx_transformer(
+        monkeypatch,
+        ff_bias=False,
+        audio_ff_bias=True,
+        use_prompt_adaln_single=False,
+        use_keyframes_abs_pos_embedding=True,
+    )
+
+    assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
+    assert not hasattr(ltx25_model, "prompt_adaln")
+    assert not hasattr(ltx25_model, "audio_prompt_adaln")
+    assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
+    assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
+    expected = torch.arange(8, dtype=torch.float32).view(1, 8)
+    loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
+
+    assert loaded == {"keyframes_abs_pos_embedding"}
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
+
+
+def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
+    class CaptureAda(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.calls = []
+
+        def forward(self, timestep, *, batch_size, hidden_dtype):
+            self.calls.append(timestep.detach().clone())
+            embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
+            return embedding, embedding
+
+    class CaptureRope(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.output_dtypes = []
+
+        def forward(self, _coords, *, device, out_dtype=None):
+            self.output_dtypes.append(out_dtype)
+            empty = torch.empty(0, device=device, dtype=out_dtype)
+            return empty, empty
+
+    model = _build_tiny_ltx_transformer(monkeypatch)
+    rope = CaptureRope()
+    model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
+    model.prompt_modulation = False
+    video_scale_shift = CaptureAda()
+    video_gate = CaptureAda()
+    audio_scale_shift = CaptureAda()
+    audio_gate = CaptureAda()
+    model.av_cross_attn_video_scale_shift = video_scale_shift
+    model.av_cross_attn_video_a2v_gate = video_gate
+    model.av_cross_attn_audio_scale_shift = audio_scale_shift
+    model.av_cross_attn_audio_v2a_gate = audio_gate
+
+    model(
+        hidden_states=torch.zeros(1, 2, 4),
+        audio_hidden_states=torch.zeros(1, 1, 4),
+        encoder_hidden_states=torch.zeros(1, 1, 8),
+        audio_encoder_hidden_states=torch.zeros(1, 1, 8),
+        timestep=torch.tensor([[0.0, 2.0]]),
+        audio_timestep=torch.tensor([[5.0]]),
+        sigma=torch.tensor([2.0]),
+        audio_sigma=torch.tensor([5.0]),
+        video_coords=torch.zeros(1, 3, 2, 2),
+        audio_coords=torch.zeros(1, 1, 1, 2),
+    )
+
+    torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
+    torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
+    assert len(rope.output_dtypes) == 4
+    assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
+
+
+def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
+    class CaptureAttention(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.encoder_hidden_states = None
+
+        def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+            if encoder_hidden_states is not None:
+                self.encoder_hidden_states = encoder_hidden_states.detach().clone()
+            return torch.zeros_like(hidden_states)
+
+    class ZeroOutput(nn.Module):
+        def forward(self, hidden_states):
+            return torch.zeros_like(hidden_states)
+
+    block = object.__new__(LTX2VideoTransformerBlock)
+    nn.Module.__init__(block)
+    block.video_cross_attn_adaln = True
+    block.audio_cross_attn_adaln = True
+    block.cross_attn_adaln = True
+    block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+    block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
+    block.norm1 = nn.Identity()
+    block.audio_norm1 = nn.Identity()
+    block.norm2 = nn.Identity()
+    block.audio_norm2 = nn.Identity()
+    block.norm3 = nn.Identity()
+    block.audio_norm3 = nn.Identity()
+    block.attn1 = CaptureAttention()
+    block.audio_attn1 = CaptureAttention()
+    block.attn2 = CaptureAttention()
+    block.audio_attn2 = CaptureAttention()
+    block.ff = ZeroOutput()
+    block.audio_ff = ZeroOutput()
+
+    hidden_states = torch.ones(1, 1, 2)
+    audio_hidden_states = torch.ones(1, 1, 2)
+    encoder_hidden_states = torch.ones(1, 1, 2)
+    audio_encoder_hidden_states = torch.ones(1, 1, 2)
+    block(
+        hidden_states,
+        audio_hidden_states,
+        encoder_hidden_states,
+        audio_encoder_hidden_states,
+        temb=torch.zeros(1, 1, 18),
+        temb_audio=torch.zeros(1, 1, 18),
+        temb_ca_scale_shift=torch.empty(0),
+        temb_ca_audio_scale_shift=torch.empty(0),
+        temb_ca_gate=torch.empty(0),
+        temb_ca_audio_gate=torch.empty(0),
+        temb_prompt=None,
+        temb_prompt_audio=None,
+        use_a2v_cross_attention=False,
+        use_v2a_cross_attention=False,
+    )
+
+    torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
+    torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
+
+
+def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
+    column_calls = []
+
+    class FakeColumnParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **kwargs):
+            column_calls.append((output_size, kwargs))
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeRowParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeParallelAttention(nn.Module):
+        def __init__(self, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+    attention = ltx2_transformer.LTX2Attention(
+        query_dim=8,
+        heads=4,
+        kv_heads=4,
+        dim_head=2,
+        qk_norm="rms_norm",
+        apply_gated_attention=True,
+        pack_qkv=False,
+    )
+
+    gate_output_size, gate_kwargs = column_calls[-1]
+    assert gate_output_size == 4
+    assert gate_kwargs["gather_output"] is False
+    assert gate_kwargs["return_bias"] is False
+    assert attention.query_num_heads == 2
@@ -5,7 +5,7 @@ import json
 import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import Mock, patch

 import numpy as np
 import pytest
@@ -363,8 +363,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
    if expected_partition == "ref2va":
        assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
    assert [source.model_or_path for source in pipeline.weights_sources] == [
-        str(tmp_path / partition_name) for partition_name in source_partitions
+        *(str(tmp_path / partition_name) for partition_name in source_partitions),
+        str(component_path),
    ]
+    assert [source.subfolder for source in pipeline.weights_sources] == [
+        *("transformer" for _ in source_partitions),
+        "text_encoder",
+    ]
+    expected_prefixes = ["transformer."]
+    if expected_dits == 2:
+        expected_prefixes.append("transformers_ref.")
+    expected_prefixes.append("text_encoder.")
+    assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
    expected_patterns = (
        pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
        if expected_partition == "combined"
@@ -936,6 +946,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
    assert list(encoder.named_parameters()) == []


+def test_global_quant_config_is_shared_by_dit_and_encoder():
+    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+    global_config = build_quant_config("fp8")
+
+    assert resolve_component_quant_config(global_config, "transformer") is global_config
+    assert resolve_component_quant_config(global_config, "text_encoder") is global_config
+
+
+def test_minimax_h3_fp8_component_selection():
+    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+    dit_only = build_quant_config({"transformer": {"method": "fp8"}})
+    encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}})
+
+    assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8"
+    assert resolve_component_quant_config(dit_only, "text_encoder") is None
+    assert resolve_component_quant_config(encoder_only, "transformer") is None
+    assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8"
+
+
+def test_text_encoder_rejects_serialized_fp8():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config
+
+    config = SimpleNamespace(
+        get_name=lambda: "fp8",
+        is_checkpoint_fp8_serialized=True,
+    )
+    with pytest.raises(ValueError, match="online FP8 only"):
+        _validate_encoder_quant_config(config)
+
+
+def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module
+
+    group = SimpleNamespace(rank_in_group=0, world_size=1)
+    method = UnquantizedLinearMethod()
+    quant_config = Mock()
+    quant_config.get_quant_method.return_value = method
+    prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj"
+
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+            group,
+            input_size=4,
+            intermediate_size=8,
+            dtype=torch.bfloat16,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    assert linear.quant_method is method
+    quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
+
+
 def test_no_offload_keeps_text_encoder_resident():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

@@ -978,15 +1048,21 @@ def test_model_offload_uses_hooked_text_encoder_call():
    pipeline.text_encoder.load_to_device.assert_not_called()


-def test_layerwise_offload_releases_text_encoder():
+@pytest.mark.parametrize(
+    "offload_flag",
+    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+)
+def test_layerwise_offload_releases_text_encoder(offload_flag):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=False,
-        enable_layerwise_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
    )
+    setattr(pipeline.od_config, offload_flag, True)
    pipeline.text_encoder = Mock()
    expected = torch.ones(2, 3)
    pipeline.text_encoder.encode_ids.return_value = expected
@@ -1988,21 +2064,25 @@ def _one_layer_text_encoder():
    group = _SingleRankEncoderGroup(0)
    layer = nn.Module()
    layer.self_attn = nn.Module()
-    layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
-        group,
-        hidden_size=_ENCODER_HIDDEN,
-        num_heads=_ENCODER_NUM_HEADS,
-        num_kv_heads=_ENCODER_NUM_KV_HEADS,
-        head_dim=_ENCODER_HEAD_DIM,
-        dtype=torch.float32,
-    )
    layer.mlp = nn.Module()
-    layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
-        group,
-        input_size=_ENCODER_HIDDEN,
-        intermediate_size=_ENCODER_INTERMEDIATE,
-        dtype=torch.float32,
-    )
+    with (
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+    ):
+        layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
+            group,
+            hidden_size=_ENCODER_HIDDEN,
+            num_heads=_ENCODER_NUM_HEADS,
+            num_kv_heads=_ENCODER_NUM_KV_HEADS,
+            head_dim=_ENCODER_HEAD_DIM,
+            dtype=torch.float32,
+        )
+        layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+            group,
+            input_size=_ENCODER_HIDDEN,
+            intermediate_size=_ENCODER_INTERMEDIATE,
+            dtype=torch.float32,
+        )
    layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)

    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
@@ -2035,6 +2115,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()):
    safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
    index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
    (path / "model.safetensors.index.json").write_text(json.dumps(index))
+    return tensors


@pytest.mark.parametrize(
@@ -2056,10 +2137,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
    tmp_path, omitted_sources, expected_detail, unreported_target
 ):
    encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
+    weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)

    with pytest.raises(RuntimeError) as excinfo:
-        encoder._load_weights(str(tmp_path))
+        encoder.load_weights(weights.items())

    message = str(excinfo.value)
    assert expected_detail in message
@@ -2068,17 +2149,17 @@ def test_encoder_load_reports_missing_fused_source_shards(

 def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
    encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
+    weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))

    with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
-        encoder._load_weights(str(tmp_path))
+        encoder.load_weights(weights.items())


 def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
    encoder = _one_layer_text_encoder()
-    _write_text_encoder_checkpoint(tmp_path)
+    weights = _write_text_encoder_checkpoint(tmp_path)

-    encoder._load_weights(str(tmp_path))
+    encoder.load_weights(weights.items())

    params = dict(encoder.named_parameters())
    q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
@@ -36,7 +36,13 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
    assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
    assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
    assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
-    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+    assert set(MiniMaxH3DiTModel._sp_plan) == {
+        "sp_prepare",
+        "local_sp_prepare",
+        "sp_gather",
+    }
+    assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
+    assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}

    model = object.__new__(MiniMaxH3DiTModel)
    nn.Module.__init__(model)
@@ -89,6 +95,279 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)


+def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        _build_rope_table,
+    )
+
+    local_freqs_half = torch.randn(2, 48)
+    local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1)
+
+    actual = _build_rope_table(local_freqs)
+    expected = torch.cat(
+        (torch.cos(local_freqs_half), torch.sin(local_freqs_half)),
+        dim=-1,
+    ).to(torch.bfloat16)
+
+    assert actual.shape == (2, 96)
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
+    from vllm_omni.diffusion import forward_context
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict")
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1)
+
+    assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2)
+
+
+@pytest.mark.parametrize(
+    ("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"),
+    [
+        (False, "strict", 16, 8, 1, 1),
+        (True, "advanced_uaa", 16, 8, 1, 1),
+        (True, "strict", 17, 8, 1, 1),
+        (True, "strict", 16, 4, 2, 1),
+        (True, "strict", 16, 1, 1, 8),
+    ],
+)
+def test_local_span_falls_back_when_local_embedding_is_unsafe(
+    monkeypatch,
+    hooks_applied,
+    mode,
+    seq_len,
+    ulysses,
+    ring,
+    allgather,
+):
+    from vllm_omni.diffusion import forward_context
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_ulysses_parallel_world_size",
+        lambda: ulysses,
+    )
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_allgather_parallel_world_size",
+        lambda: allgather,
+    )
+
+    assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == (
+        0,
+        seq_len,
+    )
+
+
+def test_local_embedding_spans_reassemble_multimodal_rows():
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    class _IdentityProjection(nn.Module):
+        def forward(self, x):
+            return x, None
+
+    class _IdentityRefiner(nn.Module):
+        def forward(self, x, **kwargs):
+            return x
+
+    model = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.hidden_size = 2
+    model.video_patch_proj = _IdentityProjection()
+    model.audio_patch_proj = _IdentityProjection()
+    model.condition_proj = _IdentityProjection()
+    model.token_refiner = _IdentityRefiner()
+    model.time_embedder = nn.Identity()
+
+    seq_len = 16
+    x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2)
+    audio_x = x + 100
+    text = torch.tensor(
+        [[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]],
+        dtype=torch.bfloat16,
+    )
+    common = {
+        "x": x,
+        "audio_x": audio_x,
+        "text_embeddings_selected": text,
+        "unique_timesteps": torch.tensor([0.2, 0.3]),
+        "img_pos": torch.tensor([1, 2, 6, 7, 11, 12]),
+        "audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]),
+        "text_pos": torch.tensor([0, 5, 10, 15]),
+        "refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32),
+        "refiner_max_seqlen": 4,
+        "seq_len": seq_len,
+        "device": torch.device("cpu"),
+        "local_span": (0, seq_len),
+    }
+
+    full, full_t_emb = model._embed(**common)
+    local_spans = []
+    for rank in range(8):
+        local_span = (rank * (seq_len // 8), seq_len // 8)
+        local, local_t_emb = model._embed(
+            **{
+                **common,
+                "local_span": local_span,
+            }
+        )
+        assert local.shape == (2, 2)
+        torch.testing.assert_close(local_t_emb, full_t_emb)
+        local_spans.append(local)
+
+    torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "local_projection",
+    [False, True],
+    ids=["fallback", "strict-local"],
+)
+def test_sp_output_boundary_preserves_logits_and_minimizes_gather(
+    monkeypatch,
+    local_projection,
+):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    class _StaticRope(nn.Module):
+        input_rows = None
+
+        def forward(self, position_ids):
+            self.input_rows = position_ids.shape[1]
+            return torch.zeros(position_ids.shape[1], 6)
+
+    class _Prepare(nn.Module):
+        def __init__(self, shard):
+            super().__init__()
+            self.shard = shard
+
+        def forward(self, hidden, rope, combined):
+            local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0]
+            return hidden[:local_len], rope[:local_len], combined[:local_len]
+
+    class _Gather(nn.Module):
+        def __init__(self, remote_rows, replace):
+            super().__init__()
+            self.remote_rows = remote_rows
+            self.replace = replace
+            self.input_width = None
+
+        def forward(self, tensor):
+            self.input_width = tensor.shape[-1]
+            if self.replace:
+                return self.remote_rows
+            return torch.cat((tensor, self.remote_rows), dim=0)
+
+    class _FinalLayer(nn.Module):
+        def forward(self, hidden, *, t_emb, inverse_indices):
+            del t_emb
+            inverse = inverse_indices.to(torch.float32).unsqueeze(-1)
+            return (
+                hidden[:, :2].to(torch.float32) + inverse,
+                hidden[:, 2:3].to(torch.float32) - inverse,
+            )
+
+    seq_len = 8
+    full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4)
+    inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0])
+    final_layer = _FinalLayer()
+    remote_video, remote_audio = final_layer(
+        full_hidden[seq_len // 2 :],
+        t_emb=torch.empty(0),
+        inverse_indices=inverse_indices[seq_len // 2 :],
+    )
+    if local_projection:
+        remote_rows = torch.cat((remote_video, remote_audio), dim=-1)
+    else:
+        remote_rows = full_hidden
+    gather = _Gather(remote_rows, replace=not local_projection)
+
+    model = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.arch = h3.MiniMaxH3DiTArchConfig(
+        latents_dim=2,
+        audio_latents_dim=1,
+        patch_size=(1, 1, 1),
+    )
+    rope = _StaticRope()
+    model.rope = rope
+    model.sp_prepare = _Prepare(shard=True)
+    model.local_sp_prepare = _Prepare(shard=False)
+    model.sp_gather = gather
+    model.blocks = nn.ModuleList()
+    model.final_layer = final_layer
+    monkeypatch.setattr(
+        model,
+        "_embed",
+        lambda **kwargs: (
+            full_hidden[: seq_len // 2] if local_projection else full_hidden,
+            torch.tensor([[0.5]], dtype=torch.float32),
+        ),
+    )
+
+    def _local_span(*args, **kwargs):
+        del args, kwargs
+        return (0, seq_len // 2) if local_projection else (0, seq_len)
+
+    monkeypatch.setattr(
+        h3,
+        "_sequence_parallel_local_span",
+        _local_span,
+    )
+
+    img_pos = torch.tensor([0, 2, 4, 6])
+    audio_pos = torch.tensor([1, 3, 5, 7])
+    video, audio = model(
+        x=torch.zeros(1, seq_len, 2),
+        audio_x=torch.zeros(1, seq_len, 1),
+        img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long),
+        unique_timesteps=torch.tensor([0.5]),
+        inverse_indices=inverse_indices,
+        update_mask=torch.ones(img_pos.numel()),
+        token_tags=torch.zeros(seq_len, dtype=torch.long),
+        prompt_embeds=torch.empty(0, 2),
+        img_pos_info={"position_ids": img_pos},
+        audio_pos_info={"position_ids": audio_pos},
+        text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)},
+        img_pos_for_infer_output_info={"position_ids": img_pos},
+        packed_seq_params={
+            "cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32),
+            "max_seqlen_q": seq_len,
+        },
+        refiner_packed_seq_params={
+            "cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32),
+            "max_seqlen_q": 0,
+        },
+    )
+
+    expected_video, expected_audio = final_layer(
+        full_hidden,
+        t_emb=torch.empty(0),
+        inverse_indices=inverse_indices,
+    )
+    torch.testing.assert_close(video, expected_video.index_select(0, img_pos))
+    torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos))
+    expected_gather_width = 3 if local_projection else full_hidden.shape[-1]
+    expected_rope_rows = seq_len // 2 if local_projection else seq_len
+    assert gather.input_width == expected_gather_width
+    assert rope.input_rows == expected_rope_rows
+
+
@pytest.mark.parametrize(
    ("tp_size", "message"),
    [
@@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
    ]


+def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
+    from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+        get_direct_mmap_adapter,
+    )
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "Attention", _FakeAttention)
+
+    arch = h3.MiniMaxH3DiTArchConfig(
+        hidden_size=1,
+        num_attention_heads=2,
+        attention_head_dim=1,
+        rope_inv_freq_len=1,
+    )
+    attention = h3.MiniMaxH3Attention(
+        arch,
+        quant_config=None,
+        prefix="blocks.0.attn",
+    )
+    transformer = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(transformer)
+    transformer.blocks = nn.ModuleList([nn.Module()])
+    transformer.blocks[0].attn = attention
+    pipeline = nn.Module()
+    pipeline.transformer = transformer
+
+    adapter = get_direct_mmap_adapter(pipeline)
+    assert adapter is not None
+    policy = adapter.policy_for(
+        "transformer.blocks.0.attn.qkv_proj.weight",
+        attention.qkv_proj.weight,
+    )
+    assert policy is not None
+    assert policy.allow_custom_loader
+    assert policy.transform is not None
+    checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+
+    assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+    assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
+
+
 def test_pipeline_resolves_transformer_component_quant_config():
    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
        _resolve_component_quant_config,
@@ -211,3 +254,18 @@ def test_pipeline_resolves_transformer_component_quant_config():
    assert transformer_config.ignored_layers == ignored_layers
    assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
    assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
+
+
+def test_pipeline_strips_prequantized_text_encoder_config():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    prequantized = Mock()
+    prequantized.get_name.return_value = "modelopt"
+    online_fp8 = Mock()
+    online_fp8.get_name.return_value = "fp8"
+    component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
+
+    assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
+    assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
+    assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -142,12 +142,6 @@ def _generate_joint_output(omni, config: QualityTestConfig):

    first = outputs[0]
    worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
-    if hasattr(first, "request_output") and isinstance(first.request_output, list):
-        first = first.request_output[0]
-        worker_peak_memory_mb = max(
-            worker_peak_memory_mb,
-            float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
-        )
    if not hasattr(first, "images") or not first.images:
        raise ValueError("Could not extract video frames from H3 output")
    frames = first.images[0]
@@ -8,15 +8,28 @@ import torch
 from PIL import Image
 from torch import nn

-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
    Wan22I2VPipeline,
+    get_wan22_i2v_post_process_func,
    get_wan22_i2v_pre_process_func,
 )
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


+def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
+    video = torch.zeros(1, 4, 1, 2, 2)
+
+    output = get_wan22_i2v_post_process_func(SimpleNamespace())(
+        video,
+        sampling_params=SimpleNamespace(output_type="latent"),
+    )
+
+    assert output is video
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
    pipeline = object.__new__(Wan22I2VPipeline)
    nn.Module.__init__(pipeline)
@@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
    return pipeline


+def _make_i2v_sampling(**overrides):
+    values = {
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "num_inference_steps": 1,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "guidance_scale_2": None,
+        "guidance_scale_2_provided": False,
+        "boundary_ratio": None,
+        "generator": None,
+        "seed": None,
+        "num_outputs_per_prompt": 1,
+        "max_sequence_length": 8,
+        "latents": None,
+        "output_type": "latent",
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
 def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
    request = SimpleNamespace(
@@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
        preprocess(missing_image)


+def _make_i2v_preprocess_request(last_image):
+    return SimpleNamespace(
+        prompt={
+            "prompt": "p",
+            "multi_modal_data": {
+                "image": Image.new("RGB", (320, 160), "red"),
+                "last_image": last_image,
+            },
+        },
+        sampling_params=SimpleNamespace(height=16, width=16),
+    )
+
+
+def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    result = preprocess(_make_i2v_preprocess_request([]))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is None
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
+
+
+def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+    last_image = Image.new("RGB", (16, 16), "blue")
+
+    result = preprocess(_make_i2v_preprocess_request([last_image]))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is last_image
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+@pytest.mark.parametrize(
+    "last_image",
+    [
+        Image.new("RGB", (16, 16), "blue"),
+        torch.zeros(3, 16, 16),
+    ],
+)
+def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    result = preprocess(_make_i2v_preprocess_request(last_image))
+
+    assert result.prompt["multi_modal_data"]["last_image"] is last_image
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+    path = tmp_path / "last.png"
+    Image.new("RGB", (16, 16), "blue").save(path)
+
+    result = preprocess(_make_i2v_preprocess_request(str(path)))
+
+    last_image = result.prompt["multi_modal_data"]["last_image"]
+    assert isinstance(last_image, Image.Image)
+    assert last_image.mode == "RGB"
+    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_rejects_multiple_last_images() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    with pytest.raises(ValueError, match="at most one last_image"):
+        preprocess(
+            _make_i2v_preprocess_request(
+                [
+                    Image.new("RGB", (16, 16), "blue"),
+                    Image.new("RGB", (16, 16), "green"),
+                ]
+            )
+        )
+
+
+def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
+    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+    with pytest.raises(TypeError, match="Unsupported last_image format"):
+        preprocess(_make_i2v_preprocess_request({"not": "an image"}))
+
+
 def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
    pipeline = _make_i2v_pipeline(expand_timesteps=True)
    latents = torch.zeros(1, 4, 2, 4, 4)
@@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
    assert first_frame_mask.shape == (1, 1, 2, 2, 2)
    assert first_frame_mask[:, :, 0].sum() == 0
    assert first_frame_mask[:, :, 1].sum() == 4
+
+
+def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    generators = [
+        torch.Generator(device="cpu").manual_seed(1),
+        torch.Generator(device="cpu").manual_seed(2),
+    ]
+
+    latents, condition, first_frame_mask = pipeline.prepare_latents(
+        image=torch.zeros(2, 3, 16, 16),
+        batch_size=2,
+        num_channels_latents=4,
+        height=16,
+        width=16,
+        num_frames=5,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=generators,
+    )
+
+    assert latents.shape == (2, 4, 2, 2, 2)
+    assert condition.shape == (2, 4, 1, 2, 2)
+    assert first_frame_mask.shape == (2, 1, 2, 2, 2)
+
+
+def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 5.0
+    pipeline.boundary_ratio = 0.875
+    pipeline.has_image_encoder = False
+    pipeline._guidance_scale = None
+    pipeline._guidance_scale_2 = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.zeros(n, 2, 3), None
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        batch_size = kwargs["batch_size"]
+        return (
+            kwargs["latents"],
+            torch.zeros(batch_size, 4, 1, 2, 2),
+            torch.ones(batch_size, 1, 2, 2, 2),
+        )
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 2, 2, 2)
+    latents_b = torch.ones(2, 4, 2, 2, 2)
+    image_a = torch.zeros(1, 3, 16, 16)
+    image_b = torch.ones(1, 3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
+                sampling_params=_make_i2v_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
+                sampling_params=_make_i2v_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    torch.testing.assert_close(
+        prepare_call["image"],
+        torch.cat([image_a, image_a, image_b, image_b]),
+    )
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 5.0
+    pipeline.boundary_ratio = 0.875
+    pipeline.has_image_encoder = False
+    pipeline._guidance_scale = None
+    pipeline._guidance_scale_2 = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="image condition"):
+        pipeline.forward(batch)
+
+
+def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
+    pipeline = _make_i2v_pipeline(expand_timesteps=True)
+    image = torch.zeros(1, 3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "last_image": image},
+                },
+                sampling_params=_make_i2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image}},
+                sampling_params=_make_i2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="mix of provided and missing last_image"):
+        pipeline.forward(batch)
@@ -93,6 +93,29 @@ def _make_pipeline() -> Wan22Pipeline:
    return pipeline


+def _make_sampling(**overrides):
+    values = {
+        "height": None,
+        "width": None,
+        "num_frames": 1,
+        "num_inference_steps": 2,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "guidance_scale_2": None,
+        "guidance_scale_2_provided": False,
+        "boundary_ratio": None,
+        "generator": None,
+        "seed": None,
+        "num_outputs_per_prompt": 1,
+        "max_sequence_length": 32,
+        "latents": None,
+        "output_type": "latent",
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
@pytest.mark.parametrize(
    ("sampling_params_kwargs", "expected_low", "expected_high"),
    [
@@ -129,9 +152,10 @@ def test_forward_delegates_denoising_to_diffuse(
    )
    batch = DiffusionRequestBatch(requests=[mock_req])

-    output = pipeline.forward(batch)
+    outputs = pipeline.forward(batch)

-    assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert len(outputs) == 1
+    assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
    assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
    assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
    assert captured["guidance_low"] == expected_low
@@ -142,6 +166,134 @@ def test_forward_delegates_denoising_to_diffuse(
    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]


+def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
+    pipeline = _make_pipeline()
+    encode_call = {}
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        encode_call.update(kwargs)
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 1, 2, 2)
+    latents_b = torch.ones(2, 4, 1, 2, 2)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "negative_prompt": "bad first"},
+                sampling_params=_make_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "negative_prompt": "bad second"},
+                sampling_params=_make_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert encode_call["prompt"] == ["first", "second"]
+    assert encode_call["negative_prompt"] == ["bad first", "bad second"]
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_forward_batches_precomputed_prompt_embeddings() -> None:
+    pipeline = _make_pipeline()
+    diffuse_call = {}
+    pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run")  # type: ignore[method-assign]
+    pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2)  # type: ignore[method-assign]
+
+    def _fake_diffuse(**kwargs):
+        diffuse_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
+    embeds_a = torch.zeros(3, 4)
+    embeds_b = torch.ones(3, 4)
+    negative_a = torch.full((3, 4), 2.0)
+    negative_b = torch.full((3, 4), 3.0)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
+                sampling_params=_make_sampling(guidance_scale=4.0),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
+                sampling_params=_make_sampling(guidance_scale=4.0),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
+    torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
+    assert len(outputs) == 2
+
+
+def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
+    pipeline = _make_pipeline()
+    kwargs = {
+        "num_channels_latents": 4,
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "dtype": torch.float32,
+        "device": torch.device("cpu"),
+    }
+
+    batched = Wan22Pipeline.prepare_latents(
+        pipeline,
+        batch_size=2,
+        generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
+        **kwargs,
+    )
+    singles = torch.cat(
+        [
+            Wan22Pipeline.prepare_latents(
+                pipeline,
+                batch_size=1,
+                generator=torch.Generator().manual_seed(seed),
+                **kwargs,
+            )
+            for seed in (1, 2)
+        ]
+    )
+
+    torch.testing.assert_close(batched, singles)
+    assert not torch.equal(batched[0], batched[1])
+
+
 def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
    pipeline = _make_pipeline()
    latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
@@ -10,6 +10,7 @@ import torch
 import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
    create_transformer_from_config,
+    get_wan22_post_process_func,
    load_transformer_config,
    retrieve_latents,
 )
@@ -26,6 +27,17 @@ class _LatentDist:
        return torch.tensor([2.0])


+def test_wan22_postprocess_honors_request_output_type() -> None:
+    video = torch.zeros(1, 4, 1, 2, 2)
+
+    output = get_wan22_post_process_func(SimpleNamespace())(
+        video,
+        sampling_params=SimpleNamespace(output_type="latent"),
+    )
+
+    assert output is video
+
+
 def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
    generator = torch.Generator(device="cpu")

@@ -1,14 +1,307 @@
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

+import numpy as np
+import PIL.Image
 import pytest
 import torch
 import torch.nn as nn

+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
+    Wan22S2VPipeline,
+    _make_clip_generators,
+)
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


+def _make_s2v_sampling(**overrides):
+    values = {
+        "height": 16,
+        "width": 16,
+        "num_frames": 8,
+        "num_inference_steps": 1,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "generator": None,
+        "seed": None,
+        "num_outputs_per_prompt": 1,
+        "max_sequence_length": 8,
+        "latents": None,
+        "output_type": "latent",
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
+    pipeline = object.__new__(Wan22S2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(dtype=torch.float32)
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.resolution_divisor = 16
+    pipeline.motion_frames = 7
+    pipeline.drop_first_motion = True
+    pipeline._DEFAULT_INFER_FRAMES = 8
+    pipeline._guidance_scale = None
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+    return pipeline
+
+
+def test_s2v_predict_noise_keeps_batch_dimension() -> None:
+    pipeline = object.__new__(Wan22S2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    transformer = MagicMock()
+    transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
+    transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
+    pipeline.transformer = transformer
+
+    result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))
+
+    assert result.shape == (1, 16, 2, 2, 2)
+
+
+def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
+    explicit_generator = torch.Generator(device="cpu").manual_seed(999)
+
+    generators = _make_clip_generators(
+        seeds=[1234, 5678, None, None],
+        request_generators=[
+            torch.Generator(device="cpu").manual_seed(1234),
+            torch.Generator(device="cpu").manual_seed(5678),
+            explicit_generator,
+            None,
+        ],
+        clip_index=2,
+        device=torch.device("cpu"),
+    )
+
+    assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
+    assert generators[2] is explicit_generator
+
+
+def test_s2v_forward_rejects_different_num_repeat() -> None:
+    pipeline = _make_s2v_validation_pipeline()
+    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16)  # type: ignore[method-assign]
+    image = PIL.Image.new("RGB", (16, 16))
+    audio = np.zeros(16000, dtype=np.float32)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "audio": audio},
+                    "additional_information": {"num_repeat": 1},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "multi_modal_data": {"image": image, "audio": audio},
+                    "additional_information": {"num_repeat": 2},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="num_repeat"):
+        pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
+    pipeline = _make_s2v_validation_pipeline()
+    pipeline.encode_audio = lambda audio, **kwargs: (  # type: ignore[method-assign]
+        torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
+        1,
+        16 if audio[0] == 0 else 24,
+    )
+    image = PIL.Image.new("RGB", (16, 16))
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="audio embedding time lengths"):
+        pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_init_first_frame() -> None:
+    pipeline = _make_s2v_validation_pipeline()
+    image = PIL.Image.new("RGB", (16, 16))
+    audio = np.zeros(16000, dtype=np.float32)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "audio": audio},
+                    "additional_information": {"init_first_frame": True},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "multi_modal_data": {"image": image, "audio": audio},
+                    "additional_information": {"init_first_frame": False},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="init_first_frame"):
+        pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
+    pipeline = _make_s2v_validation_pipeline()
+    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16)  # type: ignore[method-assign]
+    image = PIL.Image.new("RGB", (16, 16))
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
+                },
+                sampling_params=_make_s2v_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
+        pipeline.forward(batch)
+
+
+def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
+    pipeline = object.__new__(Wan22S2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = MagicMock()
+    pipeline.transformer.dtype = torch.float32
+    pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
+    pipeline.transformer.casual_audio_encoder = None
+    pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
+    pipeline.vae = MagicMock()
+    pipeline.vae.dtype = torch.float32
+    pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.resolution_divisor = 16
+    pipeline.motion_frames = 7
+    pipeline.drop_first_motion = True
+    pipeline._DEFAULT_INFER_FRAMES = 8
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
+    pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
+    pipeline.prepare_motion_latents = MagicMock(
+        side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
+    )
+    pipeline._denormalize_latents = lambda latents: latents
+    pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
+
+    image = PIL.Image.new("RGB", (16, 16))
+    audio_a = np.zeros(16000, dtype=np.float32)
+    audio_b = np.ones(16000, dtype=np.float32)
+    latents_a = torch.zeros(2, 16, 2, 2, 2)
+    latents_b = torch.ones(2, 16, 2, 2, 2)
+    generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
+    generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "negative_prompt": "negative first",
+                    "multi_modal_data": {"image": image, "audio": audio_a},
+                },
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    generator=generators_a,
+                    latents=latents_a,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "negative_prompt": "negative second",
+                    "multi_modal_data": {"image": image, "audio": audio_b},
+                },
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    generator=generators_b,
+                    latents=latents_b,
+                ),
+            ),
+        ]
+    )
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 2
+    assert outputs[0].output[0].shape[0] == 2
+    assert outputs[1].output[0].shape[0] == 2
+    assert outputs[0].output[1].shape == (16000,)
+    assert outputs[1].output[1].shape == (16000,)
+    np.testing.assert_array_equal(outputs[0].output[1], audio_a)
+    np.testing.assert_array_equal(outputs[1].output[1], audio_b)
+    assert np.shares_memory(outputs[0].output[1], audio_a)
+    assert np.shares_memory(outputs[1].output[1], audio_b)
+    torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
+    assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
+    assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
+    assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]
+
+
 def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
@@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
    pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean

    # -- Build request --
-    from vllm_omni.diffusion.request import OmniDiffusionRequest
-
-    sampling_params = MagicMock()
-    sampling_params.height = 704
-    sampling_params.width = 1024
-    sampling_params.num_frames = 80
-    sampling_params.num_inference_steps = 5
-    sampling_params.guidance_scale = 4.5
-    sampling_params.guidance_scale_provided = True
-    sampling_params.generator = None
-    sampling_params.seed = 42
+    sampling_params = _make_s2v_sampling(
+        height=704,
+        width=1024,
+        num_frames=80,
+        num_inference_steps=5,
+        guidance_scale=4.5,
+        generator=torch.Generator(device="cpu").manual_seed(42),
+        max_sequence_length=512,
+    )

    ref_image = PIL.Image.new("RGB", (1024, 704))
    audio_data = np.zeros(16000, dtype=np.float32)

-    req = MagicMock(spec=OmniDiffusionRequest)
-    req.prompts = [
-        {
-            "prompt": "test prompt",
-            "negative_prompt": "bad quality",
-            "multi_modal_data": {"image": ref_image, "audio": audio_data},
-            "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
-        }
-    ]
-    req.sampling_params = sampling_params
+    req = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="test",
+                prompt={
+                    "prompt": "test prompt",
+                    "negative_prompt": "bad quality",
+                    "multi_modal_data": {"image": ref_image, "audio": audio_data},
+                    "additional_information": {
+                        "audio_path": audio_data,
+                        "pose_video": None,
+                        "init_first_frame": False,
+                    },
+                },
+                sampling_params=sampling_params,
+            )
+        ]
+    )

    # -- Mock methods that use complex internal state --
    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
@@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
    pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)

    # Mock predict_noise_maybe_with_cfg from CFGParallelMixin
-    pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
+    pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
    pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)

    # Mock platform methods
@@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
        with patch(
            "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
        ):
-            result = Wan22S2VPipeline.forward(pipeline, req=req)
+            result = Wan22S2VPipeline.forward(pipeline, req=req)[0]

    # -- Assertions --
    # Text encoder was called
@@ -8,13 +8,14 @@ import torch
 from PIL import Image
 from torch import nn

-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
    Wan22VACEPipeline,
    create_vace_transformer_from_config,
    get_wan22_vace_pre_process_func,
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]

@@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
    return pipeline


+def _make_vace_sampling(**overrides):
+    values = {
+        "height": 16,
+        "width": 16,
+        "num_frames": 5,
+        "num_inference_steps": 1,
+        "guidance_scale_provided": True,
+        "guidance_scale": 1.0,
+        "boundary_ratio": None,
+        "generator": None,
+        "seed": None,
+        "num_outputs_per_prompt": 1,
+        "max_sequence_length": 8,
+        "latents": None,
+        "output_type": "latent",
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
 def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
    preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
    ref = Image.new("RGB", (320, 160), "green")
@@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
    assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
    assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
    torch.testing.assert_close(result, torch.ones_like(latents))
+
+
+def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = None
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    prepare_call = {}
+
+    def _fake_encode_prompt(**kwargs):
+        batch_size = len(kwargs["prompt"])
+        n = batch_size * kwargs["num_videos_per_prompt"]
+        return torch.zeros(n, 2, 3), None
+
+    def _fake_prepare_latents(**kwargs):
+        prepare_call.update(kwargs)
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
+    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
+    gen_a = torch.Generator(device="cpu").manual_seed(1)
+    gen_b = torch.Generator(device="cpu").manual_seed(2)
+    latents_a = torch.zeros(2, 4, 2, 2, 2)
+    latents_b = torch.ones(2, 4, 2, 2, 2)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first"},
+                sampling_params=_make_vace_sampling(
+                    generator=gen_a,
+                    latents=latents_a,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second"},
+                sampling_params=_make_vace_sampling(
+                    generator=gen_b,
+                    latents=latents_b,
+                    num_outputs_per_prompt=2,
+                ),
+            ),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert prepare_call["batch_size"] == 4
+    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output, latents_a)
+    torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = object()
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+
+    def _fake_preprocess(video, **kwargs):
+        del kwargs
+        mask = torch.ones_like(video)
+        return video, mask, [[]]
+
+    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
+    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
+        lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
+    )
+    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2)  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="VACE condition"):
+        pipeline.forward(batch)
+
+
+def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+        SimpleNamespace(is_available=lambda: False),
+    )
+    pipeline = _make_vace_pipeline()
+    pipeline.transformer.vace_patch_embedding = object()
+    pipeline.transformer_2 = None
+    pipeline.scheduler = StubScheduler([9])
+    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+    pipeline._sample_solver = "unipc"
+    pipeline._flow_shift = 3.0
+    pipeline.boundary_ratio = None
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline._current_timestep = None
+    pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
+
+    def _fake_preprocess(video, mask, reference_images, **kwargs):
+        del kwargs
+        return video, mask, [reference_images]
+
+    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
+    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
+        lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
+    )
+    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2)  # type: ignore[method-assign]
+    video = torch.zeros(1, 3, 5, 16, 16)
+    mask = torch.ones_like(video)
+    reference = torch.zeros(3, 16, 16)
+    batch = DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={
+                    "prompt": "first",
+                    "additional_information": {
+                        "source_video": video,
+                        "mask": mask,
+                        "reference_images": [reference],
+                    },
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={
+                    "prompt": "second",
+                    "additional_information": {
+                        "source_video": video,
+                        "mask": mask,
+                        "reference_images": [reference, reference],
+                    },
+                },
+                sampling_params=_make_vace_sampling(),
+            ),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="same number of reference images"):
+        pipeline.forward(batch)
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]
@@ -79,11 +79,11 @@ def dist_group():


@pytest.fixture
-def patched_offload_runtime(mocker):
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+def patched_offload_runtime(monkeypatch):
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)


 class TinyBlock(nn.Module):
@@ -128,6 +128,33 @@ class TestLayerwiseOffloadHook:
        assert current_block.weight.to_local().shape == torch.Size([4])
        assert not hook.is_materialized

+    def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
+        """Online-FP8 Cutlass weights must retain their transposed layout."""
+
+        class StridedBlock(nn.Module):
+            def __init__(self, start: float):
+                super().__init__()
+                base = torch.arange(start, start + 12).reshape(3, 4)
+                self.weight = nn.Parameter(base.t(), requires_grad=False)
+
+        current_block = StridedBlock(1.0)
+        next_block = StridedBlock(20.0)
+        expected = next_block.weight.detach().clone()
+        expected_stride = next_block.weight.stride()
+        assert expected_stride == (1, 4)
+
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+        hook.prefetch_layer(non_blocking=False)
+
+        assert next_block.weight.stride() == expected_stride
+        assert torch.equal(next_block.weight, expected)
+

 class _DummyBlock(nn.Module):
    def __init__(self):
@@ -16,10 +16,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
 )
 from vllm_omni.quantization.component_config import (
    ComponentQuantizationConfig,
+    resolve_component_quant_config,
 )
 from vllm_omni.quantization.inc_config import OmniINCConfig

-pytestmark = [pytest.mark.core_model]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


 # ---------------------------------------------------------------------------
@@ -347,6 +348,24 @@ class TestTalkerVisualRouting:
 # ===================================================================


+class TestResolveComponentQuantConfig:
+    def test_global_config_applies_to_each_quantization_aware_component(self):
+        config = _MockQuantConfig("fp8")
+
+        assert resolve_component_quant_config(config, "transformer") is config
+        assert resolve_component_quant_config(config, "text_encoder") is config
+
+    def test_component_map_is_the_only_scope_override(self):
+        transformer = _MockQuantConfig("fp8")
+        config = ComponentQuantizationConfig(
+            component_configs={"transformer": transformer},
+            default_config=None,
+        )
+
+        assert resolve_component_quant_config(config, "transformer") is transformer
+        assert resolve_component_quant_config(config, "text_encoder") is None
+
+
 class TestComponentResolve:
    def test_longest_prefix_match(self):
        a = _MockQuantConfig("a")
Author	SHA1	Message	Date
Hongsheng Liu	f1abb13958	[Diffusion] Register runtime cache for direct DLO H2D Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup. Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-17 11:50:32 +08:00
Hongsheng Liu	b732ad2cea	[Diffusion] Document runtime cache hardware tradeoffs Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-17 00:43:06 +08:00
Hongsheng Liu	99347e1ba7	[Diffusion] Tighten runtime cache compatibility guards Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 23:26:25 +08:00
Hongsheng Liu	db86a6a078	[Diffusion] Add node-local runtime weight cache for DLO Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 23:15:48 +08:00
Yukim1andHongsheng Liu	baba7d1e2e	[Refactor] Remove legacy stage_args YAML loader (#6200 ) Signed-off-by: zwhzzz0821 <2831474076@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 11:22:28 +08:00
Hongsheng Liu	9d2bb23ff6	[Diffusion] Add loader-owned host-weight plans for DLO (TP=1) (#6213 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 11:02:38 +08:00
Jack Chen zhangbin.327 Hongsheng Liu	3875f8ded4	[Core] Add realtime AR-Diffusion tick sessions for LingBot World 2.0 (#5491 ) Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com> Signed-off-by: Jack47 <scsvip@gmail.com> Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 09:11:20 +08:00
Yueqian Lin	9798bc3a8d	[Docs] Unify recipe serve commands on `vllm serve --omni` (#6221 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-08-15 21:18:47 +00:00
Anurag	68bf8c2c42	[Bugfix] Prevent DiffusionResultPump crash on cancelled futures (#5793 ) (#5983 ) Signed-off-by: ANURAG KANADE <anuragkanade6@gmail.com>	2026-08-15 16:19:45 +00:00
yiminghub2024andHongsheng Liu	444485650b	Add MiniMax-H3 recipe for NPU 950PR (#6120 ) Signed-off-by: yiminghub2024 <482890@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-16 00:04:38 +08:00
hurukawa	2a315e1a2d	[Perf] Support request-level batching for Wan2.2 pipelines (#5676 ) Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com>	2026-08-15 23:58:36 +08:00
SYLARandHongsheng Liu	d1e230c95b	[Diffusion][Quantization] Enable MiniMax-H3 global FP8 with DLO (#5910 ) Signed-off-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-15 23:01:26 +08:00
m0g3r	09b6a1604a	[Bugfix] Restore per-stage runtime env during launch (#6214 ) Signed-off-by: m0g3r <87276771+m0g3r@users.noreply.github.com>	2026-08-15 13:00:27 +00:00
Yueqian Lin	555336e887	[CI] Refresh Voxtral-4B-TTS perf baselines from a clean window (#6032 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-08-15 11:51:56 +00:00
Samit	136cb27b05	[Bugfix] Fix lost async diffusion outputs in request-level batching (#6023 ) Signed-off-by: samithuang <285365963@qq.com>	2026-08-15 18:27:51 +08:00
SYLAR	ce6ab0ca46	[Model] Add LTX-2.5 support (#6070 ) Signed-off-by: lishunyang12 <lishunyang12@163.com>	2026-08-15 18:27:23 +08:00
NumberWan	f39b8684eb	[BugFix][Nightly CI] Adjust Qwen-Image accuracy thresholds for stable FA-deterministic mode (#5963 ) Signed-off-by: NumberWan <wantszkin2003@gmail.com>	2026-08-15 09:38:30 +00:00
MrlixiangWE	e57574ae29	[Bugfix] Carry ec_transfer_params and num_cache_creation_tokens on OmniRequestOutput (#6152 ) Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>	2026-08-15 08:37:27 +00:00
mokeke mokashliu WeiQing Chen	7b76b6446d	[Model][Performance] Optimize MiniMax-H3 strict Ulysses boundaries (#6173 ) Signed-off-by: mokashliu <mokashliu@tencent.com> Co-authored-by: mokashliu <mokashliu@tencent.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>	2026-08-15 16:27:52 +08:00
hotTea jiangweifeng.dev Hongsheng Liu	a70d4d4b3c	[codex] Fix diffusion worker timeout for large broadcast payloads (#4845 ) Signed-off-by: jiangweifeng.dev <jiangweifeng.dev@bytedance.com> Co-authored-by: jiangweifeng.dev <jiangweifeng.dev@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-15 16:03:04 +08:00
Hongsheng Liu	1260ceaae6	[Docs] Split diffusion attention and CPU offload guides (#6075 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-15 14:41:09 +08:00
Ma1oneZhang	6f36fd3b78	[Bugfix] Preserve object storage URIs during model resolution (#5036 ) Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>	2026-08-15 14:11:48 +08:00
mjZhaoElaine Claude Opus 5 Yueqian Lin	35cacf020c	[Model] Add Gepard-1.0 native-AR FSQ/NanoCodec TTS (offline inference) (#5666 ) Signed-off-by: Mengjie Zhao <zmj0129@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-08-14 22:06:05 -07:00
wangyu	d671e6b0db	[CI] Expand E2E source_file_dependencies for shared model code (#5994 ) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>	2026-08-15 12:01:20 +08:00
MrlixiangWE	b95c73742d	[Bugfix] Stop reading the removed OmniRequestOutput.request_output accessor (#6172 ) Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>	2026-08-15 07:07:28 +08:00
Alex Brooks	fe54ecc388	[Frontend] Loosen Validation for TTS Models with No Uploaded Speakers (#5878 ) Signed-off-by: Alex Brooks <albrooks@redhat.com>	2026-08-14 17:05:11 -06:00